Skip to content

RL From Zero and PPO

The full classic RLHF stack, derived from scratch. RL vocabulary, the non-obvious step where an LLM is already a policy, why naive REINFORCE fails, the clipped surrogate, the KL penalty, and the famous four models in memory. Once you can name what each of the four models does, the rest of post-training organizes itself around what each subsequent method removes.

Rapid Recall

An LLM is already a policy: the softmax over next-token logits is \(\pi(a|s)\). Generating a response is rolling out a trajectory. The RLHF reward is sparse — one scalar at the very end. Naive policy gradient (REINFORCE) has crushing variance, off-policy issues, and destructive single-step updates. PPO layers three fixes: (1) advantage estimation with a learned critic V(s) for variance reduction; (2) clipped surrogate objective with importance ratio π_θ / π_θ_old clipped to [1-ε, 1+ε] for a trust region; (3) KL penalty to a frozen reference model to prevent reward hacking. Cost: four models in memory — policy (trained), reference (frozen SFT), reward model (frozen), value/critic (trained).

§1 Reinforcement learning, from zero

No labeled "right answer" per step — just an agent trying things and learning from outcomes. The vocabulary: an agent observes a state \(s\), picks an action \(a\) via its policy \(\pi\); the environment returns a reward \(r\) and a new state. Loop. Goal: a policy that maximizes total expected reward.

A policy is literally the agent's strategy, a function from state to action (or a distribution over actions).

§2 Mapping RL onto LLMs

The non-obvious translation that unlocks everything: an LLM is already a policy. The softmax over next-token logits is \(\pi(a \mid s)\).

RL concept policy π state sᵗ action aᵗ reward r episode LLM equivalent the LLM itself (softmax over tokens) prompt + tokens generated so far the next token RM score, at the END of the response one full generation: prompt → response
Generating a response is rolling out a trajectory. Crucially, the reward in RLHF is sparse — one scalar at the very end, not per token.

§3 Why naive REINFORCE fails

The obvious move — gradient ascent on expected reward via the policy gradient theorem — works, but suffers from:

  • High variance. A single trajectory's reward is a noisy estimate of the gradient direction.
  • Sample inefficiency. On-policy: each sample is valid for one update step.
  • No baseline. Every token in a high-reward trajectory gets reinforced equally, even the ones that did not help.
  • Destructive updates. One bad step can collapse the model into gibberish with no path back, especially with sparse rewards.

PPO exists to fix that last problem above all.

§4 The four models in PPO

Policy (actor) the LLM being trained TRAINED Reference (frozen SFT) KL penalty anchor FROZEN Reward (frozen RM) scores the full response FROZEN Value (critic) estimates V(s) for advantages TRAINED
Four models, the reason PPO for LLMs is so expensive. Two are trained (policy, critic), two are frozen (reference, reward).

§5 PPO — three machineries

5.1 Idea 1 — advantage, and why you need a critic

Instead of "was this response good?" ask "was it better than expected from this state?" That is the advantage: \(A_t = R_t - V(s_t)\). If every response scores ~+5, raw reward is uninformative; advantage centered near zero gives a clean low-variance signal.

\(V(s)\), the expected reward from state \(s\) onward, is not known analytically, so you learn it with a neural network: the critic. Same trick as the RM: SFT backbone plus a scalar value head, trained by regression on observed returns:

\[ \mathcal{L}_{\text{value}} = (V_\phi(s_t) - R_t)^2 \]

Policy plus critic training together is the actor-critic setup. This is precisely why the value model exists: the RM gives the terminal reward; the critic gives per-state baselines.

5.2 Idea 2 — the clipped objective (the "proximal" part)

Define the probability ratio:

\[ r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)} \]

PPO maximizes:

\[ \mathcal{L}^{\text{PPO}} = \mathbb{E}_t\Big[\min\big(r_t(\theta) A_t,\; \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t\big)\Big] \]

With \(\epsilon = 0.2\), the ratio is clipped to \([0.8, 1.2]\). The min creates a trust region: when \(A_t > 0\) you increase probability but stop benefiting past +20%; when \(A_t < 0\) you keep suppressing the bad action. The policy moves only modestly from \(\pi_{\text{old}}\) — no catastrophic jumps. Because the clip keeps you near the data-generating policy, you can take several gradient steps on the same batch, which is why PPO is more sample-efficient than vanilla policy gradient.

5.3 Idea 3 — KL penalty and the reference model

Even with clipping, the policy can reward-hack — drift into RM blind spots. So subtract a penalty for diverging from the frozen SFT model (the reference model):

\[ r_{\text{final}}(x, y) = R_\phi(x, y) - \beta \, D_{KL}\big(\pi_\theta(\cdot \mid x) \,\Vert\, \pi_{\text{ref}}(\cdot \mid x)\big) \]

Two different jobs

The clip ratio constrains update size per step. The KL penalty constrains total drift from the original SFT model across all of training. They are not redundant.

§6 One PPO iteration, end to end

  1. Sample prompts; generate responses with \(\pi_\theta\) (save log-probs as \(\pi_{\text{old}}\)).
  2. Score responses with the frozen RM.
  3. Build per-token rewards: KL penalty on every token, RM reward added only on the last token.
  4. Compute advantages via the critic (usually with GAE smoothing).
  5. Update policy with the clipped objective (multiple steps per batch).
  6. Update critic by MSE to observed returns. Repeat.

For the runnable rollout-plus-update implementation, see Alignment Walkthrough §5.

§7 Why RLHF is hard, in five bullets

  1. Four models in memory — expensive. Need big clusters. For a 70B base model that is crushing.
  2. RL is unstable — PPO has many hyperparameters (LR, β, ε, GAE λ), and they interact. Reward goes up, then crashes. Loops diverge.
  3. Reward hacking — the policy finds ways to game the RM. Classic example: if RM prefers longer answers, policy learns to produce rambling verbose outputs.
  4. RM quality ceiling — your final model can only be as good as your RM. If the RM misjudges a preference, the policy learns that misjudgment.
  5. Compute — one full RLHF run on a 70B model takes days on a multi-node cluster.

This is why OpenAI and Anthropic can do it (they have the infra and RL experts), but a 10-person startup usually cannot. Enter DPO.

Interview Questions

Q1: Why is the reward only at the end, not per token?

The RM was trained on complete responses, so it only knows how to score finished outputs. Per-token rewards need per-token reward models, which are harder to get from human preferences. The KL penalty is added at every token, but the RM score is added only at the final response token; GAE smooths this terminal reward backward across the trajectory.

Q2: What happens if β = 0 (no KL penalty)?

Reward hacking. The policy exploits RM blind spots and drifts into verbose, sycophantic, or repetitive outputs that score high but are useless. The KL penalty is what keeps the policy in the RM's "trust region."

Q3: What happens if β is too high?

The policy never moves; RL becomes useless because any change from SFT is penalized away. You get exactly the SFT model back, which is fine but defeats the point of running RLHF.

Q4: Why multiple gradient steps per batch (vs one for vanilla PG)?

The clip keeps the policy near the data-generating policy, so re-using the batch a few times is approximately valid. Vanilla policy gradient re-use would be off-policy and biased; PPO's clipping bounds the off-policy error so multiple updates remain trustworthy.

Q5: What is the role of the critic if you already have a reward model?

Different jobs. The RM scores complete responses after generation; the critic estimates expected future reward from any mid-generation state, used for advantages. The critic gives advantages (variance reduction); the RM gives the reward signal.

Q6: Why advantage instead of raw return?

Variance reduction. Subtracting any state-only baseline does not bias the gradient (provable from the policy gradient theorem), but sharply reduces variance. Without a baseline, all trajectories with positive reward get reinforced equally regardless of which actions were actually responsible, which makes credit assignment slow and noisy.

Q7 (Trap): Why is PPO "on-policy" if we use an importance ratio?

It is approximately on-policy. The clipping keeps us close enough to the data-generating policy that the off-policy correction (the ratio) does not blow up. Truly off-policy methods (DQN) do not need this constraint because they have explicit value functions and bootstrapping. PPO sits in the middle: re-use a batch a few times, but only if you stay near the policy that produced it.