# Lab 7: policy gradients and REINFORCE

You will build REINFORCE from the ground up. First, on a 3-armed bandit, you check by
hand and by simulation that the score-function estimator is unbiased and that a
baseline shrinks its variance. Then you train a neural policy on CartPole and
watch reward-to-go and a baseline turn a barely-learning algorithm into one
that balances the pole for the full 500 steps.

Everything you write is in **`pg.py`**. The tests are in `test_pg.py`;
`cartpole.py` is a plain-numpy CartPole (provided, nothing to do there).

## Setup

```bash
cd 07-policy-gradients
python -m pytest -q                  # all tests fail with NotImplementedError until you fill in the TODOs
RL_SOLUTIONS=1 python -m pytest -q   # the reference solution: all pass
```

Dependencies: `numpy`, `torch`, `matplotlib`, `pytest`. CPU only.

## Tasks (in order)

### Part A: the softmax bandit (numpy)

The chapter's running example: three arms with mean rewards `q = (1, 2, 4)`,
a softmax policy `pi = softmax(theta)`, and `J(theta) = sum_a pi_a q_a`.

1. **`softmax_score(theta, actions)`**: the score `grad log pi(a) = e_a - pi`,
   one row per action. Tests compare with finite differences and check that the
   score averages to zero under `pi`.
2. **`exact_gradient(theta, q)`**: `grad J = pi * (q - J)`. At the uniform policy
   this must give the chapter's `(-4/9, -1/9, 5/9)`.
3. **`reinforce_estimates(theta, actions, rewards, baseline)`**: the single-sample
   estimates `(R - b) * score`. Tests check that their mean matches `exact_gradient`
   for any baseline, and that `b = J` slashes the variance when all rewards are
   shifted by +10.

### Part B: REINFORCE on CartPole (PyTorch)

4. **`rewards_to_go(rewards, gamma)`**: `G_t = r_t + gamma * G_{t+1}` in one
   backward pass.
5. **`reinforce_loss(logits, actions, weights)`**: the surrogate loss
   `-mean(w_t * log pi(a_t | s_t))` whose gradient is the REINFORCE estimate. The
   weights must be treated as constants (`.detach()`).

`train_reinforce(variant, ...)` (provided) plays a batch of 8 episodes, turns the
rewards into per-step weights, and takes one Adam step on your loss. The three
variants are:

| variant | weight on `log pi(a_t \| s_t)` |
|---|---|
| `"total"` | the whole-episode return `G_0` (the original REINFORCE) |
| `"to_go"` | the reward-to-go `G_t` |
| `"baseline"` | `G_t - V(s_t)`, where a small value net `V` is regressed on the Monte Carlo returns `G_t` |

The baseline network is fit by plain regression to Monte Carlo returns; it never
bootstraps. Making it bootstrap (learning `V` from `r_t + gamma V(s_{t+1})`) is the step
to actor-critic, in Chapter 8.

Two honest simplifications, both standard: like nearly every implementation, the weights
drop the `gamma^t` factor that the exact discounted gradient would carry (the chapter
explains why). And CartPole's 500-step time limit is treated as a true ending.

## Run it

```bash
python run.py        # about a minute; saves reinforce.png
```

You should see two panels:

- **Left:** the variance of a single bandit estimate as a function of the constant
  baseline `b`, with rewards 11, 12, 14. It is a valley on a log scale, about 100 at
  `b = 0` and about 1 near `b = J = 12.33`.
- **Right:** CartPole learning curves, 4 seeds per variant. Typical final mean returns
  are roughly 40 (whole return), 300 (reward-to-go, with a wide spread across seeds),
  and close to 500 (baseline). Same algorithm, same learning rate, same number of
  samples: only the variance of the weights differs.

## Going further (optional)

- Add a function `optimal_constant_baseline(scores, rewards)` that returns
  `sum_i R_i ||s_i||^2 / sum_i ||s_i||^2` (Exercise 7.6 in the chapter derives it). On the
  bandit, check that its variance is at most that of `b = J`. At a non-uniform policy
  such as `theta = (2, 0, -1)` the two differ. Why are they equal at the uniform policy?
- In `train_reinforce`, normalise the weights (subtract their batch mean, divide by
  their std) for the `"to_go"` variant. How close does that get to the learned
  baseline, and why does this "batch-mean baseline" help?
- Raise the learning rate to 0.05 for `"baseline"`. Do some seeds collapse? Chapter 9
  is about exactly this fragility.