RLOO 训练器

概述

TRL 支持 RLOO Trainer，用于训练语言模型，其方法如 Arash Ahmadian、Chris Cremer、Matthias Gallé、Marzieh Fadaee、Julia Kreutzer、Ahmet Üstün 和 Sara Hooker 在论文 Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs 中所述。

该论文的摘要如下：

以基于人类反馈的强化学习（Reinforcement Learning from Human Feedback, RLHF）形式的人工智能对齐，正日益被视为高性能大语言模型的关键要素。近期文献将近端策略优化（Proximal Policy Optimization, PPO）定位为 RLHF 中强化学习部分的标准方法。然而，PPO 涉及高昂的计算成本以及对超参数调优的敏感性。我们认为，促使 PPO 开发的许多动机原则在 RLHF 中并非实际的关注点，并主张采用一种计算成本更低且能保持甚至提升性能的方法。我们在强化学习的背景下重新审视了基于人类偏好的对齐问题。以简洁性为指导原则，我们证明在 RLHF 场景中 PPO 的许多组件是不必要的，且远为简单的 REINFORCE 风格优化变体在性能上优于 PPO 以及新提出的“无 RL”方法（如 DPO 和 RAFT）。我们的研究表明，针对大语言模型对齐特性的精心适配，能够以较低的成本从在线强化学习优化中获益。

此训练后方法由 Costa Huang 贡献，随后由 Shirin Yamani 进行了重构。

快速入门

本示例演示了如何使用 RLOO 方法训练模型。我们使用 DeepMath-103K 数据集 中的提示词来训练 Qwen 0.5B Instruct 模型。您可以在这里查看该数据集的数据：

以下是用于训练模型的脚本。accuracy_reward 需要 math-verify，您可以通过 pip install trl[math_verify] 进行安装。

# train_rloo.py
from datasets import load_dataset
from trl import RLOOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = RLOOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)
trainer.train()

使用以下命令执行该脚本：

torchrun --nproc_per_node 8 train_rloo.py

深入了解 RLOO 方法

RLOO 是一种在线学习算法，这意味着它通过利用训练过程中由已训练模型自身生成的数据来迭代地提升性能。RLOO 目标背后的直觉是最大化生成补全的优势，同时确保模型保持在参考策略附近。为了理解 RLOO 的工作原理，可以将其分解为四个主要步骤：生成补全、计算优势、估计 KL 散度以及计算损失。

生成补全

在每个训练步骤中，我们采样一批提示词，并为每个提示词生成一组 $G$ 补全（记为 $o_i$).

计算奖励

在 RLOO 中，奖励由两部分组成：奖励模型（或奖励函数）提供的奖励，以及一个 KL 惩罚项，用于阻止策略偏离固定的参考策略过远

1. 对于每个 $G$ 生成的序列 $o_i = (o_{i,1}, \dots, o_{i,T})$ 在查询 $q$, 的条件下，我们使用奖励模型计算一个标量奖励 $R(o_i, q)$.
2. 同时，我们估计当前策略 $\pi_\theta$ 与固定的参考策略 $\pi_{\text{ref}}$ 在序列上的 KL 散度。序列 $o_i$ 的 KL 估计值为： $$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta\|\pi_{\mathrm{ref}}\right] = \sum_{t=1}^T \log \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}.$$

分配给序列的最终奖励 $o_i$ 为： $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

其中 $\beta > 0$ 控制 KL 惩罚的强度。

在纯在线设置（num_iterations = 1，默认值）中，数据由当前策略生成。在这种情况下，KL 惩罚直接使用当前策略计算。

在更一般的情况下（例如，每个批次执行多个梯度步骤），数据实际上是由较早的快照生成的 $\pi_{\text{old}}$. 为了使惩罚项与采样分布保持一致，KL 散度是相对于该策略定义的：

$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_{\text{old}} \,\|\, \pi_{\text{ref}}\right].$

等价地，对于采样的序列 $o$, 蒙特卡洛估计为：

$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_{\text{old}} \|\pi_{\mathrm{ref}}\right] = \sum_{t=1}^T \log \frac{\pi_{\text{old}}(o_{i,t} \mid q, o_{i,<t})}{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}.$

计算优势

在计算完每个补全的奖励后，我们计算一个基线，即同一批次中除当前样本外所有其他样本的平均奖励。该基线用于降低策略梯度估计的方差。随后，每个补全的优势值通过其自身奖励与该留一法基线之差获得。

形式上，对于一批 G 个补全，补全的基线为：$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

然后，每个补全的优势值计算为其奖励与基线之间的差值： $$A_i = r_i - b_i$$

计算损失

REINFORCE 损失函数定义如下： $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

在实践中，对同一批次执行多次梯度更新会使动作相对于当前参数实际上变为离策略（off-policy）。为了纠正这一点，我们引入了重要性采样比率。为了防止在采样和梯度更新之间策略发生变化时出现过大的更新，我们对这个比率进行裁剪： $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \min \left( \frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} \hat{A}_i, \, \text{clip}\left(\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}, 1-\epsilon, 1+\epsilon\right) \hat{A}_i \right)$$

在完全在线的单步设置（默认）中， $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ 这简化为标准的 REINFORCE。

与 GRPO 不同，RLOO 不会通过 KL 项反向传播梯度。在这里，KL 纯粹用于奖励整形：它在 no_grad 下计算，并折叠到标量奖励 $r_i$, 中，该奖励成为分离的 advantage $\hat{A}_i$.

换句话说，梯度进入策略的唯一路径是重要性比率项 $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, 而 KL 项不会反向传播到策略中。

相比之下， GRPO 在其目标函数中引入了一个显式的、可微的 KL 项，该 KL 项基于当前策略 $\pi_\theta$, 实时计算，因此其梯度会同时流经比率项和 KL 惩罚项。

记录的指标

在训练和评估过程中，我们记录以下指标：

• num_tokens：迄今为止处理的 token 总数，包括提示词和补全部分。
• step_time：每个训练步骤（包括生成）的平均耗时（以秒为单位）。
• completions/mean_length：生成的补全的平均长度。
• completions/min_length：生成补全的最小长度。
• completions/max_length：生成的补全的最大长度。
• completions/mean_terminated_length：以 EOS 终止的生成补全的平均长度。
• completions/min_terminated_length：以 EOS 终止的生成补全的最小长度。
• completions/max_terminated_length：以 EOS 终止的生成补全的最大长度。
• completions/clipped_ratio：被截断（裁剪）的补全比例。
• rewards/{reward_func_name}/mean：特定奖励函数的平均奖励。
• rewards/{reward_func_name}/std：特定奖励函数的奖励标准差。
• reward：对各个函数的奖励求和（按 reward_weights 加权）后得到的总体平均奖励。
• reward_std：跨函数求和奖励的标准差（按 reward_weights 加权），在完整批次上计算。
• frac_reward_zero_std：生成批次中奖励标准差为零的样本比例，表明该提示词下的多样性较低（所有回答均正确或均错误）。
• entropy：生成的补全中 token 预测的平均熵。（如果 mask_truncated_completions=True，则排除被掩码序列的 token。）
• aux_loss：混合专家（Mixture-of-Experts）模型的负载均衡辅助损失，在乘以 router_aux_loss_coef 并加到总损失之前。仅当模型为 MoE 模型且 router_aux_loss_coef 非零时记录。
• kl：模型与参考模型之间的平均 KL 散度，基于生成的补全内容计算。仅在 beta 非零时记录。
• clip_ratio/region_mean: RLOO 目标被裁剪以保持在信任区域内的序列概率比率： $\text{clip}\left( r_{i}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)\,, \quad r_{i}(\theta) = \frac{\pi_\theta(o_{i} \mid q)}{\pi_{\theta_{\text{old}}}(o_{i} \mid q)}$. 较高的值意味着更多的样本被裁剪，从而限制了策略 $\pi_\theta$ 可以变化的幅度。
• clip_ratio/low_mean: 信任区域下界被裁剪的序列概率的平均比率： $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: 信任域下界上每个 completion 的最小裁剪指示器： $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: 信任区域上界被裁剪的序列概率的平均比率： $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: 信任域上界上最大的逐完成裁剪指示器： $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

自定义

使用 vLLM 驱动的生成来加速训练

在使用在线方法训练时，生成过程往往是主要的性能瓶颈。为了加速生成，你可以使用 vLLM，这是一个面向 LLM 的高吞吐量、低延迟推理引擎。要启用它，首先使用以下命令安装该包

pip install trl[vllm]

我们支持在训练期间使用 vLLM 的两种方式：服务器模式和共置模式。

选项 1：Colocate 模式

在此模式下，vLLM 在训练器进程内运行，并与训练模型共享 GPU 内存。这避免了启动单独的服务器，并可能提高 GPU 利用率，但可能导致训练 GPU 上的内存争用。这是默认模式。

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

选项 2：服务器模式

在此模式下，vLLM 在独立进程中运行（并使用独立的 GPU），并通过 HTTP 与训练器进行通信。如果您拥有专门用于推理的 GPU，则此模式非常理想。

1. 启动 vLLM 服务器：

   已复制

   VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
       --weight-transfer-config '{"backend": "nccl"}' \
       --logprobs-mode processed_logprobs \
       --max-logprobs -1

2. 在训练脚本中启用服务器模式：

   已复制

   from trl import RLOOConfig
   
   training_args = RLOOConfig(
       ...,
       use_vllm=True,
       vllm_mode="server",
   )

请确保服务器使用的 GPU 与训练器不同，否则可能会遇到 NCCL 错误。您可以通过 CUDA_VISIBLE_DEVICES 环境变量来指定要使用的 GPU。

根据模型大小以及训练所需的整体 GPU 内存，您可能需要调整 RLOOConfig 中的 vllm_gpu_memory_utilization 参数，以避免资源利用率不足或内存溢出错误。

我们提供了一个 HF Space，用于根据您的模型配置和实验设置来估算推荐的 GPU 内存利用率。只需按以下方式使用，即可获得 vllm_gpu_memory_utilization 的推荐值：

如果推荐值在您的环境中不起作用，我们建议在推荐值上添加一个小缓冲（例如 +0.05 或 +0.1），以确保稳定性。

如果您仍然遇到内存不足错误，请将 vllm_enable_sleep_mode 设置为 True，这样在优化步骤期间，vllm 参数和缓存将被卸载。有关更多信息，请参阅 使用 vLLM Sleep Mode 减少内存使用。

默认情况下，RLOO 使用 MASTER_ADDR=localhost 和 MASTER_PORT=12345 作为 vLLM 的配置，但你可以通过相应地设置环境变量来覆盖这些值。

有关更多信息，请参阅 使用 vLLM 加速训练。

大规模 RLOO：在多节点上训练 70B+ 模型

在训练 Qwen2.5-72B 这样的大模型时，你需要多项关键优化，以确保训练在多 GPU 和多节点环境下高效且可扩展。这些优化包括：

• DeepSpeed ZeRO Stage 3：ZeRO 利用数据并行将模型状态（权重、梯度、优化器状态）分布在多个 GPU 和 CPU 上，从而降低每个设备的内存和计算需求。由于大模型无法装入单个 GPU，训练此类模型必须使用 ZeRO Stage 3。更多详情请参阅 DeepSpeed Integration。
• torchrun：PyTorch 的启动器会在每个节点上的每张 GPU 上启动一个进程，并通过公共的会合点将它们连接起来。更多详情，请参阅 分布式训练。
• vLLM：请参阅上一节，了解如何使用 vLLM 来加速生成。

以下是一个示例 SLURM 脚本，用于在多个节点上使用 RLOO 训练 70B 模型。该脚本在 4 个节点上训练模型，并使用第 5 个节点进行基于 vLLM 的生成。

#!/bin/bash
#SBATCH --nodes=5
#SBATCH --gres=gpu:8

# Get the list of allocated nodes
NODELIST=($(scontrol show hostnames $SLURM_JOB_NODELIST))

# Assign the first 4 nodes for training and the 5th node for vLLM
TRAIN_NODES="${NODELIST[@]:0:4}"  # Nodes 0, 1, 2, 3 for training
VLLM_NODE="${NODELIST[4]}"  # Node 4 for vLLM

# Run training on the first 4 nodes (Group 1)
srun --nodes=4 --ntasks=4 --nodelist="${NODELIST[@]:0:4}" torchrun \
     --nnodes 4 \
     --nproc_per_node 8 \
     --node_rank $SLURM_PROCID \
     --rdzv_backend c10d \
     --rdzv_endpoint ${NODELIST[0]}:29500 \
     train_rloo.py \
     --server_ip $VLLM_NODE &

# Run vLLM server on the 5th node (Group 2)
srun --nodes=1 --ntasks=1 --nodelist="${NODELIST[4]}" env VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-72B --tensor-parallel-size 8 \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1 &

wait

import argparse

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--vllm_server_host", type=str, default="", help="The server IP")
    args = parser.parse_args()

    # Example dataset from TLDR
    dataset = load_dataset("trl-lib/tldr", split="train")

    # Dummy reward function: count the number of unique characters in the completions
    def reward_num_unique_chars(completions, **kwargs):
        return [len(set(c)) for c in completions]

    training_args = RLOOConfig(
        output_dir="Qwen2.5-72B-RLOO",
        per_device_train_batch_size=4,
        bf16=True,
        deepspeed="examples/deepspeed_configs/zero3.json",
        use_vllm=True,
        vllm_mode="server",
        vllm_server_host=args.vllm_server_host.replace("ip-", "").replace("-", "."),  # from ip-X-X-X-X to X.X.X.X
    )

    trainer = RLOOTrainer(model="Qwen/Qwen2.5-72B", args=training_args, reward_funcs=reward_num_unique_chars, train_dataset=dataset)
    trainer.train()

if __name__=="__main__":
    main()

使用自定义奖励函数

RLOOTrainer 支持使用自定义奖励函数，而非稠密奖励模型。为确保兼容性，您的奖励函数必须满足以下要求：

奖励函数可以是同步的 Python 可调用对象，也可以是异步的 async def 协程。当您提供多个异步奖励函数时，它们会被并发等待（通过 asyncio.gather 并行执行），从而使它们的延迟相互重叠。

1. 输入参数：

   • 该函数必须接受以下关键字参数：

     • prompts（包含提示词），
     • completions（包含生成的补全内容），
     • completion_ids（包含分词后的补全内容），
     • trainer_state (TrainerState)：训练器的当前状态。可用于实现动态奖励函数，例如课程学习，其中奖励会根据训练进度进行调整。
     • log_extra：一个可调用对象 log_extra(column: str, values: list)，用于向 completions 表添加额外的列。参见示例 6。
     • log_metric：一个可调用对象 log_metric(name: str, value: float)，用于将标量指标记录为图表，与 kl、entropy 等并列显示。参见示例 6。
     • 数据集可能包含的所有列名（prompt 除外）。例如，如果数据集包含名为 ground_truth 的列，则函数将以 ground_truth 作为关键字参数被调用。

     满足此要求的最简单方法是在函数签名中使用 **kwargs。

   • 根据数据集格式的不同，输入会有所不同：

     • 对于标准格式，prompts和completions将是字符串列表。
     • 对于对话格式，prompts和completions将是消息字典的列表。

2. 返回值：该函数必须返回一个浮点数列表。每个浮点数代表对应单个补全的奖励。

示例 1：奖励更长的补全

以下是一个标准格式奖励函数的示例，该函数会奖励更长的补全：

def reward_func(completion_ids, **kwargs):
    """Reward function that assigns higher scores to longer completions (in terms of token count)."""
    return [float(len(ids)) for ids in completion_ids]

您可以按以下方式测试：

>>> prompts = ["The sky is", "The sun is"]  # not used in the reward function, but the trainer will pass it
>>> completions = [" blue.", " in the sky."]  # not used in the reward function, but the trainer will pass it
>>> completion_ids = [[6303, 13], [304, 279, 12884, 13]]
>>> reward_func(prompts=prompts, completions=completions, completion_ids=completion_ids)
[2.0, 4.0]

示例 1.1：奖励更长的补全（基于字符数）

与上一个示例相同，但这次奖励函数基于字符数而非 token 数。

def reward_func(completions, **kwargs):
    """Reward function that assigns higher scores to longer completions (in terms of character count)."""
    return [float(len(completion)) for completion in completions]

您可以按以下方式测试：

>>> prompts = ["The sky is", "The sun is"]
>>> completions = [" blue.", " in the sky."]
>>> completion_ids = [[6303, 13], [304, 279, 12884, 13]]  # not used in the reward function, but the trainer will pass it
>>> reward_func(prompts=prompts, completions=completions, completion_ids=completion_ids)
[6.0, 12.0]

示例 2：奖励具有特定格式的补全

以下是一个奖励函数的示例，用于检查补全内容是否具有特定格式。该示例受论文 DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning 中使用的 format reward 函数启发。 它专为对话格式设计，其中提示和补全内容均由结构化消息组成。

import re

def format_reward_func(completions, **kwargs):
    """Reward function that checks if the completion has a specific format."""
    pattern = r"^<think>.*?</think><answer>.*?</answer>$"
    completion_contents = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, content) for content in completion_contents]
    return [1.0 if match else 0.0 for match in matches]

您可以按以下方式测试此函数：

>>> prompts = [
...     [{"role": "user", "content": "What is the result of (1 + 2) * 4?"}],
...     [{"role": "user", "content": "What is the result of (3 + 1) * 2?"}],
... ]
>>> completions = [
...     [{"role": "assistant", "content": "<think>The sum of 1 and 2 is 3, which we multiply by 4 to get 12.</think><answer>(1 + 2) * 4 = 12</answer>"}],
...     [{"role": "assistant", "content": "The sum of 3 and 1 is 4, which we multiply by 2 to get 8. So (3 + 1) * 2 = 8."}],
... ]
>>> format_reward_func(prompts=prompts, completions=completions)
[1.0, 0.0]

示例 3：基于参考的奖励完成

以下是一个奖励函数的示例，用于检查补全内容是否正确。该示例受论文 DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning 中使用的 accuracy reward 函数启发。 此示例针对 standard format 设计，其中数据集包含一个名为 ground_truth 的列。

import re

def reward_func(completions, ground_truth, **kwargs):
    # Regular expression to capture content inside \boxed{}
    matches = [re.search(r"\\boxed\{(.*?)\}", completion) for completion in completions]
    contents = [match.group(1) if match else "" for match in matches]
    # Reward 1 if the content is the same as the ground truth, 0 otherwise
    return [1.0 if c == gt else 0.0 for c, gt in zip(contents, ground_truth)]

您可以按以下方式测试此函数：

>>> prompts = ["Problem: Solve the equation $2x + 3 = 7$. Solution:", "Problem: Solve the equation $3x - 5 = 10$."]
>>> completions = [r" The solution is \boxed{2}.", r" The solution is \boxed{6}."]
>>> ground_truth = ["2", "5"]
>>> reward_func(prompts=prompts, completions=completions, ground_truth=ground_truth)
[1.0, 0.0]

示例 4：多任务奖励函数

以下是在 RLOOTrainer 中使用多个奖励函数的示例。在此示例中，我们定义了两个任务特定的奖励函数：math_reward_func 和 coding_reward_func。math_reward_func 根据数学问题的正确性给予奖励，而 coding_reward_func 则根据代码解决方案是否可运行来给予奖励。

from datasets import Dataset
from trl import RLOOTrainer

# Define a dataset that contains both math and coding problems
dataset = Dataset.from_list(
    [
        {"prompt": "What is 2+2?", "task": "math"},
        {"prompt": "Write a function that returns the sum of two numbers.", "task": "code"},
        {"prompt": "What is 3*4?", "task": "math"},
        {"prompt": "Write a function that returns the product of two numbers.", "task": "code"},
    ]
)

# Math-specific reward function
def math_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "math":
            # Calculate math-specific reward
            correct = check_math_solution(prompt, completion)
            reward = 1.0 if correct else -1.0
            rewards.append(reward)
        else:
            # Return None for non-math tasks
            rewards.append(None)
    return rewards

# Coding-specific reward function
def coding_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "code":
            # Calculate coding-specific reward
            works = test_code_solution(prompt, completion)
            reward = 1.0 if works else -1.0
            rewards.append(reward)
        else:
            # Return None for non-coding tasks
            rewards.append(None)
    return rewards

# Use both task-specific reward functions
trainer = RLOOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=[math_reward_func, coding_reward_func],
    train_dataset=dataset,
)

trainer.train()

在此示例中，math_reward_func 和 coding_reward_func 旨在处理包含数学和编程问题的混合数据集。数据集中的 task 列用于确定对每个问题应用哪个奖励函数。如果数据集中的某个样本没有相关的奖励函数，该奖励函数将返回 None，RLOOTrainer 将继续使用有效的函数和任务。这使得 RLOOTrainer 能够处理具有不同适用性的多个奖励函数。

请注意，RLOOTrainer 会忽略奖励函数返回的 None 奖励，仅考虑相关函数返回的奖励。这确保了模型在相关任务上进行训练，并忽略没有相关奖励函数的任务。

示例 5：异步奖励函数

自定义奖励函数也可以定义为 async def 协程。如果您的奖励依赖于缓慢的 I/O（例如，调用远程服务），这将非常有用。当您传入多个异步奖励函数时，RLOOTrainer 会并发执行它们，从而使它们的延迟相互重叠。

以下是一个异步奖励函数的最小示例，用于模拟 I/O 密集型操作：

import asyncio

async def async_reward_func(prompts, completions, **kwargs):
    # Simulate an I/O-bound call (e.g., HTTP request, database lookup)
    await asyncio.sleep(0.01)
    # Simple toy reward: 1.0 if the completion is non-empty, else 0.0
    return [1.0 if completion else 0.0 for completion in completions]

示例 6：记录额外的列和指标

以下是一个奖励函数的示例，该函数会将额外的列记录到 completions 表中，并将标量指标记录为图表。

import re

def reward_func(completions, ground_truth, log_extra=None, log_metric=None, **kwargs):
    extracted = [re.search(r"\\boxed\{(.*?)\}", c) for c in completions]
    extracted = [m.group(1) if m else None for m in extracted]
    rewards = [1.0 if e == gt else 0.0 for e, gt in zip(extracted, ground_truth)]

    if log_extra:
        log_extra("golden_answer", list(ground_truth))
        log_extra("extracted_answer", [e or "[none]" for e in extracted])

    if log_metric:
        log_metric("accuracy", sum(rewards) / len(rewards))

    return rewards

将奖励函数传递给训练器

要将自定义奖励函数用于 RLOOTrainer，请按以下方式传入：

from trl import RLOOTrainer

trainer = RLOOTrainer(
    reward_funcs=reward_func,
    ...,
)

您可以将多个奖励函数作为列表传入；该列表可以同时包含同步函数和异步函数：

from trl import RLOOTrainer

trainer = RLOOTrainer(
    reward_funcs=[reward_func, async_reward_func1, async_reward_func2],
    ...,
)

奖励将计算为每个函数奖励的总和，或者如果配置中提供了reward_weights，则为加权总和。

请注意，RLOOTrainer 支持多种不同类型的奖励函数。有关更多详细信息，请参阅参数文档。

训练视觉语言模型

RLOO 支持在包含文本和图像的多模态数据集上训练视觉语言模型（VLM）。

支持的模型

已测试：

• Gemma3 — 例如 google/gemma-3-4b-it
• LLaVA-NeXT — 例如 llava-hf/llava-v1.6-mistral-7b-hf
• Qwen2-VL — 例如 Qwen/Qwen2-VL-2B-Instruct
• Qwen2.5-VL — 例如 Qwen/Qwen2.5-VL-3B-Instruct
• SmolVLM2 — 例如 HuggingFaceTB/SmolVLM2-2.2B-Instruct

无法保证与所有 VLM 的兼容性。如果您认为某个模型应当被支持，欢迎在 GitHub 上提交 issue——或者更好的是，提交一个包含所需更改的 pull request。

快速入门

使用 rloo_vlm.py 对 VLM 进行微调。在 lmms-lab/multimodal-open-r1-8k-verified 上训练的示例命令：

torchrun --nproc_per_node 8 examples/rloo_visual_math/rloo_visual_math.py \
  --deepspeed examples/deepspeed_configs/zero3.json \
  --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
  --output_dir rloo-Qwen2.5-VL-3B-Instruct \
  --learning_rate 1e-5 \
  --dtype bfloat16 \
  --max_completion_length 1024 \
  --use_vllm \
  --vllm_mode colocate \
  --use_peft \
  --lora_target_modules "q_proj", "v_proj" \
  --log_completions

配置技巧

• 在视觉-语言投影层上使用 LoRA
• 启用 4 位量化以降低内存使用
• VLM 对内存需求较高 — 请从较小的批处理大小开始
• 大多数模型都兼容 vLLM（server 和 colocate 模式）

数据集格式

每个训练样本应包含：

• prompt：通过处理器的聊天模板格式化的文本
• image/images：PIL Image 或 PIL Image 列表

Trainer 会通过模型的图像处理器自动处理图像到张量的转换。

RLOOTrainer

>>> from trl import RLOOTrainer
>>> from trl.rewards import accuracy_reward
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

>>> trainer = RLOOTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     reward_funcs=accuracy_reward,
...     train_dataset=dataset,
... )
>>> trainer.train()

RLOOConfig

参考

1. RLOO 论文
2. 回归基础：重新审视 LLM 中基于人类反馈学习的 REINFORCE 风格优化
3. 论文 - REINFORCE++：一种用于对齐大语言模型的简单高效方法
4. 博客文章 - 将 RL 重新引入 RLHF
5. 博客文章 - 解析 RLHF 及其变体：进展与工程实践洞察
6. Youtube - RLOO：一种用于从人类反馈中学习大语言模型的低成本优化方法

从旧版实现（0.21 及更早版本）迁移指南

随着 0.22.0 版本的发布，我们重新设计了 RLOOTrainer，使其与库中其他在线训练器（例如 GRPOTrainer）更加一致。这一新实现引入了对训练器配置参数和整体结构的若干更改。 以下是 RLOOConfig 主要更改的摘要：

以下是 RLOOTrainer 主要变更的摘要：