Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

MuZero and Value-Equivalent Models › Search Targets and Reanalysis

Michael BrenndoerferJuly 16, 20265 min read

World Models HandbookAll chapters

Part 8 of 12Decision-Centric Research Lineages

Chapter 46MuZero and Value-Equivalent Models › Search Targets and Reanalysis

Reading section 5 of 8 · Search Targets and Reanalysis · 5 minSection 5 of 8 · 5 min

Section 5 of 8 · 5 min read

MuZero and Value-Equivalent Models › Search Targets and ReanalysisOne section · Sections in this chapter5 / 8 ↓

Sections in this chapter8

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Search consumes the learned model and supplies policy targets for later training. Better targets can improve the model, but this feedback loop is not guaranteed to converge or improve control under approximation.

Write $\mathcal{S}$ for a search procedure that uses $g_\theta$ and $v_\theta$ to produce a policy distribution and a root-value estimate. MuZero trains its policy head against the search distribution, while its value target also uses observed outcomes or $n$ -step returns. Our later toy comparison directly trains on search root values as a separate experimental simplification.

Self-consistency, $v_\theta \approx \mathcal{S}(v_\theta)$ , would mean the value head matches the particular search estimates on tested states. It would not prove optimality. Classical tabular value iteration has contraction guarantees under specific assumptions; finite tree search plus a learned model and function approximation does not inherit them automatically.

MuZero actors run search while selecting real actions and store trajectories with search statistics in replay. Reusing old search-policy targets creates staleness as the network changes. Reanalyzing stored observations with newer parameters can refresh those policy targets, at additional compute cost.

ReanalysisLink Copied

The MuZero Reanalyze variant searches stored states again with a newer network to provide fresh policy targets for a reported 80% of updates. Its value target uses observed $n$ -step rewards bootstrapped from a target network's later value estimate. It does not simply replace every value label with the current root search value.

Reanalysis trades compute for more current targets. Its benefit is empirical and depends on the data regime, target network, and model quality; it is not a standalone guarantee of sample efficiency.

Training a value head directly on its own search estimates can reinforce an overestimate. Observed rewards ground the sampled transitions, while target-network bootstrapping supplies another learned estimate that may stabilize targets; neither certifies unseen imagined actions or rules out search exploiting inaccurate off-data transitions. The original MuZero objective has no explicit latent-consistency term.

The toy experiment below intentionally tests the more direct search-value target. Treat any outcome as a property of this fixed dataset and optimizer, not as a reproduction of the paper's Reanalyze value-target construction.

An experiment: same data, different targetsLink Copied

We can compare target sources while holding data and optimization fixed. Take the trained decision-target model and fine-tune it twice, on exactly the same data, for exactly the same number of epochs:

Control run. Fine-tune on Monte Carlo returns and behaviour-policy action targets.
Reanalysis run. Fine-tune on search values and search visit-count policies.

The runs start from matched model copies and use the same data, optimizer settings, gradient steps, and seed. Value and policy targets change together, so this intervention cannot identify either target's separate effect. If the search-target model does not improve, that is a useful counterexample to assuming that a more accurate raw target is always internalized by this training loop.

In[29]:

Code

rng_re = np.random.default_rng(11)
REANALYZE_IDX = rng_re.permutation(len(OBS))

SEARCH_VALUES = np.zeros(len(REANALYZE_IDX))
SEARCH_POLICIES = np.zeros((len(REANALYZE_IDX), N_ACTIONS))
for i, j in enumerate(REANALYZE_IDX):
    v, _, pi = run_mcts(VE_MODEL, OBS[j], n_sims=24, max_depth=4)
    SEARCH_VALUES[i] = v
    SEARCH_POLICIES[i] = pi

MC_TARGETS = MC_RETURNS[REANALYZE_IDX]
POLICY_TARGETS = BEHAVIOUR_POLICY[REANALYZE_IDX]
SUB_OBS = OBS[REANALYZE_IDX]
SUB_ACT = ACT[REANALYZE_IDX]
SUB_REW = REW[REANALYZE_IDX]
SUB_NXT = NXT[REANALYZE_IDX]

In[30]:

Code

MODEL_CONTROL = copy.deepcopy(VE_MODEL)
hist_control = train_model(
    MODEL_CONTROL,
    SUB_OBS,
    SUB_ACT,
    SUB_REW,
    SUB_NXT,
    MC_TARGETS,
    POLICY_TARGETS,
    epochs=80,
    seed=1,
)

MODEL_REANALYZED = copy.deepcopy(VE_MODEL)
hist_reanalysis = train_model(
    MODEL_REANALYZED,
    SUB_OBS,
    SUB_ACT,
    SUB_REW,
    SUB_NXT,
    SEARCH_VALUES,
    SEARCH_POLICIES,
    epochs=80,
    seed=1,
)

ERR_CONTROL = value_rmse(MODEL_CONTROL)
ERR_REANALYSIS = value_rmse(MODEL_REANALYZED)
ERR_SEARCH_TARGETS = float(
    np.sqrt(np.mean((SEARCH_VALUES - V_STAR[STATE_IDS[REANALYZE_IDX]]) ** 2))
)
with torch.no_grad():
    replay_observations = torch.tensor(SUB_OBS)
    _, replay_head_values = MODEL_REANALYZED.predict(
        MODEL_REANALYZED.encode(replay_observations)
    )
replay_reference = V_STAR[STATE_IDS[REANALYZE_IDX]]
ERR_HEAD_REPLAY = float(
    np.sqrt(np.mean((replay_head_values.numpy() - replay_reference) ** 2))
)
REPLAY_GOAL_COUNT = int(np.count_nonzero(STATE_IDS[REANALYZE_IDX] == 0))

Out[31]:

Visualization

Bar chart comparing value RMSE across three training conditions. — Equally weighted grid-cell mean value RMSE against the exact optimum: base 0.104, behaviour-target fine-tune 0.095, and search-target fine-tune 0.072. This plot does not compare the head with raw search values; those use a separate replay-weighted evaluation below.

Out[32]:

Console

Raw search RMSE on replay observations:   0.040
Fine-tuned head RMSE on the same replay: 0.039
Absorbing-goal replay rows:              2197/3000
Behaviour-policy cell-mean value RMSE:   0.105

Under equal weighting of the 16 grid-cell means, search-target fine-tuning lowers the direct head's error to 0.072, compared with 0.095 for behaviour-target fine-tuning and 0.104 for the base model. The raw search values were instead computed on replay observations. On those same observations, raw search RMSE is 0.040 and the fine-tuned head RMSE is 0.039. Their ordering therefore cannot be inferred from the unequal-weighted cell-mean chart. The replay metric also weights the absorbing goal heavily (2,197 of 3,000 rows), so it should not stand in for uniform state coverage. This single-run joint value-and-policy target change improves cell-mean value accuracy, but is not a general convergence or action-quality guarantee.

This distinction separates the quality of a target generator from how well a parameterized network learns its targets. The experiment cannot justify dropping search at deployment; that would require direct action-quality evaluation without search.