[Agent Memory / 强化学习] MemPO源码学习笔记 --- (7)--- 混合机制
0x00 概要
现有的基于强化学习的 Memory 管理方法往往缺乏一种有效机制针对 Memory 的更新内容进行引导优化，Memory 的内容难以保证质量。
MemPO(Self-Memory Policy Optimization)使模型对 Memory 进行自管理，并引入了基于有效信息含量的 Memory-level 的优势估计，引导 Memory 保留对解决任务更有效的信息，进而提升记忆有效性。
MemPO的独特切入点：让模型把记忆写在每轮开头()，形式上像"自我对话的草稿纸"，既是记忆又是思考链的一部分。这样，变成可训练的策略变量，用RL信号端到端地教会模型"什么值得记、怎么记”。RL 直接端到端优化这一行为，无需额外的记忆模块。
MemPO 的信息如下：
- 论文标题：MemPO: Self-Memory Policy Optimization for Long-Horizon Agents
- 论文地址：arxiv.org/abs/2603.00…
- 代码地址：github.com/TheNewBeeKi…
- 模型和数据集地址：huggingface.co/collections…
0x01 Mem Reward 机制
1.1 公式
Memory Advantage 公式(分两步)如下：
Step 1 - 计算 mem_reward (A1):mem_reward_k = P_mem_k - P_full_k
其中:
P_full_k = exp (mean (full_logp [k] × mask_k) ) ← 完整上下文→答案概率
P_mem_k = exp ( mean (mem_logp [k] × mask_k) ) ← 仅摘要→答案概率
k = 某条轨迹的某一轮
Step 2 - 归一化 (A2):mem_adv_k = (mem_reward_k - pool_mean) / (pool_std + ε)
其中:
pool = 同一 question 所有 16 条轨迹 × 所有轮次的 mem_reward(约 48 个值: 16 条 × 3 轮 / 条)
pool_mean = mean (所有 48 个值)
pool_std = std (所有 48 个值)
Step 1 含义拆解:
P_mem - P_full(Step 1): "只看笔记能答对的概率－看完整历史能答对的概率"
> 0：笔记比原文更聚焦 → 好笔记 → 正奖励
= 0：笔记无损压缩 → 中性
< 0：笔记丢了信息 → 差笔记 → 负奖励
直觉：如果你把一本书浓缩成一页笔记后，靠笔记答题，比翻全书还准 → 说明你抓住了重点
Step 2 含义拆解:
(mem_reward-pool_mean)/ pool_std (Step 2):
"我这轮的笔记质量比大家的平均水平好多少？"
pool包含同一question下所有轨迹所有轮次的mem_reward
→ 跨轨迹：不同搜索路径写出的笔记互相比较
→ 跨轮次：第2轮笔记和第4轮笔记也互相比较
→ 归一化：产生标准化的正负信号
写入位置：mem_adv[i，start_t:end_t+1]=adv
"只在<mem>...</mem>这几个token的位置赋值"
→ 其余位置 = 0
→ 只有"写笔记"的token受到笔记质量的梯度信号
1.2 P_mem 和 P_full 的完整计算细节
第一步：构建三类输入(tool_agent_loop.py)
第二步：ans_mask的构造(精确标记答案token)
第三步：一次性批量计算log_prob
# 将full_traj和mem_traj 合并成一个大 batch (减少 GPu 调用次数)
traj_input = tokenizer.pad(
{"input_ids": concat_full_traj_ids_list + concat_mem_traj_ids_list},
return_tensors="pt",
padding="longest",
return_attention_mask=True,
padding_side="left",
)
answer_input = tokenizer.pad(
{"input_ids": concat_answer_ids_list},
return_tensors="pt",
padding="longest",
return_attention_mask=True,
padding_side="left",
)
ans_mask_tensor = pad_sequence([torch.tensor(sub) for sub in concat_ans_mask_list], batch_first=True, padding_value=0, padding_side="left")
# responses(要计算 log_prob 的 token) = answer tokens，repeat 2 次对齐 full+mem
traj_batch = DataProto.from_dict(
{
"responses": answer_input["input_ids"].repeat(2, 1), #[full组+mem组】×answer
"input_ids": traj_input["input_ids"],
"attention_mask": traj_input["attention_mask"],
"position_ids": (traj_input["attention_mask"].cumsum(dim=1) - 1).clamp(min=0),
},
)
traj_log_probs = self.worker_group.compute_log_prob(traj_batch)
# 返回每个answer token在对应上下文下的 log P(token|context)
第四步：threshold 过滤(核心工程细节)，仅保留模型置信度>50%的token
第五步：计算概率并得出mem_reward
# P_full：完整上下文下，模型对答案的几何平均概率
P_full = exp(
sum(logp[full_ans_mask]) / (count(full_ans_mask) + 1e-8)
)
# P_mem：仅凭<mem>上下文，模型对答案的几何平均概率
P_mem = exp(
sum(logp[mem_ans_mask]) / (count(mem_ans_mask) + 1e-8)
)
#记忆奖励="记忆版"预测概率-"完整版"预测概率
mem_reward=P_mem -P_full #范围大致[-1，+1]
可视化一个具体例子。
题目：爱因斯坦出生城市的市长是谁？
第1轮<mem>：
"问题1：爱因斯坦出生于乌尔姆。问题2：市长未知，需继续查询。"
全文上下文 → P_full(answer="乌尔姆；某某某")=0.35
仅<mem>上下文 → P_mem(answer="乌尔姆;某某某")=0.12
mem_reward[0]=0.12 - 0.35=-0.23 ← 记忆不完整，扣分
第2轮<mem>：
"乌尔姆市长为张三，两题已解答：乌尔姆；张三。"
全文上下文 → P_full=0.68
仅<mem>上下文 → P_mem=0.71(<mem>比全文更精炼！)
mem_reward[1]=0.71-0.68=+0.03 ← 记忆充分，轻微加分
1.3 为什么不是只用P_mem？
这是一个很关键的设计决策，背后的逻辑是消除"答案本身难易度"的影响。即，P_full是基线(baseline)，把"这道题本来就难／容易"的因素剥离掉，让奖励信号只反映记忆本身的质量，而非题目难易度。这和GRPO中的"组内均值"作为基线是同一种思想一用相对分而非绝对分来评估。
1.3.1 问题：如果只用P_mem
问题A(简单)："法国首都是哪里？
答案Z="Paris"
P_mem("Paris"|<mem>法国首都是Paris</mem>)=0.95 ← 高
问题B(困难)："2018年中美贸易战中美国征收关税的具体税率是多少？
答案Z="25%"
P_mem("25%"|<mem>贸易战关税25%</mem>)=0.65 ← 低
用P_mem打分：
模型A:简单题 0.95 → 奖励高
模型B:难题 0.65 → 奖励低
但这不公平！难题即使写得很好，P_mem 也低 → 模型会学到：回答简单题比回答难题更"合算" → 产生偏好简单问题的偏差。
1.3.2 解决方案：用差值P_mem-P_full作为基线校正
P_full = P(Z|完整上下文) 这是"上界"，最好情况下答案的概率
P_mem - P_full = "仅靠记忆的能力”相对于"看完整上下文的能力”的差距 = 记忆造成的信息损失(负数表示损失，越接近0越好)
如下问题可以看出来差值P_mem-P_full的作用。
问题A(简单)：
P_full=0.95(给完整历史，答Paris很容易)
P_mem=0.93(仅凭记忆，答Paris也不难)
差值 = -0.02 → 记忆几乎没有损失，轻微负奖励
问题B(困难)：
P_full = 0.70(给完整历史，答25%还是有点不确定)
P_mem = 0.65(仅凭记忆，稍微差一点)
差值 = -0.05 → 记忆损失较小，负奖励稍大
问题B(记忆写得很差)：
P_full = 0.70(完整历史还能猜到)
P_mem = 0.20(只看记忆完全猜不到)
差值 = -0.50 → 记忆严重不足，大负奖励
1.3.3 直觉类比
就像考试打分：
单独用P_mem(绝对分)：
简单题答90分 → 奖励高
难题答65分 → 奖励低 → 难题永远得不到高奖励，模型学偏了
用P_mem - P_full(相对分)：
简单题：能力上限95分，你答了 93分 → 差2分(优秀)
难题：能力上限70分，你答了 65分 → 5分(良好)差
难题：能力上限70分，你答了 20分 → 差50分(记忆太差！)
→ 评估的是"相对于你的能力，记忆保留了多少"
1.4 P_mem - P_full 的含义
核心问题：摘要是否「足够好」，好到能替代完整的多轮对话上下文？
由这个核心问题开始，我们来看看P_mem - P_full的直觉含义。
1.4.1 直接含义
轨迹如下
- full_traj = 当前轮生成前的完整上下文(system + question + 之前所有轮次的完整对话)，长度递增
- mem_traj = mem_sys_prompt_ids(system + question) + 当前轮摘要内容的内容，非常短
两者都被要求预测同一个正确答案：
- P_full = 模型基于完整上下文预测出正确答案的概率
- P_mem = 模型仅基于摘要预测出正确答案的概率
三种情况：
| P_mem vs P_full | 含义 | mem_reward |
|---|---|---|
| P_mem > P_full | 摘要比原文更聚焦，模型反而更容易答对 | 正值 √ |
| P_mem ≈ P_full | 摘要无损压缩了关键信息 | ≈ 0 |
| P_mem < P_full | 摘要丢失了重要信息，模型答不出来了 | 负值 X |
因此，直觉含义如下：
- 如果你看了一本书的完整内容 VS 只看了别人写的摘要笔记，谁能更好地回答问题？
- 如果笔记质量很高，可能比翻全书更快更准确找到答案一这就 是P_mem>P_full 的含义。
1.4.2 更精确的数学意义
P_mem - P_full ≈ 衡量"记忆造成的信息损失"
- 如果P_mem ≈ P_fu1l→差值 ≈ 0 → 包含了回答问题所需的全部信息 → 奖励为0(中性，已经够好了)
- 如果P_mem << P_full→差值很负 → 缺失了重要信息→ 惩罚 token，促使模型写更好的记忆
- 如果P_mem > P_full → 差值为正(少见) → 比完整上下文更精准地提示了答案 → 奖励token(记忆甚至做了有用的"压缩摘要")
1.4.3 信用分配
通过把写记忆纳入GRPO的信用分配链路，MemPO驱动模型学习哪些信息值得在中保留(记忆奖励正向鼓励)，哪些无关信息应当省略(负向惩罚)。让模型知道哪些中间信息值得占用上下文预算、哪些噪声应该被主动丢弃。这种机制从根本上缓解了长程多轮交互中上下文爆炸的问题一一一模型被迫学会在有限的空间内压缩本轮关键信息，而非依赖无限增长的完整历史。
然而，MemPO目前优化的是"信息充分性"而非"压缩效率”，且跨轮次的长程信用分配仍然依赖最终结果奖励来间接传导，而非对每轮记忆进行前的多步评估。"
信用分配链路
比如：
标准 RL: 奖励 r → 反传给 /<action> token
MemPO: 奖励 r + mem_reward → 反传给 <mem> token
└─ mem_reward = P_mem - P_full
明确衡量"这段记忆的信息价值"
的生成行为确实进入了信用分配链路，不再是"附带生成的文本”。
有效信息
如何让模型知道哪些信息值得占用上下文预算”？
- 实现层面：MemPO没有显式建模"上下文预算"(token长度限制)。它的记忆奖励P_mem-P_full 只衡量"信息充分性”，不衡量"压缩效率”。
- 实际效果：如果写得很长但信息充分，P_mem高→mem_reward 正→仍然被鼓励。MemPO 鼓励的是**"信息留存”，而非"高效压缩"**。真正的"上下文预算感知"需要额外的长度惩罚项，当前代码中没有这一项。
丢弃噪声
"哪些噪声应该被主动丢弃”一实际上是隐式实现的机制：
轮次t的<mem>写了无关信息：
mem_traj[t]=[系统提示+<mem>无关内容</mem>]
P_mem预测答案的概率很低
mem_reward[t]=P_mem-P_fu1l < 0 ← 负奖励
→ 这段<mem>的token获得负梯度
→ 模型逐渐学会"写有关信息，不写无关噪声"
但这里的"丢弃"是通过负梯度间接引导，而不是模型显式地判断"这个信息不重要”。区别在于：
- MemPO的记忆行为是被动地被奖励/惩罚
- 真正的"主动丢弃决策"需要像MEM1那样有显式的keep/discard动作空间
长程交互难点
解决长程交互难点体现在：
- 上下文长度爆炸：AgentMemory.prepare_prompt()只保留最近一轮工具结果，强制模型靠传递信息
- 信用分配到中间行为(写记忆这一动作)
未完全解决的
一跨轮次的长程信用分配：第1轮写的 对第6轮答案的贡献，无法被当前 P_mem-P_full 精确衡量(只看该轮的即时答案预测能力)-训练中rollout轮次≤8，但推理时外推到更长轮次是否稳定是开放问题
1.4.4 场景3
我们接下来看看场景3 "答错但记忆好”的设计意图
答错 + 好mem → (-outcome_adv)+(+mem_adv)
<mem>区间：-outcome+(+mem_adv) → 惩罚被部分抵消
其他token：-outcome → 正常惩罚
为什么要保护"答错但记忆好"的情况？ token会被保护还是惩罚？为什么这样设计？
-
解耦能力学习：写好和最终答对是两个不同能力。模型可能：
- 很好地压缩了信息(mem好)
- 但在推理阶段犯了逻辑错误(答错)
如果mem写得好却因为别的环节答错而被惩罚，模型会"学废"记忆摘要能力。
-
长期信用分配：是在第2、3、4轮开头写的，而答案可能在第5轮才产生。没有mem_adv 的保护，好的记忆策略会被稀释在漫长的轨迹中无法获得正确归因。
-
防止退化：如果模型发现"反正答错了全被惩罚"，可能干脆写空的以避免任何风险。mem_adv 确保即使答错，好的记忆策略仍有正向激励。
-
部分保护。区间的advantage=(-outcome）+(+mem_adv），惩罚被部分抵消。设计原因：写好和最终答对是两个独立能力，如果因为推理环节答错就惩罚好的记忆策略，模型会"学废"记忆摘要能力。代码中注释掉的条件版本*(advantages.sum()>=0）选择不保护这种情况，作者最终没有采用。
注意：代码中有一个被注释的条件版本：
advantages + mem_advantages * (advantages.sum() >= 0) #含义：只对答对的轨迹施加mem_adv
这个版本选择不保护场景3，作者最终没有使用它一说明实验证明保护"答错但好mem"是更优策略。
1.5 baseline
MemPO 训练中有两个 baseline（对应两种 advantage）：
1.5.1 Outcome Advantage 的 Baseline
baseline = 同 question 16 条轨迹 score 的均值
例：question Q1 的 16 条轨迹：
scores = [1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0]
baseline = mean = 0.625
adv_i = (score_i - 0.625) /std
→ 答对的：正 advantage
→ 答错的：负 advantage
1.5.2 Memory Advantage 的 Baseline
baseline = 同 question 所有轨迹所有轮次 mem_reward 的均值
例：question Q1, 16 条轨迹，每条 3 轮有<mem>：
pool = [r_11, r_12, r_13, ← 轨迹 1 的 3 轮
r_21, r_22, r_23, ← 轨迹 2 的 3 轮
...
r_16_1, r_16_2, r_16_3] ← 约 48 个值
baseline = mean (48 个值)
adv_k = (mem_reward_k - baseline) /std
→ 比平均好的 mem：正 advantage
→ 比平均差的 mem：负 advantage
对比标准 PPO 的 Baseline
标准 PPO: baseline = V (s_t) ← Critic 网络估计的 value
GRPO/MemPO: baseline = group_mean ← 同组统计量
优劣：
- Critic: 精确到每个 token 位置，但有估计误差，需要额外训练
- group_mean: 只有一个全局值，但零误差，零额外开销
一句话：MemPO 的 baseline 就是 "同组平均水平"— 不用额外模型估计，直接从数据中统计得到。
1.6 等价性
此处的问题是：P_full 与训练时的 baseline 在数学上是否等价？因为P_full中，应该已经有部分压缩的记忆吧？
1.6.1 P_full 与 GRPO baseline 的差异
full_traj 包含如下，可以看到，P_full 中的"完整上下文"本身已经包含了前面几轮的压缩内容。它不是"纯净的原始搜索结果”。
full_traj = 完整轨迹 token 序列：
系统提示
+ 原始问题
+ Round 1 模型输出："<mem>我的记忆1</mem><think>...</think><search>..·</search>
↑ 这里已经包含了模型自己写的<mem>！
+ Round 1 搜索结果
+ Round 2 模型输出："<mem>我的记忆2</mem><think>...</think><search>..·</search>"
↑ 包含了Round2的<mem>
+ Round 2 搜索结果
+ Round T 模型输出(最终答案Z)
1.6.2 P_full 作为 baseline 与GRPO mean的对比
GRPO baseline(组内均值)：
同一道题生成16条轨迹{T1, T2, ..., T16)
baseline = mean(ri, r2,..., r16)
→ 消除"这道题本来就难"的因素 → 不同轨迹的 baseline 相同(都用组均值)
P_full baseline(本条轨迹的上界)：
baseline=P(Z丨本条轨迹的完整上下文)
→ 消除"这条轨迹本来就难"的因素(包括之前<mem>的质量) → 每条轨迹有各自的 baseline！
两者不等价，用一个例子说明：
同一道题，两条轨迹：
轨迹A:前几轮 <mem> 写得很好 → P_full=0.85（完整历史很充分）
轨迹B:前几轮 <mem> 写得很差 → P_full=0.50（完整历史也混乱）
GRPO baseline=mean(reward_A，reward_B，...) ← 两条共用同一个均值
P_full baseline ≠ 相同 ← 两条轨迹各有自己的P_full
这导致了一个微妙的问题
mem_reward=P_mem(当前轮)-P_full(当前轨迹)
P_full包含了"前几轮的质量"作为baseline
场景：前几轮很好，但当前轮写差了
P_full 高(因为完整历史包含好的早期<mem>)
P_mem 低(当前轮的<mem>质量差)
mem_reward=低-高=大负数 → 惩罚当前轮
场景：前几轮很差，当前轮稍微好一点
P_full 低(完整历史本来就乱)
P_mem略高于P_full
mem_reward=略正 → 奖励当前轮
所以：mem_reward实际衡量的是："当前轮的相对于包含前序记忆的完整上下文，是否能独立提供足够的答案信息？"
1.6.3 与GRPO 双层baseline 的关系
最终 advantage = GRPO advantage + mem advantage
GRPO advantage.(outcome baseline):
baseline=同组16条轨迹的平均奖励
目标：整条轨迹比同组平均好吗？
mem advantage(记忆baseline)：
baseline=本条轨迹的P_full(完整上下文)
目标：当前轮<mem>是否保留了足够信息？
两层 baseline正交，各自消除不同的方差来源：
GRPO baseline → 消除"问题难易度"
P_full baseline → 消除"当前轨迹的具体上下文"
1.6.4 一个潜在问题
由于 P_full 包含了前序的 ，存在这样的情况：
如果模型学会"在Round 1的里写全部信息"：
Round 1 <mem>=超级详细的笔记
Round 2~T <mem>=简短甚至空白
那么：
P_full(包含Round1 的好记忆)仍然高
P_mem(只看当轮记忆)很低
mem_reward =大负数→惩罚了这种"前置记忆"策略
MemPO实际上在每一轮都要求：
"这一轮的<mem>必须独立地包含足够回答最终问题的信息"，而不是"历史<mem>的累积"
这是MemPO记忆奖励的一个重要特性：它激励模型在每一轮都写一个独立自洽的"完整记忆摘要”，而不是把信息分散在多轮中逐步积累。这可能也是一个潜在的问题。
1.6.5 被注释的条件版本
被注释的条件版本mem_adv * (outcome_adv.sum()>=0）有什么优缺点？
优点：防止答错轨迹通过写好来抵消outcome的惩罚信号。如果答错了但mem写得好，的advantage 不会是正值，整条轨迹一致地被抑制。
缺点：答错轨迹中的完全不受mem_reward约束。模型无法学到"虽然最终答错了，但这个记忆策略是好的"，可能> 导致记忆能力学习变慢。作者最终选择无条件叠加说明实验表明保护好的mem更重要。
1.7 连续值
Memory Reward为什么设计为连续值(-1，+1）而不是离散的{0，1}？原因如下：
- 更丰富的梯度信号：离散值只能说"好/坏"，连续值能表达"好多少/坏多少”
- 更好的归一化：同组~48个连续值的mean/std统计更稳定
- 自然的定义：P_mem-P_full是概率差，天然是连续的
- 避免人为阈值：如果离散化，需要设定"什么程度算好的mem"的阔值，而连续值让GRPO归一化自动处理
1.8 mem_reward对不同轮数的奖励强度分析
1.8.2 P_full 的效果
Round 1的完全没有mem_reward 信号
原因：没有历史搜索结果，full_traj和 mem_traj 内容差不多，对比无意义
之后，各轮P_full 会自然增长
- Round 2 full_traj = 初始prompt+R1输出+search1结果 → P_full较低（只有1轮信息）
- Round 3 full_traj = 初始prompt+R1+ search1+R2+ search2 → P_full 中等（2轮信息积累）
- Round 4 full_traj = ... +R3+search3 → P_full 较高（3轮信息积累）
这导致 mem_reward= P_mem - P_full 随轮次增大而更负
1.8.2 归一化
我们再来看看归一化后的实际效果（数值模拟）
假设：16条轨迹，每条4轮，各轮raw mem_reward均值：
- Round 2：-0.10
- Round 3：-0.22
- Round 4：-0.33
池化归一化（mean=-0.22，std=0.13）后：
- Round 2: 归一化均值 ≈ +0.90 → 正优势 (鼓励改进)
- Round 3: 归一化均值 ≈ -0.03 → 中性
- Round 4: 归一化均值 ≈ -0.87 → 负优势 (惩罚相对较差)
1.8.3 核心问题
核心问题：信号强度不一致。
从模型学习角度：
- Round 2 : 系统性获得正奖励 → 梯度方向：继续写好
- Round 4 : 系统性获得负奖励 → 梯度方向：改变写法
但这不是模型真正"写得差”—而是因为 P_full 随轮数天然增大！
本质：归一化以"跨轮次的相对表现”来评判，而不是"绝对的记忆质量”
1.8.4 小结
奖励强度的系统性差异如下：
| 轮次 | 是否计算 | raw mem_reward | 归一化后 | 训练信号方向 |
|---|---|---|---|---|
| Round 1 | 不计算 | N/A | 0 | 无信号 |
| Round 2 | 计算 | 最接近 0 (负) | 正 (相对好) | 奖励改进 |
| Round 3 | 计算 | 中等负 | 中性 | 轻信号 |
| Round 4 | 计算 | 较大负 | 负 (相对差) | 惩罚 |
| Round 5 | 计算 | 最大负 | 最负 | 最强惩罚 |
根本原因：P_full 随轮次递增，而 P_mem 相对稳定，导致后期轮次系统性地获得更负的 raw reward，归一化后反而给早期轮次正信号、晚期轮次负信号。这是 MemPO 记忆奖励的一个结构性不均衡问题。
1.9 非对称过滤的微妙之处
不均衡会导致什么学习策略？
梯度信号的方向如下：
policy gradient: ∀θ log π(τ) × advantage
advantage > 0 → 增加这个 token 序列的生成概率
advantage < 0 → 降低这个 token 序列的生成概率
导致实际可能出现的学习策略
Round 2 (系统性正奖励):
梯度方向：把"在 Round 2 写出高 P_mem 的记忆"的概率拉高
→ 模型学会：第2轮必须写详细、全面的记忆摘要
Round 4 <mem>(系统性负奖励)：
梯度方向：降低"在Round 4写当前这种记忆"的概率
但没有明确方向要写什么一一一只知道不要写这个
→ 可能的结果(几种情况)：
情况A：写更短的记忆(减少token，分母变小，P_mem被提高)
情况B：抄写Round 3的记忆(避免改变，维持上一轮水平)
情况C：写更模板化、稳定的格式(降低分布，提高单词概率)
反直觉的均衡策略
理论最优(适应这个不均衡信号)：
Round2<mem>：把已知的所有信息都写进去
Round3<mem>：略微补充新发现
Round4<mem>：直接复制Round 3的内容(不"升级"记忆)
原因：
- Round2升级记忆有正信号
- Round4升级记忆被系统性惩罚(因为绝对P_mem值与Round2同，但P_full更高，差值更负，在pool中排名更低)
- 所以Round4不升级记忆反而"相对表现较好"
与outcome reward 的交互
final_advantage = outcome_adv + mem_adv
即使mem_adv 鼓励Round4写简单记忆， outcome_adv会反向激励："只有写出足够好的记忆，最终答案才能正确，outcome_reward 才能高"
两种信号的博弈：mem_ady(后期负)vs outcome_adv(答对才有正)
实际均衡：模型在第4一5轮写*足够好*但不超出的记忆
"足够好”=足以让答案正确+不会在Poo1中排太低
一个更令人担忧的可能性
如果模型学到："第2轮把所有信息写进<mem>，后面几轮直接用<answer>输出" → 利用 max_assistant_turns 限制，尽量早结束 → Round2写全，Round3直接答题
效果：逃避了后期的 mem_adv负惩罚！
但：需要 Round 2的搜索信息足够充分
这解释了为什么多跳问题(需要多次搜索)对MemPO来说比单跳问题更难学习：多跳必须经历多轮，无法逃避后期负mem_adv,单跳可以 Round2一次搞定
由于归一化的数学结构，MemPO实际上在激励："早轮写详细记忆"而非"每轮都写好记忆"。这可能导致模型在早期轮次积累信息、在晚期轮次趋于保守或复制，是一个值得在消融实验中验证的隐性偏差。
0x02 Mem Reward 实现
compute_grpo_memory_advantage 是MemPO的创新点，是MemPO专属的记忆优势。tool_agent_loop.py 用mem_rewards_idx_list 数组(0=普通token，1=开始，2=结束)标记每个 token 属于哪段记忆，从而精确定位梯度位置。
compute_grpo_memory_advantage 的内部逻辑如下
2.1 归一化
#核心思想：对同一道题的 16条轨迹，对其所有 mem_reward 做组内标准化
id2mem_rewards[uid].extend(mem_rewards[i]) #收集同题所有<mem>段的奖励
id2mean[uid] = mean(所有 mem_rewards)
id2std[uid] = std(所有 mem_rewards)
2.2 找到位置
普通GRPO：整个回答得同一个标量奖励
↓ 广播给所有token
MemPO记忆优势：
<mem>每一段</mem>都有自己独立的奖励信号
↓
mem_reward_start_idxs / mem_reward_end_idxs 标记每段 <mem> 的 token 位置
mem_rewards 是每段记忆的质量得分
→ 不同<mem>段的 token 得到不同的梯度信号
→ 模型可以更精细地学习"第几轮的记忆应该怎么写"
注意：start_idxs / end_idxs 是 response 张量中的绝对位置
2.3 写入
#仅在[start_idx，end_idx]范围(即 <mem>...</mem> token)上写入优势
for start, end, r in zip(starts, ends, mem_rewards):
adv =(r-mean)/(std +1e-6) #组内标准化
advantages[i,start:end+1] = adv #精确写入<mem>token位置
#其余 token (<think>/<search>/<answer>)的mem_advantages保持0
2.4 示例
举例如下：
关键数值关系：
假设同一题uid="Q1"，有16条轨迹，每条 2-4轮有mem_reward：
池化后约有16×348个mem_reward值
adv=(r_ij-mean_48个) / std_48个
归一化的含义：
某轮<mem>比全局均值好→正adv→鼓励
某轮<mem>比全局均值差→负adv→惩罚
2.5 完整函数
def compute_grpo_memory_advantage(
response_mask: torch.Tensor,
index: np.ndarray,
mem_reward_start_idxs: np.ndarray,
mem_reward_end_idxs: np.ndarray,
mem_rewards: np.ndarray,
) -> tuple[torch.Tensor, torch.Tensor]:
id2mem_rewards = defaultdict(list)
id2mean = {}
id2std = {}
with torch.no_grad():
bsz = response_mask.shape[0]
for i in range(bsz):
id2mem_rewards[index[i]].extend(mem_rewards[i])
for idx in id2mem_rewards:
if len(id2mem_rewards[idx]) == 0 or len(id2mem_rewards[idx]) == 1:
id2mean[idx] = torch.tensor(0.0)
id2std[idx] = torch.tensor(1.0)
elif len(id2mem_rewards[idx]) > 1:
rewards_tensor = torch.tensor(id2mem_rewards[idx])
id2mean[idx] = torch.mean(rewards_tensor)
id2std[idx] = torch.std(rewards_tensor)
else:
raise ValueError(f"no memory reward in prompt index: {idx}")
advantages = torch.zeros_like(response_mask, dtype=torch.float32)
for i in range(bsz):
for start_idx, end_idx, r in zip(
mem_reward_start_idxs[i], mem_reward_end_idxs[i], mem_rewards[i]
):
if id2std[index[i]] > 0:
adv = (r - id2mean[index[i]]) / (id2std[index[i]] + 1e-6)
else:
adv = r - id2mean[index[i]]
advantages[i, start_idx : end_idx + 1] = adv
return advantages, advantages
0x03 合并
3.1 公式
在MemPO中的完整公式
复合advantage：
final_adv = GRPO_outcome_adv + memory_adv
↑来自GRPO归一化 ↑ 来自P_mem-P_full归一化
即细化如下：
compute_grpo_outcome_advantage():
score → 广播到全序列 outcome_adv
↓
+ compute_grpo_memory_advantage():
mem_rewards → 仅写入 <mem>...</mem> 区间
↓
final_adv = outcome_adv + mem_adv
双重奖励信号的叠加效果
token:[<mem> 某某某</mem>] [<think> 推理</think>] [<search>查询</search>]...[<answer>答案</answer>]
GRPO整体优势：[+0.75+0.75+0.75][+0.75...+0.75][+0.75...]...[+0.75]
(整条轨迹答对了)
mem专属优势：[+1.20+1.20+1.20][ 0...0 ][ 0... ]... [ 0 ]
(这段<mem>信息充分，能独立预测答案)
叠加：[+1.95+1.95+1.95] [ +0.75...+0.75 ][ +0.75... ]...[ +0.75]
PPO(GRPO) loss:
ratio = exp(new_log_prob - old_log_prob)
clipped_ratio = clip(ratio,1-e,1+e)
loss = -mean( final_adv x min(ratio,clipped_ratio) x mask ) +KL_coefxKL(π || π_ref)
可以看到，复合advantage（final_adv）作为 PPO loss的 advantage 系数，作用与PPO中。
此处的关键点是：reward_tensor 是稀疏的，整条 response(可能2000 token)只有最后1个 token 有非零奖励。后续compute_grpo_outcome_advantage 会把它广播到全序列。
设计洞见：validate_format 是 MemPO 的隐式 Prompt Engineer。
- 系统提示只有"You are a helpful assistant."
- 但奖励函数强制要求格式(否则得0分)
- 等价于让RL信号代替Prompt来教会模型格式
3.2 交互点
两条路径完全解耦，唯一的交汇点是A3叠加：
④ compute_score → outcome_adv ─────┐
┼────────► ⑥ final_adv = outcome + mem
⑤ P_mem -P_full→ mem_adv ──────────┘
compute_score (B2):
- 输入:data_source, response_str, ground_truth
- 流程: extract_solution → validate_format → em_check
- 输出：score ∈ {0,1}
- 调用链：B5→B1→B2→{C2，B3，B4}
- 用途：仅供 Outcome Advantage
Memory Advantage的计算链(完全独立)：
- 输入:full_traj_list,mem_traj_list, answer_ids
- 流程:model.compute_log_prob → P_mem - P_full
- 输出:mem_reward ∈(-1,+1)
- 调用链：A4(收集)→A1(计算)→A2(归一化)
- 用途：仅供 Memory Advantage
在final_adv=outcome_adv+mem_adv中，区间的token和其他token分别接收什么信号？
- ...内的token：outcome_adv+mem_adv_t(双重信号）
- 其他token(，，等）：仅outcome_adv(单一信号）
3.3 四个阶段的数据变换
阶段一 (A5+A4): 每轮生成时收集
生成后 → 三路分支(记录 full_traj / mem_traj / idx_list(A4)):
Round 1: idx_list += [0,...,0] 不收集
Round 2 + 有<mem>: idx_list += [1,0,...,2,0,...,0]
full_traj_list.append (raw 完整上下文)
mem_traj_list.append (sys + 问题 +<mem>内容)
Round 2 + 无<mem>: idx_list += [0,...,0] 不收集
rollout 结束，从 idx_list 提取:
start_idxs = np.where (idx_list == 1) ← <mem>开始绝对位置
end_idxs = np.where (idx_list == 2) ← </mem>结束绝对位置
reward_score 写回 output.
阶段二 (A1): 一次批量前向(rollout全部完成后)
7-一次批量前向
阶段三 (A2): 池化归一化
同题所有轨迹所有轮次池化(约48值)→(r-mean)/std
写入 response张量的<mem>...</mem>位置
-------------------------------------------------
同一题的所有轨迹所有轮次 → 一个 pool (约 48 个值)
adv = (r - pool_mean) /pool_std
写入 advantages [i, start:end+1] = adv ← <mem>...</mem> 全填同一值
→ mem_advantages [bsz，seq_len]，其余位置=0
阶段四(A3)：双信号叠加
outcome_adv（来自 B1 的 0/1 score，广播全序列）
+ mem_adv(来自 A1 的 P_mem-P_full，仅 <mem>区间)
= final_adv →PPO 更新
-------------------------------------------------
<mem> token: outcome_adv+mem_adv ← 双重信号
其他 token： outcome_adv ← 单一信号
outcome_adv衡量："这条轨迹最终答对了吗？"(相对组均值)
mem_adv衡量："这轮的<mem>比同题所有<mem>写得好吗？
3.4 "信用分配"问题
EM奖励的困境：
Round2写了好的<mem> → Round5答对了,outcome_adv均匀分给所有轮次的所有token → Round 2的<mem> token不知道自己功劳最大
mem_adv的解决方案：
Round2的mem_reward高 → mem_adv正 → Round2的<mem>token额外得到正梯度 → 精确信用分配到具体轮次的具体内容
3.5 全景图
===================================================================
多轮Rollout 阶段
ToolAgentLoop._handle_generating_state()
───────────────────────────────────────────────────────────────────
每当模型生成一段<mem>...</mem>，同时记录：
full_traj_ids_list[t] = 第t轮之前的完整上下文 token ids
mem_traj_ids_list[t] = 系统提示+仅第t轮的<mem>内容token ids
mem_rewards_idx_list = token级别的位置标记数组：
[0,0,...,0，1，0,0,...，0，2，0，0,...](1=<mem>开始，2=</mem>结束)
===================================================================
│
│
▼
===================================================================
Rollout完成后：计算mem_reward
AgentLoopBase (agent_loop.py)
───────────────────────────────────────────────────────────────────
对每个<mem>片段，用Actor模型在两种上下文下分别计算答案的 log_prob：
P_full(answer)= exp(均值 log_prob) over full_traj + answer_ids
P_mem(answer) = exp(均值 log_prob) over mem_traj + answer_ids
过滤低置信度token(log_prob<log(0.5))后：
mem_reward[t] = P_mem(answer)-P_full(answer) ← prob_bias
直觉：若仅凭<mem>就能预测出正确答案→mem_reward 大(记忆质量高)
若<mem>没有全上下文有用 →mem_reward小/负(质量差)
===================================================================
│
│
▼
===================================================================
优势计算阶段：ray_trainer.py
───────────────────────────────────────────────────────────────────
# 1.标准GRPO优势(整条轨迹)
advantages, returns = compute_grpo_outcome_advantage(
token_level_rewards, response_mask, index=uid
) #→所有token 获得相同标量：(r-μ_group)/ σ_group
# 2.记忆专属优势(精确到<mem>片段)
mem_advantages, mem_returns = compute_grpo_memory_advantage(
response_mask, index=uid,
mem_reward_start_idxs, ← <mem>开始的token位置
mem_reward_end_idxs, ← </mem>结束的token位置
mem_rewards, ← P_mem -P_full
) #→仅<mem>...</mem>范围内的token 获得按组标准化的奖励
# 3.叠加两种优势
advantages = advantages + mem_advantages
===================================================================
0x04 可能的问题和改进方向
我们接下来看看这种奖励设计可能的问题和改进方向。
4.1 奖励强度是否一致
我们提出的问题是：mem_reward对不同轮数的奖励强度是否一致？
答案是：MemPO 记忆奖励确实存在一个结构性不均衡问题。
根本原因： P_full 随轮次递增，而 P_mem 相对稳定，导致后期轮次系统性地获得更负的 raw reward，归一化后反而给早期轮次正信号、晚期轮次负信号。
其本质是：归一化以"跨轮次的相对表现"来评判，而不是"绝对的记忆质量"。
Round 1：
→ Round 1的<mem>完全没有mem_reward 信号
原因：没有历史搜索结果，full_traj 和 mem_traj内容差不多，对比无意义
其它轮：
各轮P_full的自然增长：
Round 2 full_traj = 初始prompt + R1输出+ search1结果 → P_full较低(只有1轮信息)
Round 3 full_traj =初始prompt + R1 + search1 + R2+ search2 → P_full中等(2轮信息积累)
Round 4 full_traj = ... +R3 +search3 → P_full较高(3轮信息积累)
→ mem_reward=P_mem-P_full 随轮次增大而更负
因此，核心问题：信号强度不一致。
从模型学习的角度来看，具体如下：
Round 2<mem>：系统性获得正奖励→ 梯度方向：继续写好
Round 4<mem>: 系统性获得负奖励→ 梯度方向：改变写法
→ 但这不是模型真正"写得差"—而是因为P_full随轮数天然增大！
总结：奖励强度的系统性差异
| 轮次 | 是否计算 | raw mem_reward | 归一化后 | 训练信号方向 |
|---|---|---|---|---|
| Round 1 | 不计算 | N/A | 0 | 无信号 |
| Round 2 | 计算 | 最接近 0（负） | 正（相对好） | 奖励改进 |
| Round 3 | 计算 | 中等负 | 中性 | 轻微信号 |
| Round 4 | 计算 | 较大负 | 负（相对差） | 惩罚 |
| Round 5 | 计算 | 最大负 | 最负 | 最强惩罚 |
4.2 潜在问题分析
4.2.1 问题 1：P_mem 的构建天然有优势偏差
现象：
full_traj=[完整多轮上下文+杂乱信息+工具结果] → 上下文长、噪声多
mem_traj=[系统提示+精炼<mem>] → 上下文短、信息密
即使内容随机，短上下文下模型预测概率也可能系统性高于长上下文(注意力分散效应) → P_mem-P_full 可能长期为正，不能准确区分"记忆质量好"vs"上下文短带来的伪增益"。
4.2.2 问题 2：额外的 forward pass开销
每个训练步需要对所有 段做额外的 compute_log_prob：
- 若每条轨迹平均 8轮，每轮1段→batch size×8×2次额外推理
- 这些推理用的是当前 actor 权重(old_log_probs)，而非冻结的 ref model
- 随着训练进行，actor 权重不断变化，但 full_traj 是固定的 rollout 数据 → 存在 off-policy 偏差
4.2.3 问题 3：记忆奖励只看"最终答案预测能力”，忽略中间推理质量
mem_reward 衡量的是：能否直接预测答案
但的另一个功能是辅助下一轮的推理(间接价值)
可能出现：写了足够答案的余信息(直接reward高)，但浪费了宝贵的 token 预算，未来轮次更难收敛。
4.2.4 问题 6：prob_bias的值域没有归一化
mem_reward = P_mem - P_full # 范围 [-1, +1]，但实际集中在 [-0.2, +0.2]
而 GRPO 主优势的量级约为 ±1.0(标准化后)，叠加时两者量级差异大，记忆奖励的实际权重难以预测。代码中虽然有 compute_grpo_memory_advantage 内部做了组内标准化，但prob_bias 原始值的分布对效果影响很大。
4.3 改进方向
改进1：使用相对比率代替绝对差值
#当前：
mem_reward =P_mem -P_full
#改进：用对数比(更稳定，消除上下文长度偏差)
= log(P_mem + ε) - log(P_full + ε)
或者引入长度惩罚校正：
#按<mem>的token数惩罚余
mem_reward=(P_mem -P_full)-λ*len(mem_tokens)/ max_mem_len
改进2：用Ref Model计算P_full(避免off-policy偏差)
#当前：用actor当前权重计算(会随训练漂移)
traj_log_probs = self.worker_group.compute_log_prob(...)
#改进：用冻结的ref model计算P_full(对齐基线更稳定)
full_log_probs = ref_worker.compute_log_prob(full_traj_batch)
mem_log_probs = actor.compute_log_prob(mem_traj_batch)
改进3：threshold自适应
#当前：硬编码
threshold = math.log(0.5)
#改进：根据当前batch分位数动态调整
threshold =torch.quantile(al1_logps[ans_mask]，q=0.3) #过滤最低 30% 置信度token
改进4：加入记忆长度奖励
#鼓励<mem>在压缩信息的同时保持简洁
mem_compression_reward = (P_mem / P_full) / (len(mem_tokens) / context_length)
改进5：多步前瞻奖励(Memory Credit Assignment)
当前：每轮只根据当前轮次计算
改进：考虑对后续k轮的辅助效果
mem_reward[t] = sum(discount^j *(P_mem_at_t evaluated at turn t+j)) for j in 1..k
TransFormer-封面
0xFF 参考
本文使用 markdown.com.cn 排版
