MemPO源码解析，如何用强化学习端到端优化Agent记忆写入策略

在长时程任务中，大语言模型Agent常因记忆冗余或缺失而表现退化。传统方法依赖固定规则或外部模块管理记忆，缺乏对记忆内容质量的动态评估与优化能力。MemPO（Self-Memory Policy Optimization）提出一种端到端的解决方案：将记忆写入本身视为一个可学习的策略动作，通过强化学习直接优化“何时写、写什么、如何写”的决策过程。

MemPO的核心创新在于将记忆内容纳入策略梯度的优化闭环。不同于以往将记忆作为静态缓存或辅助输入，MemPO要求Agent在每轮推理开始前主动生成一段摘要式记忆（形式上置于标签内），该段落既是后续推理的上下文，也是模型自我反思的“草稿纸”。这一设计使记忆写入行为本身成为策略网络的输出变量，从而可被RL信号直接监督。

P_full_k = exp(mean(full_logp[k] × mask_k))：使用完整历史上下文生成答案的概率；
P_mem_k = exp(mean(mem_logp[k] × mask_k))：仅使用本轮生成的记忆摘要生成答案的概率。

二者之差mem_reward_k = P_mem_k − P_full_k即为原始记忆优势值。若该值为正，说明仅靠摘要就能比看全文更准确作答，表明记忆提炼有效；若为负，则说明摘要丢失关键信息，属于低质量记忆。这一设计直觉清晰：好的笔记应能“以简驭繁”，甚至超越原文的聚焦效率。

但单条轨迹的reward存在方差过大问题，因此第二步引入跨轨迹归一化：将同一问题下所有16条轨迹、每条3轮共约48个mem_reward值构成一个pool，计算其均值与标准差，得到标准化优势mem_adv_k = (mem_reward_k − pool_mean) / (pool_std + ε)。此举实现了三重比较：跨不同搜索路径的轨迹间比较、同轨迹内不同轮次间的比较，以及绝对质量与群体平均水平的对比，确保梯度信号稳定且具判别性。

关键实现细节在于梯度的精确注入位置。MemPO仅在标签所包裹的token区间内赋值mem_adv，其余位置梯度为0。这意味着只有“写记忆”这一行为受到记忆质量反馈的影响，而推理、工具调用等其他环节不受干扰。这种局部梯度注入机制解决了强化学习中典型的“信用分配”难题——避免将下游任务成败错误归因于无关模块。

具体实现中，代码通过批量处理提升效率：将完整轨迹（full_traj）与记忆轨迹（mem_traj）合并为一个batch，一次性调用model.compute_log_prob，减少GPU通信开销。ans_mask则精确标记答案部分token，确保log_prob计算仅聚焦于目标输出。整个流程在tool_agent_loop.py中完成，逻辑清晰且工程友好。

值得注意的是，MemPO并未简单采用P_mem作为唯一reward信号。若仅用P_mem，模型可能倾向于生成过于简略、回避复杂推理的记忆（例如只记“答案是A”），虽提升P_mem却损害长期能力。而P_mem − P_full的差值结构天然鼓励“保真压缩”：必须保留足以复现全文推理能力的信息，否则P_mem将低于P_full，触发负奖励。这构成了内在约束机制。

此外，该机制对非对称过滤敏感。例如，当某轮记忆遗漏了关键前提但保留了结论时，P_mem可能接近P_full（因结论正确），但实际推理链已断裂。此时mem_reward接近零，无法有效惩罚。作者在笔记中指出这是“微妙之处”，暗示未来可结合中间推理步骤的置信度进行增强。

从系统架构看，MemPO无需新增记忆模块或外部控制器，完全依托现有LLM的文本生成能力与策略梯度框架。其训练流程与标准PPO高度兼容，仅需在损失函数中加入mem_adv加权的策略熵项。论文实验显示，在需要多步推理的数学与代码任务上，MemPO显著优于基线方法，尤其在长horizon场景下记忆衰减问题明显缓解。

尽管如此，该方法仍面临若干挑战。首先是奖励强度一致性问题：不同问题难度差异导致mem_reward量级波动，归一化虽缓解但未根治；其次，当前方案假设答案可明确标注（依赖ans_mask），在开放生成任务中难以直接迁移；第三，记忆写入与推理耦合紧密，可能引发策略震荡——例如为追求高mem_reward而过度简化记忆，反而削弱后续轮次能力。

针对这些，作者提出若干改进方向：引入多粒度记忆（如分段摘要+关键事实列表）、结合内在动机信号补充外在reward、或在推理阶段动态调整记忆使用权重。长远看，MemPO代表了一种“记忆即策略”的新范式：不再将记忆视为被动存储，而是主动认知操作的一部分，其演化由任务目标直接驱动。

目前，MemPO相关代码已开源至GitHub（github.com/TheNewBeeKi…），模型与数据集亦托管于Hugging Face（huggingface.co/collections…），为社区复现与扩展提供了基础。随着Agent系统向更长时程、更复杂任务演进，此类内生式记忆优化机制有望成为提升自主智能体持续推理能力的关键组件。