[Agent Memory / 强化学习] MemPO源码学习笔记 --- (6)--- 奖励机制
0x00 概要 0x01 原理 1.1 双通路奖励机制 1.2 通俗解释优势函数 1.3 对比 1.4 Reward model 0x02 Outcome Reward 机制 2.1 Outcome Advantage公式详解 2.2 EM Check 仅用于 Outcome Advantage 2.3 EM在MemPO中的作用 2.4 校验规则 2.5 normalize_answer四步处理 2.6 思考 0x03 Outcome Reward 实现 3.1 完整调用链路 3.2 NaiveRewardManager 3.3 RewardManagerWorker 3.4 自定义 reward function 0xFF 参考
0x00 概要
现有的基于强化学习的 Memory 管理方法往往缺乏一种有效机制针对 Memory 的更新内容进行引导优化,Memory 的内容难以保证质量。
MemPO(Self-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
0x01 原理
我们回顾前文,代码具体路径上的关键点如下:
A1 _postprocess(P_mem/P_full段) MemPO核心:记忆奖励如何计算A2 compute_grpo_memory_advantage mem_adv如何归一化、作用于哪些 tokenA3 compute_advantage(mem叠加段) 两种优势如何叠加、被注释的条件版本A4 ToolAgentLoop.__init__(mem收集段) full/mem_traj 收集时机、ans_mask 构造A5 AgentMemory.prepare_prompt "倒逼记忆"机制:每轮只保留1轮工具B1 NaiveRewardManager.__call_ outcome reward计算和放置位置B2 compute_score 三种 target 类型处理、EM checkB3 validate_format 8条格式规则(隐式prompt工程)B4 compute_grpo_outcome_advantage 对比 outcome_adv vs mem_adv 的差异B5 RewardManagerWorker.compute_score Ray async 奖励计算接口B6 AgentLoopManager.generate_sequences rollout 调度+mem_rewards 收集C1 RayPPOTrainer.fit 训练主循环(宏观流程)C2 extract_solution 答案提取逻辑C3 ToolParser.register("search") <search>标签解析C4 AsearcherSearchTool,execute RAG检索调用+5次重试
后续文字中会引用这些关键点。
1.1 双通路奖励机制
MemPO 关键设计是双通路奖励机制:Outcome Reward + Memory Reward。
Outcome Reward:规则驱动,字符串精确匹配(em_check) → {0,1},信号作用于全序列,评估"最终答案是否正确"。 Memory Reward:计算驱动,模型内部概率对比(P_mem - P_full) 衡量记忆摘要质量(用actor 自身的 log_prob 对比) → {-1,+1},信号仅作用于token,评估"记忆摘要是否有效压缩信息"
前者看模型实际输出了什么,后者看模型有没有能力预测正确答案(看概率而非输出)。这使得标签内的token同时受到「答对/答错」和「记忆是否有效压缩了上下文」两个梯度信号的驱动。
reward_tensor 的形状[bsz,seq_len],全零,只在每条轨迹response 的最后一个token位置赋值score(0或1)。该信号为稀疏,这是因为outcome reward是trajectory-level的标量,不是token-level的信号。后续GRPO会将这个标量广播为outcome_adv覆盖全序列。
1.2 通俗解释优势函数
想象你在参加一个读书闯关比赛:
读一本很长的故事书,每读完一章就要做一次笔记,最后回答老师的问题。
1.2.1 Outcome Advantage
🏆 Outcome Advantage(结果分):"你最终答对了吗?"
老师看你的最终答案 — 对了加分,错了扣分。
但不是跟满分比,而是跟你同桌比(GRPO:同组 16 个同学做同一道题)。
全班答对率 60%,你答对了 → 你比平均好 → 鼓励 全班答对率 60%,你答错了 → 你比平均差 → 抑制
这个分数给你写的每个字 — 答案、推理、搜索,全部同样鼓励或抑制。
1.2.2 Memory Advantage
📝 Memory Advantage(笔记质量分):"你的笔记够好吗?好到只看笔记就能答题?"
测试方法:
让另一个你看完整本书,去答题 → 答对概率 P_full 让另一个你只看你的笔记,去答题 → 答对概率 P_mem
如果只看笔记也能答对 → 笔记写得好 → 奖励你的笔记 如果只看笔记答不出来 → 笔记漏了关键信息 → 惩罚你的笔记
这个分数只给你写笔记的那几行字 — 推理和答案不受影响。
1.2.3 Final Advantage
🎯 Final Advantage(最终综合分)
"把结果分和笔记分加在一起。"
你写的搜索内容:只看结果分
你写的笔记:结果分 + 笔记分 ← 双重关注!
你写的最终答案:只看结果分
为什么笔记要特殊对待?因为笔记写得好不好,直接影响后面几章你能不能记住前面的内容。如果笔记都不单独教,你永远学不会做好笔记。
1.3 对比
其实,Outcome Advantage 是结果奖励,Memory Advantage 类似过程奖励。
| 摘要是否有效压缩了信息 / 笔记能否替代全文 | ||
1.3.1 作用范围
Outcome Advantage和Memory Advantage的作用范围有什么区别? 、
Outcome Advantage:广播到全序列每个token(同一条轨迹内所有token相同值) Memory Advantage:仅赋值到...对应的token区间,其余位置=0
1.3.2 Memory Advantage
Memory Advantage ≠ 格式奖励:
格式奖励(如 B3 validate_format)是硬约束:格式错 → 直接score=0,整条轨迹废掉 Memory Advantage是软信号:衡量的内容质量,而非格式是否合规
更准确的类比:
格式奖励(B3):"你的输出格式合法吗?" → 不合法就没分(门槛) 结果奖励(B4):"你最终答对了吗?” → outcome_adv 记忆质量奖励(A1):"你的摘要能替代全文吗?" → mem_adv
所以mem_adv更像是一个局部过程奖励(process reward)一不是评格式,而是评「记忆摘要的信息压缩质量」,与全局的结果奖励叠加使用。
1.3.3 配合
两者如何配合?final_adv = outcome_adv + mem_adv
场景1:答对 + 好mem → (+outcome)+(+mem) = 强。正向强化 场景2:答对 + 差mem → (+outcome)+(-mem) = 正向,但被惩罚 场景3:答错 + 好mem → (-outcome)+(+mem) = 整体负,但部分被保护 场景4:答错 + 差mem → (-outcome)+(-mem)= 全面惩罚
1.3.4 不同路径
两条路径的评分方式完全不同:
Outcome路径(B系列):
评分方式:em_check(模型预测答案,标准答案)→{0,1} 信号含义:最终答案对不对
Memory路径(A系列):
评分方式:P_mem-P_full(概率差) 信号含义:记忆摘要质量 不涉及任何字符串匹配,纯粹是模型内部的 log_prob对比
Outcome Advantage 路径
B6 (AgentLoopManager.generate_sequences) — 16条轨迹并发 rollout↓ 轨迹 DONEB5 (RewardManagerWorker.compute_score) — 异步奖励入口↓B1 (NaiveRewardManager.__call__) — 解码 response → 文本↓B2 (compute_score) — 主评分入口├→ C2 (extract_solution) — 提取最后 <answer>├→ B3 (validate_format) — 8条格式规则(11项子检查)└→ B4 (em_check) — 精确匹配 → {0,1}↓B4-algo (compute_grpo_outcome_advantage) — GRPO归一化 → 全序列广播↓→ outcome_adv [bsz, seq_len]
Memory Advantage 路径
A4 (_handle_generating_state) — 每轮收集 <mem> 位置、full_traj、mem_traj↓ 所有轨迹全部完成后A1 (_postprocess) — 2N条一次前向 → P_mem - P_full↓A2 (compute_grpo_memory_advantage) — 跨轨迹跨轮次池化归一化 → 仅<mem>区间赋值↓→ mem_adv [bsz, seq_len]
汇合点
A3 (compute_advantage) — final_adv = outcome_adv + mem_adv↓PPO Update — loss.backward()
路径对比
┌──────────┬─────────────────────────────────┬─────────────────┐│ │ Outcome 路径 │ Memory 路径 │├──────────┼─────────────────────────────────┼─────────────────┤│ 标签链 │ B6→B5→B1→B2→(C2,B3,B4)→B4-algo │ A4→A1→A2 │├──────────┼─────────────────────────────────┼─────────────────┤│ 汇合 │ → A3 ← │ → A3 ← │├──────────┼─────────────────────────────────┼─────────────────┤│ 评分方式 │ 字符串匹配 │ 概率差 │├──────────┼─────────────────────────────────┼─────────────────┤│ 产出范围 │ 全序列等值 │ 仅<mem>区间 │└──────────┴─────────────────────────────────┴─────────────────┘
1.4 Reward model
MemPO 没有使用Reward Model。
1.4.1 Reward Model
传统Reward Model (如RLHF中的):
单独训练的神经网络(通常基于人类偏好数据) 输入:(prompt,response)→ 输出:标量reward 有独立的参数、独立的训练过程 MemPO没有这个
MemPO的两种"奖励机制":
Outcome Reward---- 纯规则函数(em_check + validate_format) → 不是模型,是代码逻辑 →100%确定性,无需训练 Memory Reward ---- 复用 actor 自身(P_mem-P_full =用actor 的 compute_log_prob 计算) → 不是额外的reward model → 是策略模型自身的一种"自我评估"能力 → 模型自己判断mem是否足够好 → 没有额外参数,没有额外训练,而是复用actor的前向能力
Memory Reward和PRM都是过程级奖励,不仅看最终结果一都用于优化推理过程中的中间步骤,但是,PRM需要额外训练一个reward model(通常需要人工标注中间步骤的正确性)。
1.4.2 对比
对比需要Reward Model的系统,这也是 MemPO 能用 GRPO 而不用 PPO+Critic 的另一个原因——奖励信号干净确定,不需要 value network 来降低估计方差。
┌──────────────┬────────────────────────┬─────────────────────────────┐│ │ MemPO │ 典型RLHF │├──────────────┼────────────────────────┼─────────────────────────────┤│ Reward 来源 │ 规则EM + 自身logprob │ 训练好的 Reward Model(数B参数)│├──────────────┼────────────────────────┼─────────────────────────────┤│ 额外模型 │ ❌ 不需要 │ ✅ 需要单独训练 │├──────────────┼────────────────────────┼─────────────────────────────┤│ 人类标注 │ ❌ 不需要 │ ✅ 需要偏好数据 │├──────────────┼────────────────────────┼─────────────────────────────┤│ Reward 确定性 │ 完全确定 │ 有噪声 │└──────────────┴────────────────────────┴─────────────────────────────┘
虽然都是过程级奖励,不仅看最终结果一都用于优化推理过程中的中间步骤,但是双通路奖励机制与PRM(Process Reward Model)不同。
PRM需要额外训练一个reward model(通常需要人工标注中间步骤的正确性) MemPO的mem_reward不需要额外模型,用actor自身的log_prob自我评估 PRM评估的是推理步骤的正确性 MemPO评估的是记忆摘要的信息压缩质量 MemPO的信号更客观(概率可计算) PRM依赖标注质量 关键区分: A1中的model.compute_log_prob()调用的是actor 模型本身(正在被训练的那个模型),不是一个独立的reward model。它只是用 actor 在两种不同输入条件(full context vs mem only)下的概率差异来衡量mem 质量。 所以严格来说:MemPO 没有Reward Model,只有 reward function(规则)和 reward signal(自评估概率差)。
1.4.3 代码
代码中确实有 reward_model 相关代码,但在 MemPO 的实际训练中并未启用。原因如下 :
reward_model
代码中reward_model的存在有两种含义:
作为"字段名/数据键”(不是模型):non_tensor_batch["reward_model"] = {"ground_truth": [...], "style": "rule"} → 这只是数据传递的字段名,存放 ground_truth和评分模式 → 不是真正的神经网络reward model 作为VeRL框架的通用能力(MemPO未使用): VeRL框架本身支持reward model(用于RLHF 场景) 但MemPO 的配置中use_rm =False/reward_model.enable=False 这些分支不会被执行
MemPO
MemPO实际使用的奖励路径:
if self.config.reward_model.launch_reward_fn_async:future_reward = compute_reward_async.remote(data=batch, reward_fn=self.reward_fn)else:reward_tensor = compute_reward(batch, self.reward_fn)
这里的 reward_fn 是custom_reward_function.path 指定的,即 my_reward_score.py 中的 compute_score(规则函数),不是reward model。
总结:reward_model 在代码中既是数据字段名(传递ground_truth),也是VeRL 框架保留的RLHF 能力(MemPO 未启用)。MemPO 走的是reward_fn(自定义规则函数)路径,不走reward model 路径。
路径 1:规则奖励(MemPO实际使用)
rm_executor = Nonereward_model.enable = Truereward_model.enable_resource_pool = Falseenable_async_reward = (None is not None and ..) or not True= False or False = False↓不走compute_score异步路径enable_async_reward =(self.rm_executor is not None and self.config.reward_model.enable_resource_pool) or not self.config.reward_model.enable
具体展开
情况 enable_async_reward 走哪条路───────────────────────────────────────────────────────────────────────无RM,规则奖励(MemPO):rm_executor=None, enable=True (F and ?) or False =.False→ agent 自己打分(工具的calc_reward 或rollout时计算)→ 奖励已在轨迹中,postprocess 直接用有RM,异步计算奖励:rm_executor ≠ None, enable_resource_pool=True (T and T) or ?= True→ 走 compute_score.remote(data)→ RewardManagerWorker 里调用 rm_executor (Reward Model 推理)reward_model.enable=False:(? and ?) or not False = True→ 也走 compute_score.remote,但rm_scores 为0 (相当于跳过奖励计算)
路径 2:Reward Model推理(非MemPO场景)
┌─────────────────────────────────────────────────────────┐│ BatchExecutor (合批器) ││ ││ 多条轨迹同时调用 submit_task() ││ ↓ 放入队列 ││ worker 线程持续从队列取出,攒够 micro_batch_size 后 ││ 一次性调用 rm_wg.compute_rm_score(batch) ││ ↓ ││ GPU 上的 Reward Model 打分 (Bradley-Terry 偏好模型等) ││ ↓ ││ 结果通过 Future 返回给各条轨迹 │└─────────────────────────────────────────────────────────┘
BatchExecutor的关键设计:解决"多条轨迹异步完成但RM推理需要合批"的问题,把稀疏到来的单条请求动态合批,最大化GPU 利用率。
1.4.4 决策树总结
你用的是什么奖励?│├── 规则奖励 (EM、格式检查等)│ → 配置 custom_reward_function.path = my_reward_score.py│ → rm_wg = None → rm_executor = None│ → 奖励在 rollout 阶段由 compute_score() 同步计算│ → NaiveRewardManager 最终调用 my_reward_score.compute_score()│ ✅ MemPO 就是这种│├── Reward Model (神经网络打分)│ → 配置 reward_model.enable_resource_pool = True│ → rm_wg 传入 (独立 GPU worker)│ → rm_executor = BatchExecutor 合批器│ → 奖励在 rollout 后异步打分,利用 GPU 并发│└── 混合 (先 RM 打分,再放入 NaiveRewardManager)→ reward_manager 里检查 rm_scores 是否已存在:if "rm_scores" in data.batch.keys():return data.batch["rm_scores"] # 直接用,跳过规则
0x02 Outcome Reward 机制
2.1 Outcome Advantage公式详解
公式如下:
outcome_adv_i = (score_i -group_mean) / (group_std + ε)其中:score_i = em_check的结果,∈{0,1}(第i条轨迹是否答对)group = 同一question的16条轨迹group_mean = mean([score_1, score_2,...,score_16])group_std = std([score_1,score_2,...,score_16])
含义拆解如下:
score_i-group_mean:"我比平均水平好多少?"例:16条中10条答对,mean=0.625答对的轨迹:1-0.625 = +0.375 → "我比平均好"答错的轨迹:0-0.625 = -0.625 → "我比平均差"→如果全答对:mean=1,所有adv=0(没有区分度,不更新)→如果全答错:mean=0,所有adv=0(同理)→有对有错时:产生正负信号group_std:"归一化到标准尺度"std 小 (大家表现接近): 除以小数 → adv 放大 → 微小差异也有信号std 大 (表现差异大): 除以大数 → adv 缩小 → 防止极端更新→ 不同难度的 question 产生的 adv 量级一致广播: outcome_adv [i, :] = outcome_adv_i"整条轨迹的每个 token 获得相同的 advantage"→ 不区分哪个 token 贡献了正确答案→ 全局信号: "这整条路走对了 / 走错了"
EM check(Exact Match / 精确匹配)是MemPO唯一的奖励信号 = normalize_answer + exact match。
含义:模型给出的答案,标准化后是否和标准答案完全相等?
2.2 EM Check 仅用于 Outcome Advantage
EM Check(B4)是标准答案的精确匹配评估,只出现在 B2 compute_score 中,产出稀疏的 reward_tensor[i, ast_token]=0 or 1,最终流向compute_grpo_outcome_advantage.
Memory Reward(A1)完全不看模型最终输出了什么答案,它只关心:给定摘要后,模型有没有能力预测出正确答案(看概率),而不是是否真的输出了正确答案。
em_check的完整逻辑如下:
def em_check(prediction, golden_answers: list):norm_pred=normalize("obama") # 模型预测for golden in golden_answers: # ["Barack Obama", "Obama","Barry Obama"]if normalize(golden) == norm_pred:return 1 #只要匹配其中一个别名就算对return 0 #一个都不匹配就算错
2.3 EM在MemPO中的作用
MemPO 训练脚本配置的是 algorithm.adv_estimator=grpo, 对应 compute_grpo_outcome_advantage。
模型输出如下:
<mem>上轮搜索...</mem><think>让我分析...</think><search>query</search><tool_response>文档</tool_response><answer>obama</answer>
奖励计算链如下:
1.validate_format() → 格式合法?(8条规则(11项子检查))2.extract_solution() → 提取最后-个<answer>内容→"Obama3.em_check("Obama",["Barack Obama","Obama"]) → 1(答对)
最终奖励如下:
score=1(答对)或0(答错或格式非法)写入 reward_tensor[i,最后一个token]
关键特性:
稀疏:整条response(可能1000+个token),只有最后1个位置有非零奖励 二值:只有0/1,没有部分分 隐式课程:格式不合法也是0分>迫使模型先学会格式 与mem_reward互补:EM告诉模型"答对了吗",P_mem-P_full告诉模型"记忆写得好不好"
比如,对同一道题的16个侦探(index相同的样本):
侦探得分:[1,0,1,1,0,0,1,0,1,0,1,1,0,1,0,1]均值μ=0.625标准差σ=0.496侦探i的优势=(r_i-μ)/(σ+ε)答对侦探:(1-0.625)/0.496=+0.756 → 每个输出token都获得正梯度答错侦探:(0-0.625)/0.496=-1.260 → 每个输出token都获得负梯度
关键细节:这个优势值被"广播"给该侦探的所有输出token(包括中的每个字)一这就是为什么的质量能被端到端训练:答对的侦探写的记忆被强化,答错的被弱化。
2.4 校验规则
validate_format () 校验的完整规则如下:
规则 1:<think>/ </think>必须成对 且至少出现1次
规则 2:<answer>/</answer>必须成对 且至少出现1次
规则 3: <search> / </search> 必须成对 且 至少出现 1 次
规则 4:<mem>/</mem>必须成对
规则 5:<mem>\n出现次数==assistant 轮次数 + 1(每轮开头有一个<mem>)
规则 6:<think>\n 出现次数== assistant 轮次数+1(每轮有一个<think>)
规则 7: <search>...<tool_response>...</tool_response></search> 顺序正确
规则 8:<answer> 必须在</answer>之前
!注意规则 3:校验要求至少有 1 个 <search>。这意味着直接输出 <answer> 而不搜索的轨迹格式校验失败,score=0。这强制模型必须至少搜索一次。
另外,validate_format 的规则5 是强制的关键。该规则要求每个assistant轮次(含最后一轮)都必须写。这是隐式强制出现的关键规则,而非在·system prompt中明确要求。
如果 validate_format 返回失败,OutcomeReward和MemoryReward分别受到什么影响?
OutcomeReward:score直接=0,不再做em_check。 MemoryReward:不受影响。
MemoryReward的计算完全独立于validate_format,只要rollout过程中产生了,就会计算P_mem-P_full (但如果格式错到连都没有,则mem_traj_list为空,自然无memory reward数据)
2.5 normalize_answer四步处理
目的:消除大小写、标点、冠词的影响,让"USA"和"usa.”被视为相同。
normalize_answer("The United States of America!")↓↓步骤1 lower:"the united states of america!"步骤2 rm_punc: "the united states of america" 去掉感叹号等标点步骤3 rm_art: " united states america" 去掉the/a/an步骤4 space_fix:"united states america" 多个空格合并为一个
2.6 思考
如果把的内容也加入 em_check 评分(检查mem是否包含正确答案),这种方案好不好?
不好,原因:
目标错位:好的不一定要显式包含答案文本。它需要的是"帮助模型推理出答案的关键信息",可能是间接线索 鼓励作弊:模型会学会在里直接复制答案,而不是真正压缩上下文 P_mem - P_full更优:它直接衡量"mem能否帮助模型预测答案",这才是mem质量的本质
0x03 Outcome Reward 实现
3.1 完整调用链路
Outcome Reward 的完整调用链路如下。
3.2 NaiveRewardManager
NaiveRewardManager 是Outcome Reward 路径的中间调度层(B1)。
3.2.1 特色
职责:
解码:response_ids→response_str(用 tokenizer 解码) 提取数据:从data_item 中取出 ground_truth、data_source 调用评分: self.compute_score(response_str, ground_truth) → score 写入结果:reward_tensor[i,last_token] = score (稀疏 outcome reward) 记录日志:打印前几条样本的prompt/response/score(调试用) 传递mem 信息:计算traj_avg_mem_rewards 写入 extra_info(仅用于logging)
本质:NaiveRewardManager 不"计算"奖励本身,而是:
从 DataProto 中解包数据 调用真正的评分函数(my_reward_score.compute_score) 把结果打包写回tensor 是一个"数据搬运工"/"适配器"
注意:ground_truth =data_item.non_tensor_batch["reward_model"]["ground_truth"]-这就是前面提到的 rewar d_model 作为"数据字段名"的用法,不是调用一个reward model。
3.2.2 运作机制
奖励放在哪个token上?具体如下:
reward_tensor[i,valid_response_length - 1] = reward# ↑ 最后一个有效response token,因此是稀疏奖励# 示意:# tokens: [<mem>......</mem> <think>...</think> <search>q</search> ... <answer>ans</answer> EOS]# reward: [ 0 ... 0 0 ... 0 0 ... 0 ... 0 1.0 ]# ↑ 只有这里有值# 这是Outcome Reward(结果奖励)设计,不是Dense Reward。配合GRPO,把这个标量广播给整条轨迹。
3.3 RewardManagerWorker
RewardManagerWorker 是 NatvieRewardWorker 的异步包装。NatvieRewardWorker 是同步函数,会阻塞;RewardManagerWorker 是异步封装,不阻塞 rollout。
3.3.1 问题
为什么需要这层封装?这是因为如下问题:
Agent Loop 是异步的 (asyncio):
16条轨迹并发生成 某条轨迹先完成→立即触发奖励计算 不需要等其他轨迹完成
'NaiveRewardManager.__call__()'是同步阻塞的(CPU密集)
如果直接调用,会阻塞asyncio事件循环 其他轨迹的生成也会被卡住
3.3.2 关键设计
关键设计是:与 agent loop 并行。
为什么做成异步 Ray actor:奖励计算(尤其是 reward model 推理)很耗时,异步化可以让多条轨迹的奖励计算与其他轨迹的 rollout重叠执行,提升 GPU利用率。
AgentLoopWorker (每个 trajectory):...生成下一个 token......调用搜索工具......生成结束 (<answer>← 此时异步提交奖励计算 →result = await compute_score.remote(data)↑ 非阻塞,其他 trajectory 可以同时运行output.reward_score = result["reward_score"]
3.3.3 解决方案
RewardManagerWorker 解决方案:
用 run_in_executor()把阻塞操作丢到线程池 asyncio事件循环不被阻塞 奖励计算和其他轨迹的 rollout同时进行
B5是 B1的"async adapter"—让同步的奖励计算函数能在异步 agent loop 中并行执行,不阻塞其他轨迹的 rollout。
3.4 自定义 reward function
3.4.1 异步接口层
RewardManagerWorker.compute_score 是 MemPO 中奖励计算的异步接口层。
3.4.2 compute_score
compute_score 是自定义 reward function 本身—它们是同一个东西。"自定义 reward function"是配置层面的叫法(custom_reward_function.path),compute_score 是这个 reward fu nction 的具体函数名。配置指向文件,框架从文件中加载 compute_score 函数使用。
配置链:框架加载my_reward_score.py文件 → 取其中的compute_score函数 → 注入到 NaiveRewardManager.compute_score
run_train.sh: custom_reward_function.path = "verl/utils/reward_score/my_reward_score.py"
属性调用链如下:
B5 → B1 (NaiveRewardManager)→self.compute_score →这个函数NaiveRewardManager.__init__():self.compute_score=compute_score ← 从 my_reward_score.py 导入的函数 NaiveRewardManager.__call__():score = self.compute_score(data_source, response_str, ground_truth)↑就是 my_reward_score.py 中的 compute_score
内部流程 / 调用顺序如下:extract_solution()→validate_format()→em_check()。如果格式不合法(8条规则(11项子检查)中任一条违反),直接返回score = 0,不再进行EM检查。:
extract_solution(solution_str) → 取最后一个内容(c2) validate_format(solution_str) → 8条格式规则(11项子检查)(B3) em_check(answer, ground_truth) → 精确匹配(B4) → 返回 score ∈ {0,1}
另外,compute_score 是 Outcome Advantage 专有的函数,Memory Advantage 完全不使用它。
compute_score()的三种 target 处理
target 类型 打分方式────────────────────────────────────────────────────list[list[str]] 多目标:answer按;分割,逐个EM匹配累加分list[str] 单目标多别名:answer与任意一个EM匹配得1分str 单一字符串:包装成[str]后EM匹配
奖励值放置位置
response token 序列:[t0][t1][t2]...[t_{n-2}][t_{n-1}]↑reward 放在最后一个有效token(outcome reward 设计,稀疏奖励)reward_tensor[i, valid_response_length - 1] = score
3.4.3 加载流程
自定义 reward function 的加载流程。
run_train.sh中设置:custom_reward_function.path = verl/utils/reward_score/my_reward_score.py||▼load_reward_manager() 中:compute_score = get_custom_reward_fn(config)# → 动态 import my_reward_score.py,取其 compute_score 函数||▼NaiveRewardManager(compute_score=compute_score)# → 每次调用时执行my_reward_score.compute_score()
load_reward_manager()加载逻辑
config.reward_model.custom_reward_function.path 有值?↓ 是动态 import 外部.py 文件 →加载自定义 compute_score(MemPo 用的就是my_reward_score.py)config.reward_model.reward_manager ="naive"(默认)↓注册表查找→NaiveRewardManager最终:NaiveRewardManager(tokenizer,compute_score=my_reward_score.compute_score)
我们在最后一章会介绍 Mem Reward 机制 & 混合使用。