在 Agent 系统里，记忆管理一直是工程难点：模型究竟该保留哪些信息、丢弃哪些信息，往往缺乏可训练的信号。MemPO（Self-Memory Policy Optimization）给出了一种更直接的做法：让模型在每轮开头主动写下 <mem> 内容，把记忆变成策略变量，再通过强化学习端到端优化“什么值得记、怎么记”。
围绕 MemPO 源码中的 Rollout 环节，可以看到这套方法并非简单地把记忆模块挂在推理链旁边，而是通过轨迹生成、上下文截取、概率比较和优势计算，把记忆质量纳入 PPO 更新路径。对于关注 Agent 底层机制的读者来说，这一部分恰好解释了记忆策略训练如何从“生成数据”开始落地。
Rollout 先完成轨迹生成，再收集记忆截面
从源码结构看，MemPO 的纯 Rollout 阶段主要对应 AgentLoopManager.generate_sequences()。这一阶段的核心任务是生成轨迹并收集数据。系统会将同一个 question 发送 16 次，得到 16 条独立并发轨迹。由于 LLM 采样存在随机性，这些轨迹内容并不相同，后续也将在同一 question 分组内进行比较。
单条轨迹的生成过程采用多轮形式。模型在第一轮可能生成 <mem>、<think> 和 <search> 标签，随后由工具解析和检索模块执行外部查询；第二轮继续生成新的 <mem>、思考与检索内容；直到某一轮生成 <answer> 后结束。
- 轨迹生成过程中，
_handle_generating_state会在每轮生成后收集<mem>的位置和内容。 ToolParser.parse负责解析<search>与<access>标签。AsearcherSearchTool.execute调用 RAG 检索，返回工具结果。- 轨迹完成后，
RewardManagerWorker.compute_score异步触发评分流程。
值得注意的是，Rollout 只生成一种轨迹，即完整的多轮对话轨迹。后续用于记忆奖励计算的 full_traj 和 mem_traj，并不是在 Rollout 中重新生成的两条独立轨迹，而是从这条完整轨迹中提取和构造出来的不同上下文视图。
full_traj 与 mem_traj：完整上下文与压缩记忆的对照
MemPO 在每轮生成后会保存两类数据。第一类是 full_traj，它是某一轮开始生成前的完整上下文快照。以一条 5 轮轨迹为例，第 2 轮生成前，上下文包含 system、question、第 1 轮响应和第 1 轮工具结果；第 3 轮生成前，上下文进一步加入第 2 轮响应和第 2 轮工具结果。源码中使用 deepcopy(raw_input_ids) 保存这些截面。
这里的“full”容易引起误解。它并不是指最终完整轨迹，而是指某一轮的完整上下文。它不包含当轮生成的 response，只保留到当前轮为止的对话历史。其作用是回答一个反事实问题：如果模型看到截至目前的所有对话内容，它对正确答案的预测概率有多高。
第二类是 mem_traj。它由 system prompt、question 和当轮 <mem> 摘要拼接而成，相当于一个压缩版上下文。它不保留完整历史，只保留模型自己写下的记忆内容。其作用是回答另一个问题：如果模型只看自己总结的记忆，它能否接近甚至达到完整上下文下的答案预测概率。
full_traj：某一轮生成前的完整对话历史，不包含当轮输出。mem_traj：仅包含 system、question 和当前轮<mem>内容。- 二者都服务于 Memory Reward 计算，不直接作为 PPO Update 的训练轨迹。
- 真正进入 PPO 更新的是完整轨迹，即 prompt_ids 加上 response_ids。
额外前向传播：用概率差计算记忆奖励
在常规生成之外，MemPO 还包含一次额外前向传播。这次前向并不生成新内容，而是对已经生成好的答案重新计算概率。源码中，这一步位于 A1 _postprocess，逻辑上属于 rollout 完成后的 Memory Reward 计算阶段。
具体做法是将 N 条 full_traj 和 N 条 mem_traj 拼成一个 batch，同时将对应答案重复两遍。随后一次调用 compute_log_prob，让模型分别计算：
- P_full = P(Z | full_traj)，即在完整上下文下生成答案 Z 的概率。
- P_mem = P(Z | mem_traj)，即只依靠
<mem>摘要生成答案 Z 的概率。
素材给出的示例中，P_full 为 0.72，P_mem 为 0.68，则 mem_reward = P_mem - P_full = -0.04。这意味着，仅依靠当前记忆比看完整上下文差 4%，记忆写得不够好，系统会对相应 token 施加负奖励，促使模型改进后续的记忆写入策略。
这一步的特殊之处在于，该 compute_log_prob 调用是 detached 的，不参与 PPO 的反向传播。它只用于计算 mem_reward 数值，并作为常数系数进入 advantage。换言之，额外前向传播不是为了更新模型本身，而是为了评估当前记忆内容的有效信息含量。
优势计算按 question 分组，避免难度互相干扰
MemPO 的 advantage 分为两部分：Outcome Advantage 与 Memory Advantage。前者来自最终任务得分，后者来自记忆概率差。两者叠加后形成每个 token 位置的最终 advantage，再送入 PPO loss 更新。
以一个 batch 包含 3 个 question、每个 question 生成 16 条轨迹为例，总共会有 48 条轨迹。Outcome Advantage 并不是在全 batch 上统一归一化，而是按 question 分组。同一个 question 的 16 条轨迹互相比较，独立计算 mean 和 std。素材示例中，Q1 组的 scores 为 [1, 0, 1, 1, 0, ...]，mean 为 0.6，std 为 0.5，则得分为 1 的轨迹 advantage 为 +0.8，得分为 0 的轨迹 advantage 为 -1.2。Q2、Q3 也分别独立计算。
Memory Advantage 则按同一 question 组内所有轨迹的轮次汇总 mem_reward，统一计算 mean/std，再将每条轨迹、每轮 <mem> 区间分别赋值。最终得到形状为 [48, seq_len] 的 mem_adv，每条轨迹的 <mem> 区间拥有不同值。
- Outcome Advantage：每条轨迹一个值，并广播到其所有 token。
- Memory Advantage：只在
<mem>区间赋值，不同轮次可不同。 - final_adv = outcome_adv + mem_adv，逐 token 参与 PPO 更新。
- 按 question 分组归一化，避免不同难度任务之间互相干扰。
对 Agent 记忆训练的意义：记忆从附属模块变成可优化策略
传统 Agent Memory 方案通常将记忆视为外部存储或检索组件，模型写入什么、保留什么，往往依赖规则或启发式策略。MemPO 的 Rollout 设计则把记忆写入行为放进强化学习闭环：模型先在每轮开头写 <mem>，系统再从完整轨迹中截取上下文快照，与仅含记忆的压缩上下文进行概率比较，最后把差值转化为记忆奖励。
这一设计的价值在于，它不需要额外独立的记忆评分模型，也不依赖人工定义的记忆质量标准，而是用任务答案的预测概率来衡量记忆是否保留了有效信息。如果 mem_traj 能够接近 full_traj 的答案概率，说明模型写下的记忆更接近任务所需信息；如果明显落后，则通过负奖励推动策略调整。
对于 AI 编程场景而言，这类机制尤其值得关注。编程 Agent 经常需要在多轮搜索、代码阅读、错误修复之间维护上下文，完整历史往往冗长且噪声较多。若模型能学会用更短的 <mem> 保留关键约束、接口信息和错误线索，就有机会降低上下文成本，并提升长任务中的稳定性。MemPO 的 Rollout 部分展示的正是这一目标背后的训练路径：不是简单让模型“记住更多”，而是通过奖励信号教它“记住更有用的内容”。
原创文章，作者：点点，如若转载，请注明出处：https://www.dian8dian.com/mempo-yuan-ma-chai-jie-rollout-ru-he-wei-agent-ji-yi-xun
