# Lab 9: PPO from the loss up

You will build PPO-clip piece by piece, check every piece against numbers you can compute by hand, then train it on CartPole and see what the clip buys you.

Chapter: <https://learn.xhelio.ai/rl/learn/09-ppo/>

## Files

| File | What it is |
|---|---|
| `ppo.py` | **Your file.** Four tasks marked `TODO`. The training loop `train_ppo` at the bottom is complete. |
| `common.py` | Given: the actor-critic network, `collect_rollout`, and `compute_gae` (from Lab 8). Read `collect_rollout` once. |
| `cartpole.py` | Given: CartPole in plain numpy (same as Labs 6–8). |
| `test_ppo.py` | Tests. |
| `run.py` | Trains four PPO variants and saves `ppo_cartpole.png`. |
| `solutions/ppo.py` | Reference solution. Try not to peek. |

## Tasks

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

1. **`clipped_surrogate_loss`** (Task 1). The PPO-clip objective $\min(\rho_t \hat A_t, \operatorname{clip}(\rho_t, 1-\epsilon, 1+\epsilon)\hat A_t)$, averaged and negated. The tests check the six cases from the chapter's "four cases" table, value *and* gradient, so you will know exactly which case you got wrong. One test also checks that at $\rho = 1$ your gradient equals the REINFORCE gradient.
2. **`approx_kl` and `clip_fraction`** (Task 2). The two numbers you watch while PPO trains. One test samples 200,000 actions from a 3-action policy and checks your KL estimate against the exact KL.
3. **`value_loss`** (Task 3). Half the mean squared error, with optional value clipping.
4. **`ppo_update`** (Task 4). K epochs over the batch in shuffled minibatches: advantage normalization, the three loss terms, gradient-norm clipping, and approx-KL early stopping. The docstring lists the steps in order and the PyTorch calls you need. The last test trains PPO for 15k steps (about 6 seconds) and expects an average return of at least 150.

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

## Then run the experiment

```bash
python run.py            # about 2 minutes; saves ppo_cartpole.png
```

It trains four variants on the same seed for 40k steps:

| Variant | What changes |
|---|---|
| PPO (K=10) | the default `PPOConfig`: ε = 0.2, 10 epochs, lr 1e-3 |
| PPO, aggressive | 40 epochs and lr 3e-3: squeeze each batch hard |
| no clip, aggressive | the same, with ε = 10⁹ (the clip never activates: this is the plain surrogate) |
| one epoch (K=1) | each batch used for a single pass |

**What you should see** (with the reference solution; your numbers can differ a little):

- **PPO (K=10)** reaches the 500-step ceiling in roughly 20–30k steps. Its approx KL per update stays around 0.001–0.01.
- **PPO, aggressive** learns about as fast. The clip fraction is higher (often 0.1–0.3) because many samples hit the edge of the band, but the KL stays small (around 0.01): the clip is doing its job. It is a little less steady late in training (in the reference run it dips from 500 to about 390 after 30k steps).
- **No clip, aggressive** is the cautionary tale. Its KL per update is 100 to 1000 times larger, its policy entropy collapses almost to zero after the very first update (a nearly deterministic policy), and its return is erratic and can crash (in the reference run it falls from about 390 to near 20 around 20k steps). Its clip fraction reads 0 only because ε is huge.
- **One epoch** is stable but slow: it throws each batch away after one pass, so it needs many more environment steps.

## Things to try

- Set `target_kl=0.01` on the aggressive variants. Print `hist["epochs_run"]`: how many of the 40 epochs actually run?
- Turn off advantage normalization (`norm_adv=False`). With CartPole's +1-per-step rewards the raw advantages are large; what happens to the clip fraction?
- Try ε = 0.05 and ε = 0.5 with K = 10.
- Add an entropy bonus (`ent_coef=0.01`) to the "no clip, aggressive" variant. Does it stop the entropy collapse?