Operator-on-F complements value-equivalence: a planning-time diagnostic for latent world models

Abstract: World-model evaluation for model-based reinforcement learning typically asks whether the learned model predicts reward and value well, which can leave planning-relevant errors in the model's latent rollouts unmeasured. We introduce a complementary diagnostic, operator-on-F, that compares a model's k-step latent pushforward to the environment's on an observable subset F, using the model's own predictor. On a TD-MPC2 size sweep over cheetah-run, reward-prediction error stays within [0.028, 0.091] for every model size - only about 3x variation - so an unnormalized reward-fit check has narrow resolution to distinguish them; the (unnormalized) Bellman residual and reward error themselves have weak relationships with return (Spearman -0.10 and -0.30). Operator error spans 0.28 to 2.62 over the same sizes. At 317M the operator error is 2.62 - an order of magnitude above the 0.28-0.36 cluster - and the planning return collapses to 0.9, while reward-prediction error (0.091) is the highest of the five but stays within the same small [0.028, 0.091] range as the rest of the sweep. The rank correlation between operator error and return loss is -0.90 (anchor-bootstrap 95% CI [-0.90, -0.70] at n=5 sizes; leave-one-out removal of any single size leaves it at -0.80 or stronger). The operator also returns informative, architecture-discriminating estimates in a cross-architecture comparison between TD-MPC2 and a pure-SSL latent world model. The operator diagnostic complements value-equivalence rather than replacing it.

• The paper introduces operator-on-F, a diagnostic that quantifies k-step latent rollout divergence projected onto function class F to reveal planning-critical errors.
• It demonstrates that operator-on-F error, unlike traditional reward-prediction metrics, strongly correlates with planning return, especially across diverse model sizes and architectures.
• The study validates the approach on TD-MPC2 and LeWM models, highlighting its potential for improving safety and model selection in model-based reinforcement learning.

Operator-on-$F$ as a Complementary Diagnostic for Latent World Models in Model-Based RL

Introduction

In model-based reinforcement learning (MBRL), a central question lies in the effective evaluation of learned world models. Traditional metrics, such as reward-prediction error and value-equivalence (e.g., Bellman residual), are often used as proxies for planning capability. However, these approaches may overlook critical planning-relevant discrepancies arising in the latent space rollouts performed by the model. This paper introduces operator-on-$F$, a diagnostic tailored to address such silent failure modes. The diagnostic quantifies the divergence between the model's $k$-step latent rollout and the environment’s dynamics, projected onto an observable subset $F$. The diagnostic is shown to capture failure modes missed by standard metrics and to provide additional discriminative power in comparisons across model architectures.

Operator-on-$F$: Definition and Motivation

Operator-on-$F$ measures the normalized error between the model's $k$-step latent trajectory and the environment's ground truth, mapped onto a function class $F$. Typically, $F$ covers reward, value, and projections onto the principal axes of the next-state latent geometry (via PCA). This is executed using probes (ridge regression or shallow MLPs) trained on the latent encodings, evaluated at the model’s native planning horizon (e.g., 5 steps).

Unlike value-equivalence, which collapses information onto the reward or value head, operator-on-$F$ surfaces discrepancy at the level of the latent rollouts, potentially exposing dynamics misalignments that remain silent to reward prediction.

Empirical Findings: TD-MPC2 Model Size Sweep

The method is evaluated on a sweep of TD-MPC2 world models across five sizes (1M–317M parameters) in the DeepMind Control Suite cheetah-run task. The findings are notable:

• Reward-prediction error varies narrowly ($F$0), limiting its discriminative resolution across model sizes.
• Operator-on-$F$1 error spans a much wider range ($F$2 to $F$3) and exhibits a strong negative Spearman rank correlation ($F$4) with actual planning return, robust to leave-one-out removal of any single model in the sweep.
• The largest model (317M) demonstrates an operator-on-$F$5 error an order of magnitude above the rest (2.62), coinciding with a catastrophic drop in planning return ($F$6), even as the reward-prediction error is only slightly higher than smaller sizes.

  Figure 1: The left panel displays the full-$F$7 operator-on-$F$8 error and executed return against model size, while the right panel shows reward-prediction error. Only operator error tracks the return collapse for the largest model.

These results underscore the diagnostic’s potential to flag planning-critical degradation not apparent in conventional, unnormalized value or reward checks.

Disagreement with Value-Equivalence Proxies

A key analysis contrasts operator-on-$F$9 with the Bellman residual and the unnormalized value-only error. While these two conventional value-equivalence proxies are nearly identical (Spearman $k$0), operator-on-$k$1 assigns a very different ranking to models (Spearman $k$2 with Bellman residual), and only the operator error maintains strong alignment with empirical return loss.

Figure 2: Scatter plot comparing full-$k$3 operator-on-$k$4 error with Bellman residual for TD-MPC2 model sizes. The two metrics disagree, especially at the 317M checkpoint.

This confirms that operator-on-$k$5 offers orthogonal information in model selection and evaluation relative to standard metrics.

Cross-Architecture Comparison: TD-MPC2 vs LeWM

To explore the generality of the diagnostic, the paper compares a pure self-supervised latent world model (LeWM, no reward or value head) with a single-task TD-MPC2 model, matched on cheetah-run and a 5-step rollout horizon. Operator-on-$k$6 is computed over a shared observation-level $k$7. Both models outperform a persistence baseline, but LeWM achieves lower operator error ($k$8) than TD-MPC2 ($k$9; $F$0 CI $F$1), with entirely disjoint confidence intervals across seeds.

Figure 3: Operator error on shared observation $F$2 for LeWM and TD-MPC2 at 5-step horizon, showing significant separation and both outperforming persistence.

MLP probe ablations and probe-asymmetry controls preserve the ordering and gap, confirming robustness to probe class.

Practical and Theoretical Implications

The introduction of operator-on-$F$3 as a diagnostic addresses a fundamental blind spot in prevalent MBRL metrics: the capability to surface latent dynamics mismatches that may directly undermine planning performance, even when conventional reward and value losses are low. The diagnostic:

• Provides actionable insight into when a model is unsafe for planning, complementing reward- and value-based diagnostics;
• Enables cross-architecture comparisons, including SSL-based latent models lacking task-specific decoders;
• Reveals that unnormalized reward-fit checks and Bellman-residual metrics have limited power to identify planning-relevant model failures.

The results also emphasize the importance of normalization and semantic projection in latent space evaluation—not just for interpretability but also for alignment with deployment-time metrics such as planning return.

From a methodological perspective, operator-on-$F$4 relates to Koopman operator theory as a pushforward measured over observable functionals rather than reward or policy values, aligning with recent work exploring Koopman embeddings for control.

Limitations and Future Directions

While the findings are robust within the cheetah-run environment and across the specific models analyzed, broader claims about generality require sweeps over a more diverse set of tasks, architectures, and anchoring strategies. In addition, extending the evaluation to include reconstruction-anchored models (e.g., DreamerV3) would provide insight into the positioning of operator-on-$F$5 in a wider architectural landscape.

Potential future research directions include:

• Systematic exploration of the diagnostic’s sensitivity and specificity across tasks and action spaces.
• Integration of operator-on-$F$6 as an in-training regularizer or early-stopping criterion.
• Theoretical analysis of its relationship with controllability, observability, and the information geometry of the latent dynamics in deep MBRL.

Conclusion

This work establishes operator-on-$F$7 as a critical supplement to conventional value-equivalence diagnostics in model-based RL. By offering direct visibility into latent-planning dynamics errors, it provides practitioners and researchers with a tool to diagnose and mitigate otherwise silent catastrophic failures. The operator-on-$F$8 metric is not a replacement for value-equivalence but a necessary complement, and its adoption promises sharper evaluation, cross-model comparability, and ultimately, increased safety in model-based planning frameworks.