Lesson 5 — UCB1: Directed Exploration
Lesson 4 explored with ε-greedy, which explores blindly: a fixed fraction of the time it picks a random other action — even ones it has tried thousands of times and knows are bad. Exploration is wasted, and it doesn't explore more where it's more uncertain. UCB1 fixes exactly that. It's also the last toy before search, because UCB1 inside a tree is MCTS.
The setup: a k-armed bandit
k slot machines, each with an unknown average payout. Every pull you choose one; you want maximum total reward over many pulls. The tension is the whole subject: pull the arm that looks best so far (exploit), or one you haven't tried much that might be better (explore)?
The idea: optimism under uncertainty
Don't pick the arm with the best estimate; pick the arm with the best plausible value — its average plus a bonus for how uncertain you are about it:
score(a) = mean(a) + c·√( ln N / n(a) )
where n(a) is how many times you've pulled arm a, N is total pulls so far, and c tunes how much you value exploration (≈ √2). Pick the highest score, pull, update that arm's mean and counts, repeat.
Why the bonus is shaped that way — the whole insight
- n(a) in the denominator. A barely-pulled arm has a big bonus, so you're eager to try it; each pull shrinks the bonus, so once an arm is well known only its true mean matters. Exploration self-targets toward whatever is under-explored.
- ln N in the numerator. As total time grows, every arm's bonus slowly creeps up, so even a long-ignored arm eventually gets revisited — but it's a log, so this hunger grows very slowly. Exploration never fully stops; it tapers.
- The √ together is the width of a statistical confidence interval. You score each arm by the upper end of its plausible range — hence Upper Confidence Bound. Judge each arm by its best believable case and the truth sorts itself out.
The payoff over ε-greedy: exploration is directed, not random. UCB1 pours pulls into arms that are both promising and uncertain, and quietly abandons arms it's confident are bad — no blind ε tax.
Why we care
That count-based bonus, applied at each node of a lookahead tree, is UCT —
the selection rule at the heart of Monte-Carlo Tree Search, and the bridge to AlphaZero.
The exercise: the code below already tracks actions_chosen[(s, a)] — exactly
the n(a) UCB1 needs. Replace the ε-greedy block in
get_action with a UCB1 score over those counts.
"""Lesson 5 -- UCB1: directed exploration.
UCB1 is a rule for the EXPLORE / EXPLOIT problem. Setup: k slot machines
(arms), each with an unknown average payout. Every pull you must choose one,
and you want max total reward over many pulls. The tension: pull the arm that
LOOKS best so far (exploit), or one you haven't tried much that MIGHT be
better (explore)?
WHY NOT EPSILON-GREEDY (Lesson 4)?
It explores BLINDLY. Some fixed fraction of the time it picks a random
other arm -- even arms it has already pulled 1000 times and knows are bad.
Exploration is wasted, and it doesn't explore MORE when it's MORE uncertain.
UCB1'S IDEA -- "optimism under uncertainty":
Don't pick the arm with the best ESTIMATE; pick the arm with the best
PLAUSIBLE value. For each arm, add a bonus for how uncertain you are:
score(a) = mean_reward(a) + c * sqrt( ln(N) / n(a) )
|__ exploit __| |____ explore bonus ____|
n(a) = times you've pulled arm a
N = total pulls so far (all arms)
c = exploration weight (~sqrt(2))
Pick the arm with the highest score, pull it, update its mean and counts,
repeat.
WHY THE BONUS IS SHAPED THAT WAY -- the whole insight:
- n(a) in the DENOMINATOR: a barely-pulled arm has a big bonus -> you're
eager to try it. Each pull shrinks its bonus -> once an arm is well known,
the bonus fades and only its true mean matters. Exploration self-targets
toward whatever is under-explored.
- ln(N) in the NUMERATOR: as total time grows, every arm's bonus slowly
creeps up, so even a long-ignored arm eventually gets revisited -- but
it's a LOG, so this hunger grows very slowly. Exploration never fully
stops; it tapers.
- the sqrt(...) together is the width of a statistical confidence interval.
You score each arm by the UPPER end of its plausible range -- hence
"Upper Confidence Bound." Judge each arm by its best believable case and
the truth sorts itself out.
PAYOFF over epsilon-greedy: exploration is DIRECTED, not random. UCB1 pours
pulls into arms that are both promising AND uncertain, and quietly abandons
arms it's confident are bad -- no blind epsilon tax.
WHY WE CARE: that count-based bonus, applied at each node of a lookahead tree,
is UCT -- the selection rule at the heart of MCTS, and the bridge to AlphaZero.
The body below already tracks actions_chosen[(s, a)]; that count is exactly the
n(a) UCB1 needs. The exercise: replace the epsilon-greedy block in get_action
with a UCB1 score, using those counts.
"""
import random
import statistics
from collections import defaultdict, deque
from enum import Enum, auto
import numpy as np
GAMMA = 0.9
ALPHA = 0.1
JUMP_SUCCESS_PROBABILITY = 0.60
class Action(Enum):
THRUST = 0
JUMP = 1
# --- The game, behind a wall ------------------------------------------------
# MODEL-FREE WALL: the learning loop below may call step() and NOTHING else.
# It never reads JUMP_SUCCESS_PROBABILITY and never consults a transition
# table. step() is the only window onto the dynamics, and each call returns a
# single *sampled* (s_next, reward, done) -- exactly what a real environment
# hands you. That wall is the definition of "model-free."
def step(s, a): # -> (s_next, reward, done)
if s == 10:
raise ValueError(
"state 10 is terminal; no action should be applied to it"
)
if a == Action.THRUST:
if s == 9:
return 10, 1.0, True
return s + 1, 0.0, False
# JUMP
if random.random() <= JUMP_SUCCESS_PROBABILITY:
if s >= 7: # +3 lands on or past 10 -> win
return 10, 1.0, True
return s + 3, 0.0, False
return max(0, s - 2), 0.0, False # misfire: slide back 2, floor at 0
# Ground truth from Lesson 1 value iteration -- used ONLY for scoring. The
# learner never looks at it; that would be peeking through the wall.
TRUE_V = [
0.5219919679553231,
0.5482046378586254,
0.5706240497071444,
0.6186571735660003,
0.6671991520804588,
0.7087165778019958,
0.7801916910708641,
0.8551379656365434,
0.9,
1.0,
0.0,
]
EPISODES = 50_000
WINDOW = 2_000 # rolling window for the steady-state bias/noise estimate
EPSILON = 0.10 # exploration probability
def neural_prediction() -> tuple[np.array, list[int], list[deque]]:
"""
The full write-up and stub land when this session runs. The shape of the assignment:
A tiny network from scratch. One hidden layer, implemented as plain matrix multiplies — forward pass, a squared-error loss against the TD target, and the backward pass computed by hand.
Gradient-check it. Compare your analytic gradients to numerical finite differences, so you know the backprop is right before trusting it.
Re-learn Asteroid Hop's V with the network instead of the 11-entry table, and confirm it reproduces the Lesson 1 values — same toy, known ground truth, now via function approximation.
Predictions first, as always — including what the deadly triad (bootstrapping + function approximation + off-policy data) might do to convergence.
"""
# θ ← θ + α·[ r + γ·max(Vθ(s′, a) for all a in s') − (Vθ(s, a)]·∇θ Vθ(s, a)
# intialize values to between 0 and 1
θ = np.linspace(-1, 1, 8) # correspends to a_0, a_1, a_2, a_3
# W1 = np.random.randn(3, 4) # (in, out)
# W2 = np.random.rand(4, 2)
def get_features(s, a):
s_n = s / 10
# a.value == 0 for THRUST, 1 for JUMP
return np.array(
[
1,
s_n,
s_n**2,
s_n**3,
a.value,
a.value * s_n,
a.value * s_n**2,
a.value * s_n**3,
]
)
def get_action(s, θ):
q_max = -999999
for a in [Action.THRUST, Action.JUMP]:
q = get_Q(s, a, θ)
if q > q_max:
best_action = a
q_max = q
if random.random() < EPSILON:
# explore: pick the other action
if best_action == Action.THRUST:
return Action.JUMP
else:
return Action.THRUST
else:
return best_action
def get_best_q(s, θ):
return max(get_Q(s, Action.THRUST, θ), get_Q(s, Action.JUMP, θ))
def get_Q(s, a, θ):
features = get_features(s, a)
return features @ θ
def get_gradient(θ, s, a) -> np.array:
# normalize s so that steps don't explode
return get_features(s, a)
visits = [0] * 11
actions_chosen = defaultdict(int)
def get_deque():
return deque(maxlen=WINDOW)
recent = defaultdict(get_deque)
for _ in range(EPISODES):
s = 0
while True:
a = get_action(s, θ)
visits[s] += 1
actions_chosen[(s, a)] += 1
s_next, reward, done = step(s, a)
# ∇θ Vθ(s) == grad
# ∇θ Vθ(s) means "gradient with respect to theta"
grad = get_gradient(θ, s, a)
# θ ← θ + α·[ r + γ·Vθ(s′) − Vθ(s) ]·∇θ Vθ(s)
θ = (
θ
+ ALPHA
* (
reward
+ GAMMA * (0 if done else get_best_q(s_next, θ))
- get_Q(s, a, θ)
)
* grad
)
recent[s].append(get_best_q(s, θ) - TRUE_V[s])
s = s_next
if done:
break
V = [get_best_q(s, θ) for s in range(11)]
return V, visits, recent
def report(V, visits, recent):
"""Print learned V beside ground truth, with per-state bias/noise and the
worst-fit state. `recent[s]` holds the last WINDOW signed errors at s:
its mean is the steady-state bias, its stdev the noise floor.
"""
print(
f"{'s':>2} {'visits':>8} {'V_learned':>10} {'V_true':>8} "
f"{'bias':>8} {'noise':>8}"
)
worst_s, worst_err = None, -1.0
for s in range(10):
bias = statistics.mean(recent[s]) if recent[s] else float("nan")
noise = statistics.stdev(recent[s]) if len(recent[s]) > 1 else 0.0
err = abs(V[s] - TRUE_V[s])
if err > worst_err:
worst_s, worst_err = s, err
print(
f"{s:>2} {visits[s]:>8} {V[s]:>10.3f} {TRUE_V[s]:>8.3f} "
f"{bias:>+8.3f} {noise:>8.3f}"
)
print(
f"\nworst-fit state: {worst_s} (|error| = {worst_err:.3f}, "
f"{visits[worst_s]} visits)"
)
if __name__ == "__main__":
random.seed(0)
V, visits, recent = neural_prediction()
report(V, visits, recent)
# ============================================================================
# BONUS (Lesson 3.2) — turn this into a real two-layer neural network
# ============================================================================
#
# PROMPT:
# The version above is *linear* function approximation — V_θ(s) = θ·φ(s),
# which is just one layer (and the hand-built cubic features do all the
# work). Make it a genuine two-layer net: one hidden layer with a
# nonlinearity, so the network learns its OWN features instead of you
# picking s, s², s³ by hand. Same Asteroid Hop, same TD target, same δ —
# only ∇θ V_θ(s) changes (now it comes from backprop). Gradient-check the
# backward pass against finite differences before trusting it.
#
# CHEAT SHEET — neural networks in NumPy (scalar-output value net):
#
# # --- params (shapes) ---
# W1 = np.random.randn(n_in, n_hid) * 0.1 # b1, W2 likewise
# b1 = np.zeros(n_hid) # bias = per-neuron additive offset
# W2 = np.random.randn(n_hid) * 0.1 # output is scalar -> 1-D
# b2 = 0.0
#
# # --- forward ---
# z = x @ W1 + b1 # bias just adds; lets each neuron shift its threshold
# h = np.tanh(z) # nonlinearity (the whole reason for a hidden layer)
# V = h @ W2 + b2 # scalar
#
# # --- backward: grads of V w.r.t. each param (chain rule) ---
# dV_db2 = 1.0
# dV_dW2 = h # V = h·W2 -> ∂/∂W2 = h
# dV_dz = W2 * (1 - h**2) # through tanh: d tanh = 1 - tanh²
# dV_db1 = dV_dz # bias grad = the incoming delta (∂(·+b)/∂b = 1)
# dV_dW1 = np.outer(x, dV_dz) # ∂(x·W1)/∂W1 = outer(input, delta)
#
# # --- semi-gradient TD update: θ ← θ + α·δ·∇θV ---
# W2 += ALPHA * δ * dV_dW2; b2 += ALPHA * δ * dV_db2
# W1 += ALPHA * δ * dV_dW1; b1 += ALPHA * δ * dV_db1
#
# Three things to anchor:
# - BIAS is a learnable constant added after the matmul. Its gradient is
# always just the delta flowing into that layer (∂(stuff + b)/∂b = 1), so
# dV_db = dV_dz.
# - BACKPROP is the chain rule applied right-to-left: start with ∂V/∂V = 1,
# push it back through each op, multiplying by that op's local derivative
# (W2 for the matmul, 1−h² for tanh).
# - Same δ as before — backprop only changes the ∇θV part; the TD error out
# front is unchanged.
#
# Web write-up: web/templates/lesson_3_bonus.html (/lessons/3/bonus)