PPO for RLHF in Practice

OpenAI's InstructGPT demonstrated that a 1.3 B parameter model fine-tuned with RLHF was preferred over the raw 175 B GPT-3 by human evaluators. The mechanism behind that jump is not magic: it is Proximal Policy Optimisation applied to a learned reward signal. Understanding why PPO is used here, rather than simpler gradient methods, and how it fits into the four-model apparatus, is the minimum prerequisite for diagnosing alignment training runs in practice.

The four-model apparatus

RLHF with PPO keeps four distinct models in memory simultaneously. Conflating them is the single most common source of confusion.

The policy is the model you care about. The reference policy is its initialisation point, kept frozen so you can measure how far the policy has drifted. The reward model was trained on human preference comparisons; it scores any (prompt, completion) pair with a scalar. The value function estimates expected future reward from any intermediate token position and is required to compute generalised advantage estimates (GAE).

All four models are forward-passed on every training step, which is why RLHF is GPU-memory intensive. Techniques like LoRA on the policy (keeping the reference as the frozen backbone) or sharing the backbone between policy and value head reduce cost, but complicate the training loop.

The shaped reward and why it matters

The reward signal the policy actually optimises is not \(r_\phi\) alone. It is:

where \(x\) is the prompt, \(y\) is the generated response, and \(\beta\) is a tunable coefficient (InstructGPT used \(\beta \approx 0.02\)).

The KL term penalises the policy for generating token distributions that diverge from the reference. Without it, the policy rapidly exploits the reward model's blind spots: finding short, formulaic, or linguistically bizarre completions that score high on \(r_\phi\) but look nothing like coherent text. This is reward hacking, and it happens within hundreds of gradient steps if \(\beta = 0\).

Anthropic's 2022 training analysis found a roughly linear relationship between RL reward and \(\sqrt{D_\text{KL}}\), which suggests a natural operating point exists and that returns diminish sharply past a certain KL budget.

Why PPO and not a simpler policy gradient

Vanilla REINFORCE updates the policy with:

Two problems make this impractical for LLMs. First, it is extremely high-variance on long sequences where the reward is sparse (a single scalar at the end of hundreds of tokens). Second, a large gradient step can collapse the policy irreversibly; there is no upper bound on how far parameters move.

PPO addresses both. The core clipped surrogate objective is:

where \(r_t(\theta) = \pi_\theta(a_t|s_t) / \pi_{\theta_\text{old}}(a_t|s_t)\) is the probability ratio and \(\hat{A}_t\) is the advantage estimate. The clip (typically \(\epsilon = 0.2\)) prevents the ratio from moving too far from 1, bounding the per-step policy shift without the expensive second-order computation that TRPO requires.

For RLHF, \(a_t\) is a token and \(s_t\) is the context up to that token. The advantage \(\hat{A}_t\) is computed via GAE over the value function, which propagates reward signal backwards through the sequence. This is what allows the method to assign partial credit to individual token choices rather than treating the entire sequence as a single action.

A typical RLHF-PPO training loop looks like this:

for each iteration:
    1. Sample prompts from prompt dataset
    2. Roll out π_θ to produce completions y
    3. Score (x, y) with reward model r_φ
    4. Compute per-token KL penalty against π_ref
    5. Add KL penalties to scalar reward to get R(x, y)
    6. Run value function V_ψ over sequence; compute GAE advantages
    7. Run PPO update for K mini-epochs on the collected batch
    8. Update V_ψ with MSE loss against R-bootstrapped targets

Steps 7-8 can be interleaved or separated. Running K > 1 epochs on the same rollout batch is safe because of the clipping; it would be unsafe under REINFORCE.

Practical hyperparameters and stability

Getting PPO to converge on an LLM requires far more care than the RL literature suggests.

Batch size and rollout length. Generating on-policy completions is slow. Practitioners typically collect a large rollout buffer (hundreds to thousands of prompts) before each PPO update phase, then run several mini-epochs over it. Rollout batch size and the number of mini-epochs are jointly the most sensitive hyperparameters.

Value function initialisation. Starting the critic from the same checkpoint as the policy is common, but it means the value head is initially uninformative. A warm-up phase where you freeze the policy and train only the critic on a supervised regression signal (using SFT log-probabilities as proxy rewards) can accelerate early training.

KL coefficient scheduling. Fixed \(\beta\) is simple but brittle. Adaptive KL scheduling (increase \(\beta\) if \(D_\text{KL}\) exceeds a target, decrease if below) is more robust and is what InstructGPT used in practice.

Reward normalisation. Raw reward model scores can have arbitrary scale and shift. Normalising reward statistics (running mean and variance) per batch stabilises value function learning.

Token-level vs. sequence-level reward. \(r_\phi\) produces one scalar per completion. That scalar is typically assigned to the final token; all earlier tokens receive only the KL penalty term. Some implementations distribute a fraction of the reward to each token, but this introduces its own biases.

When it falls down

Reward hacking. The reward model is trained on a finite, biased sample of human preferences. The policy will eventually find completions that score high on \(r_\phi\) but violate the implicit preferences the annotators had in mind. Common failure modes: excessive length (annotators often prefer longer responses, so the model inflates), sycophantic agreement, and repetitive formatting that pattern-matches what raters rewarded historically.

KL collapse. If \(\beta\) is too high, the KL penalty dominates and the policy never moves from the reference. You get a well-behaved model that is no better than the SFT checkpoint. If \(\beta\) is too low, the policy diverges rapidly. The operating range is narrow and varies with model scale.

Critic overfitting. When the rollout batch is small relative to the mini-epoch count, the value function overfits the current batch's noise. The resulting advantage estimates are biased, gradient variance increases, and training destabilises. Reducing mini-epochs or increasing rollout batch size usually fixes this.

Distribution shift between reward model and policy. The reward model was trained on comparisons generated by the SFT checkpoint. As the policy drifts away from that distribution, out-of-distribution inputs to the reward model produce unreliable scores. This is sometimes called "reward model overoptimisation" and is distinct from reward hacking, though the two compound each other.

Memory pressure. Holding four models (or three models plus a shared backbone with two heads) in fp16 or bf16 while generating long sequences exhausts GPU memory faster than inference alone. This forces smaller batch sizes, which increases variance. Gradient checkpointing, LoRA adapters, and offloaded reference/reward models are the standard mitigations.

Further reading

• Ouyang et al., "Training language models to follow instructions with human feedback" (InstructGPT): https://arxiv.org/abs/2203.02155
• Schulman et al., "Proximal Policy Optimization Algorithms": https://arxiv.org/abs/1707.06347
• Stiennon et al., "Learning to summarize from human feedback" (NeurIPS 2020): https://arxiv.org/abs/2009.01325
• Bai et al., "Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback" (Anthropic, 2022): https://arxiv.org/abs/2204.05862