DL S063 · RLHF with PPO — Intuition + Minimal Code

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 63 / 80 · Module M10 · ~2 hours

The story

The algorithm that made ChatGPT — and that almost nobody uses anymore

RLHF with PPO is the algorithm that turned GPT-3.5 into ChatGPT. It's also, by 2024, largely being replaced by DPO (S064). So why spend a whole session on it?

Because (a) it's the historically correct answer to the alignment problem and understanding it makes every subsequent technique legible; (b) the PPO objective, KL penalty, and value baseline concepts show up in RLAIF, GRPO, RLHF-lite, and every future RL-style alignment method; (c) frontier labs still use PPO for the final polish — Anthropic's constitutional AI stack, OpenAI's later GPT-4 iterations, DeepMind's Gemini all have PPO in the loop somewhere.

Here's the puzzle: you have an SFT'd model (a policy π_θ) and a reward model r_φ. You want to update the policy to score higher under r_φ. Naive approach: just do policy-gradient RL — sample from π_θ, get rewards, ascend the gradient. Two things go wrong: (1) you catastrophically forget everything the SFT model knew, and (2) you reward-hack the RM into oblivion. PPO fixes both with two extra terms: a KL penalty against a frozen reference model and a clipped surrogate objective that bounds how far the policy moves per step.

That's the whole idea. Everything else is bookkeeping.

Prerequisites

• Session 062 (reward modeling) — the RM is the whole engine of RLHF. If you don't have one, you can't do this.
• Session 058 (SFT) — the policy π_θ is initialized from an SFT'd model, and the reference π_ref is a frozen copy of the same model.
• A little RL background helps but isn't required — I'll re-derive the policy gradient below.

1 · The setup — four models loaded at once

Here's the whole PPO-for-LLMs picture in one image:

Four models in GPU memory simultaneously:

1. Policy π_θ — trainable. This is what we're optimizing. Init from SFT model.
2. Reference π_ref — frozen copy of the SFT model. Used to compute the KL penalty that keeps π_θ from drifting away.
3. Reward model r_φ — frozen, from S062. Provides the scalar reward per response.
4. Value model V_ψ — trainable, small. Predicts expected return from each token position. Used as a baseline for variance reduction (this is the "V" in advantage-based policy gradient).

For a 7B policy in fp16, you're looking at ~14GB × 4 = 56GB before activations. This is why PPO with 70B models needs 8 × A100 minimum. This is also why LoRA-PPO exists (only the policy and value get LoRA-adapted; ref and RM stay frozen — memory drops to ~30GB).

1.1 The generation → scoring → update cycle

Each PPO iteration:

1. Sample a batch of prompts from the training set.
2. Generate responses from the policy π_θ.
3. Score each response with the reward model → per-response scalar r.
4. Compute per-token log-probs under π_θ and π_ref.
5. Compute per-token KL: KL(π_θ(·|x, y_<t) || π_ref(·|x, y_<t)).
6. Combine per-token KL and terminal reward into a shaped per-token reward: r_t = -β · KL_t for t < T, and r_T = r(x, y) - β · KL_T at the final token.
7. Compute advantages A_t using the value model V_ψ (GAE — Generalized Advantage Estimation).
8. Update π_θ and V_ψ for 2–4 mini-epochs on this batch with the PPO clipped objective.
9. Go to 1.

That's it. The whole algorithm.

2 · The PPO clipped objective, unpacked

For each token position t, we have:

• Old log-prob under the policy at generation time: log π_θ_old(y_t | x, y_<t)
• New log-prob after some update steps: log π_θ(y_t | x, y_<t)
• Advantage A_t (how much better this token turned out than expected)

Define the probability ratio:

$\rho_t(\theta) = \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\theta_{old}}(y_t | x, y_{<t})} = \exp\bigl(\log\pi_\theta - \log\pi_{\theta_{old}}\bigr)$

The vanilla policy gradient objective is E[ρ_t · A_t]. Ascend this gradient → policy shifts toward tokens with high advantage. Problem: nothing stops the update from making ρ_t huge, which can catastrophically overshoot.

PPO's fix — the clipped surrogate objective:

$\mathcal{L}^{CLIP}(\theta) = \mathbb{E}_t\Bigl[\min\bigl(\rho_t A_t,\ \text{clip}(\rho_t,\ 1-\epsilon,\ 1+\epsilon) A_t\bigr)\Bigr]$

with ε typically 0.2. Reading this:

• If A_t > 0 (token was good): the objective is bounded above by (1+ε)·A_t. We can't gain more from this token past ρ = 1.2.
• If A_t < 0 (token was bad): the objective is bounded below by (1-ε)·A_t. We can't move too aggressively away either.

In effect, PPO says: "each policy update can only shift a per-token probability by at most a factor of 1±0.2 within an iteration." That trust-region-like bound is what makes PPO stable where vanilla policy gradient is chaos.

2.1 The KL penalty term

On top of the clipped objective, we subtract a KL penalty against the reference:

$\mathcal{L}^{PPO-KL}(\theta) = \mathcal{L}^{CLIP}(\theta) - \beta \cdot \mathbb{E}_t\bigl[\text{KL}(\pi_\theta \|\ \pi_{ref})\bigr]$

β is a scalar around 0.01–0.1. If β is too small, the policy drifts arbitrarily and reward-hacks. If β is too big, the policy can't move away from the SFT distribution at all — reward doesn't improve.

Note: some PPO implementations bake the KL penalty into the reward (r_shaped = r − β·KL) rather than the loss. Mathematically equivalent under GAE with γ=1. TRL does it in the reward.

3 · Where the value model comes in

The value model V_ψ(x, y_<t) predicts the expected shaped reward from position t onward. Its job is variance reduction: instead of pushing the policy in proportion to raw reward R, we push it in proportion to the advantage A_t = R_t − V_ψ(x, y_<t) — "how much better this trajectory turned out than the value model expected."

Advantages have mean ≈ 0 and lower variance than raw returns, which drastically stabilizes training. Without a value baseline, PPO on LLMs is nearly impossible to tune.

GAE (Generalized Advantage Estimation) is a specific formula for computing A_t from the sequence of one-step temporal-difference errors, blended with a λ knob (typically 0.95). Full formula is out of scope for this session — TRL computes it automatically. Just know: it's a low-variance estimator of the advantage.

The value model itself is trained by regressing its output against the observed shaped returns:

$\mathcal{L}^{V}(\psi) = \mathbb{E}_t\bigl[(V_\psi(x, y_{<t}) - R_t)^2\bigr]$

Standard MSE. Trained jointly with the policy on the same batches.

3.1 Architectural note

The value model is usually implemented as a head on top of the policy — you take the same transformer backbone, add a Linear(d, 1) head parallel to the LM head, and train the value head. This shares almost all parameters with the policy, saving memory and warm-starting the value estimates from the policy's representations.

In TRL:

from trl import AutoModelForCausalLMWithValueHead
model = AutoModelForCausalLMWithValueHead.from_pretrained("mistralai/Mistral-7B-v0.3")

This one line gives you a model with both the LM head (for generating and computing log-probs) and a value head (for GAE). Backprop flows through both.

4 · A skeletal PPO iteration in code

I won't write a full working PPO — TRL's PPOTrainer is 800 lines and depends on a rollout buffer, GAE machinery, gradient accumulation, etc. But here's the shape of one iteration, in ~50 lines of pseudo-torch:

# Given: policy (with value head), ref_model (frozen), reward_model (frozen)
# Given: batch of prompts
 
# ---- 1. Rollout ----
with torch.no_grad():
    responses     = policy.generate(prompts, max_new_tokens=256, do_sample=True)
    log_probs_old = policy.compute_log_probs(prompts, responses)         # (B, T)
    log_probs_ref = ref_model.compute_log_probs(prompts, responses)      # (B, T)
    values_old    = policy.value_head(prompts, responses)                # (B, T)
    rewards_end   = reward_model(prompts, responses)                     # (B,) — terminal reward

Prepare per-token shaped rewards and advantages:

# ---- 2. Shape rewards + compute GAE ----
kl_per_token = log_probs_old - log_probs_ref               # (B, T) — approx per-token KL
shaped_rewards = -beta * kl_per_token                      # penalty at every token
shaped_rewards[:, -1] += rewards_end                       # terminal reward at last token
 
advantages, returns = compute_gae(                          # standard GAE(λ=0.95)
    rewards=shaped_rewards, values=values_old, gamma=1.0, lam=0.95
)
# normalize advantages for stability
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

The core PPO update — usually 2-4 mini-epochs over the same rollout batch:

# ---- 3. PPO update (2-4 mini-epochs) ----
for _ in range(ppo_epochs):
    for minibatch in shuffle(...):
        log_probs_new = policy.compute_log_probs(mb.prompts, mb.responses)
        values_new    = policy.value_head(mb.prompts, mb.responses)
 
        # Policy loss (clipped surrogate)
        ratio = torch.exp(log_probs_new - mb.log_probs_old)
        obj1  = ratio * mb.advantages
        obj2  = torch.clamp(ratio, 1 - eps, 1 + eps) * mb.advantages
        policy_loss = -torch.min(obj1, obj2).mean()
 
        # Value loss (MSE against returns, sometimes also clipped)
        value_loss = 0.5 * (values_new - mb.returns).pow(2).mean()
 
        # Entropy bonus (optional, encourages exploration)
        entropy_bonus = -policy.entropy(mb.prompts, mb.responses).mean()
 
        loss = policy_loss + vf_coef * value_loss + ent_coef * entropy_bonus
        loss.backward(); optimizer.step(); optimizer.zero_grad()

That's the whole loop. Repeat for 100–1000 iterations, tune β, ε, LR, batch size — and pray.

4.1 In practice: TRL PPOTrainer

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer
 
cfg = PPOConfig(
    learning_rate=1.4e-5,
    batch_size=64, mini_batch_size=8, ppo_epochs=4,
    cliprange=0.2, cliprange_value=0.2,
    vf_coef=0.1, init_kl_coef=0.05, target_kl=6.0,
)
 
policy = AutoModelForCausalLMWithValueHead.from_pretrained("./sft-mistral")
ref    = AutoModelForCausalLMWithValueHead.from_pretrained("./sft-mistral")
rm     = AutoModelForSequenceClassification.from_pretrained("./rm-mistral")
 
trainer = PPOTrainer(cfg, policy, ref, tokenizer, reward_model=rm)
 
for batch in dataloader:
    responses = trainer.generate(batch["prompt_ids"])
    rewards   = rm(batch["prompt_ids"], responses)
    stats     = trainer.step(batch["prompt_ids"], responses, rewards)
    trainer.log_stats(stats, batch, rewards)

That's the whole user-facing API. About 30 lines. All the ugly PPO plumbing is internal to .step().

5 · Reading a PPO training log (survival guide)

TRL logs many metrics per step. The three that matter:

• ppo/mean_scores: mean reward. Should trend UP over training. If flat, β might be too high (policy can't move). If wildly spiking then crashing, LR too high.
• ppo/policy/policykl: mean KL between current policy and rollout policy (within one mini-epoch). Should stay small — <0.02 per step. Big → clipping isn't containing the update, drop LR or increase clip strength.
• objective/kl: KL between current policy and reference. Should grow slowly over training. If it explodes past target (default 6.0), TRL's adaptive KL controller will crank β up automatically. If β saturates high and KL still grows, you're reward-hacking.

Rules of thumb for a healthy PPO run:

• Reward grows monotonically for 100-500 steps then plateaus.
• KL(π || π_ref) grows to 5-15 and stabilizes.
• Policy loss oscillates around zero (positive means "we made things worse in the mini-epoch").
• Value loss decreases smoothly to a floor.

6 · Why the field is moving to DPO

PPO on LLMs is famously fiddly:

• Four models in memory simultaneously (56GB for a 7B fp16 policy).
• 15+ hyperparameters (LR, β, ε, γ, λ, clip ranges, mini-epochs, batch sizes, target KL, adaptive KL coefficients, entropy bonus, value function coefficient, ...).
• Rollouts are slow (generation dominates), and you need many iterations.
• Reward-hacking failures happen quietly — you notice on your eval set, not in the training loss.
• Debugging is opaque: a 5% degradation could be RM, KL, LR, or bad rollouts.

Rafailov et al. (Stanford, 2023) noticed something startling: you can derive a closed-form loss that turns the preference dataset directly into a training objective on the policy, with no RM, no rollouts, no value head, no PPO. That's DPO. Same theoretical foundation (Bradley-Terry + KL constraint), but algebraically manipulated into a supervised loss on preference pairs.

For most SFT teams in 2024, DPO gives 90-100% of the quality of PPO with 20% of the complexity. Which is why every open-source instruct model since Zephyr uses DPO, not PPO.

Where PPO still wins:

• Iterative RLHF — when you keep collecting fresh preferences during training. DPO's dataset is fixed at start-of-training.
• Reward shaping — if your reward is more than a single Bradley-Terry pref score (e.g., safety+helpfulness+conciseness combined), PPO handles it natively; DPO needs contortions.
• Frontier labs — anecdotally, at the very best-model tier, PPO's extra flexibility still buys measurable quality. It's why GPT-4-turbo and Claude 3 still have PPO in the recipe somewhere.

For your first alignment run: use DPO (S064). For your fifth: consider whether PPO is worth the complexity for your task.

7 · Diagram — one PPO iteration

8 · GRPO, RLVR, and the 2024–2025 revival of RL for LLMs

Section 6 ("Why the field is moving to DPO") was the 2023 conventional wisdom. Then in early 2025 the pendulum swung back toward RL — hard — because of two ideas: GRPO and verifiable rewards. Here's the update.

8.1 The DPO ceiling and why RL came back

By late 2024 the open community had thoroughly explored DPO/SimPO/ORPO and hit a ceiling: preference-tuned models were polished but plateaued on hard reasoning benchmarks (MATH, GPQA, competition code). Ivison et al. (2024) showed DPO couldn't push past ~40% on MATH even with unlimited preference data.

Meanwhile, OpenAI's o1 (Sept 2024) demonstrated that large-scale RL with chain-of-thought could crush MATH (94%) and Codeforces (89th percentile). The recipe wasn't public but the signal was clear: RL wasn't dead, it had been misapplied.

January 2025 changed the game: DeepSeek-R1 dropped, and its methods paper was public.

8.2 GRPO — Group Relative Policy Optimization (Shao et al., DeepSeek, 2024–2025)

GRPO is PPO without the value model. Here's the swap:

PPO advantage: $\hat{A}_t = R_t - V_\phi(s_t)$ where $V_\phi$ is a learned critic (a whole second model).

GRPO advantage: for each prompt, sample $G$ completions (say $G=64$). Compute reward $R_i$ for each. Set $\hat{A}_i = (R_i - \text{mean}(R)) / \text{std}(R)$ within the group. That's the advantage. No value model.

Why this works:

• The group mean is an unbiased baseline for the prompt (same as value model would estimate) but comes for free from sampling.
• Halves the memory footprint — you drop the value model, which was as big as the policy.
• Better with sparse/binary rewards — the value model in PPO struggled when rewards were mostly 0 with occasional 1; group normalization handles this natively.

DeepSeek-R1 used GRPO with $G=64$ completions per prompt. On DeepSeek-Math-7B, GRPO added +5.5 pts on GSM8K and +6.4 pts on MATH vs PPO with the same reward model. Faster wall-clock too because no critic training.

Further reading:

• DeepSeek-Math paper (GRPO intro): https://arxiv.org/abs/2402.03300
• DeepSeek-R1 report: https://arxiv.org/abs/2501.12948

8.3 RLVR — Reinforcement Learning with Verifiable Rewards

The other DeepSeek-R1 innovation: skip the reward model entirely for domains where the answer can be checked programmatically.

• Math: extract final numeric answer, compare to ground truth. Reward = 1 if match, 0 otherwise.
• Code: run generated code against unit tests. Reward = pass rate.
• Format: regex-check that reasoning is wrapped in <think>...</think> tags. Reward = 1 if formatted correctly.

No human labels, no RM, no reward hacking (the verifier is unhackable in the way an RM isn't). Combined with GRPO this is what let DeepSeek-R1 match o1 on MATH and code with an open-source recipe.

Allen AI adopted the same recipe for Tulu-3-RLVR (the third stage after SFT + DPO). +3.2 pts on GSM8K, +7.1 pts on MATH beyond the DPO checkpoint. This is the current SOTA open recipe as of mid-2025.

8.4 R1-Zero — SFT-free RL

DeepSeek also showed R1-Zero: starting from a raw base model (no SFT!), you can jump straight to GRPO+RLVR and the model develops chain-of-thought reasoning entirely through RL. Emergent behaviors: self-verification ("wait, let me check that"), backtracking, considering alternatives.

R1-Zero was unreadable (mixed languages, weird formatting) so R1-proper added a small SFT warm-start. But the finding is deep: SFT teaches format, RL teaches capability. If you have verifiable rewards, RL can bootstrap capability without ever seeing a human demonstration.

8.5 PPO isn't dead — where it still wins

Even with GRPO's rise, classic PPO still dominates in three regimes:

1. Dense rewards from a scalar RM. If your RM outputs continuous scores (not binary) and you have plenty of memory, PPO's value model gives cleaner variance reduction than group normalization.
2. Very long trajectories (agents, multi-turn tool use). PPO's step-level advantage is more informative than GRPO's episode-level normalization when episodes are 100+ steps.
3. When you've already invested in a PPO codebase. GRPO is easier to run but PPO is more mature; frontier labs (Anthropic, OpenAI) reportedly still use PPO variants for their main alignment loops.

8.6 The 2025 alignment stack

Current frontier open recipe (as of Tulu-3, DeepSeek-R1, Llama-3.1):

Steps 1–3 are what you'll ship for a chat product. Steps 4–5 are what frontier labs add to top leaderboards. The recipes are open. The compute is the moat.

Further reading:

• Tulu-3 RLVR: https://arxiv.org/abs/2411.15124
• Nathan Lambert's Interconnects post-mortem of 2024 alignment: https://www.interconnects.ai/p/2024-alignment-year-in-review

9 · War story — the KL coefficient that ate my weekend

Early 2024. I was reproducing the TRL PPO tutorial on a 1.4B Pythia model for a talk. Followed the recipe. Everything looked normal for 200 steps — reward climbing from 0.2 to 0.8, KL creeping up to 8 (init KL coef 0.2). Then step 210, reward jumped to 4.5 (max scale was 5.0). Suspiciously good.

Inspected the samples. The policy had converged to generating a single response for every prompt: "Certainly! I would be happy to help you with that. As an AI assistant, my role is to provide helpful, accurate, and safe information. Here is what I would recommend: I would strongly recommend that you consult a qualified professional for this matter."

That string had reward 4.9 from my RM (which was trained on UltraFeedback — which is dominated by GPT-4 responses — which love that exact preamble). The policy found the sycophancy attractor and camped there. KL was 24 by that point but the reward gain was outweighing the KL penalty.

Fix: I bumped the initial KL coef from 0.2 to 0.5 and set target_kl=6 (adaptive KL controller in TRL). Also added a length penalty in the reward: R ← R - 0.001 * (len - 50).clamp(0). Retrained. Reward plateaued at 2.8 (lower nominal) but sample diversity survived. Human eval score on a held-out set: 2× the sycophantic-collapsed policy.

Lessons:

1. Watch KL like a hawk. If KL > 15 sustained, you're overfitting to the RM's quirks.
2. target_kl with adaptive coefficient is the difference between working PPO and non-working PPO.
3. Length penalties are cheap insurance against verbose collapse.
4. Always sample and read outputs during training. Reward numbers lie; sample text doesn't.

10 · Retention scaffold

Recall

Q1. Write down the PPO-KL objective for LLMs.

Q2. What role does each of the four models play?

Q3. Why do we use reverse KL rather than forward KL in the PPO penalty?

Q4. In a healthy PPO training log, which direction should reward and KL each move?

Q5. Why is the field moving from PPO to DPO?

Stretch prompt

You're setting up a PPO run and want to prevent reward hacking as much as possible. Design a "sanity harness" that runs every N steps during training: generates from the policy on a fixed set of held-out prompts, computes RM score, and flags if either (a) mean length is drifting, (b) any characteristic n-gram is dominating outputs, or (c) RM score has grown past a threshold consistent with hacking. Sketch pseudo-code.

In your own words

Explain PPO to a colleague who understands SFT and reward modeling but has never touched RL. In 5–6 sentences, address: (a) what PPO adds beyond just "gradient-ascend the reward", (b) what each of the four models does, (c) what the failure modes look like.

Spaced review

• S062 — Reward modeling. PPO is the algorithm that consumes the RM. Without a good RM, PPO can only fail.
• S058 — SFT. The policy π_θ and reference π_ref both start as your SFT'd model. PPO is downstream of SFT.
• S022 — Cross-entropy loss. Compare the PPO clipped objective (weighted log-prob) to CE loss — same core object, different weighting scheme.

Next session

S064 — DPO: RLHF without the RL. Tomorrow we derive DPO from PPO's own KL-regularized RL objective, show that the optimal policy can be written in closed form as a function of the reference policy and reward, invert that to solve for reward as a function of log-ratios, plug into the Bradley-Terry loss, and end up with a supervised loss that trains the policy directly on preference pairs. Same underlying theory, 10× less code.

Bring back tomorrow

• Four models: policy, reference, reward, value.
• PPO clipped objective bounds per-token probability ratios to 1 ± ε.
• Reward hacking is the #1 failure mode; watch samples, not just numbers.

Previous: ← DL S062 · Next: DL S064 →