读一个陌生的强化学习仓库最怕的不是代码难而是不知道作者到底想把梯度信号灌到哪个变量上。MemPO 这个仓库我刚拿到手的时候就有这种感觉目录里既有记忆库的增删改查又有一整套 rollout 和组内归一化的训练逻辑第一眼看过去像是把 RAG 和 GRPO 硬拼在一起。Agent Memory 和强化学习的结合点其实非常窄窄到只有记忆的写入与检索这两个动作到底该不该被当成策略来训这一件事。这篇笔记就把 MemPO 当成一个普通的开源项目来拆从整体结构讲到关键实现把我在读源码过程中踩过的坑和想通的关节都摊开讲。如果你正在做长程任务的智能体、正在给 Agent 加长期记忆、或者手上已经有一套 RAG 记忆但效果始终靠调参玄学那这套思路值得花时间看一遍。内容偏源码向需要一点强化学习基础但我会尽量把每个概念落回到代码里具体哪一行在干活。1. 先搞清楚 MemPO 在优化什么在动手读任何一行代码之前我习惯先把这个问题问清楚这个项目训练出来的模型输入是什么、输出是什么、优化目标是什么。对 MemPO 来说如果这三个问题回答错误后面看的每一行代码都会理解偏。1.1 从规则驱动的记忆到策略驱动的记忆绝大多数 Agent 记忆系统的做法是固定的对话告一段落调一次大模型做记忆抽取把摘要写进向量库下一轮任务开始时用当前问题做一次向量召回取相似度最高的 top-k 条塞回上下文。这套流程里所有的关键决策都是超参数——什么时候写、写多细、检索几条、要不要丢弃旧的。而这些超参数一旦定死就再也没法根据任务反馈调整。问题出在哪出在一条记忆到底有没有用这件事往往要等到几十步之后才知道。你在第 3 步写了一条关于用户偏好的记忆第 40 步这个偏好才真正影响任务成败。这种延迟反馈用监督学习是学不了的因为你没法给第 3 步该不该写这条记忆标注一个正确标签。但强化学习恰好擅长处理这种只有终局奖励、中间过程没有标签的场景。MemPO 的核心主张就是把记忆的写、读、合并、丢弃这四个动作看成策略的动作空间把任务的最终成败当成奖励让模型自己学会什么时候该记住、什么时候该忘掉。这个视角转换带来的直接后果是记忆模块从一个被调用的工具函数变成了策略网络输出的一部分。原来记忆是写死在环境里的现在记忆操作要参与梯度回传。这是读这个仓库时最需要绷紧的一根弦凡是涉及记忆操作的 token它的 loss mask 处理方式都和普通生成 token 不一样。1.2 和 MemGPT 类分层记忆的分界线很多人第一次看到 MemPO 会联想到分层记忆那套思路主上下文装不下了就往外部存储里挪需要了再读回来。两者表面上有相似之处都会做写入外部存储和检索回上下文这两件事但本质差别在两个地方。第一是策略的来源。分层记忆系统里的换页规则通常写在提示词里比如当上下文超过 80% 时把最早的三轮对话摘要写入长期存储。这是人工规则模型只是执行者。MemPO 里这些决策是模型自己生成的规则只负责解析和容错。第二是优化目标。分层记忆系统的目标函数是隐式的——别丢信息别爆上下文。MemPO 的目标函数是显式的任务奖励甚至会把记忆效率比如平均记忆条数、检索命中率作为辅助项加进奖励里。这个差别决定了前者只能靠人工调参后者可以端到端训练。注意不要指望 MemPO 能直接替换掉你现有的向量库。它训练的是记忆操作的策略模型检索本身依然可以是普通的向量检索加轻量重排改动的是决策层而不是存储层。1.3 仓库的整体骨架与数据流把目录铺开看这类仓库的骨架基本是稳定的几块配置层负责模型路径、组大小、回合上限这些memory 目录放记忆条目的数据结构、索引、合并逻辑env 或 task 目录放任务环境和工具调用rollout 目录负责采样和轨迹组装trainer 目录是训练循环和 loss 计算reward 目录定义奖励函数。数据流可以按一条轨迹的生命周期来理解。任务实例进来环境初始化给一份初始记忆库策略模型开始多轮交互每一轮可能产生一次记忆写入或一次检索轨迹跑完拿到终局奖励同一个任务再采样若干条轨迹组成一组组内做归一化得到优势优势回传到记忆操作对应的 token 上更新策略参数。整个链路里组这个概念只存在于训练阶段推理阶段每条轨迹都是独立跑自己的记忆库。我建议第一次读的时候先别碰 trainer把 rollout 那条路径手动跟一遍用一个最小的任务跑出三条轨迹把每条轨迹里的 memory_ops 打印出来看看长什么样。这一步做完后面看 loss 计算会顺畅很多。2. 核心建模记忆策略是怎么变成 MDP 的强化学习的代码难读通常是因为作者把 MDP 的定义散落在好几个文件里。状态在 env 里、动作在 parser 里、奖励在 reward 函数里谁也不挨着谁。所以读之前先在脑子里把这三件事对齐会省掉大量来回翻文件的时间。2.1 状态、动作、奖励的落地定义状态这块比较直观当前对话历史、当前记忆库的一个摘要视图通常是条数和最近几条的标题、已经用掉的回合数。这里有个工程细节值得留意很多实现不会把整个记忆库塞进状态因为那样状态空间太大一般只给模型看记忆库里有什么主题这种粗粒度信息真正的内容等模型发起检索时再取。动作空间的设计是这套方法里最巧的一环。MemPO 没有额外接一个 policy head 去输出离散动作而是把记忆操作做成特殊的标记让模型在正常生成的过程中顺带吐出来比如用一对标签包住要写的内容或者用一个带参数的检索指令。这么做的好处是不改动模型结构直接复用现有的对话模板和推理框架成本极低。代价也很明显动作空间变成了自然语言空间的一个子集模型完全可能生成格式错误的动作。所以 parser 层必须写得很厚要处理标签不闭合、参数缺失、检索条数超限、写入内容为空这些情况。我做类似项目时一般会给非法动作一个小的负奖励而不是直接抛异常让模型在训练中学会规避而不是让整个训练进程崩掉。奖励通常由三块组成终局的任务得分、过程性惩罚非法动作、重复写入高度相似的记忆、上下文超长、以及可选的记忆效率项。最后这块要小心如果对记忆条数惩罚太重模型会学成什么都不记因为不记永远不会错。我的经验是效率项的权重不要超过终局奖励的五分之一。2.2 延迟奖励怎么摊到具体的记忆操作上这是整个方法里最难也最核心的部分。终局奖励是一个标量但一条轨迹里可能有十几次记忆操作梯度怎么分最朴素的做法是折扣回报第 t 步的动作拿到的回报是终局奖励乘以一个折扣因子的幂次越早的操作折扣越狠。这个做法实现简单但有个明显缺陷——它对所有记忆操作一视同仁没有区分哪条记忆后来真的被用上了。更细的做法是引入记忆归属的概念一条记忆被写入之后如果它在后续某个回合被检索出来并进入了上下文那这个区间内的任务表现就和它有关。实现上就是给每条记忆打一个区间标记把这个区间内的奖励变化作为这条记忆的近似贡献。这个近似很粗糙但它把信用分配的粒度从整条轨迹缩到了记忆存活区间训练稳定性会有肉眼可见的提升。还有个更省事的折中方案只在记忆操作对应的 token 位置上算 loss其他 token 全部 mask 掉。这样梯度不会污染普通的对话生成能力同时也天然地把哪些 token 应该为结果负责这件事限制在了记忆操作上。缺点是如果一条轨迹里记忆操作很少样本利用率会很低。实操心得这三种方案我都试过早期建议直接用第三种加上折扣回报实现最简单先跑通再优化。等 baseline 稳定了再上归属区间那套否则一旦训练不收敛你分不清是归属算错了还是别的环节有问题。2.3 为什么用组内归一化而不用 PPO 那套选 GRPO 这类组内归一化的方案理由其实很务实。PPO 需要一个和策略模型同规模的 critic 来估计状态价值对于 Agent 这种序列长度动辄几千 token 的场景critic 的价值估计噪声极大而且显存直接翻倍。组内归一化把 baseline 换成了同一个任务采样出来的若干条轨迹的平均奖励完全省掉了 critic 这一整块。这个替换成立的前提是同一个任务可以被反复采样出多条轨迹。Agent 任务恰好满足这个条件因为任务集是固定的同一个任务换个随机种子就能再跑一遍。相比之下那些一次性的、无法重复的环境就不适合这套方法。至于为什么不直接用偏好优化那类方法原因更简单这里的监督信号是标量奖励不是成对偏好。你要硬套偏好优化得先把标量奖励转成两两比较信息损失很大而且采样成本翻倍。所以在这类标量终局奖励 可重复任务的场景下组内归一化几乎是默认答案。3. 源码主干从入口到损失函数把建模思路理清之后读代码就有地图了。我按执行顺序把主干拆成四段配置与入口、记忆库数据结构、采样与轨迹组装、训练循环。每一段我只讲读的时候真正需要停下来的地方。3.1 配置层与入口脚本这类项目的配置一般是分层的一份默认配置打底实验配置做覆盖。入口脚本做的事情无非三件——初始化分布式环境、构建任务环境和奖励函数、实例化 trainer 并启动。需要重点确认的字段有这么几个策略模型和参考模型的路径、组大小、单条轨迹的最大回合数、检索条数上限、奖励函数名、KL 系数、裁剪范围、学习率、采样温度。这里最容易被忽略的是采样温度Agent 任务里温度设得太低组内几条轨迹几乎一模一样组内方差接近零优势全为零训练直接原地踏步。我一般会把它设在 0.8 到 1.0 之间保证组内有足够多样性。# 典型的训练配置片段 policy_model: ./ckpt/policy ref_model: ./ckpt/ref group_size: 8 max_turns: 30 retrieval_top_k: 6 kl_coef: 0.02 clip_ratio: 0.2 temperature: 0.9 memory_merge_threshold: 0.86组大小这个参数值得单独说。理论上越大越好因为 baseline 估计更准但显存和采样时间线性增长。实际跑下来8 到 16 是一个比较舒服的区间。小于 4 的时候组内方差抖动很大训练曲线会很难看。3.2 记忆条目的数据结构与索引记忆库这一块是整个项目里最工程的部分也是我建议第一个读透的地方。一条记忆条目通常包含这些字段唯一 id、文本内容、向量表示、创建时的步数、最近一次被访问的步数、累计访问次数、来源轨迹 id、状态标记。状态标记这个字段容易被忽略但很关键。记忆通常有几种状态活跃、已合并、已删除。删除并不是物理删除而是打个标记因为训练过程中需要回溯这条记忆在当时的决策下是被保留还是被丢弃物理删掉就没法算了。检索环节一般分两步先用向量召回一批候选再用最大边际相关性之类的策略做去冗余重排。去冗余这一步在 Agent 场景里比在普通问答里重要得多因为记忆库里很容易积累大量语义重复的条目如果直接取相似度最高的前几条很可能全是同一件事的不同表述白白占掉上下文预算。写入环节的去重逻辑通常有两道一道是向量相似度阈值超过阈值就认为是重复另一道用模型判断这条信息是否真的带来了增量。两道都要过才算新记忆。3.3 采样与轨迹组装采样部分的实现有一个容易踩的坑同一个任务采样出来的多条轨迹记忆库必须互相独立。如果你图省事共用了一个记忆库那后面写入的记忆会污染其他轨迹的状态组内比较就失去了意义训练出来的策略会朝奇怪的方向跑。我在别的项目里见过有人为了省内存做成共享的结果训练曲线看着像收敛了实际推理时完全不是那个表现。每条轨迹需要记录的东西比普通强化学习多。除了常规的输入 id、输出 id、动作 mask还要额外记一份记忆操作日志里面写清楚每一步做了哪类记忆操作、影响了哪条记忆。这份日志是后面算信用分配的依据。动作 mask 的处理有个细节环境返回的工具调用结果、检索回来的记忆内容这些都不是模型生成的绝对不能算 loss。有些人为了简单整个序列一起算 loss结果模型学会了复读检索到的内容因为它以为那些也是它该负责的输出。3.4 训练循环与损失计算训练循环的骨架和常规的组内归一化实现差不多但有两处改动需要留意。第一是优势计算的位置。标准做法是对每个输出 token 用同一个序列级优势但如果你要做记忆归属就要在序列级优势的基础上叠加一个位置相关的权重。实现上通常是准备一个和序列等长的权重向量记忆操作位置权重为 1其他位置为 0 或者一个很小的值。第二是 KL 项的系数。KL 约束在这里的作用不只是防止策略跑偏还有一个很实际的功能防止模型为了拿奖励而退化成一味输出格式化的记忆操作。如果 KL 系数设得太小训练几百步之后你会看到模型不管什么情况都往外吐写记忆的标签因为它发现写一次就能拿到过程奖励。这个现象我踩过一次排查了大半天才反应过来是 KL 太松。# 组内优势计算的核心逻辑示意 rewards rewards.view(group_size, -1).mean(dim-1) # 每条轨迹的终局奖励 mean_r rewards.mean() std_r rewards.std(unbiasedFalse) advantages (rewards - mean_r) / (std_r 1e-4) # 只对记忆操作位置的 token 赋予权重 token_adv per_token_weight * advantages.unsqueeze(-1)上面这段是简化版真实实现里还要处理组内奖励全相同、权重全为零这些边界情况。等优势算完剩下的就是常规的裁剪替代目标和反向传播了。4. 关键参数与实现细节读到这一步主干已经通了剩下的问题是为什么这些数字是这些值。这一节挑四个最影响结果的点展开。4.1 记忆写入的触发条件与合并策略写入触发有两种流派。一种是定时触发每跑够若干回合就让模型判断一次要不要写另一种是模型主动触发模型觉得信息重要就自己吐出写入标签。MemPO 走的是后者为主、前者兜底的路子。定时兜底是必要的。如果没有兜底模型在训练早期几乎不会主动写记忆因为它还没学会这个动作能带来收益而主动写入是需要多花 token 的早期策略更倾向于少说话。设一个兜底的定时检查能保证训练初期记忆库里至少有东西模型才有机会观察到记忆带来的收益。合并策略里最敏感的参数是相似度阈值。设得太高语义相近的记忆堆积检索时互相挤占名额设得太低本来不该合并的两条信息被糊在一起细节丢失。参考实现一般在 0.85 到 0.9 之间。我的建议是先用 0.86 跑一遍然后统计一下合并前后的记忆条数比例如果合并率超过三成说明阈值偏低往上调。合并之后文本怎么处理也有讲究。直接把两段拼起来最简单但会让记忆条目越来越长。用模型重新做一次摘要效果更好代价是每次合并都要多调一次模型。我的折中做法是只在两条记忆的文本总长度超过阈值时才触发摘要否则直接拼接。4.2 组内归一化的数值稳定性组内归一化看起来只有一行代码实际藏着不少坑。最常见的问题是组内奖励全部相同标准差为零优势全为零这一组样本就白采了。在任务成功率很低或者很高的阶段这种情况会批量出现。处理方式有几种。一是在除以标准差时加一个足够大的 epsilon让结果接近零而不是除爆。二是做动态采样丢弃那些组内方差低于阈值的组重新采直到满足条件。三是把奖励设计得更细不要只有 0 和 1 两个取值中间的状态可以给部分分。我自己更倾向第二种加第三种组合。动态采样会增加采样开销但样本效率提升明显奖励细化则从根子上减少全同组出现的概率。第三种的实现要点是别把奖励维度堆得太高否则模型会去钻某一个小项的漏洞我一般控制在三到四个维度以内。还有一个隐蔽的坑如果奖励里有负值标准差的计算方式会影响结果的符号分布。用有偏估计还是无偏估计在小样本组里差别不小。参考实现通常用有偏估计你如果换成无偏可能需要重新调学习率。4.3 KL 约束与参考模型的选择参考模型一般是策略模型的初始检查点冻结不动。这里有个成本问题参考模型要和策略模型一样大推理时显存占用几乎翻倍。想省显存的话可以用更小的参考模型做蒸馏式约束但效果会打折而且约束的尺度不好对齐。KL 系数的设置我摸索出来一个经验值从零开始 warmup前十步线性升到目标值之后保持不变。原因是训练最开始的几步策略分布和参考分布本来就接近KL 项没什么作用直接给大值反而会压制探索。另外要注意 KL 项的计算位置。有些实现只在记忆操作 token 上算 KL有些在整个序列上算。前者更聚焦能保护对话生成能力不被带跑后者更保守整体分布不会漂。如果发现训练后期模型的正常对话能力退化得厉害就把 KL 改成全序列计算。4.4 超参数速查把上面提到的关键参数整理成一张表方便对照。参数建议取值作用调大后的影响组大小8 ~ 16控制 baseline 估计精度更稳但更耗显存采样温度0.8 ~ 1.0控制组内多样性多样性增加可能过散检索条数上限4 ~ 8控制上下文预算信息更全但噪声增加合并相似度阈值0.85 ~ 0.90控制记忆库冗余度合并变少库更臃肿KL 系数0.01 ~ 0.05约束策略漂移更保守探索受限裁剪范围0.2限制单步更新幅度更保守收敛变慢最大回合数20 ~ 40控制轨迹长度更长但信用分配更难这张表里的数值只是起点任务不同差异很大。我建议每次只动一个参数并且记录下改动前后组内奖励方差的变化那个指标比总奖励更早反映问题。5. 常见问题与排查实录前面讲的是应该怎么做这一节讲实际跑起来会怎么坏。这几个问题都是我在复现类似项目时真实撞上的排查过程也一并写上。5.1 训练信号全零模型完全不学最常见的症状是训练几十步之后奖励曲线纹丝不动打印出来的优势全是零。这时候按顺序排查三件事。先看组内奖励是不是全相同。直接打印每个组的奖励列表如果有超过一半的组所有元素相等就是采样多样性不够。把温度调高或者在提示词里加一点随机扰动。再看采样是否真的走了不同的路径。有时候温度看着调高了但模型在某个特殊 token 上的概率极度集中导致几条轨迹在这个位置上还是选了一样的动作。可以统计一下组内轨迹的编辑距离如果普遍很小说明探索不足。最后看奖励函数返回的数值范围。有些实现会做奖励归一化如果归一化写错了所有奖励都被压到同一个值上优势自然全零。这个最好在奖励函数里加一行日志确认。踩坑记录我曾经因为奖励函数里一个整数除法写成了整除导致所有部分分都被截断成同一个整数排查了整整一个下午。教训是奖励函数一定要打印原始值和归一化后的值两套。5.2 记忆库膨胀检索越来越慢训练跑了几千步之后记忆库条目数爆炸式增长检索耗时从几十毫秒涨到几秒。这个问题的根源通常是写入的收益远大于删除的收益模型学到的是多写少删。短期止血的办法是给记忆库设硬上限超过就按访问次数和存活时间做一个淘汰。中期的做法是给写入加成本每写一条记忆扣一点过程奖励具体扣多少要做实验扣太多模型就不写了。长期来看还是要靠训练本身。我观察到一个规律当任务奖励的区分度足够高时模型会自己学会清理无用记忆因为无用的记忆会挤占检索名额间接拉低任务得分。所以如果你的记忆库一直膨胀先怀疑奖励信号是不是太粗糙。还有一个性能层面的优化容易被忽略向量检索的索引结构。如果条目数上万了还在用暴力检索那就是纯粹的实现问题换一个近似最近邻索引能省掉大部分时间。5.3 复现时最容易踩的几个坑第一个是环境版本。这类项目对推理框架的版本敏感尤其是涉及自定义采样和 logprob 计算的接口版本不一致会导致算出来的概率和实际采样的概率对不上训练直接跑偏。跑之前把依赖锁死别用浮动版本号。第二个是参考模型和策略模型的初始化必须完全一致。有人为了省事用了一个近似版本当参考模型结果 KL 项从一开始就是个大值训练全程被压制。这个错误很隐蔽因为曲线看着只是收敛得慢不会报错。第三个是随机种子。Agent 任务里有工具调用的随机性、检索的随机性、采样的随机性三者叠加之后复现性很差。调试阶段建议把工具调用固定成确定性的先把训练逻辑调通再放开随机性。第四个是显存碎片。长序列加上组采样显存占用波动很大跑一段时间容易出现碎片化导致的分配失败。开一个合适的显存分配策略或者定期做一次显存整理能缓解不少。5.4 问题速查表症状最可能的原因优先排查动作优势全为零组内奖励全同提高采样温度检查奖励归一化奖励不涨但 loss 在降KL 系数过小调大 KL检查是否退化成只会写记忆记忆库暴涨写入奖励高于删除加写入成本设库容量上限检索耗时长索引结构低效换近似最近邻索引限制返回条数复现结果差异大随机源太多固定种子工具调用改为确定性正常对话能力退化KL 只算在记忆 token 上改为全序列 KL训练后期崩溃学习率过大或显存碎片降学习率检查显存分配这张表可以直接当排查清单用按症状找原因比从头捋一遍代码快得多。我在实际做类似项目的过程中最深的一个体会是记忆策略的训练难点从来不在算法本身而在怎么把记忆操作和它的后果关联起来这件工程活上。组内归一化那套数学是现成的抄过来就能跑真正决定成败的是你对每条记忆的生命周期有没有清晰的追踪是你在奖励函数里有没有把无用记忆的隐性成本体现出来。这两个地方做扎实了模型会给你惊喜做马虎了再漂亮的算法也只是在原地打转。