1 Introduction↩︎

World-model evaluation in model-based reinforcement learning often treats two questions as equivalent: does the learned model predict reward and value well, and does it model the task’s dynamics well? Value-equivalence [1]–[3] argues that for planning, what a model gets right about reward and value is what matters. When the answers to the two questions agree, training a model to predict reward is enough. When they disagree, a reward-prediction check can be silent at the failure mode that matters for planning.

This paper proposes a complementary diagnostic, operator-on-\(F\), that compares the model’s \(k\)-step latent rollout against the environment’s at the level of an observable subset \(F\) rather than at the reward head. We then report results on released TD-MPC2 [4] checkpoints in which the two metrics disagree at the largest scale of the released mt80 sweep: at \(317\)M the operator-on-\(F\) error is an order of magnitude above the \(0.28\)–\(0.36\) cluster (\(2.62\) vs.next-highest \(0.91\), at \(1\)M) while reward-prediction error (\(0.091\)) stays within the same small \([0.028, 0.091]\) range as the rest of the sweep, and the planning return collapses to \(0.9\). We also report the same diagnostic on a pure-SSL world model trained from scratch (LeWM, [5]) used as a cross-architecture comparison; both of these are single-environment observations, not a general claim.

Our contributions are: (i) the operator-on-\(F\) definition and a probe-based estimator that runs natively on a model’s own latent predictor at its planning-time horizon; (ii) on the released TD-MPC2 mt80 size sweep, reward-prediction error stays within the narrow \([0.028, 0.091]\) band across all five sizes (\(\sim 3\times\) variation) while operator error spans \(0.28\) to \(2.62\), with the \(317\)M model showing the highest operator error (an order of magnitude above the \(0.28\)–\(0.36\) cluster) and a return collapse to \(0.9\) (rank correlation between operator error and return \(-0.90\), anchor-bootstrap \(95\%\) CI \([-0.90, -0.70]\), with leave-one-out drop of any single size leaving Spearman at \(-0.80\) or stronger); (iii) a re-ranking of the same five models that disagrees with the Bellman residual (Spearman \(+0.30\), anchor-bootstrap \(95\%\) CI \([+0.30, +0.60]\)), contrasted with the per-size near-identity between the unnormalized value-only error and the Bellman residual (Spearman \(+1.00\)); and (iv) a cross-architecture sanity check between single-task TD-MPC2 and pure-SSL LeWM on cheetah-run with disjoint \(95\%\) CIs that are preserved under a \(1\)-hidden-layer MLP-probe ablation.

2 Background and method↩︎

2.1 Operator-on-\(F\)↩︎

Let \(s_t\) be a state, \(a_t\) an action, and \(z_t = \mathrm{enc}(s_t)\) a model’s latent encoding. We write the model’s \(k\)-step rollout as \(\hat{z}_{t+k} = O_k(z_t,\, a_{t:t+k-1})\) and the environment’s encoded next-state as \(z_{t+k}' = \mathrm{enc}(s_{t+k})\). For TD-MPC2 [4], \(O_k\) is its latent dynamics rolled \(k\) times open-loop; for LeWM [5], \(O_k\) is its native predictor at a \(k\)-step frameskip. A probe \(\phi_F : Z \to F\) maps latents to a chosen observable subset \(F\) [6], so a fair cross-model comparison can happen on a surface both encoders share.

For each functional \(\phi \in F\) and each anchor \((s_t,\, a_{t:t+k-1},\, s_{t+k})\), we compute the per-anchor operator error \(|\phi(\hat{z}_{t+k}) - \phi(z_{t+k}')|\) normalized by the across-anchor spread of the true pushforward, and aggregate by RMS over \((\phi, \text{anchor})\). The value slice sets \(F = \{r, V\}\) using the model’s own reward and value heads. The full-\(F\) aggregate adds a held-out per-anchor PCA basis on the encoded next-state geometry (top-\(16\) directions, whitened by their singular values); the basis is fit on a held-out half of the anchors and evaluated on the other half so it cannot adapt to operator error. Appendix 6 gives the full expression.

2.2 Epistemic class and probe↩︎

We report all operator-on-\(F\) values as native: \(O_k\) is each model’s own predictor at the same \(k\) a planner would use at deployment, and the comparison is to encoded ground-truth next-states. We never substitute a synthetic predictor. LeWM has no reward or value head, so for the cross-architecture comparison we report the operator on a shared observation-level \(F\) (cheetah-run observation plus reward) probed identically on both models.

The probe is a ridge regression with feature standardization and \(\lambda\) chosen by grid search on a held-out sub-split, reported with per-observable \(R^2\) so probe quality is visible. Unregularized probes degrade sharply on high-dimensional latents (in a pilot, unregularized ridge on the TD-MPC2 mt80-\(19\)M \(768\)-dim latent gave strongly negative \(R^2\) on some observables), an instance of the capacity-vs-information confounds documented in the probing-classifier literature [7], which makes unregularized probes an unfair surface for cross-model comparison. Appendix 7 gives the protocol in full, including the commonly-well-probed restriction used as the probe-asymmetry control in Section 4.3.

3 Experimental setup↩︎

We evaluate three sources of world models on the DMC cheetah-run task [8]: the released TD-MPC2 [4] mt80 multitask checkpoints at five parameter counts (\(1\)M, \(5\)M, \(19\)M, \(48\)M, \(317\)M), used as the size-sweep axis; the released TD-MPC2 single-task cheetah-run checkpoint, used as the single-task control; and three LeWM (LeWorldModel, [5]) seeds trained from pixels for \(100\) epochs each with the SIGReg regularizer of [9] active. LeWM has no reward or value head, so the cross-architecture comparison reports the operator on observation-level \(F\) only.

Operator error is measured at a matched horizon of \(5\) collect-steps. LeWM predicts the jump in one forward pass (frameskip \(5\)); TD-MPC2 rolls its latent dynamics \(5\times\) open-loop under the recorded actions. The cross-architecture comparison (Section 4.3) pools \(1552\) evaluation anchors over \(16\) held-out episodes; the TD-MPC2 size-sweep (Sections 4.1 and 4.2) uses the \(1500\) pooled anchors from the size-sweep rollouts. Configurations, analysis scripts, and figure scripts to reproduce the operator-error, metric-disagreement, and cross-architecture results reported in Section 4 are available at https://github.com/DonnaVakalis/operator-on-F.

4 Results↩︎

4.1 Value-equivalence can be silent when the operator is wrong↩︎

On the released TD-MPC2 mt80 sweep over five parameter counts (Figure 1), reward-prediction error stays within \([0.028, 0.091]\) at every size — only \(\sim 3\times\) variation, narrow resolution for an unnormalized reward-fit check to distinguish them — while the full-\(F\) operator-on-\(F\) error spans \(0.28\) to \(2.62\). The \(317\)M checkpoint is the clearest case: operator error is \(2.62\), an order of magnitude above the \(0.28\)–\(0.36\) cluster (next-highest \(0.91\), at \(1\)M), and the planning return collapses to \(0.9\), while reward-prediction error (\(0.091\)) remains within the same small \([0.028, 0.091]\) range as the rest of the sweep. Neither value-equivalence proxy, as conventionally (unnormalized) reported, orders the models by return: Spearman \(-0.10\) for the Bellman residual and \(-0.30\) for reward-prediction error (normalizing each channel by its target spread, as operator-on-\(F\) does, recovers the return ordering; Appendix 8). Per-size operator-on-\(F\) values have tight anchor-bootstrap CIs (Appendix 8): the across-size rank correlation with return is \(-0.90\) (anchor-bootstrap \(95\%\) CI \([-0.90, -0.70]\)), and leave-one-out removal of any single size leaves the Spearman at \(-0.80\) or stronger, so the rank order is stable to single-point outliers, though \(n=5\) is inherently small. Because reward-prediction error has limited variance across the sweep, a partial correlation controlling for it reproduces the marginal value (\(-0.89\)); we report this as consistency, not an independent control.

4.2 Operator-on-\(F\) disagrees with value-equivalence proxies↩︎

On the same five TD-MPC2 sizes, the full-\(F\) operator-on-\(F\) rank-disagrees with the Bellman residual: Spearman \(+0.30\) (anchor-bootstrap \(95\%\) CI \([+0.30, +0.60]\)), and \(+0.10\) (\([+0.10, +0.30]\)) against the kernel-divergence operator. The paper’s value slice \(F = \{r, V\}\) (normalized per-functional aggregate) also disagrees with the Bellman residual at Spearman \(+0.10\) (\([+0.10, +0.20]\)); it is the held-out PCA basis that lets full-\(F\) pick up return-tracking structure neither metric sees.

To anchor the comparison: a naive value-equivalence proxy, the per-anchor unnormalized \(|V(\hat{z}) - V(z')|\) averaged across anchors, is numerically near-identical to the Bellman residual on this sweep (per size: \(2.66 / 1.33 / 2.80 / 2.70 / 6.78\) vs.Bellman residual \(2.65 / 1.33 / 2.78 / 2.67 / 6.73\); Spearman \(+1.00\)). This is the \(\gamma\)-scaled equivalence that TD-MPC2’s small reward-prediction error (\(\mathrm{error} \in [0.028, 0.091]\), nearly two orders of magnitude below the value-error component) produces between the two quantities. The normalized value slice and the full-\(F\) operator both depart from this equivalence; the full-\(F\) operator is the one whose ranking tracks return loss (Section 4.1).

4.3 Cross-architecture sanity check↩︎

To check whether the diagnostic returns informative estimates across architectures, we compare pure-SSL LeWM (\(3\) seeds, no anchor) against the released single-task TD-MPC2 on cheetah-run, both at the matched \(5\)-step horizon and on a shared observation-level \(F\) probed identically. LeWM’s operator error on shared \(F\) across seeds is \(0.384 \pm 0.009\); single-task TD-MPC2’s is \(0.840\) (\(95\%\) CI \([0.79, 0.89]\)); both beat the persistence baseline (\(1.33\) for LeWM, \(1.40\) for TD-MPC2), and the \(95\%\) CIs are disjoint (Figure 3 in Appendix 9). A re-run with a \(1\)-hidden-layer MLP probe (matched protocol, \(\alpha\) grid-searched on the same sub-split) preserves the ordering with disjoint CIs (LeWM \(0.494 \pm 0.006\), TD-MPC2 \(0.997\), ratio \(0.495 \pm 0.006\); Appendix ¿sec:app:probe-mlp?), so the gap is not a ridge-probe-capacity artifact. Probe-asymmetry and prediction-mode controls (Appendix 9) also preserve the gap. We do not draw a broader claim about which family of model the operator metric favors from a single environment with a two-point architecture spectrum.

5 Limitations and discussion↩︎

The cross-architecture comparison is on a single environment (cheetah-run), and the architecture spectrum is two points (TD-MPC2 and LeWM) with no reconstruction-anchored midpoint such as DreamerV3 [10]; the TD-MPC2 size-sweep correlation is \(n=5\) and reported as correlational. The \(317\)M return collapse that anchors that correlation is specific to cheetah-run: on four sibling mt80 tasks the same \(317\)M checkpoint plans competently (walker-walk and reacher-easy return \(\sim 800\)–\(930\)) with correspondingly low operator error, so it is one operating point of a broader operator–return association (pooled partial Spearman \(-0.49\) across the \((\text{task}, \text{size})\) cells) rather than a globally degenerate checkpoint. The probe is a confound, since a shared \(F\) across different latent geometries must be read out through a probe; we report probe \(R^2\) with each operator-on-\(F\), cross-check on the commonly-well-probed restriction (Appendix 9), and verify that a \(1\)-hidden-layer MLP-probe variant preserves the cross-architecture ordering with disjoint CIs (Appendix ¿sec:app:probe-mlp?).

Operator-on-\(F\) is a Koopman-style pushforward comparison on a chosen observable subset [11]; [12] train a task-anchored Koopman latent for control and [13] examine Koopman invariants inside JEPA-style models. We use operator-on-\(F\) as a diagnostic on existing models, not as a training objective. We do not claim that value-equivalence [1], [2] is wrong, only that on this sweep and this environment a reward- or Bellman-residual check, as conventionally reported, is silent at the failure mode the operator metric detects; we recommend reporting both rather than one in place of the other.

Acknowledgments↩︎

We gratefully acknowledge financial support from the Canada CIFAR AI Chairs program, and compute resources from Mila (mila.quebec).

6 Operator-on-\(F\): full definition↩︎

For an anchor \((s_t,\, a_{t:t+k-1},\, s_{t+k})\) and a functional \(\phi : Z \to \mathbb{R}\) in \(F\), the per-anchor operator error is \[e_\phi(t) \;=\; \frac{\bigl|\,\phi(\hat{z}_{t+k}) \,-\, \phi(z_{t+k}')\,\bigr|}{\sigma_\phi} , \qquad \hat{z}_{t+k} = O_k(z_t,\, a_{t:t+k-1}), \quad z_{t+k}' = \mathrm{enc}(s_{t+k}),\] where \(\sigma_\phi\) is the standard deviation of \(\phi(z')\) across the anchor pool. Per-functional aggregation is the root-mean-square over anchors. The full-\(F\) value is the RMS of \(\{e_\phi\}_{\phi \in F}\), with \(F\) split into two reported groups:

• Value slice: \(F = \{r,\, V\}\), where \(r\) is the model’s reward head and \(V\) its value head.

• Full-\(F\): the value slice plus a per-anchor PCA basis on the encoded next-state geometry. We use the top-\(16\) principal directions and whiten by their singular values so each direction contributes a unit-variance functional. The PCA is fit on a held-out half of the anchor pool; the operator-on-\(F\) is evaluated on the other half. This prevents the basis from being chosen to minimize operator error on the evaluation half.

The choice of \(k\) matches each model’s native planning step: \(k = 5\) collect-steps throughout (one frameskip-\(5\) predictor call for LeWM, five open-loop dynamics steps under the recorded actions for TD-MPC2).

7 Probe protocol↩︎

The probe \(\phi_F : Z \to F\) for each observable is a ridge regression trained on a feature-standardized latent, with the regularizer \(\lambda\) chosen by grid search on a held-out sub-split. We report per-observable \(R^2\) alongside each operator-on-\(F\) number. For the cross-architecture comparison in Section 4.3 we also restrict \(F\) to the commonly-well-probed subset — those observables for which both LeWM and TD-MPC2 achieve \(R^2 \geq 0.7\), giving \(n = 7\) observations on cheetah-run. The restricted-\(F\) ratio \(\mathrm{LeWM} / \mathrm{TD\text{-}MPC2}\) is \(0.395 \pm 0.012\) across the three LeWM seeds, compared with \(0.457 \pm 0.011\) on the shared \(F\); the gap grows on the restriction, so probe asymmetry is not the source of the disjoint CIs reported in Section 4.3.

The unregularized-probe symptom noted in Section 2.2 — strongly negative \(R^2\) on some observables for the \(768\)-dim mt80-\(19\)M latent before regularization — is what motivates the \(\lambda\) grid search; the chosen \(\lambda\) values are recorded per model in the released configs.

MLP-probe ablation for the cross-architecture comparison↩︎

The reviewer-relevant question for the cross-architecture comparison in Section 4.3 is whether the LeWM/TD-MPC2 ordering depends on the probe family. We re-ran the protocol with a \(1\)-hidden-layer MLP probe (\(64\) ReLU units, weight decay \(\alpha\) grid-searched on an \(80/20\) sub-split of the train half, identical otherwise to the ridge protocol). The probe family is the only thing that changed: same compare data, same \(50/50\) train/eval split, same target \(F\).

Compared to the ridge baseline (ratio \(0.457 \pm 0.011\), Table 3), the MLP-probe ratio of \(0.495 \pm 0.006\) is slightly closer to one but the LeWM advantage and disjoint \(95\%\) CIs hold. The cross-architecture gap in Section 4.3 is therefore not a ridge-probe-capacity artifact.

8 Per-size rows for the TD-MPC2 size sweep↩︎

Table 2 reports the five-row source data behind Figure 1 and the metric-disagreement numbers cited in Section 4.2.

Across-size Spearman (full-\(F\) operator, return) \(= -0.90\) (anchor-bootstrap \(95\%\) CI \([-0.90, -0.70]\), \(P(\rho < 0) = 1.000\) across \(1000\) resamples). Leave-one-out values: dropping \(1\)M \(\to -0.80\), \(5\)M \(\to -1.00\), \(19\)M \(\to -0.80\), \(48\)M \(\to -1.00\), \(317\)M \(\to -0.80\), so the rank order is stable to single-point removal but the \(n=5\) sample remains small. Reward-prediction error has range \([0.028, 0.091]\), so a partial correlation that conditions on it reproduces the marginal value (reported \(-0.89\)); we treat this as consistency, not as an independent control. Other across-size Spearmans (anchor-bootstrap CIs over the full-\(F\) operator, comparator held at its point estimate): vs.Bellman residual \(+0.30\) (\([+0.30, +0.60]\)); vs.kernel divergence \(+0.10\) (\([+0.10, +0.30]\)); for the normalized value-slice \(F = \{r, V\}\) vs.Bellman residual \(+0.10\) (\([+0.10, +0.20]\)). Bellman residual vs. return: \(-0.10\); reward error vs.return: \(-0.30\).

Under operator-on-\(F\)’s own per-functional normalization (each channel divided by the across-anchor spread of its true pushforward), the reward slice and the value slice \(F = \{r, V\}\) track return at Spearman \(-0.90\) and \(-1.00\) respectively, as strong as the full-\(F\) operator (\(-0.90\)); their conventional unnormalized forms (raw reward error, Bellman residual) instead give \(-0.30\) and \(-0.10\). On this sweep the resolution gap between operator-on-\(F\) and the value-equivalence proxies is thus largely a normalization effect, and operator-on-\(F\)’s distinct contribution is to carry this normalized pushforward comparison onto a shared observable basis and onto models without reward or value heads (Section 4.3).

Figure 2 plots the metric-disagreement scatter (full-\(F\) operator vs.Bellman residual) referenced from Section 4.2.

9 Cross-architecture controls↩︎

Table 3 reports the three LeWM seeds against the fixed single-task TD-MPC2 checkpoint. Each cell is generated by re-running the shared protocol on a checkpoint trained from scratch with a different random seed.

As an additional control, we replace the mt80 multitask TD-MPC2 with the released single-task checkpoint and re-run the comparison; the TD-MPC2 number is essentially identical (\(0.840\) vs.\(0.833\) on mt80), so the cross-architecture gap is not a multitask handicap.

References↩︎