# Lab 12: RLHF from the reward model up

You will run the whole InstructGPT recipe in miniature: collect rankings from a (simulated) labeler, fit a Bradley–Terry reward model, then optimize TinyLM against it with PPO and a per-token KL penalty. Then you will watch the reward model get hacked when the KL leash is off.

Chapter: <https://learn.xhelio.ai/rl/learn/12-rlhf/>

## The toy world

TinyLM (from `tinylm.py`, the same model as the other Part III labs) answers prompts like `37+48=` with a few digits and `<eos>`. After a short SFT it gets about half of the sums right.

A **hidden labeler** (in `prefs.py`) ranks answers. They don't check the arithmetic. They have a style preference:

| Answer | Hidden score `gold` |
|---|---|
| well formed, last digit even (`84`, `130`) | +1 |
| well formed, last digit odd (`85`) | 0 |
| malformed: no `<eos>` within the 4-token budget (`4444`), empty, or a stray `+`/`=` | −1 |

Shown K = 4 answers to the same prompt, the labeler sorts them by $\kappa\cdot$gold + Gumbel noise ($\kappa = 2$). Inside any one ranking, each pair follows Bradley–Terry: $P(y_1 \text{ above } y_2) = \sigma(\kappa(\text{gold}_1 - \text{gold}_2))$. The labeler is noisy on purpose: answers with the same gold score are ordered by a coin flip, just as real annotators disagree.

The **reward model** is deliberately tiny: one weight per digit, summed over the answer's digits, $r(y) = \sum_{d \in y} w_d$. It cannot see where a digit sits or whether the answer ended properly. On the answers it was trained on (well formed, from the SFT model) that hardly matters. Off that distribution it has a blind spot, and PPO will find it.

## Files

| File | What it is |
|---|---|
| `rlhf.py` | **Your file.** Five tasks marked `TODO`. The training loops at the bottom (`train_reward_model`, `train_rlhf`) are complete and call your functions. |
| `prefs.py` | Given: the hidden labeler, `collect_rankings`, the `DigitRM` reward model, the `ValueModel` critic, masked helpers. Read it once. |
| `tinylm.py` | Given: the tiny transformer, sampling, per-token log-probs, SFT (shared by all Part III labs). |
| `test_rlhf.py` | Tests. |
| `run.py` | The full pipeline and four RLHF runs; saves `rlhf_runs.png`. |
| `solutions/rlhf.py` | Reference solution. Try not to peek. |

## Tasks

Run one task's tests at a time with `-k`, e.g. `python -m pytest -q -k task1`.

**Part A: the reward model**

1. **`bt_loss`**: the Bradley–Terry loss $-\log\sigma(r_w - r_l)$, averaged over pairs. Tests check three hand-computed values, the gradient $-(1-\sigma(\Delta))/N$, shift invariance, and numerical stability at $\Delta = -100$.
2. **`ranking_to_pairs`**: a ranking of K answers becomes $\binom{K}{2}$ (winner, loser) pairs. A test then trains the given `train_reward_model` (which uses both of your functions) and checks that the RM learns "even digits good" and reaches the accuracy ceiling set by the labeler's noise.

**Part B: PPO with a KL penalty**

3. **`kl_penalized_rewards`**: per-token rewards $r_t = -\beta\log\frac{\pi_\theta(y_t\mid\cdot)}{\pi_\text{ref}(y_t\mid\cdot)}$, plus the RM score on the last real token. One test is the chapter's worked example ("Paris", ".", EOS with $r_\phi = 1.5$, $\beta = 0.1$), so you can check your output against the table in the chapter.
4. **`gae_per_token`**: GAE over the tokens of each response with $\gamma = 1$, treating the state after the last real token as terminal (value 0) and ignoring padding. One test continues the chapter's "Paris" example (advantages 0.47, 0.32, 0.10 with $\lambda = 1$).
5. **`ppo_rlhf_loss`**: the clipped policy loss and the value loss, each averaged over real tokens only. Tests check values, the clip fraction, and that clipped and padding tokens get exactly zero gradient. The last test runs 12 PPO iterations (a few seconds) and checks that the RM score rises.

```bash
python -m pytest -q                   # all tests; they fail until you fill in the TODOs
python -m pytest -q -k task3          # just Task 3
RL_SOLUTIONS=1 python -m pytest -q    # run the tests against the reference solution
```

The first run trains and caches a small SFT checkpoint, `sft_tinylm.pt` (about 10 seconds).

## Then run the experiment

```bash
python run.py            # about 1.5-2 minutes; saves rlhf_runs.png
```

It fits the RM on 600 ranked prompts (3,600 pairs), prints its weights and held-out accuracy, then runs PPO four times for 120 iterations:

| Run | What changes |
|---|---|
| β = 0 (no leash) | no KL penalty at all |
| β = 0.1 | a light leash |
| β = 1 | a strong leash |
| β = 0.1 + PPO-ptx | light leash, plus a supervised loss on correct additions mixed into every PPO step (InstructGPT's fix for the "alignment tax") |

**What you should see** (reference solution, default seeds; your numbers may differ a little):

- **The reward model** reaches about 0.63 held-out accuracy against a ceiling of about 0.64. That sounds low, but most pairs of SFT answers have the same gold score, so the labeler flips a coin and nobody could do better. The weights come out positive for even digits (about +0.4) and negative for odd ones (about −0.5).
- **β = 0** is the cautionary tale. For about 30 iterations the RM score and the hidden labeler's score rise together. Then the policy discovers the blind spot: it stops emitting `<eos>` and fills the budget with even digits (`4444`). The RM score keeps climbing (to about +2.5), the hidden score crashes to −1, the KL to the reference passes 40 nats, and arithmetic accuracy goes to 0. This is reward hacking, and the gold curve has the rise-then-fall shape that [Chapter 19](https://learn.xhelio.ai/rl/learn/19-failure-modes/) studies.
- **β = 0.1** gets the hidden labeler's score to about +1 (nearly every answer ends in an even digit) at a KL of about 2.6 nats. But accuracy falls from 0.5 to under 0.2: the labeler never rewarded correct sums, so the policy happily answers `37+48=` with `84`.
- **β = 1** stays close to the reference (KL ≈ 0.25 nats). It gains less on the labeler's score (about +0.7) and keeps most of the arithmetic (accuracy ≈ 0.44).
- **β = 0.1 + PPO-ptx** keeps accuracy near the SFT level (≈ 0.46) while still raising the labeler's score (≈ +0.8). The supervised term protects the skill that the reward doesn't measure.

## Things to try

- **Adaptive KL.** Set `target_kl=1.0` (with `beta=0.1` as the starting value) and plot `hist["beta"]`. The controller is Ziegler et al.'s: $\beta \leftarrow \beta\,(1 + 0.1\,\mathrm{clip}((\mathrm{KL} - \mathrm{KL}_\text{target})/\mathrm{KL}_\text{target}, -0.2, 0.2))$.
- **Gold vs. √KL.** For the β = 0 run, plot `hist["gold"]` against `sqrt(hist["kl"])`. Gao et al. (2022) fit the RL curve with $d(\alpha - \beta\log d)$, $d = \sqrt{\mathrm{KL}}$ (their $\alpha, \beta$ are fitted constants, not the KL coefficient). Does yours rise and then fall?
- **Close the blind spot.** In `prefs.py`, add an 11th feature to `digit_counts` (1 if the answer contains `<eos>`) and an 11th weight to `DigitRM`. Collect the rankings at `temperature=1.5` so that some malformed answers appear in the data. Does the β = 0 run still collapse to `4444`, or does it find a different hack?
- **More or fewer labels.** Collect 100 or 2,000 prompts instead of 600, or use K = 2 instead of 4. How do the RM's held-out accuracy and the β = 0.1 run change?
- **No critic.** Replace the GAE advantages with `score - score.mean()` broadcast to every token (keep the KL part per token). That is most of the way to the critic-free methods of [Chapter 15](https://learn.xhelio.ai/rl/learn/15-rlvr-grpo/).

## What this lab simplifies

- The reward model is a 10-parameter linear model on digit counts, not a language model with a scalar head. Its loss and training loop are exactly the real ones.
- The critic is a separate TinyLM with a scalar value head. It is initialized from the SFT model; InstructGPT initializes it from the reward model.
- One response per prompt, 64 prompts per batch, 4 full-batch PPO epochs, and no EOS trick or length limit beyond TinyLM's 4-token budget.