跳转至

reward.py — 该文件负责奖励函数的加载和管理

文件概述

模块路径: verl.trainer.ppo.reward

该文件负责奖励函数的加载和管理。它提供了从外部文件动态加载自定义奖励函数的能力,以及初始化奖励管理器(Reward Manager)的功能。

在训练流程中的位置

在训练和评估过程中,奖励函数用于给模型生成的回复打分。这是 RLHF 的关键环节——奖励信号指导策略的优化方向。

关键代码讲解

1. 动态加载自定义奖励函数

def get_custom_reward_fn(config) -> Optional[RawRewardFn]:
    """从外部文件动态加载奖励函数"""
    reward_fn_config = config.reward.get("custom_reward_function") or {}
    module_path = reward_fn_config.get("path")
    if not module_path:
        return None

    fn_name = reward_fn_config.get("name")
    assert fn_name is not None

    # 动态加载外部模块中的函数
    from verl.utils.import_utils import load_extern_object
    raw_fn = load_extern_object(module_path=module_path, object_name=fn_name)

    # 将额外参数绑定到函数上
    reward_kwargs = dict(reward_fn_config.get("reward_kwargs", {}))
    if not inspect.iscoroutinefunction(raw_fn):
        return partial(_call_with_kwargs, raw_fn, reward_kwargs)
    else:
        return partial(_call_with_kwargs_async, raw_fn, reward_kwargs)

这个设计非常灵活:用户只需在配置文件中指定奖励函数的文件路径和函数名,verl 就能自动加载和使用。同时支持同步和异步奖励函数。

2. 参数合并辅助函数

def _call_with_kwargs(raw_fn, extra_kwargs, *args, **kwargs):
    """将额外参数合并到函数调用中"""
    merged_kwargs = {**kwargs, **extra_kwargs}
    return raw_fn(*args, **merged_kwargs)

async def _call_with_kwargs_async(raw_fn, extra_kwargs, *args, **kwargs):
    """异步版本"""
    merged_kwargs = {**kwargs, **extra_kwargs}
    return await raw_fn(*args, **merged_kwargs)

extra_kwargs 来自配置文件中的 reward_kwargs,优先级高于运行时传入的 kwargs。

3. 奖励管理器加载

def load_reward_manager(config, tokenizer, **reward_kwargs) -> RewardManagerBase:
    """加载并初始化奖励管理器"""

    # 1. 尝试加载自定义奖励函数
    compute_score = get_custom_reward_fn(config)
    final_compute_score = compute_score

    # 2. 根据配置选择奖励管理器类
    reward_manager_cfg = config.reward.reward_manager
    if reward_manager_cfg.source == "register":
        # 从注册表获取
        reward_manager_cls = get_reward_manager_cls(reward_manager_cfg.name)
    elif reward_manager_cfg.source == "importlib":
        # 从外部模块导入
        reward_manager_cls = load_extern_object(
            module_path=module_cfg.path,
            object_name=reward_manager_cls_name
        )

    # 3. 如果没有自定义奖励函数,使用默认的
    if compute_score is None:
        sandbox_config = config.reward.get("sandbox_fusion")
        sandbox_url = sandbox_config.get("url") if sandbox_config else None
        if sandbox_url:
            # 使用沙箱执行(安全的代码执行环境)
            _concurrent_semaphore = sandbox_manager.Semaphore(max_concurrent)
            final_compute_score = partial(
                default_compute_score,
                sandbox_fusion_url=sandbox_url,
                concurrent_semaphore=_concurrent_semaphore,
                memory_limit_mb=memory_limit_mb,
            )
        else:
            final_compute_score = default_compute_score

    # 4. 实例化奖励管理器
    return reward_manager_cls(
        config=config,
        tokenizer=tokenizer,
        compute_score=final_compute_score,
        **reward_kwargs,
    )

奖励管理器支持三种模式: - 自定义奖励函数:用户提供的 Python 函数 - 沙箱执行:在安全环境中执行代码(如代码正确性验证) - 默认评分:使用内置的默认评分函数

4. 奖励提取

def extract_reward(batch: DataProto):
    """从 batch 中提取奖励张量和额外信息"""
    reward_tensor = batch.batch["rm_scores"]
    reward_extra_keys = batch.meta_info.get("reward_extra_keys", [])
    reward_extra_infos_dict = {key: batch.non_tensor_batch[key] for key in reward_extra_keys}
    return reward_tensor, reward_extra_infos_dict

这个函数在 ray_trainer.py 的 fit() 中被调用,将奖励模型的输出提取出来供后续计算使用。

核心类/函数列表

名称 类型 作用
get_custom_reward_fn(config) function 从外部文件加载自定义奖励函数
load_reward_manager(config, tokenizer) function 加载并初始化完整的奖励管理器
extract_reward(batch) function 从 batch 中提取奖励张量
_call_with_kwargs() function 合并额外参数的函数调用包装器
_call_with_kwargs_async() function 异步版本的参数合并包装器

数据流和调用关系

配置文件 (reward.custom_reward_function.path/name)
      |
      v
get_custom_reward_fn(config)
      |
      v
load_reward_manager(config, tokenizer)
      |
      +-- 自定义函数? --> partial(_call_with_kwargs, raw_fn, kwargs)
      +-- 沙箱模式? --> partial(default_compute_score, sandbox_url=...)
      +-- 默认模式? --> default_compute_score
      |
      v
RewardManagerBase(compute_score=final_compute_score, ...)

ray_trainer.py:
      +-- extract_reward(batch) --> reward_tensor, reward_extra_infos_dict

小结

reward.py 是奖励系统的入口,提供了:

  1. 灵活的奖励函数加载:支持动态导入用户定义的奖励函数
  2. 多种评分模式:自定义函数、沙箱执行、默认评分
  3. 同步/异步支持:自动检测并适配同步或异步奖励函数
  4. 参数注入:通过配置文件传递额外参数给奖励函数

这种设计使得用户可以很方便地定义自己的奖励逻辑,而不需要修改 verl 的源代码。