# Lab 2: Bandits

You will write the four pieces of a bandit agent toolkit from [Chapter 2](https://learn.xhelio.ai/rl/learn/02-bandits/) and race them on the 10-armed testbed: greedy, ε-greedy, optimistic greedy, UCB1 and Thompson sampling. numpy only.

## Files

| File | What it is |
|---|---|
| `bandits.py` | **Your file.** Four tasks, each marked `# TODO`. |
| `envs.py` | Bandit environments and the experiment runner (complete; read it once). |
| `test_bandits.py` | Tests, grouped by task. |
| `run.py` | Races the agents and saves a plot. |
| `solutions/bandits.py` | Reference solution. Peek only when stuck. |

## Setup and running

```bash
cd 02-bandits
python -m pytest -q                    # all tests (they fail until you fill in the TODOs)
python -m pytest -q -k task1           # only Task 1 (also: task2, task3, task4)
python run.py                          # race on the 10-armed testbed -> bandits.png (~20 s)
python run.py --switch                 # nonstationary version -> bandits_switch.png (~20 s)

RL_SOLUTIONS=1 python -m pytest -q     # the same tests against the reference solution
RL_SOLUTIONS=1 python run.py           # the plots you should get
```

## Tasks (in order)

1. **`argmax_tiebreak(values, rng)`**: argmax with random tie-breaking. Two lines. It matters more than it looks: at the start all estimates are 0, and `np.argmax` would always pick arm 0.
2. **`EpsilonGreedy`**: `select_action` (explore with probability ε, otherwise greedy) and `update` (the incremental update $Q \leftarrow Q + \text{step}\,(r - Q)$, with step $1/N$ or a constant α). The one class covers greedy (ε = 0), ε-greedy, and optimistic greedy (large `initial_value`, constant `step_size`).
3. **`ucb_scores` and `UCB1`**: the score $Q(a) + c\sqrt{\ln t / N(a)}$, with $+\infty$ for untried arms, and an agent that pulls the arm with the highest score. The test uses the chapter's worked example (scores 1.29, 1.50, 1.73).
4. **`ThompsonBernoulli`**: keep a Beta(α, β) belief per arm, sample one value from each, pull the argmax; update α on a win and β on a loss. One test checks Exercise 2.4's answer: with Beta(2, 1) vs Beta(1, 1) you should pick the first arm 2/3 of the time.

The last two tests check that everything learns: on 20 random 10-armed Bernoulli problems, ε-greedy, UCB1 and Thompson must beat random play by a wide margin, Thompson must beat ε-greedy, and greedy must lock onto the worse arm of a (0.9, 0.5) bandit about 36% of the time (Exercise 2.2).

## What you should see

`python run.py` (200 runs, 1000 steps; numbers from the reference solution):

- **Greedy** flatlines: about 16% optimal action and regret growing in a straight line (≈ 250 by step 1000). In most runs the first arm that pays anything wins forever.
- **ε-greedy (ε = 0.1)** learns fast but keeps paying for exploration: regret ≈ 70, still rising in a straight line.
- **ε-greedy (ε = 0.01)** explores so rarely that it is slow to find the best arm (regret ≈ 135).
- **Optimistic greedy** (start every estimate at 1, step 0.1) does surprisingly well here (regret ≈ 40).
- **UCB1 with the textbook c = √2** is the slowest starter of the good methods (regret ≈ 130 at step 1000). Its bonus is conservative: it is built to be safe for any rewards in [0, 1]. Its regret does grow only logarithmically in the long run, but at this horizon the constant dominates: even at 5,000 steps it is still slightly behind ε = 0.1 (about 274 vs 257 over 100 runs), and it overtakes somewhere between 5,000 and 7,000 steps (`python run.py --steps 10000`, about 2 minutes). Or edit `agents_stationary()` in `run.py` to use `B.UCB1(k, rng, c=0.5)`: regret ≈ 34 at step 1000, close to Thompson.
- **Thompson sampling** wins clearly: regret ≈ 27 and still flattening.

`python run.py --switch` (arms shuffled at step 1000): everyone crashes at the switch. The sample-average ε-greedy agent recovers slowly, because 1000 old rewards outvote every new one; the constant-α version forgets in about 1/α = 10 pulls per arm and recovers fast. UCB1 recovers because its bonus keeps growing for neglected arms. Thompson sampling recovers too, but only after the old best arm's sharp posterior has been worn down.

## Stretch goals

- Add a `ThompsonGaussian` agent for the Gaussian testbed (`envs.gaussian_testbed`): with prior $\mathcal N(0, 1)$ and unit reward noise, the posterior of arm $a$ after $n$ pulls with reward sum $S$ is $\mathcal N\big(S/(n+1),\, 1/(n+1)\big)$.
- Plot regret against $\ln t$ instead of $t$. Logarithmic-regret algorithms become straight lines.
- Implement ε-greedy with a decaying schedule $\varepsilon_t = \min(1, c/t)$ and compare.