[Agent Memory / 强化学习] MemPO源码学习笔记 —4— Rollout实现细节文章目录[Agent Memory / 强化学习] MemPO源码学习笔记 ---4--- Rollout实现细节0x00 概要0x01 回顾0x02 mem_sys_prompt_ids2.1 mem_sys_prompt_ids 定义2.2 mem_sys_prompt_ids作用为什么需要它关键特点小结0x03 ans_mask3.1 位置Outcome AdvantageMemory Advantage3.2 ans_mask 的精确构造过程Step1构造完整的答案序列Step 2: 单独 token 化 core_response_strStep 3: 计算 ans_mask完整示例(具体数字)0x04 threshold4.1 作用4.2 位置4.3 过滤的实际意义4.4 示例4.5 作用域4.6 问题0x05 Misc5.1 16个样本5.2 mem_rewards_idx_list5.3 首尾0xFF 参考0x00 概要现有的基于强化学习的 Memory 管理方法往往缺乏一种有效机制针对 Memory 的更新内容进行引导优化Memory 的内容难以保证质量。MemPOSelf-Memory Policy Optimization使模型对 Memory 进行自管理并引入了基于有效信息含量的 Memory-level 的优势估计引导 Memory 保留对解决任务更有效的信息进而提升记忆有效性。MemPO的独特切入点让模型把记忆写在每轮开头()形式上像“自我对话的草稿纸“既是记忆又是思考链的一部分。这样变成可训练的策略变量用RL信号端到端地教会模型“什么值得记、怎么记。RL 直接端到端优化这一行为无需额外的记忆模块。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本篇看一些实现的细节主要是rollout方面的实现细节。0x01 回顾我们首先回顾下。Rollout只生成一种轨迹一一一完整的多轮对话轨迹。full_traj和mem_traj是从这条轨迹中提取/构造出来的。构建 mem_traj 的方案即mem_traj 的最终token 组成如下ans_mask 会对答案 token 做 mask其含义只关注核心答案内容的 log_prob忽略/ 标签token。answer_idstokenize(\nthink...\n/think\nanswer\nKathryn Bigelow\n/answer)ans_mask:[0000,11,111,1110,0,0,0]↑Kathryn Bigelow 的token ↑\n/answer的4个tokenthreshold 会过滤mask其含义忽略模型完全没把握的tokenfull_ans_maskans_mask AND(full_logplog(0.5))mem_ans_maskans_mask AND(mem_logplog(0.5))具体示例(Round 3)如下|im_start|system You are a helpful assistant.|im_end||im_start|user 乔布斯在哪所大学读书|im_end||im_start|assistantmem乔布斯曾就读于俄勒冈州里德学院1972年入学6个月后学但仍在校旁听书法课。 最终答案应该是里德学院(Reed College)/mem|im_end|这就是P_mem的计算上下文系统提示 原始问题 当前轮内容然后用这个上下文计算模型生成正确答案的概率。接下来我们分析下具体细节。0x02 mem_sys_prompt_ids我们来看看 mem_sys_prompt_ids 包含哪些内容(具体 token 组成)。2.1 mem_sys_prompt_ids 定义mem_sys_prompt_ids第一轮生成前的初始prompt_ids的深拷贝即system prompt 用户问题(不含任何多轮对话历史)。具体内容是mem_sys_prompt_idstokenize(apply_chat_template([{role:system,content:You are a helpful assistant...}{role:usercontent:Who directed the 2o1o Best Picture?}]))# 用于构建mem_traj 时作为干净前缀与mem 摘要拼接token组成如下|部分|内容|────────────────────────────────────────────────────────────||im_start|system\n|系统角色起始||You are a helpful assistant|系统提示文本|||im_end|\n|系统段结束|||im_start|user\n|用户起始||原始问题文本|训练数据中的question|||im_end|\n|用户段结束|||im_start|assistant\n|助手起始(generation prompt)|注意prompt_ids 是第一轮开始时的初始 prompt此时 messages 只包含系统提示 用户问题没有任何历史搜索结果或 内容。2.2 mem_sys_prompt_ids作用作用构建mem_traj时作为干净的前缀与摘要拼接# tool_agent_loop.pymem_traj_ids_list.append(agent_data.mem_sys_prompt_idsresponse_mem_ids)mem_traj[systemquestion][mem摘要内容/mem]↑ mem_sys_prompt_ids ↑ response_mem_ids为什么需要它在A1中要对比 full_traj 和 mem_traj两者必须有相同的起点(systemquestion)才能公平比较。而 mem_sys_prompt_ids 提供这个相同起点。full_traj[systemquestion多轮完整对话历史 → P_full ← 完整历史 mem_traj[systemquestionmem摘要/mem]→ P_mem ← 仅摘要关键特点每一轮的 mem_traj 都共享同一个mem_sys_prompt_ids(第一轮的prompt)→ 无论到了第几轮上下文起点永远是系统提示原始问题→ mem_traj不累积历史每轮都重新从原点开始评估这使得P_mem真正测量的是“仅凭这一条模型能从原始问题出发回答正确吗”随着轮次推进prompt_ids会不断增长(加入工具结果等)但mem_traj 需要的始终是最初的 system question → 必须在第一轮就deepcopy保存。小结简言之mem_sys_prompt_ids 是Memory Reward计算中 “如果模型只看问题摘要” 这个假设条件的实现。这样可以让模型在只看问题摘要” vs 看完整历史两种条件下预测答案比较概率差异。我们接下来介绍 ans_mask 和 threshold。0x03 ans_mask3.1 位置ans_mask 和 threshold 都不作用于 Outcome Advantage。它们仅作用于 Memory Advantage 路径。Outcome AdvantageOutcome Advantage 路径中的mask 作用如下:只用 response_mask[bsz seq_len]→ 区分 prompt token(0)vs response token(1)→ 在 PPO loss 中loss-mean(adv × ratio × response_mask)→ 不涉及 ans_mask 或 thresholdMemory AdvantageMemory Advantage 路径中的 mask 和 threshold (A1):ans_mask标记answer_ids 中核心答案token的位置thresholdlog(0.5)过滤低置信度tokenfull_ans_maskans_mask(full_logpthreshold)mem_ans_maskans_mask(mem_logpthreshold)→ 用于计算P_full和P_mem → 产出mem_reward两条路径对比Outcome: response_str → em_check → {0,1} (无mask / threshold)Memory: log_prob → ans_mask × threshold过滤 → P_mem - P_full3.2 ans_mask 的精确构造过程ground_truth “里德学院”core_response_ids [里德学院]→ len 4Step1构造完整的答案序列ground_truth_text里德学院#从数据集取第一个答案answer_response_str(\nthink\nI have sufficient information to provide the final answers.\n/think\nanswer\n里德学院\n# ground_truth_text/answer)Step 2: 单独 token 化 core_response_strcore_response_str里德学院# 只有纯答案文本无 XML 标签core_response_idstokenizer(里德学院).input_ids# 假设[里德学院] 4 个 tokenlen 4Step 3: 计算 ans_maskans_masknp.zeros_like(answer_response_ids)# 全零ans_mask[-1*(len(core_response_ids)4):-4]1↑ ↑ 从倒数第(core_len4)个位置 到倒数第4个位置(不含)为什么 4 和 -4我们看 answer_response_str 的末尾结构...\n 里 德 学 院 \n/answer↑ ↑ core 开始前面有\n 末尾4个tokenn/answer这4个不应计入答案末尾4个token(Qwen tokenizer)对应\n即[‘\n’‘/’‘answer’‘’]。实际上\n被token化后恰好是4个token(硬编码假设)\n 是1token/answer是3tokens(或tokenizer可能分不同方式)这些是格式标签token不是答案内容本身排除它们可以确保只评估模型对核心答案内容的预测能力。因此得到具体标记结果如下answer_response_ids:[\nthink\n I.../think\nanswer\n 里 德 学 院 \n/answer]01..N N1..M M1M2M3M4M5M6M7末尾4个 ans_mask:[00...00...000011110000]↑ 从-8到-5↑末尾4个保持0(以4个core token为例)完整示例(具体数字)ground_truth “里德学院”core_response_ids [里德学院]→ len 4answer_response_str token 序列 (假设共18个 token)如下位置0123...1314151617\nthi nk \n \nans wer\n 里 德 学 院 \n/ans wer↑ 这里开始 ↑ core4个 ↑末尾4个 ans_mask[-1*(44):-4]ans_mask[-8:-4]1index:012...91011121314151617mask:000...0000011110000↑里德学院 ↑ \nanswer这4个是1这4个是0关键约束与潜在风险如下要素内容-4 硬编码假设\n恰好4个token适用条件Qwen tokenizer 中 \n 为 1 token, 为 3 token (可能是/answer)风险不同 tokenizer 可能 分词结果不同导致 mask 偏移即如果 tokenizer 对 n 的分词不是恰好 4个 token(不同 tokenizer、不同语言)答案mask 会错位奖励计算错误。正确效果只有纯答案文本(无 XML 标签)的 token 参与概率计算这样设计的原因计算 P(答案丨上下文)时不希望 \n 这些格式 token 干扰概率估算只关注实际答案词的预测概率。0x04 threshold4.1 作用threshold 在Memory Reward路径(A1)中使用(对full_logp和mem_logp各自独立过滤)。作用是过滤掉模型完全没信心的answer token(如人名的中间子词)避免噪声token拉低P_mem和P_full的区分度。注意threshold 与 Outcome 路径无关——Outcome 路径(B 系列)的 em_check 是字符串匹配不涉及任何概率计算或 threshold。threshold 的特点如下方面内容主要目的过滤掉模型完全不懂的token避免随机噪声效果让mem_reward 聚焦于有意义的答案token副作用两边过滤不同token→P_mem可能被高估硬编码风险prob50%是拍脑袋的阈值没有消融实验支撑改进方向可以改为min(full_logpmem_logp) threshold确保同一token 才对比以下面为例因为 “ryn” 和 “elow” 无论给什么上下文都难预测(子词特性)。如果不过滤这些噪声 token 会拉低 P_mem 和 P_full导致 P_mem - P_full ≈ 0(两边都被噪声淹没)。例如答案 “Kathryn Bigelow”我们得到tokenize 为[Kath,ryn, Big,elow]log_prob:[-0.2,-3.5,-0.1,-2.8]threshold:-0.693过滤后[-0.2,x,-0.1,x]← 只保留Kath和 Big丢弃 丢弃4.2 位置threshold 在 A1_postprocess 中使用属于 Memory Reward 路径(A 路径)。调用位置如下文件verl/experimental/agent_loop/agent_loop.py 函数AgentLoopManager.generate_sequences()的后处理段(即 A1)路径A4(收集)→[A1]_postprocess → A2(归一化)→ A3(叠加)↑ threshold 在这里调用链如下① rollout 完成 → 收集到 full_traj_list, mem_traj_list② A1: compute_log_prob(2N条) → 得到 full_logp, mem_logp③ threshold math.log(0.5) ← 这一步④ 过滤 计算 P_mem - P_full⑤ 结果存入 mem_rewards → 流向 A2, A34.3 过滤的实际意义情景A模型认识这个答案token(prob50%)full_logp-0.3→ KEEP(full_traj 能预测)mem_logp-0.4→ KEEP(mem_traj也能预测)→ 两边都参与计算正常对比情景B:full_traj 能预测但mem_traj不能full_logp-0.3→ KEEP mem_logp-2.0→ FILTER(过滤掉)→ 只有P_full的分子增大P_mem的分子不增大 → 实际效果P_mem的均值被计算为跳过这个token情景C两边都不认识这个tokenfull_logp-5.0→ FILTER mem_logp-6.0→ FILTER → 这个token在两边的概率计算中都被排除 → 对比差异0(不干扰信号)4.4 示例比如假设答案 亚硫酸盐沉淀反应中间体”(罕见术语)没有thresholdfull_logp(亚)-8.0prob0.0003mem_logp(亚)-9.0prob0.0001P_full 均值 ≈ exp(-8.0)0.0003P_mem均值 ≈ exp(-9.0)0.0001mem_reward0.0001-0.0003-0.0002◄─── 惩罚仅0.02%信号极弱且来自无意义的随机猜测差异有threshold(过滤掉prob50%的token)full_logp(亚)-8.0-0.693→ FILTER mem_logp(亚)-9.0-0.693→ FILTER → 这条轨迹的答案token 全被过滤有效 mask 数0→ P_memexp(0/(01e-8))exp(0)1.0→ P_full同上 ≈1.0→ mem_reward1.0-1.00(中性不产生信号)我们再对threshold log(0.5)的过滤效果分析。过滤规则如下thresholdlog(0.5)≈-0.693只有logp-0.693(即概率50%)的token才参与计算 full_ans_maskans_mask AND(full_logpthreshold)mem_ans_maskans_mask AND(mem_logpthreshold)4.5 作用域threshold会作用于 full_logpmem_logp。但是两者会各自独立过滤。thresholdmath.log(0.5)# -0.693full_logp_mask_bool(full_logpthreshold)# 过滤full中低概率tokenmem_logp_mask_bool(mem_logpthreshold)# 过滤mem中低概率tokenfull_ans_mask_boolans_mask_boolfull_logp_mask_bool# 交集full 的最终 maskmem_ans_mask_boolans_mask_boolmem_logp_mask_bool# 交集mem 的最终mask# 只对通过过滤的token计算平均log_prob → 再expP_fullexp(sum(logp*mask)/sum(mask))P_memexp(sum(logp*mask)/sum(mask))# 注意两者的 mask是独立的可能不同 → 各自用自己有信心的token来估算概率样例如下这意味着P_full和 P_mem用的是各自的 logp 来过滤两边可能过滤掉不同的.token一一一一这是设计意图每个条件下模型对不同 token的置信度可能不同各自用自己有信心的token来估算概率。答案Reed College(3tokens:Re,ed,GCollege)情景full_traj 对GCollege很确信mem_traj不确信 full_logp:[Re-0.2,ed-0.3,GCollege-0.1]→ 全部 KEEP mem_logp:[Re-0.5,ed-0.4,GCollege-1.5]→ GCollege 被 FILTER P_full基于3个tokentoken012的均值 P_mem基于2个tokentoken01的均值(跳过了GCollege)P_mem的分母减小(仅2个有效token)→ P_mem被拉高(分母变小)减轻了惩罚 这是一个潜在问题 当mem_traj 对某些 token 没有把握时这些 token 被排除 导致P_mem计算基于更容易预测的子集可能虚高4.6 问题thresholdlog(0.5)是硬编码超参 完全没有配置化对于不同大小的模型(7B vs 70B)合理值差异很大训练初期模型很弱大部分 token 被过滤P_mem 分子为 0→奖励无意义训练后期模型强了几乎不过滤→值失效0x05 Misc此处介绍其它细节。5.1 16个样本16是actor_rollout_ref.rollout.n的配置值一每个question生成16条独立的rollout轨迹其含义是同一个question送入LLM16次 → 每次用不同的随机采样samplingtemperature0→ 得到16条内容不同的多轮对话轨迹16是GRPO的group size(actor_rollout_ref.rollout.n16。GRPO用组内均值和标准差归一化advantage组太小(如2条→ 均值/方差估计不准信号噪声大组太大(如64条→ 计算开销大rollout时间长。为什么MemPO每个question要生成16条rollout轨迹16是常见的平衡点。同时MemoryAdvantage也受益于大组每个question约有16x3 48个。mem_reward值用于归一化统计更稳定。GRPO需要同一个question的多条轨迹来计算组内统计量 group_meanmean([score_1,score_2,...,score_16])group_stdstd([score_1,score_2,...,score_16])adv_i(score_i-mean)/std 如果只有1条→无法归一化16条→ 统计量估计相对稳定 例子 Question:Who directed Inception?轨迹1 searchInception→答对→score1轨迹2 search(2010 film→答错→score0轨迹3 searchInception director→ 答对→ score1......轨迹16:searchNolan movies→答对 → score1mean0.75,std0.43轨迹1:adv1-0.75/0.430.58鼓励 轨迹2:adv0-0.75)/0.43-1.74抑制 这个值是可配置的在 run_train.sh 中通过 actor_rollout_ref.rollout.n16设置。5.2 mem_rewards_idx_listmem_rewards_idx_list 中0、1、2分别代表什么0无关token(不在区间内)1开始位置-2结束位置5.3 首尾第1轮(Round1)为什么不收集mem数据第1轮是模型第一次生成没有之前的多轮历史需要总结因此不会(也不应该)产生摘要。此时mem_rewards_idx_list 全部填 0。如果rollout被截断最后一个没有系统如何处理丢弃该轮的mem 数据。检测方式start_idxs比end_idxs多一个→删除最后一个start。0xFF 参考