Agent Memory / 强化学习 MemPO源码学习笔记 --- (6)--- 奖励机制
目录
- [Agent Memory / 强化学习 MemPO源码学习笔记 --- (6)--- 奖励机制](#[Agent Memory / 强化学习] MemPO源码学习笔记 --- (6)--- 奖励机制)
- [0x00 概要](#0x00 概要)
- [0x01 原理](#0x01 原理)
- [1.1 双通路奖励机制](#1.1 双通路奖励机制)
- [1.2 通俗解释优势函数](#1.2 通俗解释优势函数)
- [Outcome Advantage](#Outcome Advantage)
- [Memory Advantage](#Memory Advantage)
- [Final Advantage](#Final Advantage)
- [1.3 对比](#1.3 对比)
- [1.4 Reward model](#1.4 Reward model)
- [Reward Model](#Reward Model)
- 对比
- 代码
- reward_model
- MemPO
- [路径 1:规则奖励(MemPO实际使用)](#路径 1:规则奖励(MemPO实际使用))
- [路径 2:Reward Model推理(非MemPO场景)](#路径 2:Reward Model推理(非MemPO场景))
- 决策树总结
- [0x02 Outcome Reward 机制](#0x02 Outcome Reward 机制)
- [2.1 Outcome Advantage公式详解](#2.1 Outcome Advantage公式详解)
- [2.2 EM Check 仅用于 Outcome Advantage](#2.2 EM Check 仅用于 Outcome Advantage)
- [2.3 EM在MemPO中的作用](#2.3 EM在MemPO中的作用)
- [2.4 校验规则](#2.4 校验规则)
- [2.5 normalize_answer四步处理](#2.5 normalize_answer四步处理)
- [2.6 思考](#2.6 思考)
- [0x03 Outcome Reward 实现](#0x03 Outcome Reward 实现)
- [0xFF 参考](#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://huggingface.co/collections/NewBeeKing/mempo
0x01 原理
我们回顾前文,代码具体路径上的关键点如下:
python
A1 _postprocess(P_mem/P_full段) MemPO核心:记忆奖励如何计算
A2 compute_grpo_memory_advantage mem_adv如何归一化、作用于哪些 token
A3 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 check
B3 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 通俗解释优势函数
想象你在参加一个读书闯关比赛:
读一本很长的故事书,每读完一章就要做一次笔记,最后回答老师的问题。
Outcome Advantage
🏆 Outcome Advantage(结果分):"你最终答对了吗?"
老师看你的最终答案 --- 对了加分,错了扣分。
但不是跟满分比,而是跟你同桌比(GRPO:同组 16 个同学做同一道题)。
- 全班答对率 60%,你答对了 → 你比平均好 → 鼓励
- 全班答对率 60%,你答错了 → 你比平均差 → 抑制
这个分数给你写的每个字 --- 答案、推理、搜索,全部同样鼓励或抑制。
Memory Advantage
📝 Memory Advantage(笔记质量分):"你的笔记够好吗?好到只看笔记就能答题?"
测试方法:
- 让另一个你看完整本书,去答题 → 答对概率 P_full
- 让另一个你只看你的笔记,去答题 → 答对概率 P_mem
如果只看笔记也能答对 → 笔记写得好 → 奖励你的笔记 如果只看笔记答不出来 → 笔记漏了关键信息 → 惩罚你的笔记
这个分数只给你写笔记的那几行字 --- 推理和答案不受影响。
Final Advantage
🎯 Final Advantage(最终综合分)
"把结果分和笔记分加在一起。"
-
你写的搜索内容:只看结果分
-
你写的笔记:结果分 + 笔记分 ← 双重关注!
-
你写的最终答案:只看结果分
为什么笔记要特殊对待?因为笔记写得好不好,直接影响后面几章你能不能记住前面的内容。如果笔记都不单独教,你永远学不会做好笔记。
1.3 对比
其实,Outcome Advantage 是结果奖励,Memory Advantage 类似过程奖励。
| Outcome Advantage | Memory Advantage | |
|---|---|---|
| 公式 | (score - mean) / std | (P_mem-P_full - mean) / std |
| 类似于 | ✔ 结果奖励 (Result Reward) | 类似过程奖励 |
| 评什么 | 最终答案是否正确 (EM match) | 摘要是否有效压缩了信息 / 笔记能否替代全文 |
| 奖励信号来源 | 离散:em_check / score ∈ | log_prob 概率差:P_mem - P_full ∈ (-1, +1) (连续值) |
| 归一化池 | 同question同batch(16条 轨迹) | 同question所有轮次池化(16xN轮),比如~48个值(16轨迹×3轮) |
| 作用范围 | 全序列所有 token | 仅... 区间的 token |
| 广播方式 | outcome reward 广播到全序列 | 逐轮写入,不重叠 |
| 本质 | 你答对了吗?/ 这条路走对了吗 | "你的记忆摘要能否替代完整上下文来预测正确答案?/ 这个笔记写得好吗 |
作用范围
Outcome Advantage和Memory Advantage的作用范围有什么区别? 、
- Outcome Advantage:广播到全序列每个token(同一条轨迹内所有token相同值)
- Memory Advantage:仅赋值到...对应的token区间,其余位置=0
Memory Advantage
Memory Advantage ≠ 格式奖励:
- 格式奖励(如 B3 validate_format)是硬约束:格式错 → 直接score=0,整条轨迹废掉
- Memory Advantage是软信号:衡量的内容质量,而非格式是否合规
更准确的类比:
- 格式奖励(B3):"你的输出格式合法吗?" → 不合法就没分(门槛)
- 结果奖励(B4):"你最终答对了吗?" → outcome_adv
- 记忆质量奖励(A1):"你的摘要能替代全文吗?" → mem_adv
所以mem_adv更像是一个局部过程奖励(process reward)一不是评格式,而是评「记忆摘要的信息压缩质量」,与全局的结果奖励叠加使用。
配合
两者如何配合?final_adv = outcome_adv + mem_adv
- 场景1:答对 + 好mem → (+outcome)+(+mem) = 强。正向强化
- 场景2:答对 + 差mem → (+outcome)+(-mem) = 正向,但被惩罚
- 场景3:答错 + 好mem → (-outcome)+(+mem) = 整体负,但部分被保护
- 场景4:答错 + 差mem → (-outcome)+(-mem)= 全面惩罚
不同路径
两条路径的评分方式完全不同:
Outcome路径(B系列):
- 评分方式:em_check(模型预测答案,标准答案)→
- 信号含义:最终答案对不对
Memory路径(A系列):
- 评分方式:P_mem-P_full(概率差)
- 信号含义:记忆摘要质量
- 不涉及任何字符串匹配,纯粹是模型内部的 log_prob对比
Outcome Advantage 路径
python
B6 (AgentLoopManager.generate_sequences) --- 16条轨迹并发 rollout
↓ 轨迹 DONE
B5 (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 路径
python
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]
汇合点
python
A3 (compute_advantage) --- final_adv = outcome_adv + mem_adv
↓
PPO Update --- loss.backward()
路径对比
python
┌──────────┬─────────────────────────────────┬─────────────────┐
│ │ Outcome 路径 │ Memory 路径 │
├──────────┼─────────────────────────────────┼─────────────────┤
│ 标签链 │ B6→B5→B1→B2→(C2,B3,B4)→B4-algo │ A4→A1→A2 │
├──────────┼─────────────────────────────────┼─────────────────┤
│ 汇合 │ → A3 ← │ → A3 ← │
├──────────┼─────────────────────────────────┼─────────────────┤
│ 评分方式 │ 字符串匹配 │ 概率差 │
├──────────┼─────────────────────────────────┼─────────────────┤
│ 产出范围 │ 全序列等值 │ 仅<mem>区间 │
└──────────┴─────────────────────────────────┴─────────────────┘
1.4 Reward model
MemPO 没有使用Reward Model。
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(通常需要人工标注中间步骤的正确性)。
对比
对比需要Reward Model的系统,这也是 MemPO 能用 GRPO 而不用 PPO+Critic 的另一个原因------奖励信号干净确定,不需要 value network 来降低估计方差。
python
┌──────────────┬────────────────────────┬─────────────────────────────┐
│ │ 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(自评估概率差)。
代码
代码中确实有 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实际使用的奖励路径:
python
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实际使用)
python
rm_executor = None
reward_model.enable = True
reward_model.enable_resource_pool = False
enable_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
具体展开
python
情况 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场景)
python
┌─────────────────────────────────────────────────────────┐
│ BatchExecutor (合批器) │
│ │
│ 多条轨迹同时调用 submit_task() │
│ ↓ 放入队列 │
│ worker 线程持续从队列取出,攒够 micro_batch_size 后 │
│ 一次性调用 rm_wg.compute_rm_score(batch) │
│ ↓ │
│ GPU 上的 Reward Model 打分 (Bradley-Terry 偏好模型等) │
│ ↓ │
│ 结果通过 Future 返回给各条轨迹 │
└─────────────────────────────────────────────────────────┘
BatchExecutor的关键设计:解决"多条轨迹异步完成但RM推理需要合批"的问题,把稀疏到来的单条请求动态合批,最大化GPU 利用率。
决策树总结
python
你用的是什么奖励?
│
├── 规则奖励 (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公式详解
公式如下:
python
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])
含义拆解如下:
python
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的完整逻辑如下:
python
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。
模型输出如下:
python
<mem>上轮搜索...</mem>
<think>让我分析...</think>
<search>query</search>
<tool_response>文档</tool_response>
<answer>obama</answer>
奖励计算链如下:
python
1.validate_format() → 格式合法?(8条规则(11项子检查))
2.extract_solution() → 提取最后-个<answer>内容→"Obama
3.em_check("Obama",["Barack Obama","Obama"]) → 1(答对)
最终奖励如下:
python
score=1(答对)或0(答错或格式非法)
写入 reward_tensor[i,最后一个token]
关键特性:
- 稀疏:整条response(可能1000+个token),只有最后1个位置有非零奖励
- 二值:只有0/1,没有部分分
- 隐式课程:格式不合法也是0分>迫使模型先学会格式
- 与mem_reward互补:EM告诉模型"答对了吗",P_mem-P_full告诉模型"记忆写得好不好"
比如,对同一道题的16个侦探(index相同的样本):
python
侦探得分:[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 () 校验的完整规则如下:
python
- 规则 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."被视为相同。
python
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)。
特色
职责:
- 解码:response_ids→response_str(用 tokenizer 解码)
- 提取数据:从data_item 中取出 ground_truth、data_source
- 调用评分: self.compute_score(response_str, ground_truth) → score
- 写入结果:reward_tensori,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。
运作机制
奖励放在哪个token上?具体如下:
python
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。
问题
为什么需要这层封装?这是因为如下问题:
Agent Loop 是异步的 (asyncio):
- 16条轨迹并发生成
- 某条轨迹先完成→立即触发奖励计算
- 不需要等其他轨迹完成
'NaiveRewardManager.__call__()'是同步阻塞的(CPU密集)
- 如果直接调用,会阻塞asyncio事件循环
- 其他轨迹的生成也会被卡住
关键设计
关键设计是:与 agent loop 并行。
为什么做成异步 Ray actor:奖励计算(尤其是 reward model 推理)很耗时,异步化可以让多条轨迹的奖励计算与其他轨迹的 rollout重叠执行,提升 GPU利用率。
python
AgentLoopWorker (每个 trajectory):
...生成下一个 token...
...调用搜索工具...
...生成结束 (<answer>
← 此时异步提交奖励计算 →
result = await compute_score.remote(data)
↑ 非阻塞,其他 trajectory 可以同时运行
output.reward_score = result["reward_score"]
解决方案
RewardManagerWorker 解决方案:
- 用 run_in_executor()把阻塞操作丢到线程池
- asyncio事件循环不被阻塞
- 奖励计算和其他轨迹的 rollout同时进行
B5是 B1的"async adapter"---让同步的奖励计算函数能在异步 agent loop 中并行执行,不阻塞其他轨迹的 rollout。
3.4 自定义 reward function
异步接口层
RewardManagerWorker.compute_score 是 MemPO 中奖励计算的异步接口层。
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
python
run_train.sh: custom_reward_function.path = "verl/utils/reward_score/my_reward_score.py"
属性调用链如下:
python
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 ∈
另外,compute_score 是 Outcome Advantage 专有的函数,Memory Advantage 完全不使用它。
compute_score()的三种 target 处理
python
target 类型 打分方式
────────────────────────────────────────────────────
list[list[str]] 多目标:answer按;分割,逐个EM匹配累加分
list[str] 单目标多别名:answer与任意一个EM匹配得1分
str 单一字符串:包装成[str]后EM匹配
奖励值放置位置
python
response token 序列:
[t0][t1][t2]...[t_{n-2}][t_{n-1}]
↑
reward 放在最后一个有效token
(outcome reward 设计,稀疏奖励)
reward_tensor[i, valid_response_length - 1] = score
加载流程
自定义 reward function 的加载流程。
python
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()加载逻辑
python
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)