MemPO自记忆策略优化源码学习笔记

一、项目概述

当前基于强化学习的智能体记忆管理方案普遍存在核心痛点：缺乏针对记忆内容的有效引导优化机制，导致生成的记忆质量参差不齐，难以真正服务于任务目标。MemPO，即Self-Memory Policy Optimization，正是为解决这一问题提出的全新框架。它的核心创新在于让智能体自主生成每一轮交互的记忆摘要，将记忆内容转化为可训练的策略变量，通过强化学习信号端到端地教会模型“应该记住什么、如何记录”，无需依赖额外的外部记忆模块。

该方案让模型在每轮交互开头生成记忆内容，形式上如同“自我对话的草稿纸”，既是历史信息的浓缩，也是推理过程的一部分。以下是MemPO的相关资源信息：

• 论文标题：MemPO: Self-Memory Policy Optimization for Long-Horizon Agents
• 论文链接：https://arxiv.org/abs/2603.00680
• 代码仓库：https://github.com/TheNewBeeKing/MemPO
• 模型与数据集：https://huggingface.co/collections/NewBeeKing/mempo

二、强化学习记忆系统的背景与挑战

MemPO的核心目标不止于用强化学习训练智能体，更在于为记忆本身设计可学习、可归因的优化信号，让智能体在交互过程中主动压缩、组织并保留最有助于任务完成的信息。在深入讲解MemPO之前，我们先梳理用强化学习训练记忆系统的核心要点与常见难点。

1. 训练记忆系统的核心要点

现有基于强化学习的记忆管理方法普遍缺少对记忆更新内容的引导优化机制，导致记忆质量难以保证。在长周期交互场景中，奖励信号往往稀疏且延迟，信用分配难度极大，智能体很难明确哪一轮的记忆决策最终影响了任务成败，最终容易出现记忆冗余、关键信息遗漏等问题。

2. 主要挑战

3. 业界主流思路

如何让智能体在交互过程中自主组织历史信息，通过强化学习优化记忆质量并对齐最终任务目标？目前行业内主要有以下几类解决方案：

• 隐式记忆压缩：包括MemPO和MEM1的思路。MemPO让模型自主生成每轮的记忆摘要，通过强化学习奖励反向驱动模型学习“哪些信息值得保留”；MIT在2025年提出的MEM1方案则维持固定大小的内部状态，每步整合并丢弃无关记忆，在问答任务上相比全上下文方案节省3.7倍内存，性能提升3.5倍。
• 外部检索增强（RAG-in-the-loop）：智能体在推理过程中主动调用搜索工具，如Search-R1、ASearcher等，将外部知识动态注入上下文，并通过强化学习训练“何时搜索、搜索什么”的策略。
• 分层强化学习+大语言模型规划：由大语言模型负责长时序抽象规划（高层决策），小型强化学习策略负责原子动作（低层执行），两者分离训练。
• 多粒度记忆系统：结合短期工作记忆与长期向量数据库/知识图谱，通过注意力机制动态检索所需信息。

4. 假想的强化学习记忆训练方案

我们可以通过一个标准的强化学习记忆训练框架，展示如何设计并训练智能体记忆系统，使其在交互中学会利用关键历史信息完成任务。该方案包含以下核心部分：

4.1 方案目标

训练一个强化学习智能体，通过交互学习获得实用的记忆能力，在部分可观察的环境中利用历史信息更好地完成任务。

4.2 任务与环境设计

选择迷宫导航任务作为训练场景：

• 环境：网格世界迷宫，智能体仅能观测当前网格周围的局部状态，属于部分可观测场景。
• 任务目标：智能体需要记住已访问区域，避免重复探索，最终找到目标位置。
• 核心挑战：当前局部观测不足以推断全局状态，智能体必须依赖内部记忆完成任务。

4.3 模型架构

智能体核心由策略网络与记忆模块组成：

• 策略网络：基于Actor-Critic框架，如PPO或A3C，其中Actor负责输出动作分布，Critic负责评估当前状态价值。
• 输入：当前观测与记忆模块生成的隐藏状态。
• 记忆模块：可根据实际场景选择RNN、Transformer或外部记忆系统实现。

4.4 训练流程

采用多步强化学习算法如PPO进行训练，核心流程包括：

• 智能体每一步获取当前观测与奖励，输入记忆模块生成状态表示。
• 策略网络基于当前观测与记忆状态生成动作分布，智能体执行动作后环境返回下一状态与奖励。
• 奖励信号设计包含两部分：稀疏任务完成奖励，以及记忆相关奖励，如在正确存储或提取关键信息时给予额外奖励。

4.5 优化策略

为提升训练效果，可采用以下优化手段：

• 记忆消耗正则化：限制智能体对记忆模块的使用，防止存储过多无关数据。
• 长时目标分层学习：将复杂长期任务分解为多个阶段，分别学习记忆与决策逻辑。
• 注意力优化记忆：通过调整Transformer中的注意力权重，使模型更关注与高奖励关联的观测信息。

三、MemPO核心论文设计

MemPO的核心洞察在于无需依赖外部记忆模块，通过在构建下一轮提示时，让模型在每个助手轮次的开头主动生成历史上下文的记忆摘要。系统仅保留最近一轮的工具结果（经过short_text截断处理），之前的所有历史信息都需要由模型自行写入记忆中，通过强化学习奖励信号端到端地驱动模型学习“什么值得记忆、如何记录记忆”，这就是“倒逼”模型学会记忆的核心机制。

MemPO采用多轮强化学习框架，在rollout采样阶段，模型与外界进行多轮交互，每一轮交互都会生成当前轮次的记忆内容。在优势计算环节，MemPO结合了两类优势估计方式：

• 全局轨迹优势：用于衡量整个轨迹的准确性，通过答案准确性与格式准确性计算奖励，该部分的优势估计会覆盖整个响应序列。
• 信息记忆优势：用于衡量每一段生成的记忆内容中保留了多少对解决问题有效的信息，其奖励计算通过在已知记忆内容的情况下，生成最终正确答案的后验概率来表示。

四、系统总体架构

1. 核心代码路径解析

MemPO的代码核心路径可以分为三个部分：

• A路径（记忆相关）：负责计算记忆奖励，即模型log概率的差值，该信号仅作用于记忆标签对应的token区间。例如通过流程可以了解记忆奖励的计算逻辑与作用范围，其中A5为评估专用路径，会限制上下文窗口，训练阶段不会使用。
• B路径（结果相关）：对应论文中的全局轨迹级奖励信号，通过字符串EM匹配计算，属于稀疏奖励，归一化后会广播到整个响应序列，为所有token提供梯度信号。
• C路径（系统主循环）：涵盖整个系统的骨架，包括训练循环、工具解析、RAG检索等核心功能。

2. 关键算法细节速查表

3. 基于VeRL的扩展设计

MemPO基于VeRL框架开发，复用了其GPU通信、FSDP、SGLang集成、Ray调度等基础设施，未对核心底层代码进行修改。MemPO的贡献主要集中在“如何定义和计算记忆奖励”这一层面，具体通过4处外科手术式修改实现：标记记忆token位置、计算记忆信息价值、组装记忆奖励信号、将记忆奖励叠加到GRPO优势估计中，同时新增了独立的训练配置与评估代码。

VeRL的可扩展点从易到难分为多个层级：

1. 配置文件修改：无需编写代码，仅调整yaml或shell配置即可，例如替换奖励函数是最常用的操作
2. 外部文件扩展：仅编写新的Python文件即可，例如替换奖励管理器
3. 继承注册类：编写新类并注册到框架中，例如添加自定义工具
4. 修改数据收集逻辑：调整tool_agent_loop.py文件，修改轨迹数据收集逻辑
5. 修改优势函数：调整core_algos.py文件，修改强化学习优势计算逻辑
6. 修改训练主循环：调整ray_trainer.py文件，修改整体训练流程

MemPO新增的核心代码包括：

• core_algos.py：新增记忆优势计算函数
• ray_trainer.py：添加记忆优势与结果优势的叠加逻辑
• tool_agent_loop.py：添加记忆token标记与完整/记忆轨迹收集功能
• agent_loop.py：实现完整上下文与记忆上下文的概率计算逻辑
• my_reward_score.py：新增EM奖励与格式校验函数文件
• asearcher_*_tool.py：新增搜索与访问工具文件

五、设计思路详解

1. 通俗解释

我们可以用一个侦探办案的例子来通俗解释MemPO的设计思路：假设侦探需要回答“爱因斯坦出生在哪个城市，该城市的市长是谁？”这个问题。

• 传统办案方式：每次查完一条线索，就把所有笔记都带进下一个房间。随着调查深入，需要携带的笔记越来越多，书包越来越重，最终根本无法负担。
• MemPO办案方式：每进入一个新房间，侦探只能携带一张自己写的小纸条。例如，他在小纸条上写下“目前已知：爱因斯坦生于乌尔姆市，我还需要查询乌尔姆市的市长信息”。在下一个房间中，他只需要查看这张小纸条和刚查到的最新信息，就可以继续推理。

侦探每次推理的格式是固定的：

<mem> [本轮对之前所有信息的压缩摘要] </mem>
<think> [推理过程] </think>
<search> 查询词 </search>

2. 详细设计解析

MemPO的核心设计是双通路奖励机制：

• 结果奖励：通过EM匹配正确答案计算，奖励信号作用于整个响应序列。
• 记忆奖励：通过完整上下文与记忆上下文的模型输出概率差值衡量记忆摘要的质量，该信号仅作用于记忆标签对应的token区间。

六、训练与推理流程

1. 单轮交互格式

无论是训练还是推理阶段，每一轮助手的输出都必须遵循固定格式：

<mem> [本轮对之前所有信息的压缩摘要] </mem>
<think> [推理链] </think>
<search> 查询词 </search>

2. 核心数据流串联

MemPO的核心数据流可以概括为：通过强制要求生成记忆内容制造学习压力，记录模型生成的记忆内容，测量记忆内容的质量差异，最终将记忆奖励与结果奖励结合，实现端到端的训练优化。

七、环境实现细节

1. 环境设计特点

MemPO的环境并非传统的Gym环境，而是基于VeRL框架与SGLang推理引擎定制的异步多轮智能体循环，其核心特点包括：

• 智能体与环境合一：传统强化学习中，智能体与环境是分离的，但在MemPO中，智能体循环本身就是环境。其中“动作”是模型生成的token序列，“观测”是工具返回的搜索结果，“奖励”仅在轨迹结束后才会计算。
• 异步状态机设计：ToolAgentLoop是一个异步状态机，每个请求独立运行，包含等待编码、内容生成、工具处理、完成四个核心状态。
• 真实外部RAG服务：工具并非模拟的，而是调用真实的稠密检索服务，返回维基百科的检索结果。
• 仅轨迹结束奖励：与传统强化学习每步给出奖励不同，MemPO仅在整个交互轨迹结束后给出一次奖励，这也是为什么需要GRPO来进行信用分配的原因。
• 停止词机制替代动作空间：通过定义停止词列表，当模型生成对应token时停止生成，相当于环境接管控制权。
• 区分训练与非训练token：通过掩码标记哪些token参与梯度计算，模型生成的内容参与梯度，工具返回或用户输入的内容不参与梯度。
• 记忆收集嵌入循环：环境会在每轮生成后自动收集完整轨迹与记忆轨迹，这一功能嵌入在状态机处理逻辑中。

2. 两套独立环境设计

MemPO实际上包含两套完全独立的环境，分别用于训练与评估，具体差异如下：

3. RAG服务部署

MemPO包含两套独立的RAG服务，分别对应训练与评估场景：

• 训练RAG服务（端口8013）：基于retriever conda环境，使用e5-base-v2模型，基于维基百科语料构建索引，启动脚本为sglang_multiturn/retrieval_launch.sh。
• 评估RAG服务（端口8002）：基于mempo-eval conda环境，使用e5-base-v2模型，基于自定义知识库构建索引，启动脚本为eval/scripts/launch_local_server.sh。

八、参考资料

本文基于MemPO相关论文与开源代码整理，完整资源信息可参考项目官方仓库与论文页面。