The RLHF Pipeline: Reward Models and PPO

The previous post trained a model on demonstrated responses. Those demonstrations can encode human preferences, but they do not directly label comparisons among the model’s alternative answers.

Preference data adds that comparison: for the same prompt, which response did the annotator favor? RLHF uses those judgments to guide further training.

A reward-model-based RLHF pipeline turns preference judgments into scores, then uses those scores to optimize a language model. It is one way to use comparative feedback; direct preference methods are another.

Prerequisites

This post assumes you've read:

• Next-Token Prediction for the base-model setup and next-token loss.
• Supervised Fine-Tuning for the SFT stage we build on top of.
• Attention From Scratch for what "policy" and "logits" mean in transformer terms.
• Cross-entropy and softmax. The policy-gradient calculation below uses $\nabla \log \pi$.

The shape of the pipeline

Take an SFT model, sample responses from it, ask humans which is better, fit a reward model to those preferences, then use the reward model to steer the language model with reinforcement learning.

Four moving pieces:

1. SFT model ($\pi_{\text{SFT}}$). Frozen as the reference for this policy-training stage.
2. Preference dataset. Pairs $(x, y_w, y_l)$ where $y_w$ is the chosen response and $y_l$ is the rejected one, given the same prompt $x$.
3. Reward model ($r_\phi$). For example, a pretrained transformer with the language-modeling head swapped out for a single scalar. Trained to rank the chosen response above the rejected one.
4. Policy ($\pi_\theta$). Starts as a copy of $\pi_{\text{SFT}}$, and PPO gradually updates it to produce responses the reward model likes. The KL penalty tethers it back to $\pi_{\text{SFT}}$ to discourage drift.

A learned critic often estimates values during PPO as well. PPO’s old rollout policy and the frozen reference serve different roles: the first appears in the surrogate ratio, the second in the KL penalty.

Reward model: turning "A > B" into a scalar

Pairwise judgments can be easier to collect consistently than absolute ratings, though agreement depends on the task, rubric, and annotators.

Given a prompt $x$ and two candidate responses $y_a$ and $y_b$, a labeler picks one. That's the entire data format: $(x, y_w, y_l)$, where $w$ means "won" and $l$ means "lost."

The reward model learns to assign a scalar $r_\phi(x, y)$ to any (prompt, response) pair, such that the chosen response scores higher than the rejected one. The objective comes from the Bradley–Terry preference model, which has been the workhorse for pairwise ranking for over seventy years. The model says: the probability that $y_w$ is preferred over $y_l$ is

where $\sigma$ is the logistic sigmoid. It converts the reward gap to a binary preference probability.

The training loss is the negative log-likelihood of the observed preferences:

A reward model can start from a pretrained or instruction-tuned transformer with a scalar readout on the final token. It need not match the policy’s size or architecture. Its ability to generalize preference judgments affects which errors policy optimization may exploit.

In this setup, the reward model scores a full response, not each token. It reads the whole thing, then outputs one number. In PPO we'll assign that terminal reward to the last token of the sequence and treat every token before it as having reward zero.

Pairwise training identifies score differences. Adding the same prompt-dependent constant to both scores leaves this loss unchanged. A raw score such as $+1.3$ therefore has no intrinsic quality meaning, and preference calibration still needs evaluation.

PPO for language models

Proximal Policy Optimization (PPO) uses the reward model’s scores to update the policy initialized from SFT.

A direct policy-gradient estimate can have high variance, and several updates on the same rollout can move the policy far from the one that generated it. PPO’s clipped surrogate reduces the incentive for large changes on individual samples.

PPO-Clip uses a clipped surrogate objective. At each update, we compute the probability ratio

between the current policy and the policy we used to sample the data. If the ratio is one, nothing has changed. If it's 1.5, we've made the action 50% more likely. PPO limits the objective’s incentive for certain changes in that ratio; it does not impose a hard bound on the ratio itself. Specifically, the clipped surrogate objective is

where $\hat{A}_t$ is the advantage (how much better this action was than the baseline) and $\varepsilon$ is typically $0.2$.

For positive advantage, the surrogate becomes flat when the ratio exceeds $1+\varepsilon$. For negative advantage, it becomes flat below $1-\varepsilon$. In the opposite directions it remains sloped. Shared parameters and other samples can still move a clipped sample’s probability, so clipping is not a hard trust-region constraint.

The distinction matters when debugging: a ratio outside the interval does not by itself tell you whether that sample’s surrogate gradient is zero. You also need the advantage sign.

Reward, minus a baseline, minus KL

The PPO loss needs an advantage estimate $\hat{A}_t$. In the actor-critic recipe here, that estimate combines rewards, a value baseline, and a reference penalty:

1. Terminal reward. A response-level reward is placed at the last valid response token; intermediate tokens can still have KL penalties.
2. Returns and baseline. Discounted returns or generalized advantage estimation propagate later rewards to earlier actions. A value function commonly supplies the baseline. Subtracting a batch mean from immediate token rewards is not an equivalent calculation.
3. Reference penalty. A sampled log ratio to a frozen reference estimates KL under the policy that sampled that token. Its sum has the sequence-KL interpretation when sampling and probability accounting match. Temperature changes or truncated sampling require care.

For the rollout-return construction illustrated here, freeze the sampled log probabilities and use

where $r^{\text{RM}}_t$ is zero except at the last valid response token. A sampled log ratio can be negative; the corresponding exact KL is nonnegative after expectation under the sampling policy.

The reference penalty discourages exploiting a reward model by moving far from the reference distribution. Its strength trades reward against divergence, but it cannot guarantee correctness or prevent all reward hacking.

Why RLHF is hard: the reward model is the bottleneck

If the reward model favors length, confident tone, or plausible errors, policy optimization can amplify those patterns. Check them against independent evaluations rather than assuming a higher reward implies better answers.

The usual mitigations are all imperfect: regularize harder with KL (limits drift but also improvement), re-train the RM periodically on fresh policy samples (expensive), use multiple RMs and take a conservative aggregate, or cap training steps with humans watching for drift. Their effectiveness depends on the reward model and evaluation coverage.

The Helpful and Harmless paper and Casper et al.’s review examine these limitations.

Worked example

Consider one PPO update for a single prompt, using constructed rewards and token probabilities.

Prompt: "Explain gradient descent briefly."

Step 1: rollout. Sample four completions from the rollout policy. This toy assumes sampling from the same policy distribution whose log probabilities are recorded.

Step 2: score. We run each $(x, y_i)$ through the reward model.

Step 3: toy advantage. To isolate clipping, omit KL and use each completion’s reward minus the mean $0.4$ as a fixed advantage for every shown token. This is a constructed sequence-level baseline example, not a full GAE implementation.

Step 4: compute ratios. The update runs for a couple of inner epochs on this same rollout, so by the time we compute gradients, $\pi_\theta$ has drifted from $\pi_{\theta_{\text{old}}}$. We compute the per-token ratio $r_t = \pi_\theta(y_t \mid \cdot) / \pi_{\theta_{\text{old}}}(y_t \mid \cdot)$. The token-level ratio goes into the token-level clipped surrogate; do not replace it with a product across the whole response.

Step 5: clip. Positive-advantage terms flatten above 1.2; negative-advantage terms flatten below 0.8. Elsewhere the unclipped term supplies the slope.

The advantage signs explain the clipping in the table:

• Completion 1 has advantage $+2.0$ (positive). Its per-token ratios hover just above 1 (the policy now mildly prefers it). All ratios stay inside $[0.80, 1.20]$, so the full gradient flows and the policy gets nudged further toward completion 1.
• Completion 4 has advantage $-2.0$ (strongly negative). Its ratios have already dropped to $[0.36, 0.71]$, meaning the policy has already moved substantially away. Every one of those ratios is outside $1 - \varepsilon = 0.80$. PPO clips the gradient contribution to zero. The policy already moved; PPO stops further churn on this particular sample.
• Completion 3 also has negative advantage, but its ratios are only at $[0.79, 0.91]$, mostly inside the clipping interval. The 0.79 ratio is on the flat side for a negative advantage; the other four still contribute gradients.

Implementation sketch

The following tensor functions implement reward-model loss, rollout-return construction, and the clipped policy surrogate. They can run independently, but a complete PPO trainer also needs sampling, a value estimator, and an optimizer loop.

import torch
import torch.nn.functional as F
 
def rm_loss(chosen_scores, rejected_scores):
    return F.softplus(-(chosen_scores - rejected_scores)).mean()
 
@torch.no_grad()
def rollout_returns(logp_old, logp_ref, rm_rewards, response_mask, beta=0.02):
    # Right-padded response tokens. Log-probs come from the rollout policy.
    mask = response_mask.bool()
    lengths = mask.sum(-1)
    if (lengths == 0).any():
        raise ValueError("each response must contain a valid token")
    expected = torch.arange(mask.size(1), device=mask.device)[None, :] < lengths[:, None]
    if not torch.equal(mask, expected):
        raise ValueError("response mask must be right-padded")
    rewards = -beta * torch.where(mask, logp_old - logp_ref, 0.0)
    rewards[torch.arange(mask.size(0), device=mask.device), lengths-1] += rm_rewards
    # Undiscounted returns: each earlier action sees subsequent rewards.
    returns = rewards.flip(-1).cumsum(-1).flip(-1)
    return returns * mask
 
def ppo_loss(logp_new, logp_old, advantages, response_mask, eps=0.2):
    # Advantages are fixed rollout estimates, e.g. returns minus a baseline.
    mask = response_mask.bool()
    if not mask.any():
        raise ValueError("no response tokens to score")
    # Select valid entries before arithmetic: padded sentinels may be infinite.
    ratio = (logp_new[mask] - logp_old.detach()[mask]).exp()
    advantage = advantages.detach()[mask]
    objective = torch.minimum(ratio * advantage,
                              ratio.clamp(1-eps, 1+eps) * advantage)
    return -objective.mean()
 
 

softplus(-(r_w − r_l)) computes $-\log \sigma(r_w - r_l)$ in a numerically stable form.

The rollout helper places the terminal reward at the last valid token, then computes backward cumulative returns. A production PPO trainer commonly uses a learned value function and GAE. The surrogate accepts detached advantages and masks padding; it does not differentiate through reward construction.

The reference and recorded rollout log probabilities are fixed during a surrogate update. Recomputing new-policy log probabilities supplies the gradient path for the policy loss; a critic has a separate value loss.

Misconceptions

"The reward model directly measures human values." It learns from judgments over a particular set of samples and instructions. Its generalization to new languages, tasks, or policy outputs requires evaluation, especially when the policy is optimized against its scores.

"PPO keeps every ratio inside the clip interval." It does not. The surrogate becomes flat only on one side for each advantage sign. Other terms and parameter sharing can change those ratios further.

"The KL penalty guarantees the policy stays useful." It penalizes distribution shift from a chosen reference. A useful setting depends on reward scale and training recipe, and some recipes use different regularization.

"RLHF makes the model honest / aligned / safe." It optimizes feedback-derived objectives. The data, reward model, sampling distribution, optimizer, and evaluation all affect the result. Preferences can miss factual errors or reward style over substance, so an improved preference score is not a guarantee of any broader property.

What's next

The next post covers DPO and its variants, the supervised shortcut that replaces the entire RM-plus-PPO pipeline with a single closed-form loss on preference pairs, and the assumptions behind that replacement.

Additional reading (and watching)

• Ouyang, L., et al. (2022). Training language models to follow instructions with human feedback. arXiv:2203.02155 (InstructGPT). The paper that made the modern RLHF pipeline legible — SFT model, preference data, reward model, PPO with KL.

• Bradley, R. A., & Terry, M. E. (1952). Rank Analysis of Incomplete Block Designs: I. The Method of Paired Comparisons. Biometrika, 39(3/4), 324-345. The original formulation of the Bradley–Terry preference model, which the RM loss is a direct application of.

• Schulman, J., et al. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347. The clipped-surrogate objective itself. Includes the clipping equation and comparisons with TRPO and vanilla policy gradients.

• Ziegler, D. M., et al. (2019). Fine-Tuning Language Models from Human Preferences. arXiv:1909.08593. The earlier OpenAI paper that laid out the per-token KL penalty against the SFT reference. Describes an early language-model preference-training recipe.

• Christiano, P., et al. (2017). Deep Reinforcement Learning from Human Preferences. arXiv:1706.03741. The foundational paper showing that a reward model trained from pairwise preferences can steer a policy. Predates RLHF-for-LLMs but the template is identical.

• Bai, Y., et al. (2022). Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback. arXiv:2204.05862. Studies reward-model and policy training, including reward hacking and the helpfulness–harmlessness tradeoff.

• Casper, S., et al. (2023). Open Problems and Fundamental Limitations of Reinforcement Learning from Human Feedback. arXiv:2307.15217. A taxonomy of problems in feedback collection, reward modeling, and policy optimization.

• Rafailov, R., et al. (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. NeurIPS 2023. Derives a preference objective that avoids fitting a separate reward model; covered in the next post.

• Bai, Y., et al. (2022). Constitutional AI: Harmlessness from AI Feedback. arXiv:2212.08073. The RLAIF (RL from AI Feedback) approach that replaces human preference labelers with a second LM guided by written principles.

• Lightman, H., et al. (2023). Let's Verify Step by Step. arXiv:2305.20050. The process-reward-model paper — score every step of a reasoning chain rather than just the final answer. Compares process supervision and outcome supervision on mathematical reasoning.