# Lab 15: GRPO from scratch on TinyLM

You will build GRPO piece by piece, check each piece against numbers from the chapter, and then use it to train a tiny transformer to add two-digit numbers better, with nothing but a verifier as the reward. This is the centerpiece lab of the course: everything from REINFORCE (Lab 7) and PPO (Lab 9) comes together here.

Chapter: <https://learn.xhelio.ai/rl/learn/15-rlvr-grpo/>

## Files

| File | What it is |
|---|---|
| `grpo.py` | **Your file.** Five tasks marked `TODO`. The training loop `train_grpo` and the helper `token_entropy` at the bottom are complete. |
| `tinylm.py` | Given: a ~100k-parameter causal transformer, the `"37+48="` task, its verifiable reward, sampling, per-token log-probs, SFT and evaluation. Read its docstring first. |
| `test_grpo.py` | Tests, grouped by task. |
| `run.py` | Trains GRPO from the SFT checkpoint and saves `grpo_tinylm.png` (about a minute). Takes flags for experiments. |
| `passk.py` | Compares pass@k of the SFT and GRPO models (run after `run.py`). |
| `colab_grpo_trl.md` | Optional, needs a GPU: the same algorithm on a real 0.5B model with Hugging Face TRL. Not run by the tests. |
| `solutions/grpo.py` | Reference solution. Try not to peek. |

The first time anything needs the starting policy, `tinylm.load_or_train_sft` trains it by supervised fine-tuning (about 5 seconds, seeded) and caches it as `tinylm_sft.pt`. It answers about 51% of problems correctly when sampling at temperature 1, and about 76% greedily.

## Tasks

Run the tests for one task at a time with `-k`:

```bash
python -m pytest -q -k task1           # just Task 1
python -m pytest -q                    # everything (about 25 s once all tasks work)
RL_SOLUTIONS=1 python -m pytest -q     # the reference solution: all pass
```

1. **`sample_group`** (Task 1). Repeat each of the P prompts G times, consecutively, sample all P·G responses in one batch, decode them, and score them with the verifiable reward `tinylm.reward`. Everything downstream relies on rows `i*G … i*G+G-1` belonging to prompt `i`.
2. **`group_advantages`** (Task 2). $\hat A^{(i)} = (r^{(i)} - \text{mean}) / (\text{std} + \varepsilon)$ per group, or mean-only. The tests use the chapter's worked examples: rewards (1, 0, 0, 1) must give (+1, −1, −1, +1), and all-equal groups must give exactly 0.
3. **`k3_kl`** (Task 3). The per-token estimator $k_3 = u - \log u - 1$ with $u = \pi_\text{ref}/\pi_\theta$. One test samples 200,000 tokens and checks that it is unbiased for KL.
4. **`grpo_loss`** (Task 4). The DeepSeekMath objective: per-token clipped surrogate minus $\beta k_3$, masked, averaged over each response's own tokens, then over responses, negated. The tests check the loss *and its gradient* on a hand-computed batch, so you will know exactly which case you got wrong. One test checks that a zero-signal group gives exactly zero gradient through a real model.
5. **`grpo_step`** (Task 5). One iteration: sample groups, advantages, frozen `old_logp` and `ref_logp`, then `inner_steps` gradient steps. The last test runs 100 iterations (about 15 s) from the SFT checkpoint and expects sampled accuracy to rise by at least 6 points (the reference gains 12 to 18).

## Then run the experiment

```bash
python run.py        # about a minute; saves grpo_tinylm.png and tinylm_grpo.pt
python passk.py      # about 15 s; saves passk.png
```

**What you should see** with the defaults (G = 8, 16 prompts per iteration, 2 inner steps, lr 3e-4, ε = 0.2, β = 0.04, 300 iterations). Your numbers will differ a little:

- **Accuracy:** sampled accuracy on 500 held-out problems rises from about 0.51 to about 0.8; greedy from about 0.76 to about 0.9.
- **Zero-signal groups:** the fraction of groups whose 8 answers are all right (or all wrong) climbs from about 10% to about 40%. Those samples were generated and then contributed nothing.
- **KL to the reference:** climbs to roughly 0.1–0.2 nats per token, with spikes. $k_3$ on single samples is heavy-tailed; the smoothed curve is the one to read.
- **Entropy:** falls steadily from about 0.34 to about 0.14 nats per token. The policy is sharpening onto its preferred answers.
- **pass@k** (`passk.py`): GRPO wins clearly at k = 1 and 2, the curves meet around k = 16, and at k = 64 the SFT model is slightly ahead. That is the "RL sharpens rather than expands" finding from the chapter, in miniature. The large-k gap is small (a few problems out of hundreds), so treat it as suggestive.

## Things to try

`run.py` takes flags; use `--out name.png` to keep plots side by side.

| Command | Question |
|---|---|
| `python run.py --inner-steps 8 --out clip.png` vs `python run.py --inner-steps 8 --clip 1000 --out noclip.png` | What does the clip buy you when you reuse a batch 8 times? Watch KL and greedy accuracy. |
| `python run.py --G 2 --prompts 64 --out g2.png` | Same samples per iteration, tiny groups. What fraction of groups carries signal? Does it still learn? |
| `python run.py --no-std --out meanonly.png` | Mean-centred advantages only (Dr. GRPO style). Any difference at this scale? |
| `python run.py --beta 0 --out nokl.png` and `--beta 0.5` | How far does the policy drift with no leash? With a strong one? |
| `python run.py --temperature 0.5 --out t05.png` | Colder sampling: fewer mistakes per group, but what happens to the zero-signal fraction? |

## Common bugs

- Computing `old_logp` *after* the first optimizer step. Then the ratio is not 1 on the first inner step, and `test_task5_one_inner_step_never_clips` fails.
- Averaging over all tokens of the batch at once, instead of per response first. That is a different loss (the "token-level" aggregation of [Chapter 16](https://learn.xhelio.ai/rl/learn/16-policy-gradient-zoo/)), and the hand-computed gradient test catches it.
- Forgetting the mask, so padding positions leak into the averages.
- Letting gradients flow into `old_logp` or `ref_logp`. Compute both under `torch.no_grad()`.
- Using the sample std (`unbiased=True`). GRPO implementations use the population std over the group; the test for rewards (1, 0, 0, 0) expects +1.73, not +1.5.

Before you zip or share the folder you can delete the generated `*.png` and `*.pt` files; they are recreated on demand.