Lab 03 — Tabular Reinforcement Learning: Q-Learning + Bandits

Phase 32 · Lab 03 · Phase README · Warmup

The problem

Supervised learning gets told the right answer; reinforcement learning only gets a reward signal — and often a delayed one. The agent must learn which actions were responsible for a reward that arrives many steps later (the credit-assignment problem) while balancing acting on what it knows against trying what it doesn't (the exploration-exploitation dilemma). Those two problems, and the Bellman equation that solves the first, are the entire conceptual core of RL — everything from AlphaGo to RLHF is scaffolding around them.

You build the two smallest systems that exhibit both problems honestly. A GridWorld MDP — states, four actions, deterministic transitions, a per-step cost and a terminal goal reward — trained with tabular Q-learning: the ε-greedy behavior policy explores (seeded RNG, reproducible), and every observed transition applies the Bellman update Q(s,a) ← Q(s,a) + α·[r + γ·max Q(s′,·) − Q(s,a)] until the table satisfies the Bellman optimality equation and the greedy policy walks the shortest path. And a multi-armed bandit — RL with one state, where only exploration matters — solved twice: ε-greedy (explore by coin flip, forever) and UCB1 (explore by confidence bound, provably less waste, zero randomness).

What you build

File map

Run it

python3 -m pytest test_lab.py -q                     # your lab.py (red until you implement)
LAB_MODULE=solution python3 -m pytest test_lab.py -q # the reference (green)
python3 solution.py                                   # solve the grid + run both bandits

Success criteria

• GridWorld.step moves correctly, clamps off-grid moves in place (still charging the step), pays the goal reward with done=True, and rejects unknown actions / terminal-state steps.
• update from an all-zero table equals α·r; a terminal transition ignores future value; a non-terminal one bootstraps from max next-Q; the TD error is returned.
• ε=0 is pure greedy with a deterministic tie-break; ε=1 visits all four actions.
• Trained on the 4×4 grid, the greedy path reaches the goal in exactly 6 moves (the Manhattan distance — no wasted steps).
• With ε=1.0 (pure random behavior!) the Bellman residual is driven to ~0 and the extracted greedy policy is still optimal — the off-policy property, demonstrated.
• State values increase along the path to the goal, and γ=0.95 values the start state higher than γ=0.5 does.
• ε-greedy bandit finds the best arm (most pulls, accurate estimate); UCB1 pulls every arm once, then concentrates >50% of pulls on the best arm; the confidence bonus forces a pull of an undersampled arm even when its mean estimate is lower.
• Everything is bit-for-bit reproducible from a seed (UCB needs no seed at all).
• All 28 tests pass under both lab and solution.

How this maps to the real stack

• GridWorld is a Gymnasium environment in miniature: env.step(action) returning (observation, reward, terminated, ...) is the exact API contract of gymnasium.Env — ours takes the state explicitly instead of holding it, which makes tests stateless and is the only difference in shape.
• Q-learning is the algorithm under DQN (Mnih et al. 2015): replace the table with a neural network Q(s,a;θ), add a replay buffer and a target network for stability, and you have the agent that played Atari. The Bellman target r + γ·max Q(s′,·) is literally the same line of code; everything else in DQN exists because a function approximator, unlike our table, can diverge.
• The off-policy property you prove in the tests (learning optimal values from ε=1 random behavior) is why replay buffers work at all: DQN trains on transitions collected by old policies. SARSA — one symbol different, using the actually-taken next action — is on-policy and can't do that; the Warmup covers the contrast.
• ε-greedy and UCB run production systems today: recommender and ad systems allocate traffic with bandits, and UCB's "optimism in the face of uncertainty" is the core of Monte-Carlo Tree Search (AlphaGo's UCT is UCB applied to tree nodes). A/B testing with adaptive allocation is a bandit.
• Policy-gradient methods (REINFORCE → PPO) are the other family — optimize the policy directly instead of a value table — and PPO is what RLHF uses to fine-tune LLMs against a reward model. The Warmup surveys them; the value-based mechanics you built here are the contrast that makes them make sense.

Limits of the miniature. Real environments are stochastic (transition probabilities) — Q-learning's update is unchanged but converges in expectation, needing a decaying α. Real state spaces don't fit in a table — that's the jump to function approximation (DQN), with its own instabilities. Our bandit rewards are Gaussian with known-ish noise; real bandits face non-stationary arms (yesterday's best ad isn't today's) and use sliding windows or discounted estimates. And γ, α, ε here are hand-picked; at scale they're tuned like any hyperparameter (this phase's Warmup §15).

Extensions (your own machine)

• Implement SARSA (use the actually-chosen next action in the target) and compare the learned path near a "cliff" row of large negative rewards — the classic cliff-walking experiment where SARSA learns the safe path and Q-learning the risky-optimal one.
• Add stochastic transitions (10% chance of a random slip) and watch why a decaying α becomes necessary.
• Add ε decay (start 1.0, anneal to 0.05) and compare episodes-to-optimal against fixed ε.
• Implement value iteration directly on the known MDP and confirm Q-learning's table matches it — model-based vs model-free, same fixed point.
• Add a Thompson sampling bandit (Beta posteriors over Bernoulli arms) and compare regret curves against ε-greedy and UCB1.

Interview / resume signal

"Built tabular Q-learning and multi-armed bandits from scratch in pure Python: a GridWorld MDP, the Bellman update with ε-greedy exploration, convergence verified against the Bellman optimality equation (residual → 0) and the known optimal path — including a demonstration of the off-policy property (optimal values learned from a fully random behavior policy) — plus ε-greedy and UCB1 bandits with incremental-mean estimates. I can explain what DQN adds to Q-learning and why, and where bandits ship in production (adaptive experiments, recommenders, MCTS)."