[Agent Memory / 强化学习] MemPO源码学习笔记 --- (6)--- 奖励机制

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 的信息如下:

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)

0xFF 参考