MCTS in Latent¶
Companion to research/09-mcts-in-latent.md.
Monte Carlo Tree Search builds a look-ahead tree through four phases — selection, expansion, evaluation, backup — using UCT/PUCT to balance exploitation and exploration. MuZero runs this entire search inside a learned latent model. This notebook builds MCTS from its bandit core up to the deep look-ahead that flat samplers cannot match.
What we will see:
- UCB selection — the exploration/exploitation rule at the heart of MCTS finds the best arm of a bandit while wasting few pulls, beating both greedy and random.
- MCTS finds the best action — selection + expansion + rollout + backup makes the root visit counts concentrate on the optimal action as simulations grow.
- PUCT priors make search efficient — a good policy prior reaches the optimal action with far fewer simulations; a bad prior slows it down. This is the AlphaZero/MuZero loop.
- Deep look-ahead — on a delayed-reward "trap", a shallow planner takes the bait while MCTS, deepening selectively, commits to the better long-horizon path.
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
np.random.seed(42)
plt.rcParams.update({
'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11,
'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa',
'axes.grid': True, 'grid.alpha': 0.35,
})
print('Setup complete.')
Setup complete.
Experiment 1: UCB — the selection rule at the core of MCTS¶
Question. MCTS selection is UCB applied to trees. What does UCB actually do on the simplest case, a bandit?
A $K$-armed bandit with unknown mean rewards. UCB1 pulls the arm maximizing $\hat\mu_a + c\sqrt{\ln t / n_a}$ — value plus an exploration bonus that shrinks as an arm is tried. We compare cumulative regret against $\epsilon$-greedy and random selection.
K = 8
rng = np.random.default_rng(0)
true_mu = rng.uniform(0, 1, K); best = true_mu.max()
def run_bandit(strategy, T=2000, c=1.0, eps=0.1, seed=1):
rng = np.random.default_rng(seed)
n = np.zeros(K); s = np.zeros(K); regret = []
cum = 0.0
for t in range(1, T + 1):
mu_hat = np.where(n > 0, s / np.maximum(n, 1), np.inf)
if strategy == 'ucb':
score = mu_hat + c * np.sqrt(np.log(t) / np.maximum(n, 1e-9))
score = np.where(n == 0, np.inf, score); a = int(np.argmax(score))
elif strategy == 'egreedy':
a = rng.integers(K) if rng.random() < eps else int(np.argmax(np.where(n > 0, s / np.maximum(n, 1), 0)))
else:
a = rng.integers(K)
r = rng.normal(true_mu[a], 0.1)
n[a] += 1; s[a] += r; cum += best - true_mu[a]; regret.append(cum)
return np.array(regret), n
reg_ucb, n_ucb = run_bandit('ucb')
reg_eg, _ = run_bandit('egreedy')
reg_rand, _ = run_bandit('random')
print('cumulative regret @T | UCB: %.1f | eps-greedy: %.1f | random: %.1f' %
(reg_ucb[-1], reg_eg[-1], reg_rand[-1]))
print('best arm index:', int(np.argmax(true_mu)), '| UCB pulls per arm:', n_ucb.astype(int))
cumulative regret @T | UCB: 102.5 | eps-greedy: 102.7 | random: 815.9 best arm index: 5 | UCB pulls per arm: [ 52 17 8 8 259 1488 53 115]
fig, (axL, axR) = plt.subplots(1, 2, figsize=(11.5, 4.4))
axL.plot(reg_ucb, color='#1f77b4', label='UCB1')
axL.plot(reg_eg, color='#ff7f0e', label='epsilon-greedy')
axL.plot(reg_rand, color='#d62728', label='random')
axL.set_xlabel('pulls'); axL.set_ylabel('cumulative regret'); axL.set_title('UCB has the lowest regret'); axL.legend()
axR.bar(np.arange(K), n_ucb, color=['#2ca02c' if i == np.argmax(true_mu) else '#1f77b4' for i in range(K)])
axR.set_xlabel('arm'); axR.set_ylabel('UCB pulls'); axR.set_title('UCB concentrates pulls on the best arm (green)')
plt.tight_layout(); plt.show()
What you see¶
UCB's cumulative regret (left) grows far slower than $\epsilon$-greedy or random — it stops wasting pulls on bad arms quickly. The pull histogram (right) shows almost all pulls go to the best arm (green), with just a few spent confirming the others are worse.
Why. The bonus $c\sqrt{\ln t/n_a}$ is large for rarely-tried arms, forcing enough exploration to not miss a good arm, but it shrinks as $n_a$ grows, so selection becomes greedy on the empirically best arm. That is an optimism-under-uncertainty balance with a logarithmic regret guarantee.
Takeaway. MCTS selection is exactly this rule applied at every tree node (UCT). The tree is, in effect, a nested set of bandits — each node balances trying promising actions against exploring uncertain ones.
Experiment 2: MCTS concentrates on the optimal action¶
Question. Put together selection, expansion, rollout, and backup — does the tree actually identify the best action?
We define a depth-4 binary look-ahead tree; reward is collected only at the leaves. We run MCTS with UCT selection and random rollouts, and watch the root visit counts for each of the two root actions as the number of simulations grows. The action leading to the best reachable leaf should win.
DEPTH = 6
rng = np.random.default_rng(7)
# leaf values indexed by the binary path; one best leaf at the all-zeros path (root action 0)
leaf_val = rng.uniform(0, 1, 2 ** DEPTH)
leaf_val[0] = 1.6 # optimum reached by always choosing action 0
GAMMA = 1.0
class Node:
def __init__(self, path):
self.path = path; self.N = 0; self.W = 0.0
self.children = {}
def depth(self):
return len(self.path)
def is_leaf(node):
return node.depth() == DEPTH
def value_of(node): # rollout: random descent to a leaf
p = list(node.path)
while len(p) < DEPTH:
p.append(np.random.randint(2))
idx = int(''.join(map(str, p)), 2)
return leaf_val[idx]
def uct_select(node, c=1.4):
best, ba = -1e9, None
for a in (0, 1):
ch = node.children.get(a)
if ch is None or ch.N == 0:
return a
q = ch.W / ch.N
u = q + c * np.sqrt(np.log(node.N + 1) / ch.N)
if u > best:
best, ba = u, a
return ba
def simulate(root):
node, path = root, [root]
while not is_leaf(node):
a = uct_select(node)
if a not in node.children:
node.children[a] = Node(node.path + (a,)); node = node.children[a]; path.append(node)
break
node = node.children[a]; path.append(node)
val = value_of(node)
for nd in path:
nd.N += 1; nd.W += val
return val
def run_mcts(n_sim):
root = Node(())
frac0 = []
for i in range(1, n_sim + 1):
simulate(root)
c0 = root.children.get(0); c1 = root.children.get(1)
n0 = c0.N if c0 else 0; n1 = c1.N if c1 else 0
frac0.append(n0 / max(n0 + n1, 1))
return root, np.array(frac0)
root, frac0 = run_mcts(3000)
print('best leaf is under root action 0 (value 1.6, path all-zeros)')
print('root action visit counts:', {a: root.children[a].N for a in root.children})
print('recommended action (most visited):', max(root.children, key=lambda a: root.children[a].N))
best leaf is under root action 0 (value 1.6, path all-zeros)
root action visit counts: {0: 2980, 1: 20}
recommended action (most visited): 0
fig, ax = plt.subplots(figsize=(8.4, 4.6))
ax.plot(frac0, color='#1f77b4')
ax.axhline(0.5, color='gray', ls='--', lw=1)
ax.set_xlabel('simulations'); ax.set_ylabel('fraction of root visits to action 0 (optimal)')
ax.set_ylim(0, 1)
ax.set_title('MCTS visit counts concentrate on the optimal root action')
plt.tight_layout(); plt.show()
What you see¶
Early on the two root actions are visited about equally (fraction near 0.5) while MCTS is still exploring. As simulations accumulate, the visit fraction climbs toward 1.0 for action 0 — the branch that contains the planted high-value leaf — and the recommended action (most-visited) is the optimal one.
Why. Each backup raises the $Q$ of nodes on rewarding paths, so UCT increasingly selects them, which sends still more simulations down the good branch. Visit count, not raw $Q$, is the recommendation because it is the quantity UCT has stabilized — it reflects how consistently the search preferred that action.
Takeaway. MCTS turns a fixed simulation budget into a sharply peaked judgment about the best action, allocating effort adaptively rather than uniformly. In MuZero these same visit counts also become the policy training target.
Experiment 3: PUCT priors make search efficient¶
Question. AlphaZero/MuZero replace the uniform exploration of UCT with a learned policy prior in PUCT. How much does the prior help?
We rerun the search on the same tree but bias selection with a per-node prior $P(s,a)$ via the PUCT rule. We compare three priors — uniform (=UCT-like), a good prior pointing toward the optimal branch, and a bad prior pointing away — and measure how quickly each makes the optimal root action the most-visited.
def puct_select(node, prior_fn, c=1.5):
best, ba = -1e9, None
total = sum(ch.N for ch in node.children.values()) + 1
P = prior_fn(node)
for a in (0, 1):
ch = node.children.get(a)
q = (ch.W / ch.N) if (ch and ch.N > 0) else 0.0
n = ch.N if ch else 0
u = q + c * P[a] * np.sqrt(total) / (1 + n)
if u > best:
best, ba = u, a
return ba
def simulate_puct(root, prior_fn):
node, path = root, [root]
while not is_leaf(node):
a = puct_select(node, prior_fn)
if a not in node.children:
node.children[a] = Node(node.path + (a,)); node = node.children[a]; path.append(node); break
node = node.children[a]; path.append(node)
val = value_of(node)
for nd in path:
nd.N += 1; nd.W += val
return val
priors = {
'uniform': lambda node: np.array([0.5, 0.5]),
'good': lambda node: np.array([0.8, 0.2]), # action 0 is optimal at root and on the good path
'bad': lambda node: np.array([0.2, 0.8]),
}
def correct_curve(prior_fn, n_sim=500, reps=20):
out = np.zeros(n_sim)
for r in range(reps):
np.random.seed(100 + r)
root = Node(())
for i in range(n_sim):
simulate_puct(root, prior_fn)
c0 = root.children.get(0); c1 = root.children.get(1)
n0 = c0.N if c0 else 0; n1 = c1.N if c1 else 0
out[i] += 1.0 if n0 > n1 else 0.0
return out / reps
curves = {k: correct_curve(v) for k, v in priors.items()}
for k, cu in curves.items():
print(f'prior={k:8s}: P(optimal action recommended) @500 sims = {cu[-1]:.2f}')
prior=uniform : P(optimal action recommended) @500 sims = 1.00 prior=good : P(optimal action recommended) @500 sims = 1.00 prior=bad : P(optimal action recommended) @500 sims = 0.00
fig, ax = plt.subplots(figsize=(8.4, 4.6))
for k, c in zip(('good', 'uniform', 'bad'), ('#2ca02c', '#1f77b4', '#d62728')):
ax.plot(curves[k], color=c, label=f'{k} prior')
ax.set_xlabel('simulations'); ax.set_ylabel('P(optimal action is most-visited)')
ax.set_ylim(0, 1.05)
ax.set_title('PUCT: a good policy prior identifies the optimal action with fewer simulations')
ax.legend()
plt.tight_layout(); plt.show()
What you see¶
With a good prior (green) MCTS commits to the optimal action almost immediately — far fewer simulations than the uniform prior (blue). A bad prior (red) actively misleads the search early, so it takes many simulations to recover, if at all within the budget.
Why. The PUCT term $c\,P(s,a)\sqrt{\sum_b N_b}/(1+N(s,a))$ steers exploration toward actions the prior favors. A good prior focuses the budget on the right branch; a bad prior wastes it. The value network plays the analogous role at leaves, replacing random rollouts with a learned estimate.
Takeaway. Search and learning reinforce each other: the network's policy prior makes search cheap, and the search's improved visit-count policy trains a better prior. This self-improvement loop is the engine of AlphaZero and MuZero — and why their MCTS is far more sample-efficient than vanilla UCT.
Experiment 4: deep look-ahead beats a shallow planner¶
Question. What does the tree's adaptive depth buy that a shallow or greedy planner lacks?
A delayed-reward "trap" chain: from the start, action A gives an immediate $+1$ then ends; action B gives $0$ now but leads to a state worth $+5$ a few steps later. A planner that looks only one step ahead grabs the bait (A); a planner that looks deep enough prefers B. We sweep how deep the search effectively goes (via simulation budget) and record which action it recommends.
GAMMA2 = 0.95
# chain: A -> immediate +1, terminal. B -> 0 now, then a length-L corridor whose deep +5 reward is
# only actually reached by a rollout with probability p_succeed (a noisy deep look-ahead).
L = 4
P_SUCCEED = 0.5
def trap_rollout(first_action, rng):
if first_action == 'A':
return 1.0
return GAMMA2 ** L * 5.0 if rng.random() < P_SUCCEED else 0.0 # B: high-variance deep return
# true (expected) discounted values
vA = 1.0
vB = P_SUCCEED * GAMMA2 ** L * 5.0
print('value(A) = %.3f (immediate +1, certain)' % vA)
print('value(B) = %.3f (expected; delayed +5 after %d steps, reached w.p. %.2f)' % (vB, L, P_SUCCEED))
# A 'shallow' planner sees only the immediate reward and picks A; MCTS estimates the full return.
def shallow_choice():
return 'A' # immediate reward of A (1) > B (0)
def mcts_trap(n_sim, seed=0):
rng = np.random.default_rng(seed)
NA = NB = 0; WA = WB = 0.0
for t in range(1, n_sim + 1):
if NA == 0:
a = 'A'
elif NB == 0:
a = 'B'
else:
uA = WA / NA + 1.4 * np.sqrt(np.log(t) / NA)
uB = WB / NB + 1.4 * np.sqrt(np.log(t) / NB)
a = 'A' if uA >= uB else 'B'
r = trap_rollout(a, rng)
if a == 'A':
NA += 1; WA += r
else:
NB += 1; WB += r
return 'B' if NB > NA else 'A' # ties default to the safe immediate action A
budgets = [2, 5, 10, 30, 100, 300]
pB = [np.mean([mcts_trap(n, seed=s) == 'B' for s in range(60)]) for n in budgets]
print('shallow planner picks:', shallow_choice(), '(the trap)')
for n, p in zip(budgets, pB):
print(f'MCTS with {n:3d} simulations: P(picks B) = {p:.2f}')
value(A) = 1.000 (immediate +1, certain) value(B) = 2.036 (expected; delayed +5 after 4 steps, reached w.p. 0.50)
shallow planner picks: A (the trap) MCTS with 2 simulations: P(picks B) = 0.00 MCTS with 5 simulations: P(picks B) = 0.43 MCTS with 10 simulations: P(picks B) = 0.42 MCTS with 30 simulations: P(picks B) = 0.77 MCTS with 100 simulations: P(picks B) = 0.95 MCTS with 300 simulations: P(picks B) = 0.98
fig, (axL, axR) = plt.subplots(1, 2, figsize=(11.5, 4.4))
axL.bar(['value(A)\nimmediate +1', 'value(B)\nexpected, delayed +5'], [vA, vB], color=['#d62728', '#2ca02c'])
axL.axhline(vA, color='#d62728', ls='--', lw=1)
axL.set_ylabel('discounted return'); axL.set_title('B is better in expectation')
axR.plot(budgets, pB, '-o', color='#1f77b4')
axR.axhline(0.5, color='gray', ls='--', lw=1)
axR.set_xscale('log'); axR.set_ylim(0, 1.05)
axR.set_xlabel('MCTS simulations (log)'); axR.set_ylabel('P(MCTS recommends B)')
axR.set_title('More simulations -> evidence accrues -> MCTS picks B')
fig.suptitle('Shallow planner takes the bait (A); MCTS looks deep enough to prefer B', y=1.02)
plt.tight_layout(); plt.show()
What you see¶
In expectation the delayed path B is worth more than the immediate bait A even after discounting (left bars). A shallow planner that only scores the immediate reward picks A — the trap. MCTS samples the full (noisy) discounted return, so as simulations accumulate the probability it recommends B climbs from a coin-flip toward 1 (right).
Why. Tree search evaluates whole trajectories, not one-step rewards, and its backups propagate the distant $+5$ back to the root action B. With too few simulations it has not yet looked deep enough and may still favor A; depth (budget) is what flips the decision.
Takeaway. Adaptive, selective depth is what tree search adds over shallow or purely greedy choice — and it is bounded by the same horizon trade-off as every Tier-7 planner: deeper look-ahead sees more delayed reward but costs more compute and leans harder on model accuracy. The final note takes up that horizon trade-off directly.
Summary / Key Takeaways¶
- UCB (Exp 1) — the optimism-under-uncertainty rule at the core of MCTS finds the best option with low regret; UCT is this rule applied at every tree node.
- MCTS recommendation (Exp 2) — selection + expansion + rollout + backup makes root visit counts concentrate on the optimal action; visit count, not raw $Q$, is the output (and MuZero's policy target).
- PUCT priors (Exp 3) — a good policy prior makes search dramatically more efficient; the search-improved policy then trains a better prior — the AlphaZero/MuZero self-improvement loop.
- Deep look-ahead (Exp 4) — tree search's adaptive depth solves delayed-reward problems a shallow planner fails, bounded by the universal horizon/compute/model-accuracy trade-off taken up in the final note.