Reward Trainer

RewardTrainer trains a reward model: a model with a scalar output head that scores how good a response is. Reward models are used by online RL algorithms (PPO) and for best-of-n sampling / rejection sampling.

Dataset format

The same preference format as DPO — prompt, chosen, rejected. The reward model learns to assign a higher score to chosen than to rejected.

Example

from datasets import load_dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardConfig, RewardTrainer

model = AutoModelForSequenceClassification.from_pretrained(
    "Qwen/Qwen2.5-0.5B-Instruct", num_labels=1
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

trainer = RewardTrainer(
    model=model,
    args=RewardConfig(output_dir="Qwen2.5-0.5B-Reward"),
    train_dataset=dataset,
    processing_class=tokenizer,
)
trainer.train()

The trained reward model can then be plugged into PPOTrainer as the reward_model, or used to rank candidate generations at inference time.