MemPO源码解析，GRPO强化学习在长时序Agent中的应用

在强化学习领域，智能体（Agent）的记忆管理一直是研究热点。传统的基于强化学习的记忆方法往往缺乏对记忆内容质量的有效引导，导致记忆信息冗余或关键细节丢失。为解决这一问题，MemPO（Self-Memory Policy Optimization）提出了一种全新的记忆自管理机制，通过引入基于有效信息含量的Memory-level优势估计，显著提升了记忆的有效性和任务解决能力。

MemPO的核心创新在于将记忆视为可训练的策略变量，使其不仅作为存储信息的载体，同时兼具思考链的功能。这种设计使得记忆生成过程可以直接接受强化学习信号的指导，无需额外的记忆模块支持。论文指出，MemPO通过"自我对话的草稿纸"形式，将记忆写在每轮交互的开头，既记录了重要信息，又构成了决策过程的一部分。

在具体实现上，MemPO采用了GRPO（Group Relative Policy Optimization）作为其核心算法框架。GRPO是PPO的一个变体，其主要区别在于优势函数（advantage）的计算方式。标准PPO需要训练一个Critic网络来估计价值函数，而GRPO则通过组内统计（均值和标准差）替代了Critic网络，从而简化了模型结构并降低了训练复杂度。

图1清晰展示了PPO与GRPO在Rollout、奖励计算、优势估计及更新流程上的差异。在GRPO中，每个问题（question）会生成16条轨迹（rollout），并通过组内归一化计算优势值。这种方式避免了对额外7B模型的需求，同时也消除了单独训练Critic网络的必要性。具体而言，GRPO的优势计算公式为：adv = (reward - mean) / std，其中mean和std分别表示组内16个score的均值和标准差。

MemPO进一步在GRPO的基础上进行了改进，提出了附加的优势值计算方法。除了传统的outcome_adv（基于最终答案正确性的奖励），MemPO还额外计算了一个mem_adv（记忆信息含量优势）。这一创新使得记忆生成过程能够更精确地指导哪些信息值得保留，从而确保记忆在保持简洁的同时最大化任务相关性。

从技术实现角度看，MemPO的GRPO模块主要包括以下几个关键部分：阶段划分、模型架构、优势函数计算、前向传播、损失函数设计、梯度计算以及KL约束处理。这些组件共同构成了一个完整的强化学习训练流程，使得模型能够在端到端的框架下优化记忆生成策略。

通过对比分析，可以发现GRPO相较于传统PPO的主要优势在于：1）去掉了需要额外训练的Critic网络；2）利用组内统计替代了价值基线估计；3）简化了模型结构，降低了计算复杂度；4）通过组内归一化提高了奖励信号的稳定性。这些改进使得GRPO特别适合应用于长时序任务中的记忆管理场景。

总的来说，MemPO通过引入GRPO框架，并结合记忆信息含量优势估计，为强化学习中的记忆管理提供了一个高效且实用的解决方案。这一创新不仅提升了记忆的有效性，也为长时序Agent的设计提供了新的思路。