How RLHF trains a reward model

Training a reward model is the critical middle step in RLHF — it translates human preferences into a scalar signal that reinforcement learning can optimize. If you’re searching for how RLHF reward model works in practice, this guide walks through the full pipeline: collecting comparison data, architecting the model head, training with pairwise ranking loss, and validating before you plug it into PPO. The code assumes PyTorch, Hugging Face Transformers, and the TRL library, but the principles transfer to any stack.

Step 1: Prepare preference comparison data

Reward models learn from pairwise comparisons: given a prompt and two completions, which is better? You need a dataset of (prompt, chosen, rejected) triples. Sources include human annotators, synthetic generation from a stronger model, or existing datasets like Anthropic’s HH-RLHF or OpenAssistant.

# data/preferences.jsonl — one record per line
{"prompt": "Explain quantum entanglement simply.", "chosen": "Quantum entanglement links two particles so measuring one instantly determines the state of the other, no matter the distance.", "rejected": "It's when particles are connected magically."}
{"prompt": "Write a Python decorator that retries a function.", "chosen": "import functools\nimport time\n\ndef retry(max_attempts=3, delay=1):\n    def decorator(fn):\n        @functools.wraps(fn)\n        def wrapper(*args, **kwargs):\n            for attempt in range(max_attempts):\n                try:\n                    return fn(*args, **kwargs)\n                except Exception as e:\n                    if attempt == max_attempts - 1:\n                        raise\n                    time.sleep(delay)\n        return wrapper\n    return decorator", "rejected": "def retry(fn):\n    def wrapper():\n        try:\n            return fn()\n        except:\n            return fn()\n    return wrapper"}

Clean the data: deduplicate prompts, filter non-English or off-topic pairs, and ensure chosen/rejected are meaningfully different. Aim for 10k–100k comparisons for a 7B–13B base model; more is better but diminishing returns appear around 50k for most domains.

Verify: Load a sample and confirm chosen responses are genuinely preferred. Compute length statistics — if chosen is consistently longer, the model may learn a length proxy instead of quality.

import json, statistics
lengths_chosen = []
lengths_rejected = []
with open("data/preferences.jsonl") as f:
    for line in f:
        d = json.loads(line)
        lengths_chosen.append(len(d["chosen"].split()))
        lengths_rejected.append(len(d["rejected"].split()))
print(f"Chosen avg tokens: {statistics.mean(lengths_chosen):.1f}")
print(f"Rejected avg tokens: {statistics.mean(lengths_rejected):.1f}")

Step 2: Choose base model and tokenizer

Start from a pretrained LLM — typically the same model you’ll later fine-tune with PPO, or a smaller variant if compute is constrained. Common choices: Llama-2-7B, Mistral-7B, or their instruct-tuned versions. The reward model adds a scalar head on top of the final hidden state.

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

MODEL_ID = "mistralai/Mistral-7B-v0.1"  # or your base model
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"  # critical for batching variable-length pairs

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_ID,
    num_labels=1,  # scalar reward
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model.config.pad_token_id = tokenizer.pad_token_id

The num_labels=1 replaces the classification head with a single linear projection from hidden size → 1. This is the reward head. Freeze the base model initially if you want faster, cheaper training (LoRA on the base + full head), but full fine-tuning typically yields better reward quality.

Step 3: Format pairs for pairwise ranking loss

The standard loss is Bradley-Terry: maximize log σ(r_chosen - r_rejected). Each batch contains concatenated chosen and rejected sequences for the same prompt. TRL’s RewardTrainer handles this, but understanding the data collator helps debug.

from torch.utils.data import Dataset
from dataclasses import dataclass
from typing import Dict, List
import torch

class PreferenceDataset(Dataset):
    def __init__(self, path: str, tokenizer, max_length: int = 2048):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = []
        with open(path) as f:
            for line in f:
                self.data.append(json.loads(line))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        prompt = item["prompt"]
        chosen = item["chosen"]
        rejected = item["rejected"]

        # Format: prompt + completion (no special tokens between)
        chosen_text = prompt + "\n" + chosen
        rejected_text = prompt + "\n" + rejected

        chosen_enc = self.tokenizer(
            chosen_text, truncation=True, max_length=self.max_length,
            padding="max_length", return_tensors="pt"
        )
        rejected_enc = self.tokenizer(
            rejected_text, truncation=True, max_length=self.max_length,
            padding="max_length", return_tensors="pt"
        )

        return {
            "input_ids_chosen": chosen_enc["input_ids"].squeeze(0),
            "attention_mask_chosen": chosen_enc["attention_mask"].squeeze(0),
            "input_ids_rejected": rejected_enc["input_ids"].squeeze(0),
            "attention_mask_rejected": rejected_enc["attention_mask"].squeeze(0),
        }

@dataclass
class RewardDataCollator:
    tokenizer: any
    def __call__(self, features: List[Dict]) -> Dict[str, torch.Tensor]:
        batch = {}
        for key in ["input_ids_chosen", "attention_mask_chosen",
                    "input_ids_rejected", "attention_mask_rejected"]:
            batch[key] = torch.stack([f[key] for f in features])
        return batch

Verify: Decode a batch and confirm prompt/completion boundaries are correct. Check that no truncation silently drops the completion.

ds = PreferenceDataset("data/preferences.jsonl", tokenizer)
collator = RewardDataCollator(tokenizer)
batch = collator([ds[0], ds[1]])
print(tokenizer.decode(batch["input_ids_chosen"][0]))
print("---")
print(tokenizer.decode(batch["input_ids_rejected"][0]))

Step 4: Train with pairwise ranking loss

TRL’s RewardTrainer wraps the loss and logging. Configure learning rate, batch size, and gradient accumulation for your GPU memory. A typical run: 1–3 epochs, LR 1e-5 to 5e-6, global batch size 32–64.

from trl import RewardTrainer, RewardConfig
from transformers import EarlyStoppingCallback

training_args = RewardConfig(
    output_dir="reward-model-output",
    per_device_train_batch_size=2,      # adjust for VRAM
    gradient_accumulation_steps=16,     # effective batch = 2*16*gpus
    num_train_epochs=2,
    learning_rate=2e-5,
    weight_decay=0.01,
    bf16=True,
    logging_steps=10,
    evaluation_strategy="steps",
    eval_steps=200,
    save_steps=200,
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    report_to="tensorboard",
    remove_unused_columns=False,  # keep our custom columns
)

trainer = RewardTrainer(
    model=model,
    args=training_args,
    train_dataset=ds,
    eval_dataset=ds.select(range(min(1000, len(ds)))),  # small eval split
    data_collator=collator,
    tokenizer=tokenizer,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)

trainer.train()

The loss computed internally is:

loss = -log(sigmoid(reward_chosen - reward_rejected))

Monitor two metrics in TensorBoard: train/loss (should decrease) and eval/accuracy (fraction of pairs where reward_chosen > reward_rejected). Accuracy above 70% is typical for clean data; above 80% is strong. If accuracy stalls near 50%, check data quality or increase model capacity.

Verify: After training, run a quick sanity check on held-out pairs.

model.eval()
with torch.no_grad():
    for i in range(5):
        item = ds[i]
        for key in ["input_ids_chosen", "attention_mask_chosen",
                    "input_ids_rejected", "attention_mask_rejected"]:
            item[key] = item[key].unsqueeze(0).to(model.device)
        r_chosen = model(input_ids=item["input_ids_chosen"],
                         attention_mask=item["attention_mask_chosen"]).logits.item()
        r_rejected = model(input_ids=item["input_ids_rejected"],
                           attention_mask=item["attention_mask_rejected"]).logits.item()
        print(f"Pair {i}: chosen={r_chosen:.3f} rejected={r_rejected:.3f} diff={r_chosen-r_rejected:.3f}")

All diffs should be positive.

Step 5: Calibrate and validate the reward scale

Raw reward logits are uncalibrated — their magnitude depends on initialization and training dynamics. Before using in PPO, you need a sense of the distribution and, optionally, a normalization layer.

import numpy as np
rewards = []
with torch.no_grad():
    for i in range(min(2000, len(ds))):
        item = ds[i]
        for key in ["input_ids_chosen", "attention_mask_chosen"]:
            item[key] = item[key].unsqueeze(0).to(model.device)
        r = model(input_ids=item["input_ids_chosen"],
                  attention_mask=item["attention_mask_chosen"]).logits.item()
        rewards.append(r)
rewards = np.array(rewards)
print(f"Mean: {rewards.mean():.3f}, Std: {rewards.std():.3f}")
print(f"5th/95th percentile: {np.percentile(rewards, 5):.3f} / {np.percentile(rewards, 95):.3f}")

Typical output: mean ~0.5–2.0, std ~0.5–1.5. If the scale is extreme (mean > 10 or std > 5), add a learnable temperature parameter or normalize at inference:

class NormalizedRewardModel(torch.nn.Module):
    def __init__(self, base_model, mean=0.0, std=1.0):
        super().__init__()
        self.base = base_model
        self.register_buffer("mean", torch.tensor(mean))
        self.register_buffer("std", torch.tensor(std))

    def forward(self, input_ids, attention_mask):
        logits = self.base(input_ids=input_ids, attention_mask=attention_mask).logits
        return (logits - self.mean) / self.std

Validation beyond accuracy: Test on edge cases — adversarial prompts, out-of-distribution topics, and known-good/bad completions. Build a small eval suite (50–100 handcrafted pairs) covering your deployment domain. The reward model should rank them correctly.

eval_pairs = [
    ("Write a secure password generator.", "import secrets\nprint(secrets.token_urlsafe(16))", "password = '123456'"),
    ("Summarize the paper in one sentence.", "The paper proposes a new attention mechanism reducing complexity to O(n).", "This paper is about AI."),
]
for prompt, good, bad in eval_pairs:
    good_text = prompt + "\n" + good
    bad_text = prompt + "\n" + bad
    g = tokenizer(good_text, return_tensors="pt").to(model.device)
    b = tokenizer(bad_text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        rg = model(**g).logits.item()
        rb = model(**b).logits.item()
    print(f"Good: {rg:.3f}  Bad: {rb:.3f}  {'✓' if rg > rb else '✗'}")

Step 6: Export for PPO / RLHF pipeline

Save the model in a format your PPO trainer expects. Most frameworks (TRL, OpenRLHF, custom) load a AutoModelForSequenceClassification with num_labels=1. Push to Hub or save locally.

model.save_pretrained("reward-model-final")
tokenizer.save_pretrained("reward-model-final")
# Optional: push to hub
# model.push_to_hub("your-org/your-reward-model")
# tokenizer.push_to_hub("your-org/your-reward-model")

If your PPO loop runs on a different machine or framework, export to ONNX or TorchScript for lower-latency inference:

import torch
model.eval()
dummy_input = tokenizer("test prompt\ncompletion", return_tensors="pt")
torch.onnx.export(
    model,
    (dummy_input["input_ids"], dummy_input["attention_mask"]),
    "reward_model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["reward"],
    dynamic_axes={"input_ids": {0: "batch", 1: "seq"},
                  "attention_mask": {0: "batch", 1: "seq"},
                  "reward": {0: "batch"}},
    opset_version=17,
)

Verify the export: Load the ONNX model and compare outputs on a batch.

import onnxruntime as ort
session = ort.InferenceSession("reward_model.onnx")
ort_inputs = {k: v.numpy() for k, v in dummy_input.items()}
ort_out = session.run(None, ort_inputs)[0]
with torch.no_grad():
    torch_out = model(**dummy_input).logits.numpy()
print(f"Max diff: {np.abs(ort_out - torch_out).max():.6f}")  # should be ~1e-5

Step 7: Monitor reward hacking in production

A reward model is only as good as its coverage. Once deployed in PPO, the policy will exploit blind spots — generating high-reward but low-quality outputs (reward hacking). Mitigations:

1. KL penalty: Keep the policy close to the reference model. Typical kl_coef=0.05–0.2.
2. Reward clipping: Clip rewards at ±3σ or a fixed range (e.g., [-10, 10]) before feeding to PPO.
3. Ensemble: Train 3–5 reward models with different seeds/data shuffles; use mean reward and variance as uncertainty signal. Reject high-variance samples.
4. Periodic refresh: Collect new human comparisons on policy outputs monthly; retrain or fine-tune the reward model.

# Example: ensemble inference
ensemble_models = [model1, model2, model3]  # loaded separately
def ensemble_reward(input_ids, attention_mask):
    rewards = []
    for m in ensemble_models:
        with torch.no_grad():
            rewards.append(m(input_ids=input_ids, attention_mask=attention_mask).logits)
    rewards = torch.stack(rewards)  # [ensemble, batch, 1]
    return rewards.mean(0), rewards.std(0)

Verify in PPO: Track reward_mean, reward_std, and kl_div per batch. If reward_mean climbs while human eval scores drop, you’re hacking. Pause training, expand the comparison dataset with the hacked examples labeled as rejected, and retrain.

Training a reward model is straightforward mechanically — the difficulty is data quality and ongoing maintenance. Start with a clean 10k-pair dataset, train a 7B base for 2 epochs at 2e-5 LR, validate on a held-out suite, and deploy with clipping and KL control. Treat the reward model as a living component: budget for quarterly retraining as your policy distribution shifts. That’s how RLHF reward model works in production — not a one-off artifact, but a feedback loop you operate.