MuZero — Visual Experiments¶

Goal: See MuZero's pillars — search that improves a policy, and a decoder-free latent trained only on reward/value/policy.

What we'll explore¶

#	Experiment	Question answered
1	Search improves policy	Does MCTS produce a better policy than the network's raw prior?
2	PUCT selection	How does PUCT trade the prior against learned value when allocating visits?
3	Value-equivalent latent	Does a reward/value-trained latent merge dynamics-equivalent observations?
4	Unrolled model	Can the model predict reward and value over a rollout without reconstructing observations?

---

Linked theory: research/05-muzero.md

In [1]:

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.')

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: Search improves policy¶

Question: MuZero acts using the MCTS visit-count policy, not the network's raw prior. Does running the learned model through search actually yield a better policy?

We set up actions with known true values and a misleading prior, then run PUCT-MCTS using a value estimate at the leaves. The visit-count policy should shift onto the truly best action as simulations grow.

In [2]:

np.random.seed(1)

Q_true = np.array([0.2, 0.9, 0.4, 0.1])     # action 1 is best
prior = np.array([0.5, 0.1, 0.3, 0.1])      # prior wrongly favours action 0
c_puct = 1.5


def run_mcts(n_sims):
    A = len(Q_true)
    N = np.zeros(A); W = np.zeros(A)
    for _ in range(n_sims):
        Q = np.where(N > 0, W / np.maximum(N, 1), 0.0)
        u = c_puct * prior * np.sqrt(N.sum() + 1) / (1 + N)   # PUCT exploration term
        a = int(np.argmax(Q + u))
        leaf_value = Q_true[a] + 0.05 * np.random.randn()     # noisy value estimate at leaf
        N[a] += 1; W[a] += leaf_value
    return N / N.sum()


sim_counts = [4, 8, 16, 32, 64, 128, 256, 512]
prior_value = prior @ Q_true
mcts_values = [run_mcts(n) @ Q_true for n in sim_counts]
print(f'prior policy value: {prior_value:.3f}   best possible: {Q_true.max():.3f}')
print('mcts policy value by sims:', np.round(mcts_values, 3))

np.random.seed(1) Q_true = np.array([0.2, 0.9, 0.4, 0.1]) # action 1 is best prior = np.array([0.5, 0.1, 0.3, 0.1]) # prior wrongly favours action 0 c_puct = 1.5 def run_mcts(n_sims): A = len(Q_true) N = np.zeros(A); W = np.zeros(A) for _ in range(n_sims): Q = np.where(N > 0, W / np.maximum(N, 1), 0.0) u = c_puct * prior * np.sqrt(N.sum() + 1) / (1 + N) # PUCT exploration term a = int(np.argmax(Q + u)) leaf_value = Q_true[a] + 0.05 * np.random.randn() # noisy value estimate at leaf N[a] += 1; W[a] += leaf_value return N / N.sum() sim_counts = [4, 8, 16, 32, 64, 128, 256, 512] prior_value = prior @ Q_true mcts_values = [run_mcts(n) @ Q_true for n in sim_counts] print(f'prior policy value: {prior_value:.3f} best possible: {Q_true.max():.3f}') print('mcts policy value by sims:', np.round(mcts_values, 3))

prior policy value: 0.320   best possible: 0.900
mcts policy value by sims: [0.3   0.3   0.344 0.622 0.73  0.792 0.823 0.841]

In [3]:

fig, ax = plt.subplots(figsize=(8.7, 5.2))
fig.suptitle('MCTS turns a misleading prior into a near-optimal policy', fontsize=13, fontweight='bold')

ax.plot(sim_counts, mcts_values, '-o', color='seagreen', lw=2, label='MCTS visit-count policy')
ax.axhline(prior_value, color='crimson', ls='--', lw=2, label=f'raw prior value ({prior_value:.2f})')
ax.axhline(Q_true.max(), color='black', ls=':', lw=2, label=f'optimal ({Q_true.max():.2f})')
ax.set_xscale('log', base=2)
ax.set_xlabel('number of simulations'); ax.set_ylabel('policy value  $\\sum_a \\pi(a)Q^*(a)$')
ax.set_title('More search => better policy')
ax.legend(loc='center right')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(8.7, 5.2)) fig.suptitle('MCTS turns a misleading prior into a near-optimal policy', fontsize=13, fontweight='bold') ax.plot(sim_counts, mcts_values, '-o', color='seagreen', lw=2, label='MCTS visit-count policy') ax.axhline(prior_value, color='crimson', ls='--', lw=2, label=f'raw prior value ({prior_value:.2f})') ax.axhline(Q_true.max(), color='black', ls=':', lw=2, label=f'optimal ({Q_true.max():.2f})') ax.set_xscale('log', base=2) ax.set_xlabel('number of simulations'); ax.set_ylabel('policy value $\\sum_a \\pi(a)Q^*(a)$') ax.set_title('More search => better policy') ax.legend(loc='center right') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The MCTS policy value (green) starts near the weak prior and climbs toward the optimal dashed line as simulations increase, overtaking the raw prior (red) quickly.
2. Why it looks this way: Even though the prior favours the wrong action, PUCT keeps trying alternatives; once the value estimates reveal action 1 is best, visits pile onto it, so the visit-count policy concentrates on the optimal action.
3. Key takeaway: Search is a policy improvement operator — this is why MuZero trains its policy network toward the MCTS policy, not the other way around.

---

Experiment 2: PUCT selection¶

Question: PUCT picks $\arg\max_a [Q(a) + c\,P(a)\sqrt{\sum N}/(1+N_a)]$. How does this balance following the prior early against exploiting value later?

We log how visits accumulate across actions during one MCTS run and compare the final visit distribution to the prior and the true values.

In [4]:

np.random.seed(2)

A = len(Q_true)
N = np.zeros(A); W = np.zeros(A)
n_sims = 400
visit_curve = np.zeros((n_sims, A))
for t in range(n_sims):
    Q = np.where(N > 0, W / np.maximum(N, 1), 0.0)
    u = c_puct * prior * np.sqrt(N.sum() + 1) / (1 + N)
    a = int(np.argmax(Q + u))
    N[a] += 1; W[a] += Q_true[a] + 0.05 * np.random.randn()
    visit_curve[t] = N / N.sum()
final_visits = N / N.sum()
print('final visit policy:', np.round(final_visits, 3))

np.random.seed(2) A = len(Q_true) N = np.zeros(A); W = np.zeros(A) n_sims = 400 visit_curve = np.zeros((n_sims, A)) for t in range(n_sims): Q = np.where(N > 0, W / np.maximum(N, 1), 0.0) u = c_puct * prior * np.sqrt(N.sum() + 1) / (1 + N) a = int(np.argmax(Q + u)) N[a] += 1; W[a] += Q_true[a] + 0.05 * np.random.randn() visit_curve[t] = N / N.sum() final_visits = N / N.sum() print('final visit policy:', np.round(final_visits, 3))

final visit policy: [0.05  0.9   0.042 0.008]

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('PUCT: explore via the prior, then concentrate on high value', fontsize=13, fontweight='bold')

ax = axes[0]
for a in range(A):
    ax.plot(visit_curve[:, a], lw=2, label=f'action {a} (Q*={Q_true[a]})')
ax.set_xlabel('simulation'); ax.set_ylabel('visit fraction')
ax.set_title('Visit fractions over the search')
ax.legend(fontsize=9)

ax2 = axes[1]
x = np.arange(A); w = 0.27
ax2.bar(x - w, prior, w, color='crimson', label='prior P(a)')
ax2.bar(x, Q_true / Q_true.sum(), w, color='gray', label='normalized Q*(a)')
ax2.bar(x + w, final_visits, w, color='seagreen', label='final visit policy')
ax2.set_xlabel('action'); ax2.set_ylabel('probability')
ax2.set_title('Prior vs value vs resulting policy')
ax2.legend(fontsize=9)

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('PUCT: explore via the prior, then concentrate on high value', fontsize=13, fontweight='bold') ax = axes[0] for a in range(A): ax.plot(visit_curve[:, a], lw=2, label=f'action {a} (Q*={Q_true[a]})') ax.set_xlabel('simulation'); ax.set_ylabel('visit fraction') ax.set_title('Visit fractions over the search') ax.legend(fontsize=9) ax2 = axes[1] x = np.arange(A); w = 0.27 ax2.bar(x - w, prior, w, color='crimson', label='prior P(a)') ax2.bar(x, Q_true / Q_true.sum(), w, color='gray', label='normalized Q*(a)') ax2.bar(x + w, final_visits, w, color='seagreen', label='final visit policy') ax2.set_xlabel('action'); ax2.set_ylabel('probability') ax2.set_title('Prior vs value vs resulting policy') ax2.legend(fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Early on, visits spread according to the prior (action 0 gets attention), but as value estimates firm up, the visit fraction for the high-value action 1 takes over. The right bars show the final visit policy ignoring the misleading prior and tracking the true values.
2. Why it looks this way: The PUCT exploration term $c P(a)\sqrt{\sum N}/(1+N_a)$ dominates when an action has few visits (prior-driven exploration) but decays as $N_a$ grows, letting the value term $Q(a)$ decide (exploitation).
3. Key takeaway: PUCT uses the prior to explore efficiently but lets the searched value override it — the mechanism that makes MuZero's search both directed and self-correcting.

---

Experiment 3: Value-equivalent latent¶

Question: MuZero's latent is trained only to predict reward/value/policy. Does that make it merge observations that differ in appearance but are identical for control?

We build observations with an extra irrelevant bit (states come in dynamics-equivalent pairs). A latent trained to predict value should collapse each pair; a reconstruction latent keeps them apart.

In [6]:

np.random.seed(3)

# 4 'true' states with values; each appears as 2 observations differing in an irrelevant bit
true_value = np.array([0.1, 0.4, 0.7, 1.0])
obs, val, cls = [], [], []
for s in range(4):
    for irrelevant in (0, 1):
        feat = np.zeros(5)
        feat[s] = 1.0                      # identity of the true state
        feat[4] = irrelevant * 3.0         # irrelevant high-variance bit
        obs.append(feat); val.append(true_value[s]); cls.append(s)
obs = np.array(obs); val = np.array(val); cls = np.array(cls)

# value-equivalent latent: 1-D projection trained to predict value
w = np.linalg.lstsq(obs, val, rcond=None)[0]
z_ve = obs @ w
# reconstruction latent: top PCA direction (keeps the high-variance irrelevant bit)
oc = obs - obs.mean(0)
z_rec = oc @ np.linalg.svd(oc, full_matrices=False)[2][0]
print('value-equivalent latent merges the pairs:', np.round(z_ve, 2))

np.random.seed(3) # 4 'true' states with values; each appears as 2 observations differing in an irrelevant bit true_value = np.array([0.1, 0.4, 0.7, 1.0]) obs, val, cls = [], [], [] for s in range(4): for irrelevant in (0, 1): feat = np.zeros(5) feat[s] = 1.0 # identity of the true state feat[4] = irrelevant * 3.0 # irrelevant high-variance bit obs.append(feat); val.append(true_value[s]); cls.append(s) obs = np.array(obs); val = np.array(val); cls = np.array(cls) # value-equivalent latent: 1-D projection trained to predict value w = np.linalg.lstsq(obs, val, rcond=None)[0] z_ve = obs @ w # reconstruction latent: top PCA direction (keeps the high-variance irrelevant bit) oc = obs - obs.mean(0) z_rec = oc @ np.linalg.svd(oc, full_matrices=False)[2][0] print('value-equivalent latent merges the pairs:', np.round(z_ve, 2))

value-equivalent latent merges the pairs: [0.1 0.1 0.4 0.4 0.7 0.7 1.  1. ]

In [7]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('Value-equivalent latent collapses dynamics-equivalent observations', fontsize=13, fontweight='bold')

ax = axes[0]
ax.scatter(z_ve, val, c=cls, cmap='viridis', s=120, edgecolor='k')
ax.set_xlabel('value-equivalent latent'); ax.set_ylabel('true value')
ax.set_title('Each pair maps to one latent (aligned with value)')

ax2 = axes[1]
ax2.scatter(z_rec, val, c=cls, cmap='viridis', s=120, edgecolor='k')
ax2.set_xlabel('reconstruction latent (PCA)'); ax2.set_ylabel('true value')
ax2.set_title('Pairs split by the irrelevant bit')

sm = plt.cm.ScalarMappable(cmap='viridis', norm=plt.Normalize(0, 3))
fig.colorbar(sm, ax=axes, label='true state id', fraction=0.025)
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('Value-equivalent latent collapses dynamics-equivalent observations', fontsize=13, fontweight='bold') ax = axes[0] ax.scatter(z_ve, val, c=cls, cmap='viridis', s=120, edgecolor='k') ax.set_xlabel('value-equivalent latent'); ax.set_ylabel('true value') ax.set_title('Each pair maps to one latent (aligned with value)') ax2 = axes[1] ax2.scatter(z_rec, val, c=cls, cmap='viridis', s=120, edgecolor='k') ax2.set_xlabel('reconstruction latent (PCA)'); ax2.set_ylabel('true value') ax2.set_title('Pairs split by the irrelevant bit') sm = plt.cm.ScalarMappable(cmap='viridis', norm=plt.Normalize(0, 3)) fig.colorbar(sm, ax=axes, label='true state id', fraction=0.025) plt.show()

What you see¶

1. What the plot shows: In the value-equivalent latent (left), the two observations of each true state land at the same place and the latent lines up with value. In the reconstruction latent (right), the irrelevant bit splits each pair into two separated points.
2. Why it looks this way: Training the latent to predict value gives it no reason to encode the irrelevant bit, so equivalent observations collapse together. A reconstruction objective must preserve everything, including the irrelevant high-variance bit.
3. Key takeaway: Value equivalence yields an abstract latent that keeps only control-relevant structure — the foundation MuZero shares with TD-MPC2 and the reason it ignores observation detail.

---

Experiment 4: Unrolled model¶

Question: MuZero unrolls its dynamics to predict future rewards and values, never observations. Does an abstract latent rollout track the true reward/value along a trajectory?

We learn an abstract latent on a chain MDP (states map to a 1-D latent purely by value), unroll the dynamics, and compare predicted reward and value to ground truth.

In [8]:

np.random.seed(4)

L, gamma = 8, 0.9
rewards = np.array([0, 0, 0, 1, 0, 0, 2, 0], dtype=float)   # reward per chain state
# true value (sum of discounted future rewards) along the chain
true_value = np.zeros(L)
for s in range(L - 1, -1, -1):
    true_value[s] = rewards[s] + (gamma * true_value[s + 1] if s + 1 < L else 0)

# abstract latent = state's value (decoder-free: latent carries value, not appearance)
latent = true_value.copy()
# learned dynamics/reward/value heads here are exact tabular maps (the point is the rollout)
pred_reward = rewards.copy()
pred_value = true_value.copy()

# unroll from state 0
cum_pred, cum_true = [], []
cp = ct = 0.0
for s in range(L):
    cp += gamma ** s * pred_reward[s]; ct += gamma ** s * rewards[s]
    cum_pred.append(cp); cum_true.append(ct)
print('root predicted value:', round(pred_value[0], 3), ' true:', round(true_value[0], 3))

np.random.seed(4) L, gamma = 8, 0.9 rewards = np.array([0, 0, 0, 1, 0, 0, 2, 0], dtype=float) # reward per chain state # true value (sum of discounted future rewards) along the chain true_value = np.zeros(L) for s in range(L - 1, -1, -1): true_value[s] = rewards[s] + (gamma * true_value[s + 1] if s + 1 < L else 0) # abstract latent = state's value (decoder-free: latent carries value, not appearance) latent = true_value.copy() # learned dynamics/reward/value heads here are exact tabular maps (the point is the rollout) pred_reward = rewards.copy() pred_value = true_value.copy() # unroll from state 0 cum_pred, cum_true = [], [] cp = ct = 0.0 for s in range(L): cp += gamma ** s * pred_reward[s]; ct += gamma ** s * rewards[s] cum_pred.append(cp); cum_true.append(ct) print('root predicted value:', round(pred_value[0], 3), ' true:', round(true_value[0], 3))

root predicted value: 1.792  true: 1.792

In [9]:

fig, axes = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('The unrolled model predicts reward/value, not observations', fontsize=13, fontweight='bold')

ax = axes[0]
ax.plot(range(L), cum_true, '-o', color='black', lw=2, label='true discounted return')
ax.plot(range(L), cum_pred, '--s', color='seagreen', lw=2, label='unrolled prediction')
ax.set_xlabel('unroll step'); ax.set_ylabel('cumulative discounted reward')
ax.set_title('Predicted return tracks truth along the rollout')
ax.legend()

ax2 = axes[1]
ax2.plot(range(L), true_value, '-o', color='black', lw=2, label='true value V(s)')
ax2.plot(range(L), pred_value, '--s', color='royalblue', lw=2, label='predicted value')
ax2.bar(range(L), rewards, color='orange', alpha=0.4, label='reward r(s)')
ax2.set_xlabel('chain state'); ax2.set_ylabel('value / reward')
ax2.set_title('Latent carries value, not appearance')
ax2.legend(fontsize=9)

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13, 5)) fig.suptitle('The unrolled model predicts reward/value, not observations', fontsize=13, fontweight='bold') ax = axes[0] ax.plot(range(L), cum_true, '-o', color='black', lw=2, label='true discounted return') ax.plot(range(L), cum_pred, '--s', color='seagreen', lw=2, label='unrolled prediction') ax.set_xlabel('unroll step'); ax.set_ylabel('cumulative discounted reward') ax.set_title('Predicted return tracks truth along the rollout') ax.legend() ax2 = axes[1] ax2.plot(range(L), true_value, '-o', color='black', lw=2, label='true value V(s)') ax2.plot(range(L), pred_value, '--s', color='royalblue', lw=2, label='predicted value') ax2.bar(range(L), rewards, color='orange', alpha=0.4, label='reward r(s)') ax2.set_xlabel('chain state'); ax2.set_ylabel('value / reward') ax2.set_title('Latent carries value, not appearance') ax2.legend(fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The unrolled prediction of cumulative discounted reward sits on top of the true curve (left), and the predicted value matches the true value at every chain state (right), with the rewards shown as bars.
2. Why it looks this way: The model only needs to reproduce reward and value along the rollout — which it does — without ever representing what the states look like. The latent here literally is the value, an abstract quantity.
3. Key takeaway: A decoder-free model can plan accurately by predicting only the quantities planning consumes (reward, value, policy), exactly MuZero's design.

---

Summary¶

Experiment	Key insight
1. Search improves policy	The MCTS visit-count policy overtakes a misleading prior — search is policy improvement.
2. PUCT selection	The prior drives early exploration; the searched value takes over as visits accumulate.
3. Value-equivalent latent	A value-trained latent merges dynamics-equivalent observations, dropping irrelevant detail.
4. Unrolled model	The model predicts reward/value over a rollout without reconstructing observations.

Connection to the broader theory¶

MuZero is the canonical decoder-free planning agent: three functions (representation, dynamics, prediction) plus MCTS in latent, trained purely on reward/value/policy. TD-MPC2 carries the same value-equivalence idea to continuous control. The tier now turns to VLA systems — OpenVLA and π0 — where the latent must not only plan but act in the physical world.