Lab 01 — Preference Optimization: Reward Model, DPO & PPO-with-KL

Phase: 07 — Alignment: RLHF, PPO & DPO Difficulty: ⭐⭐⭐☆☆ (the math is clean; the insight — why DPO == BT with an implicit reward, why KL is load-bearing — is ⭐⭐⭐⭐⭐) Time: 3–4 hours

Every aligned chat model turns human preferences into better behavior through one of three mechanisms — a Bradley-Terry reward model, DPO, or PPO with a KL leash. This lab builds all three from scratch in pure stdlib so the math is visible: numerically stable sigmoid/log_sigmoid, the Bradley-Terry preference loss, the DPO loss and its implicit reward margin, a stable kl_divergence, the PPO clipped surrogate and a PPO step loss that combines it with the load-bearing KL-to-reference penalty, and a tiny deterministic reward-model trainer that proves the preference is actually being learned.

What you build

• sigmoid / log_sigmoid — the logistic function and its log, computed stably so ±1000 doesn't overflow or hit log(0). log_sigmoid(x) = min(x,0) - log1p(e^-|x|). This stability is the lesson, not an afterthought.
• bradley_terry_loss / reward_accuracy — the preference loss -log σ(r_w − r_l) (mean over a batch), and the fraction of pairs ranked correctly. The loss → 0 as the margin → +∞, grows linearly as it → −∞.
• dpo_loss / dpo_implicit_reward_margin — DPO: -log σ(β·[(logπ_w − logπ_ref,w) − (logπ_l − logπ_ref,l)]). No reward model, no rollouts — the implicit reward is β·log(π/π_ref), and the margin is the thing inside the σ.
• kl_divergence — Σ p log(p/q), stable and guarded (skip p=0, reject q=0 where p>0, clamp round-off). The leash PPO uses to stay near the reference.
• ppo_clipped_objective / ppo_step_loss — the clipped surrogate min(ρA, clip(ρ,1±ε)A) (we maximize it), and the step loss -surrogate + kl_coeff·KL that prevents reward hacking.
• train_reward_model — gradient descent on the Bradley-Terry loss for a 1-D linear scorer; deterministic (seeded), and the loss provably decreases.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why log_sigmoid is computed in log-space and what log(sigmoid(-1000)) does to a naive implementation.
• You can write the Bradley-Terry loss and the DPO loss from memory and explain why DPO is "BT with an implicit, reference-relative reward."
• You can explain why test_dpo_loss_decreases_when_policy_widens_chosen_margin is the soul of the lab — it is the entire DPO learning signal in one assertion.
• You can explain why ppo_clipped_objective(1.5, 1.0, 0.2) returns 1.2 and why the min (not max) makes the update proximal.
• You can explain why test_ppo_step_kl_penalty_increases_loss_as_policy_diverges shows the KL term doing its job, and why removing it causes reward hacking.

How this maps to the real stack

Limits of the miniature (be honest in the interview): the reward "model" is a 1-D linear scorer, not a transformer with a scalar head; we use scalar log-probs, where a real DPO/PPO step uses per-token sequence log-probs summed over a response; PPO here is a single clipped step, not the full loop with rollouts, a value head, and GAE advantages; and the KL is over a tiny discrete distribution, not a token-level KL over the vocabulary. The algorithms are exact; the scale and plumbing are what the frameworks add.

Extensions (build these toward the real thing)

• Replace the scalar log-probs with per-token sequence log-probs: sum log π(y_t | x, y_<t) over a (tiny, hand-rolled) response and feed those into dpo_loss.
• Add IPO: swap the -log σ objective for a squared loss targeting a finite margin, and show it over-optimizes less on near-deterministic preference pairs.
• Add KTO-style unpaired learning: a loss over individual good/bad responses instead of pairs.
• Build a one-step PPO loop: a toy policy over a few tokens, a fixed reward, advantages from a baseline, and watch the KL term cap the drift over iterations.
• Plot the gold-vs-proxy over-optimization curve: train against a noisy proxy reward and show held-out "true" reward peak then decline.

Interview / resume

• Talking points: "Write the DPO loss and derive it from the RLHF objective." "Why is the KL-to-reference penalty load-bearing — what happens without it?" "DPO vs PPO vs GRPO — when each?" "What does β control in DPO?"
• Resume bullet: Implemented the core preference-optimization stack from scratch — numerically stable Bradley-Terry reward loss, DPO with its implicit reference-relative reward, and the PPO clipped surrogate with a KL-to-reference penalty — with a test suite proving the DPO learning signal, the PPO clip boundaries, and KL-driven anti-reward-hacking.