MuZero and Value-Equivalent Models › Value Equivalence and Task Sufficiency
World Models HandbookAll chapters MuZero and Value-Equivalent Models › Value Equivalence and Task SufficiencyOne section · Sections in this chapter4 / 8 ↓
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
The tutorial's decision-target model illustrates the motivation for value equivalence, but we have not proven that it satisfies a formal value-equivalence criterion. Here is the criterion and the distinction it makes.
In Grimm et al.'s definition, two models and are value equivalent relative to a selected policy set and function set when their Bellman operators agree on every selected pair: for all and . This is a statement about one-step Bellman updates on chosen test functions, not automatically equality of every policy value. Full model equivalence additionally requires agreement on transition kernels and reward functions.
Model equivalence asks for a faithful transition/reward description. Value equivalence asks for agreement only on selected decision-relevant Bellman tests. A model can differ in transitions to states that a chosen function set assigns the same value, yet agree on those Bellman updates. Whether that suffices for planning depends on the policies and functions included; a narrow test set gives a weaker guarantee.
Here is the expected immediate reward plus discounted expected at the next state under policy in model . If the selected function set is sufficiently rich and closed under relevant Bellman iterations, agreement can imply agreement of corresponding policy values. Equality of all is not the paper's Definition 1.
Grimm and colleagues use this criterion to study model classes sufficient for selected planning calculations. The lesson is not that prediction accuracy is irrelevant, but that the quantities needed by the planner determine which predictive distinctions matter. A model that reconstructs pixels accurately can still misrank actions; one that omits pixels can still support good decisions on its tested distribution.
Reconstruction versus decision-target representations
Consider two models of our gridworld:
- Model A has an observation decoder and is trained to reconstruct all 20 observation coordinates, including the four unpredictable bits.
- Model B produces an 8-dimensional latent with no interpretation whatsoever, and no decoder.
For a decoder of the current encoded observation, retaining the noise bits can reduce reconstruction loss. For a predictor of future independent bits, no deterministic estimate can eliminate expected squared error. These are distinct tasks; the plot below tests how much current-bit information each encoder retains, not whether future noise is forecast.
Model B has no direct reconstruction incentive, but its latent can still encode distractor bits incidentally. Consistency may reduce that information without guaranteeing its removal. We compare trained probes rather than infer what happened from the objectives alone.
This decoder comparison does not test formal value equivalence. For a Bellman example, imagine two gridworld models with the same position transitions and rewards but different distributions for the next observation's independent noise bits. Their full observation-transition kernels differ. Yet for policies and value functions that depend only on position, both models give the same immediate reward and expected next-position value, hence the same selected Bellman updates. That restricted agreement illustrates value equivalence without claiming that either trained network above satisfies the formal criterion.
Task sufficiency
The practical goal is task sufficiency: the latent should carry enough information for the policies and values the agent needs. In this fixed-goal gridworld, position matters, while the four independent distractor bits do not affect optimal actions. The goal location is constant, so it need not be separately encoded in every observation.
Sufficiency is relative to a task, state distribution, and decision rule. A representation adequate for navigation may omit information needed for contact-rich manipulation. A compact decision-centric latent can be efficient for its training objective, but transfer to a new reward or environment must be tested.
There is a relationship to bisimulation, discussed in Part III: Representing Agents and Worlds. Under usual MDP assumptions, bisimilar states have equal immediate rewards for each action and matching transition probabilities over equivalence classes. Such a quotient preserves values for policies compatible with the abstraction. Value equivalence asks only for selected Bellman tests and can therefore be weaker.
Measuring what the latent keeps
We can test all of this empirically. We train a second model, identical in architecture except that it also has a decoder and is trained with a reconstruction loss instead of the consistency loss. Then we ask two questions of each model:
- How large is the sampled between-position share of latent variance? A noise-invariant latent would vary little within a position, but eight sampled noise masks per cell also put sampling variation into this statistic. Decision sufficiency alone does not require noise invariance.
- How accurately can a linear probe read the noise bits out of the latent? If the latent is noise-invariant, a linear probe should be no better than chance.
These two diagnostics are complementary. The finite-sample variance share describes one aspect of latent organization, while the linear probe tests recoverable current-noise information. A model can score differently on the two measures, so we inspect both without treating either as a sufficiency test.
class ReconstructiveNet(MuZeroNet):
def __init__(
self,
obs_dim=OBS_DIM,
n_actions=N_ACTIONS,
latent_dim=LATENT_DIM,
hidden=64,
):
super().__init__(obs_dim, n_actions, latent_dim, hidden)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, obs_dim),
)
def decode(self, z):
return self.decoder(z)
REC_MODEL = ReconstructiveNet()
hist_rec = train_model(
REC_MODEL,
OBS,
ACT,
REW,
NXT,
MC_RETURNS,
BEHAVIOUR_POLICY,
epochs=200,
use_consistency=False,
use_reconstruction=True,
)@torch.no_grad()
def latent_position_variance(model, n_noise=8, seed=7):
"""Finite-sample between-position share of latent variance."""
rng_local = np.random.default_rng(seed)
Z = np.zeros((N_STATES, n_noise, model.latent_dim))
for s in range(N_STATES):
p = (s // N, s % N)
for k in range(n_noise):
o = observe(p, rng_local)
Z[s, k] = model.encode(torch.tensor(o).unsqueeze(0)).numpy()[0]
grand_mean = Z.reshape(-1, model.latent_dim).mean(axis=0)
between = ((Z.mean(axis=1) - grand_mean) ** 2).sum(axis=1).mean()
within = ((Z - Z.mean(axis=1, keepdims=True)) ** 2).sum(axis=2).mean()
return float(between / (between + within))
def distractor_probe_accuracy(model, obs, n_train=2000, seed=0):
"""Linear probe accuracy for recovering the noise bits from the latent."""
with torch.no_grad():
Z = model.encode(torch.tensor(obs)).numpy()
Y = obs[:, N_STATES:]
perm = np.random.default_rng(seed).permutation(len(Z))
tr, te = perm[:n_train], perm[n_train:]
A = np.concatenate([Z[tr], np.ones((len(tr), 1))], axis=1)
W, *_ = np.linalg.lstsq(A, Y[tr], rcond=None)
A_te = np.concatenate([Z[te], np.ones((len(te), 1))], axis=1)
pred = (A_te @ W) > 0.5
return float((pred == (Y[te] > 0.5)).mean())
VAR_VE = latent_position_variance(VE_MODEL)
VAR_REC = latent_position_variance(REC_MODEL)
ACC_VE = distractor_probe_accuracy(VE_MODEL, OBS)
ACC_REC = distractor_probe_accuracy(REC_MODEL, OBS)
The decision-target model has a higher sampled between-position variance share, but its distractor probe remains above chance in this run. We drew eight independent noise masks per cell from the 16 possibilities, with repeats possible, so random mask variation can enter the between-cell term; this is not an unbiased estimate of variance caused by position alone. The reconstructive model retains still more current-bit information, consistent with its decoder objective. The comparator also changes decoder architecture and independently initializes the shared layers, so the gap cannot be attributed to the auxiliary objective alone. Probe test observations are held out from the linear probe fit, not from representation training. This comparison is suggestive, not a proof that consistency eliminates noise or that one objective dominates for every task.
This is not an argument that reconstruction is bad. It provides supervision when rewards are sparse and can preserve information for tasks not known during training. The two tutorial objectives emphasize different signals; whether either yields better control or transfer requires a matched downstream evaluation that this probe does not provide.