Skip to content

RLHF with PPO

Goal

Optimize the SFT policy to produce responses that the reward model prefers, while staying close to the original SFT model.

The "stay close" part is important. Without it, the model may exploit the reward model and drift away from fluent language.

Main components

  • Policy model — the trainable model. Initialized from the SFT checkpoint.
  • Reference model — frozen copy of the SFT model. Provides the KL anchor.
  • Reward model — frozen, trained on preferences. Provides the scalar reward signal.
  • Value model — separate critic that estimates expected reward; used to compute advantages. Often initialized from the reward model.
  • KL penalty — penalizes the policy for diverging from the reference model.
  • PPO update — clipped policy gradient update.

Simplified loop

  1. Sample prompts from a distribution (often the SFT prompt set).
  2. Policy generates responses.
  3. Reward model scores responses.
  4. Add per-token KL penalty against the reference model: r_t' = r_t - β · KL(π || π_ref).
  5. PPO updates the policy with the clipped objective: $$ \mathcal{L}_{\text{PPO}} = \mathbb{E}\left[\min(\rho_t \hat{A}_t,\, \text{clip}(\rho_t, 1-\epsilon, 1+\epsilon)\hat{A}_t)\right] $$ where ρ_t = π_θ(a|s) / π_old(a|s).
  6. Monitor reward, KL, length, downstream evals, safety metrics.

Why PPO is hard

PPO-based RLHF is powerful but finicky.

Challenges:

  • Expensive sampling loop — every step requires generating responses, which is the most compute-heavy part.
  • Reward hacking — the policy finds outputs that score high but are bad.
  • KL tuning — too low → drift and reward hacking; too high → no learning.
  • Instability — value-function divergence, advantage explosion.
  • Mode collapse — policy converges to a narrow set of phrasings.
  • Degraded benchmarks — capabilities the SFT model had may regress.
  • Complex infrastructure — running 4 models (policy + ref + reward + value) at scale.

InstructGPT reported that RLHF can regress on some public NLP datasets and used a pretraining-mix term ("PPO-ptx") — adding pretraining loss back into the PPO objective — to reduce regressions.

Knobs that matter

Knob Effect
KL coefficient β Higher = stay close to SFT, less drift, less learning. Start ~0.05.
Clip range ε Standard 0.2. Lower for stability.
Rollout batch size Bigger = more stable gradient estimates.
Number of PPO epochs per rollout 2–4 is typical.
Learning rate Much lower than SFT — often 1e-6 to 1e-7.
Sampling temperature High enough for diversity, low enough for quality.

When PPO is worth the complexity

PPO outperforms DPO when:

  • You have a good reward model and can afford to keep refreshing it.
  • Reward must be evaluated during generation (e.g., environment rewards from an agent task).
  • You have tasks with verifiable rewards (math checkers, unit tests) where you can replace the RM with a deterministic checker — this is RLVR, the modern replacement for classical RM-based PPO. See DPO and friends for RLVR.
  • You can iterate the cycle (sample → label → retrain RM → PPO) several times.

For most chat-alignment use cases, DPO is now preferred because it gets most of the benefit at a fraction of the engineering cost.

Practical tips

  • Validate the reward model first. Run RewardBench and held-out preference accuracy. A bad RM guarantees bad PPO.
  • Watch KL like a hawk. Plot KL per token over training. If it climbs unboundedly, β is too low or the RM has reward-hacking exploits.
  • Hold out broad evals. Run MMLU/HumanEval/IFEval/safety evals during PPO. RLHF can silently regress capabilities.
  • Add pretraining-mix loss for capability preservation (PPO-ptx).
  • Iterate. Single-shot PPO is rarely the best result. Re-collect preferences from the new policy, retrain the RM, run PPO again.
  • Consider GRPO instead for reasoning tasks: it removes the value model and is much simpler. See DPO and friends and Reasoning.

Further reading