自定义 Reward 函数与数据

RLHF 训练的"灵魂"是 reward。OpenRLHF 把 reward 抽象得很灵活——可以是 RM 模型、可以是 Python 函数、 可以是远端 HTTP 服务、甚至可以是一个"Agent function"和环境交互。本章教你怎么写自定义 reward + 数据格式， 以及 v0.10 新加的"agent / VLM"功能。

10.1四种 reward 来源

10.2写一个本地规则 reward

仓库自带 examples/python/reward_func.py 是最小例子：

import torch

def reward_func(queries, prompts, labels, **kwargs):
    """
    Args:
        queries (List[str]):  prompt + response 的完整文本
        prompts (List[str]):  仅 prompt
        labels (List[str]):   ground truth（数据里 label_key 字段）
    Returns:
        dict: 包含 rewards / scores / extra_logs 三个 key
    """
    reward = torch.randint(0, 2, (len(queries),)).float()

    return {
        "rewards": reward,                       # 用于 advantage 计算
        "scores": reward,                        # 用于 dynamic filtering（0-1 之间）
        "extra_logs": {"dummy_scores": reward},  # 写到 wandb
    }

三个 key 的含义：

10.3数学题 reward 函数（实战范本）

examples/python/math_reward_func.py：

from openrlhf.utils import extract_boxed_answer, grade_answer

def reward_func(queries, prompts, labels, **kwargs) -> dict:
    rewards = []

    for query, prompt, label in zip(queries, prompts, labels):
        # 抠出 response 部分
        if isinstance(prompt, str) and prompt in query:
            response = query[len(prompt):]
        else:
            response = query

        # 抠 \boxed{} 内的答案
        pred_answer = extract_boxed_answer(response)
        # 用 math_verify 库判断对错
        is_correct = grade_answer(pred_answer, label)
        rewards.append(1.0 if is_correct else 0.0)

    rewards_tensor = torch.tensor(rewards, dtype=torch.float)
    accuracy = rewards_tensor.mean()

    return {
        "rewards": rewards_tensor,
        "scores": rewards_tensor,
        "extra_logs": {"math_accuracy": accuracy},
    }

这种"答案对就 +1，错就 0"是 R1 风格训练的标准模板。

10.4多 reward 加权组合

常见诉求：

$r = w_1 \cdot r_\text{correct} + w_2 \cdot r_\text{format} + w_3 \cdot r_\text{length}$

def reward_func(queries, prompts, labels, **kwargs):
    rewards = []
    correct_log, format_log, length_log = [], [], []

    for query, prompt, label in zip(queries, prompts, labels):
        response = query[len(prompt):]
        pred = extract_boxed_answer(response)

        r_correct = 1.0 if grade_answer(pred, label) else 0.0
        r_format = 0.1 if has_proper_format(response) else -0.1
        r_length = -0.001 * max(0, len(response) - 2048)

        rewards.append(r_correct + r_format + r_length)
        correct_log.append(r_correct)
        format_log.append(r_format)
        length_log.append(r_length)

    return {
        "rewards": torch.tensor(rewards),
        "scores": torch.tensor(correct_log),       # filtering 只看正确性
        "extra_logs": {
            "r_correct": torch.tensor(correct_log).mean(),
            "r_format":  torch.tensor(format_log).mean(),
            "r_length":  torch.tensor(length_log).mean(),
        },
    }

10.5数据格式

OpenRLHF 用 HuggingFace datasets / parquet / jsonl 都行。关键是 --data.input_key 和 --data.label_key 指向哪个字段。

偏好对齐（PPO + RM）的 prompt 数据

// 一条样本
{
  "context_messages": [
    {"role": "user", "content": "请解释什么是相对论"}
  ]
}

# 启动时
--data.input_key context_messages \
--data.apply_chat_template               # 自动拼模板

数学题数据（GRPO + 规则 reward）

// 一条样本
{
  "prompt": "求方程 x² + 2x - 3 = 0 的解",
  "label": "x=1 或 x=-3"
}

# 启动时
--data.input_key prompt \
--data.label_key label \
--data.apply_chat_template

偏好对（RM 训练用）

// 一条样本
{
  "prompt": "...",
  "chosen": "...",     // 被偏好的回复
  "rejected": "..."    // 被拒绝的回复
}

# 用 train_rm 训练 RM
python -m openrlhf.cli.train_rm \
  --pretrain <SFT 模型> \
  --dataset <偏好数据> \
  --apply_chat_template \
  --max_len 4096 \
  --save_path ./rm_output

10.6Remote URL reward

把 reward 计算放外部服务，训练时调 HTTP：

启动远端 RM 服务

python -m openrlhf.cli.serve_rm \
  --reward_pretrain <RM 模型> \
  --port 8000 \
  --max_len 4096

训练时连过去

--reward.remote_url http://reward-server:8000/score

用途：

• 把 RM 部署到独立机器，省训练机器显存；
• 多 trainer 共享一个 RM；
• 用 GPT-4 / Claude 当 judge，写一个 HTTP wrapper。

10.7Agent 函数：环境交互式 RL

v0.10 引入。每次"调 reward" 都让模型与某个环境进行多步交互，最后给一个 reward。 仓库自带 examples/python/agent_func.py 是范本。

典型 agent_func 签名

async def step(observation, action, **kwargs):
    """
    每一步交互。
    Args:
        observation: 当前环境状态（字符串 / 图像）
        action:      模型刚生成的 action（string）
    Returns:
        next_observation: 下一步环境给模型看的内容
        reward:           当前步 reward（可以是 0，最后一步给总分）
        done:             bool
        info:             extra logs
    """
    # 例如：在数学题环境里
    if is_final_answer(action):
        is_correct = grade(action, ground_truth)
        return None, 1.0 if is_correct else 0.0, True, {}
    else:
        # 多轮思考链
        return f"<refine>{action}</refine>", 0.0, False, {}

启用方式

--train.agent_func_path examples/python/agent_func.py \
--rollout.max_new_tokens 32000 \
--algo.advantage_estimator group_norm \
--rollout.n_samples_per_prompt 8

OpenRLHF 把 agent_func 装到 vLLM 周围，每个 prompt 对应一个 agent rollout，中间多轮交互、最后一步反馈 reward。 适合：

• 多轮工具调用（function calling）；
• 多轮反思 / 自我纠错；
• 代码执行环境（写代码 → 跑测试 → 修改）。

10.8多轮 VLM Agent

v0.10 进一步扩展到 VLM（视觉语言）。examples/python/vlm_multiturn_agent.py：每次交互 observation 可以是 新的图像（如截图、相机帧），action 是模型生成的文本指令。

使用脚本：

bash examples/scripts/train_vlm_math_hybrid_engine.sh

这是当前最前沿的 RLHF 形态——把 LLM Agent 引入"图像 → 推理 → 行动"完整闭环。

10.9dataset 配置常用 flag

10.10常用公开数据集

10.11reward 调试的 3 个技巧

1. 用 max_samples=8 dry run：先确认 reward_func 能正常返回，shape 对，类型对。
2. 把 extra_logs 打满：所有中间 reward 信号都写到 wandb，方便排查 reward hacking。
3. 跑一次 SFT 模型的 baseline：用 SFT 模型直接采样、过 reward_func，看分数分布。如果都 0 或都 1 说明数据 / 函数有问题。

10.12reward hacking 的早期信号

遇到这些现象，第一时间降低 lr + 升 KL coef，把 policy 拉回锚定。

10.13Agent 接口源码锚（openrlhf/utils/agent.py）

ch06.14 提到了 AgentExecutor。"我要写一个多轮 agent" 时实际改的是这两个基类：

典型 instance 实现（参考 examples/python/agent_func.py）：

class MathAgentInstance(AgentInstanceBase):
    def reset(self, prompt, **kw):
        self.problem = prompt
        return {"prompt": prompt}

    def step(self, action):
        # action 是模型刚生成的 token / 字符串
        # 返回工具调用结果作为下一轮 observation
        if "<tool_call>" in action:
            code = extract_code(action)
            result = sandbox.run(code)
            return {"prompt": result}, 0.0, False, {}
        else:
            answer_correct = check_answer(self.problem, action)
            return {}, float(answer_correct), True, {"correct": answer_correct}

启用：--reward.remote_url examples/python/agent_func.py --train.agent_async_enable。OpenRLHF 自动给每条 prompt 起一个 instance，跑完整条 trajectory 后才进 RL。

10.14VLM 多轮 agent（2026-04 新）

2026-04 之后 OpenRLHF 把 agent 框架扩到了 VLM。看 vlm_multiturn_agent.py 与默认文本 agent 的差异：

用法：--vlm_enable --reward.remote_url examples/python/vlm_multiturn_agent.py。典型场景：网页 agent / GUI agent / 视觉推理。

10.15Remote HTTP reward server 模式

"reward 在另一台机器上跑"这个需求很常见（外部沙箱 / 大型 RM）。OpenRLHF 通过 --reward.remote_url http://... 支持：

"为什么不直接 RPC？" → HTTP 跨语言、跨集群最方便。reward server 可以是 Python / Rust / Go，OpenRLHF 端只看 JSON。

10.16这章你需要带走的

• OpenRLHF reward 来源四种：RM 模型 / 本地函数 / Remote HTTP / Agent 函数；
• 函数签名固定：reward_func(queries, prompts, labels, **kw) -> {rewards, scores, extra_logs}；
• 多 reward 加权组合在函数内做；
• v0.10 新增的 Agent / VLM 模式让 RLHF 能做多轮交互；
• 调试 reward 必走 dry run + baseline + extra_logs 三连。