[LLM 3/10] RLHF and PPO: Training a Model on a Reward You Cannot Differentiate

In chapter 2 we taught the model by "imitating the answer key" one token at a time. But the qualities that make an AI assistant genuinely usable — correct, polite, not making things up, not sliding into English — have no answer key to imitate, and cannot be written directly as a loss function. This chapter is the most classical answer to that problem: RLHF (Reinforcement Learning from Human Feedback) with PPO. We will train a real reward model from Thai preference pairs, then write the PPO loop ourselves from scratch in about 120 lines, and close with my favorite experiment in the series: taking off the KL leash and watching the model cheat the reward live. This is deliberately the heaviest chapter of the series, because chapter 4 (DPO) and chapter 5 (GRPO) both start from this chapter's equations and each choose a different piece to delete.

1. The Problem​

SFT in chapter 2 carries one hidden assumption: there must be an answer key to imitate. But think about what we actually want — say, "solve the math problem correctly, and explain it in readable Thai." That sentence has no single answer key: good answers come in a hundred shapes, and "readable" cannot be written as an equation.

Try to optimize these things directly and you always run into two walls:

Wall one — quality cannot be written as a loss. "Better" cannot be defined as a function, but humans are very good at comparing: show them two answers and ask which one they prefer, and they can do it immediately and fairly consistently. So the data you can actually collect comes in threes: a prompt $x$, the chosen answer $y_w$, the rejected answer $y_l$.

Wall two — even with a score, you cannot backprop. Suppose a magical function $r(x,y)$ scored every answer. You still couldn't train supervised, because the answer $y$ comes from sampling tokens one at a time. The score arrives after the sampling is done, and derivatives cannot travel back through sampling — the path from $r$ back to the weights $\theta$ breaks exactly there.

Hence the chapter's title: we are about to optimize a reward we cannot differentiate. The tool that can do that is called reinforcement learning.

2. What We're Going to Do​

RLHF gets past both walls with a two-step walk:

• Stage A — Reward Model: train a model $r_\phi(x,y)$ to imitate human comparisons from preference pairs (defeats wall one, and stands in for humans who can't score fast enough)
• Stage B — PPO: use policy-gradient RL to push the policy toward $r_\phi$'s scores without differentiating through the sampling (defeats wall two), with a KL leash holding it back from fleeing the starting model

The price you pay is complexity: during training there are four models resident in VRAM at once — the policy $\pi_\theta$ (the one being trained), the reference $\pi_{\text{ref}}$ (the frozen starting model), the reward model $r_\phi$, and a value network $V_\psi$ that hasn't introduced itself yet (wait for section 3.4).

And one more sentence to hold on to for the whole series: this chapter's objective is the master equation of the second half of the series. Chapter 4 (DPO) solves it in closed form until the reward model and the RL loop cancel out. Chapter 5 (GRPO) changes how the advantage is estimated until the value network disappears. Understand this one chapter, and the next two become "deleting parts" — readable at sight.

3. The Equations​

3.1 The reward model: Bradley–Terry​

Stage A trains $r_\phi$ with one short loss:

• $r_\phi(x,y)$ = a single scalar score per text — in practice a language model whose head is swapped for a single linear layer (num_labels=1)
• $\sigma$ = the sigmoid, turning the score difference into the probability a human picks $y_w$ (the Bradley–Terry model)
• the wider the gap $r_\phi(x,y_w) - r_\phi(x,y_l)$, the lower the loss

The point people overlook and pay for later: this loss sees only the difference of the scores. Replace $r_\phi$ with $r_\phi + c$ for any constant $c$ — the loss doesn't change at all. Which means the absolute scale of a reward model is meaningless and is not pinned down by training. Two runs can yield mean scores of 3.7 and −12.4 that rank identically. This is why you must always standardize rewards before feeding them into PPO (subtract the mean, divide by the std) — hold onto this; it returns in sections 7 and 9.

3.2 The master equation: the RLHF objective​

If you memorize a single equation from this whole series, memorize this one:

In plain language: "collect as much reward as you can, but every step you walk away from the starting model costs a fine."

• $\pi_\theta$ = the policy, the model being trained — note that $y$ is sampled from $\pi_\theta$ itself; this is the structural difference from SFT, which learns from static data sitting in a file
• $\pi_{\text{ref}}$ = the reference, the starting model (the post-SFT model from chapter 2), frozen throughout training
• $\beta$ = the price per nat of wandering off — the tightness of the leash
• $\mathbb{D}_{\text{KL}}$ = the distributional distance between the policy and the reference

3.3 The PPO clipped surrogate: Stage B's engine​

Raw policy gradient (REINFORCE) can use a batch of rollouts for exactly one update before throwing it away — very expensive, because generation is the bottleneck. PPO wants to squeeze the same rollouts for several epochs, so it needs a correction factor (the importance sampling ratio):

• $s_t$ = the state at position $t$: the prompt plus every token sampled so far
• $a_t$ = the "action": the next token, already sampled during the rollout
• $\pi_{\theta_{\text{old}}}$ = a snapshot of the policy at rollout time — computed once and frozen

Then $\rho_t$ gets pinched with a clip:

• $\hat A_t$ = the advantage: "how much better than expected was this token" (defined in the next subsection)
• $\epsilon$ = the width of the trust region (standard value 0.2)

The heart is that min + clip work together with deliberate pessimism: if $\hat A_t$ is positive (a good token), the payoff from pushing $\rho_t$ is capped at $1+\epsilon$ — pushing beyond that earns nothing, gradient zero. But if $\hat A_t$ is negative (a bad token), the min always picks the worse branch — the penalty has no floor. One-sentence summary: gains capped, losses uncapped. The policy therefore moves in small steps that stay close to where it already is.

3.4 GAE: computing the advantage without drowning in noise​

The advantage is built from the TD error of the value network $V_\psi$:

• $V_\psi(s_t)$ = the value network's prediction of "from here to the end, how much more reward will be collected" — this is model number four
• $r_t$ = the per-token reward (in our task: a KL penalty at every position, plus the task score at the final token)
• $\gamma$ = the discount factor (LLM work usually uses 1.0)
• $\lambda$ = the bias–variance dial: $\lambda = 0$ trusts $V_\psi$ wholeheartedly (high bias if $V_\psi$ mispredicts), $\lambda = 1$ doesn't trust it at all and waits for the real outcome to the very end (high variance, carrying the noise of the whole trajectory); the popular value is 0.95

3.5 The full PPO loss: three terms, two models​

Assemble every piece into the one loss the optimizer actually sees (written as a minimization):

• first term = the clipped surrogate from 3.3 (negated because we want to maximize it)
• second term = the value loss teaching $V_\psi$ to predict close to the actual return $\hat R_t$; $c_1$ is usually 0.5
• third term = the entropy bonus $\mathcal{H}$, keeping the distribution from collapsing too fast; $c_2$ is usually 0.01
• as for the KL leash from equation 3.2, standard practice folds it into the per-token reward: $r_t \leftarrow r_t - \beta\,(\log\pi_\theta - \log\pi_{\text{ref}})$, which is the approach we take in section 7

Count all the toys that need tuning: 4 models, plus $\epsilon, \beta, \gamma, \lambda, c_1, c_2$, and two separate learning rates. This is why PPO is famous for "run it twice with different seeds, get two completely different stories" — and it is the reason for chapter 4's entire existence.

4. Seeing the Equations​

Bradley–Terry: the gradient piles onto the pairs still ranked wrong​

Pick the Bradley-Terry mode in the tool below and drag the margin: pairs the reward model already ranks confidently right (large positive margin) have almost no gradient left — Stage A training automatically spends its budget on the pairs it still ranks wrong.

If this plot reminds you of DPO's gradient in chapter 4 — that is no coincidence. DPO lifts this Bradley–Terry model wholesale, merely changing what plays the role of the reward.

The trust region made visible: min + clip​

Study the asymmetry closely — it is PPO's entire personality: moving toward something good earns at most 20% per step, but hand too much probability to a bad token and you are hauled back at full force, every time. The "flat region" in the figure is the trust region that lets PPO reuse old rollouts for several epochs without exploding.

GAE's bias–variance dial​

Read the lines from the bottom up: at $\lambda = 0$ (green) the advantage hugs zero nearly everywhere — the end-of-episode score never reaches the early tokens, because everything is filtered through a $V_\psi$ that mispredicts. At $\lambda = 1$ (red), every token gets full credit from the ending but carries the accumulated noise of the whole trajectory with it. $\lambda = 0.95$ (blue) is the middle ground the whole field has settled on — the signal travels far, but the noise is damped.

Before moving on, build your intuition for the word advantage — "how much better than expected" — by hand: this tool uses the group mean as its baseline instead of $V_\psi$ (a full-blown spoiler of chapter 5). Press the All correct preset and watch what happens when every answer earns the same reward:

Advantage zero across the whole group = no signal to learn from — hold onto that feeling for section 9.

The KL leash: equation 3.2 drawn as a picture​

The green line's stopping point is not a guess — it is the mathematics of equation 3.2: optimization stops exactly where the reward gained per nat equals β; walking further loses money. When $\beta = 0$ that stopping condition doesn't exist — every nat of wandering that buys even a sliver of reward is "profit," so the model keeps running out of natural language for as long as the reward number still twitches upward.

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier is enough).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

Why RLHF is expensive, in one picture​

The numbers in the figure are just the weights — before activations, the KV cache during generation, gradients, and optimizer state. And at 0.6B everything still looks tiny, but the ×4 multiplier goes nowhere as you scale: at 7B it is 56 GB before you've done anything at all.

6. Preparing the Data​

Stage A — preference pairs for the reward model​

We use iapp/dpo_thai_tutorial (100 pairs, Apache-2.0) — a Thai preference dataset I built myself for this series and released for free reuse. Each row has prompt, chosen, rejected, hand-curated with an emphasis on politeness and natural-sounding Thai.

Split 80/20: train on 80 pairs, hold 20 pairs out — never touched during training. Stage A's passing bar is pairwise ranking accuracy on the 20 held-out pairs, with a Wilson 95% CI that does not straddle 0.5 — "significantly better than a coin flip" already proves the mechanism, because 20 pairs genuinely cannot make the CI any tighter.

Stage B — problems a rule can grade​

For the PPO loop we use 64 math problems from VISAI-AI/gsm8k-thai (a Thai translation of GSM8K) and score them with a verifiable rule instead of the Stage A reward model:

import re

def rule_reward(response: str, gold: int) -> float:
    """+1.0 if the last integer in the answer is correct, +0.2 if the answer is actually in Thai"""
    nums = re.findall(r"-?\d+", response.replace(",", ""))
    correct = 1.0 if nums and int(nums[-1]) == gold else 0.0
    thai = sum("\u0e01" <= ch <= "\u0e5b" for ch in response)   # the Thai Unicode block, U+0E01..U+0E5B
    thai_bonus = 0.2 if thai / max(len(response), 1) >= 0.5 else 0.0
    return correct + thai_bonus

7. The Main Code​

7.1 Stage A — training a real reward model (~6 minutes)​

Turn a language model into a scoring machine: the LM head is replaced with a single linear layer returning one scalar.

import torch, torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from peft import LoraConfig, get_peft_model

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
rm = AutoModelForSequenceClassification.from_pretrained(
    "Qwen/Qwen3-0.6B",
    num_labels=1,                          # scalar head: one score per text
    torch_dtype=torch.float16,             # the T4 has no bf16
    attn_implementation="sdpa",
).cuda()
rm.config.pad_token_id = tok.pad_token_id

rm = get_peft_model(rm, LoraConfig(
    task_type="SEQ_CLS", r=8, lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    modules_to_save=["score"],             # the score head starts random — must be trained in full
))
for p in rm.parameters():
    if p.requires_grad:
        p.data = p.data.float()            # the fp16 lesson from chapter 1: train in fp32

def rm_loss(chosen, rejected):
    s_w = rm(**chosen).logits.squeeze(-1)     # r_φ(x, y_w)
    s_l = rm(**rejected).logits.squeeze(-1)   # r_φ(x, y_l)
    return -F.logsigmoid(s_w - s_l).mean()    # equation 3.1, literally

Train 3 epochs on 80 pairs, then measure pairwise accuracy on the 20 held-out pairs with a Wilson CI. If the CI clears 0.5 — you have just trained the first reward model of your life, on 80 rows of data.

7.2 Stage B — PPO written from scratch, ~120 lines (~10 minutes)​

We will not use TRL's PPOTrainer, and this is a deliberate decision, not stubbornness: TRL has moved PPOTrainer into trl.experimental and announced plans to remove it in 0.29.0 — code taught against that library will stop running within a few months. A self-written PPO loop of about 120 lines will run for as long as PyTorch exists. And more important still: write it yourself and you will know what every line does, just like the 25-line DPO loss in chapter 4.

The pieces on the board: policy = base + LoRA adapter (from chapter 2, is_trainable=True), reference = the same base with the adapter off, and a value head — a two-layer MLP plugged into the final hidden state:

value_head = torch.nn.Sequential(
    torch.nn.Linear(1024, 1024), torch.nn.Tanh(),
    torch.nn.Linear(1024, 1),
).float().cuda()                           # ~1M parameters — tiny next to the other three

The rollout: sample answers 8 prompts at a time (max_new_tokens=200, do_sample=True), score with rule_reward, standardize the scores within the batch, and store the rollout-time log-probs, detached, immediately. Then enter this update loop — the 40 lines that are the heart of the whole chapter (the full loop is in the notebook):

def compute_gae(rewards, values, gamma=1.0, lam=0.95):
    """rewards: [T] for one response, values: [T+1] (last slot = 0 after the end)"""
    adv, acc = torch.zeros_like(rewards), 0.0
    for t in reversed(range(rewards.shape[0])):
        delta = rewards[t] + gamma * values[t + 1] - values[t]   # equation 3.4
        acc = delta + gamma * lam * acc                          # Â_t = δ_t + γλ Â_{t+1}
        adv[t] = acc
    return adv

def ppo_update(rollout, eps=0.2, beta=0.05):
    out = policy(rollout.ids, output_hidden_states=True)         # one forward, two outputs
    logp = gather_logprobs(out.logits, rollout.ids)              # [B, T], carries gradient
    with torch.no_grad(), policy.disable_adapter():
        logp_ref = gather_logprobs(policy(rollout.ids).logits, rollout.ids)

    # per-token reward = KL penalty at every position + task score at the last token (equation 3.2)
    rew = -beta * (logp.detach() - logp_ref)
    rew[:, -1] += rollout.scores_std                 # the standardized rule score

    hidden = out.hidden_states[-1].detach().float()  # detach = cut the gradient into the trunk
    values = value_head(hidden).squeeze(-1)
    adv = torch.stack([compute_gae(r, F.pad(v, (0, 1))) for r, v in zip(rew, values)])
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)    # standardize the advantage once more

    log_ratio = (logp - rollout.logp_old).clamp(-10, 10)   # prevents fp16 overflow!
    ratio = log_ratio.exp()                                # ρ_t (equation 3.3)
    surr = torch.min(ratio * adv.detach(),
                     ratio.clamp(1 - eps, 1 + eps) * adv.detach())

    returns = (adv + values).detach()                # the value head's target
    m = rollout.resp_mask                            # count response tokens only
    loss = (-(surr * m).sum() + 0.5 * ((values - returns).pow(2) * m).sum()) / m.sum()
    return loss                                      # (the entropy term is in the notebook)

Settings: $\epsilon = 0.2$, $\beta = 0.05$, $\gamma = 1.0$, $\lambda = 0.95$, 4 epochs per rollout on 64 prompts, adapter LR 1e-5, value head LR 1e-4. Roughly 10 minutes per run on a T4 (the notebook runs twice: $\beta = 0.05$ and $\beta = 0$).

8. Results​

The notebook tracks three curves simultaneously at every update and writes them to results.json:

1. Mean reward per rollout — should climb (this is what we're buying)
2. Per-token KL against $\pi_{\text{ref}}$ — should grow and then saturate under the ceiling that $\beta$ sets (this is the price we pay)
3. Mean answer length — the disease detector: length spiking or cratering abnormally is the first sign of a policy going strange

Catching reward hacking red-handed​

The notebook's second run sets $\beta = 0$ and touches absolutely nothing else — no leash. The pattern expected: reward climbs as fast or faster, but KL soars with no ceiling, and the answers start to degenerate — repetitive, oddly short, or collapsing into a fixed formula with a number stuffed at the end of the sentence. Because rule_reward sees only the final number and the fraction of Thai characters, everything it cannot see is free for the model to throw away.

Samples of degenerate answers from the $\beta = 0$ run are printed by the notebook's final cell:

[This space is filled only from the notebook's real run — I will not invent degenerate
 examples myself, because this whole series stands on the rule that no number or output
 is ever made up. Run the notebook and the "hacking exhibits" cell will show 2-3 real
 answers along with their KL.]

9. Comparison​

The notebook runs the same battery over three systems on held-out Thai math problems (TH-MATH):

Next chapter: DPO — DPO deletes both the reward model and the RL loop with pure algebra, starting from master equation 3.2, the one you just memorized.

References​

1. Schulman et al. (2017). Proximal Policy Optimization Algorithms — the original PPO paper: the clipped surrogate in section 3
2. Schulman et al. (2015). High-Dimensional Continuous Control Using Generalized Advantage Estimation — GAE: the advantage estimator PPO uses
3. Christiano et al. (2017). Deep reinforcement learning from human preferences — the paper that started RL from human preferences
4. Stiennon et al. (2020). Learning to summarize from human feedback — the first convincingly working RLHF, on summarization
5. Ouyang et al. (2022). Training language models to follow instructions with human feedback — InstructGPT: the origin of the whole SFT -> RM -> PPO pipeline
6. Bai et al. (2022). Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback — the helpful/harmless dataset and its lessons about KL
7. Zheng et al. (2023). Secrets of RLHF in Large Language Models Part I: PPO — the PPO implementation details other papers omit
8. Bradley & Terry (1952). Rank Analysis of Incomplete Block Designs: I. The Method of Paired Comparisons — the Bradley-Terry model every reward model rests on

The writing, code and notebooks in this series are licensed under CC BY-NC-SA 4.0 — reuse and adapt them freely with attribution, for non-commercial purposes, and share your adaptations under the same terms. Third-party models and datasets referenced here keep their own licences.