# Lab 8: actor-critic and GAE

You will write the four pieces that turn REINFORCE into an actor-critic method, then
train A2C on CartPole and see how the GAE parameter λ changes learning.

Chapter: [Actor-critic and advantage estimation](https://learn.xhelio.ai/rl/learn/08-actor-critic/).

## Files

| File | What |
|---|---|
| `actor_critic.py` | **Your file.** Four tasks marked `TODO`, then complete A2C scaffolding (network, vectorized CartPole, rollout collection, training loop). |
| `test_actor_critic.py` | Tests. Exact values on small hand-computable cases, plus one "does it learn" run. |
| `run.py` | Trains A2C for λ ∈ {0, 0.5, 0.95, 1} × 3 seeds and saves `a2c_lambda.png` (about 1–2 min). |
| `cartpole.py` | CartPole in numpy (copied from the course's shared helpers). |
| `solutions/actor_critic.py` | Reference solution. Try first. |

Requirements: Python 3.10+, `numpy`, `torch`, `matplotlib`, `pytest`.

## Tasks (in order)

1. **`td_errors(rewards, values, last_value, dones, gamma)`**
   The one-step TD error $\delta_t = r_t + \gamma V(s_{t+1})(1 - d_t) - V(s_t)$.
   The only subtlety is the data layout: a rollout is a block of `T` consecutive steps that can
   cross episode boundaries. `values[t + 1]` is the value of the next state *unless* the
   episode ended at step `t` (`dones[t] = 1`), in which case the next row belongs to a
   brand-new episode. The value of the state after the last row is passed in as `last_value`.
   Make it work for shape `(T,)` and `(T, N)` (N parallel environments): if you only
   index along the first axis, it will.

2. **`compute_gae(rewards, values, last_value, dones, gamma, lam)`**
   $\hat A_t = \delta_t + \gamma\lambda(1 - d_t)\hat A_{t+1}$, computed backward from the end.
   Return the advantages and the λ-returns `advantages + values` (the critic's targets).
   The first test is the chapter's worked example: rewards (0, 0, 0, 1), values
   (0.5, 0.4, 0.8, 0.6), γ = 1, λ = 0.5 gives (0.1, 0.4, 0.0, 0.4).

3. **`normalize_advantages(adv, eps=1e-8)`**
   Zero mean and unit (population) standard deviation over the batch.

4. **`actor_critic_loss(logits, values, actions, advantages, returns, value_coef, entropy_coef)`**
   policy loss $-\text{mean}[\log\pi(a\mid s)\,\hat A]$, value loss $\text{mean}[(V - R)^2]$,
   entropy bonus, combined as `policy + value_coef * value - entropy_coef * entropy`.
   Advantages and returns are targets: detach them.

Run the tests for one task at a time with `-k`, e.g. `python -m pytest -q -k td_errors`,
`-k gae`, `-k normalize`, `-k loss`. The last test (`-k learns`) trains A2C for 80k steps
(a few seconds) and needs all four tasks.

## Run

```bash
python -m pytest -q                    # your code
RL_SOLUTIONS=1 python -m pytest -q     # reference solution: all pass
python run.py                          # train and plot (uses your code)
```

## What you should see

`run.py` prints a one-line summary and saves `a2c_lambda.png`. With the reference
solution (3 seeds, 200k steps each), roughly:

- **λ = 0.95 and λ = 1** climb to an average return of about 200–300 within 50k steps and
  reach 300–400 by the end (500 is the maximum). The curves are jumpy: A2C takes one
  gradient step per batch, and a single bad step can drop the return for a while
  (Chapter 9's PPO is the fix).
- **λ = 0** (pure one-step TD advantages) is far slower: it stays below 100 for the first
  100k steps or more. It even starts *below* the random policy (about 10 instead of about 22):
  run `train_a2c(total_steps=30_000, lam=0.0, verbose=True)` and watch the entropy collapse
  from 0.69 to about 0.1 as the actor obeys a critic that has not learned anything yet.
  **λ = 0.5** is in between.

Why does high λ win here, when the demo in the chapter showed that a good critic makes
small λ best? On CartPole the critic **starts** badly wrong: true values are in the tens
(γ = 0.99, +1 per step) while a fresh network outputs about 0. One-step advantages lean
entirely on that wrong critic (bias), while λ near 1 uses the real rewards. CartPole
rewards are not noisy, so the extra variance of λ = 1 costs little. Exercise 8.8 in the
chapter asks you to explain this, including the early collapse at λ = 0.

## Things to try

- `train_a2c(..., normalize=False)`: does advantage normalization matter here? (Exercise 8.6)
- `entropy_coef=0.0` vs `0.05`: watch the `entropy` printed with `verbose=True`.
- `shared=True`: one network trunk with two heads. Try `value_coef` of 0.5 vs 0.01 and
  think about why the value loss's scale now matters for the policy.
- In `collect_rollout`, delete the time-limit bootstrap (`r = r + gamma * v_final * trunc`)
  and compare. Episodes rarely reach 500 steps early in training, so look at the end of
  training (Exercise 8.9).