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 是奖励系统的入口,提供了:
- 灵活的奖励函数加载:支持动态导入用户定义的奖励函数
- 多种评分模式:自定义函数、沙箱执行、默认评分
- 同步/异步支持:自动检测并适配同步或异步奖励函数
- 参数注入:通过配置文件传递额外参数给奖励函数
这种设计使得用户可以很方便地定义自己的奖励逻辑,而不需要修改 verl 的源代码。