MemPO 源码学习笔记 GRPO 强化学习

MemPO(Self-Memory Policy Optimization)的论文标题是 MemPO: Self-Memory Policy Optimization for Long-Horizon Agents,论文地址在 arXiv:2603.00680。这篇源码学习笔记系列聚焦在 GRPO(Group Relative Policy Optimization)在 MemPO 中的具体应用与原理。

把这件事的几个关键数字列一下:MemPO 的独特切入点,是让模型把记忆写在每轮开头的 mem 标签区间内,形式上像"自我对话的草稿纸",既是记忆又是思考链的一部分。MemPO 的核心机制是叠加 Outcome Advantage 与 Memory Advantage 两种优势函数,形成最终用于 PPO 更新的复合 Advantage。代码实现中使用了 Schulman 提出的 k3 低方差估计器(ratio_ref - log(ratio_ref) - 1)计算 KL 约束。MemPO 相比原版 GRPO 多了 1 次 extra forward pass,但该次同时批量处理了 full_traj 和 mem_traj,实际吞吐开销约是标准 old_log_prob 的 1.5 到 2 倍。

文章想拆开说三件事:第一,MemPO 选用 GRPO 而非 PPO+Critic 的工程原因;第二,双重优势函数叠加机制到底怎么工作;第三,这套机制对长程 Agent 训练的实际意义。

MemPO 不选用标准 PPO+Critic 的架构,主要基于四个工程理由。

第一,Critic 在多轮长序列中极难训练。标准 PPO 的 V(s_t) 需要为序列中每个 token 位置预测未来累计回报,预测"未来能否答对"。但是 MemPO 的序列结构为多轮交互,长度可达数千 token,奖励仅在最末尾(sparse reward)。这导致 Critic 面对三个挑战:序列极长、需要巨大容量的 value 网络;奖励极稀疏、V(s) 几乎处处为 0;多轮工具交互、状态空间复杂、value estimation 噪声大。

第二,GRPO 比较适合 outcome-based 稀疏奖励。GRPO 用同组轨迹均值替代 value baseline,比较适合 trajectory-level 离散奖励。GRPO 的假设是奖励是 trajectory-level 的标量,完美匹配 EM check 的 0、1 评分。不需要学习 V(s_t),baseline 等于同 question 16 条轨迹的均值,零额外参数,零额外训练,无 value estimation 误差。

第三,计算资源节约。PPO+Critic 需要额外一个与 actor 同规模的 Critic 网络(7B 参数),Critic 需要额外前向加反向,显存翻倍:actor(7B)+ ref(7B)+ critic(7B)=21B 参数。GRPO 仅 actor(7B)+ ref(7B)=14B 参数,省下的资源用于更多并发 rollout(16 条/question)。

第四,Memory Reward 本身的特殊性。Memory Reward 本身自带 baseline(P_mem - P_full),无需 Critic 估计。如果用 Critic,还需要为 mem 区间单独训练 value head,但 mem 的好坏取决于未来能否答对(极长时间依赖),Critic 几乎不可能准确估计这个 value。GRPO 方案:直接跨轨迹归一化 mem_reward,简单有效。

双重优势函数:MemPO 的核心创新

MemPO 并未采用两个独立的 Loss 进行多目标优化,而是将反映最终答案正确性的 Outcome Advantage 与反映记忆信息含量的 Memory Advantage 直接相加,构成 final_adv。这种叠加方式既避免了独立 Loss 可能导致的梯度冲突,又天然兼容 PPO 的 Clip 约束,确保更新幅度可控。

具体来说,Outcome Advantage 用 GRPO 标准流程计算:分组是同一 question 的 16 条轨迹,归一化公式是 adv = (score - group_mean) / group_std。这就是 GRPO 的核心——用组内相对排名替代 Critic。

Memory Advantage 用 GRPO 风格但维度不同:分组是同一 question 的所有轨迹乘以所有轮次(约 48 个值),归一化公式是 adv = (mem_reward - pool_mean) / pool_std。借鉴了 GRPO 的"组内归一化"思想,但池化范围更大(跨轨迹加跨轮次)。这种设计使得梯度信号能精确作用于 mem 标签内的记忆生成行为,同时保持整体策略的稳定性。

最终优化是 PPO clipped surrogate loss(GRPO 只是 advantage 计算方式,优化器仍是 PPO)。PPO loss 公式为 -mean(final_adv × clip(ratio, 1-ε, 1+ε) × response_mask) + KL_coef × KL(π || π_ref)。

我的评价

把 MemPO 这套机制放到长程 Agent 训练的更大视野里看,会发现它解决的是一个非常具体的工程问题:在稀疏奖励、长序列、多轮交互的场景下,如何让 Agent 的记忆管理变得可学习。

过去几年,长程 Agent 的记忆管理一直是工程难题。一种思路是设计外挂的记忆模块(向量数据库、长期记忆系统),Agent 在每轮交互中查询和更新这些模块。但这种外挂式的设计,记忆的更新质量很难直接优化,因为记忆模块和 Agent 主模型之间存在梯度断层。

MemPO 的思路完全不同:把记忆当成 Agent 自己的一部分,让 Agent 把记忆写在每轮开头的 mem 标签区间内,然后用强化学习直接优化"记忆写得好不好"。这种"内化"的设计,让记忆管理和 Agent 决策在同一个优化目标下共同进化。

论文里给出的 Memory Reward 设计(P_mem - P_full)非常巧妙:它衡量的是"模型在只看到 mem 区间的情况下,对最终答案的预测概率"减去"模型在看到完整上下文的预测概率"的差值。这个差值越大,说明 mem 区间承载的信息量越高。这种设计天然适合稀疏奖励场景,因为它不需要外部标注,直接从模型自身的行为里提取信号。

反面观点:看上去很美,但几个现实问题需要看清

MemPO 的设计看上去很美,但有几个现实问题需要看清。

第一,16 条 trajectory 的成本不可忽视。GRPO 方案要求每个 question 生成 16 条轨迹,在长程 Agent 任务里,每条轨迹的生成成本已经很高。16 条并发的 rollout,对算力的要求是单条场景的 16 倍。这种成本不是所有团队都能承担。

第二,Outcome Advantage 的稀疏性。MemPO 的最终优化目标仍然是 outcome-based,奖励信号只在轨迹末尾出现。这意味着模型在前面轮次做的所有操作,只能通过 PPO 的 advantage 反向传播回去。对于一个 5 轮交互的任务,前面 4 轮的操作对最终奖励的贡献很难被精确归因。

第三,Memory Advantage 的局限性。Memory Advantage 衡量的是 mem 区间承载的信息量,但"信息量高"不等于"对任务有用"。某些记忆可能信息量大但相关性低,某些记忆可能简洁但关键。P_mem - P_full 这个指标,无法区分这两种情况。

第四,KL 约束的强度难以拿捏。代码实现中使用了 k3 低方差估计器计算 KL 惩罚,但 KL 系数 KL_coef 的设定直接影响训练稳定性。设置过松,模型可能偏离参考策略太远,出现 reward hacking;设置过严,模型可能学不到足够的新策略。

第五,泛化性问题。MemPO 的实验主要在特定类型的长程 Agent 任务上验证,对于其他类型的 Agent 任务(比如开放域对话、多人博弈),效果如何还需要进一步验证。论文本身也没有给出跨任务的系统比较。

关键素材

• • MemPO 论文地址 arXiv:2603.00680,代码地址 github.com/TheNewBeeKing/MemPO
• • 模型和数据集地址 huggingface.co/collections/NewBeeKing/mempo
• • GRPO 的 baseline 计算使用同 question 的 16 条轨迹的均值
• • 显存占用对比:PPO+Critic 需要 actor(7B)+ref(7B)+critic(7B)=21B 参数,GRPO 只需 14B 参数
• • Memory Advantage 分组约 48 个值(跨轨迹加跨轮次),归一化公式 adv = (mem_reward - pool_mean) / pool_std
• • KL 约束使用 k3 低方差估计器 ratio_ref - log(ratio_ref) - 1
• • MemPO 相比原版 GRPO 多了 1 次 extra forward pass,吞吐开销约是标准 old_log_prob 的 1.5 到 2 倍
• • actor 在 run_train.sh 中加载 NewBeeKing/MemPo_Qwen2.5-SFT,ref 加载同一个模型但参数冻结