World Models with Predictable Long-Horizon Marginals

Yuhao Du; Shunian Chen

cs.AISep 26, 2026

World Models with Predictable Long-Horizon Marginals

Authors: Yuhao Du, Shunian Chen

Organizations: The Chinese University of Hong Kong, Shenzhen · Shenzhen Loop Area Institute

Abstract

Accurate one-step predictions do not ensure that a world model's rollouts retain the data distribution. We make the model's decoded stationary law explicit by learning a decoder of a fixed Gaussian reference and constraining the behaviour-averaged transition to preserve that reference. For controlled systems, a joint transition uses a conditional action chart to preserve behaviour occupancy without requiring invariance at each fixed action. Joint state--action rotations and parallel Gaussian noise give an exactly preserving transition with a tractable conditional density. We derive an absolute convergence bound from finite initialization banks and control departure from the reference through conditional action-space divergence. Across $216$ fitted pixel checkpoints on twelve control tasks, the occupancy model with a reference mixture retains every evaluated chain at $10^5$ steps in all $36$ task--seed cells, with a rollout-minus-reference energy-statistic difference of $-0.0002\pm0.0003$ (training-seed standard error). Each of the four nonpreserving comparison arms loses chains, although the Gaussian arm is more accurate at ten steps. An offline DreamerV3 reference also achieves better short-horizon accuracy. These results distinguish three properties of a world model: the distribution it approaches, the rate of approach, and the conditional dynamics it learns.

Figures & tables

Figure 1: A transition with an exactly preserved Gaussian reference. Latent state

z

and charted action

\alpha

feed joint rotations under behaviour

\pi_{b}

, a parallel Gaussian noise stage with

s_{i}^{2}=1-c_{i}^{2}

, and a preserving post-map, so the Gaussian reference holds at every parameter value. Decoding the reference gives the static observation law

p_{\theta}

. The rollout discrepancy from data equals the static discrepancy plus a bounded signed transient

r_{T}

(Section 4 ).

Figure 2: Static reference samples predict model-specific long-horizon discrepancies. (a) Seven Lorenz models span a

31\times

range of predicted energy statistics;

w

is the marginal-loss weight. The inset enlarges the two smallest values; an open marker denotes a value below evaluation resolution. (b) One capped Reacher checkpoint, using Euclidean distance in 256-dimensional frozen-autoencoder coordinates: the static estimate (

8

Monte Carlo replicates), resampling (eight 2,048-point clouds per reference block; SE over eight blocks), and data- and prior-initialized

10^{5}

-step rollouts. The offset is consistent with static-estimate sampling variability.

Figure 3: Finite-start and policy-shift checks for the parallel transition. (a) A 16-point initialization bank supplies the moment in equation 9 . (b) A known policy shift has conditional action KL

\kappa=0.0334

. Dots are Monte Carlo gaps for

f(z)=\exp(-\|z\|^{2}/16)

, with 8,192 independent chains per horizon. Solid curves are population bounds; dashed curves add a simultaneous

99\%

Hoeffding allowance. (c) Median forward/backward and sampling times per batch of 32, on four CPU threads, with min–max whiskers over five repetitions (prior only). These are randomized-kernel diagnostics without dynamics fitting (Appendix G ).

Transition / objective	Survival	Obs. NLL $\downarrow$	10-step MSE $\downarrow$	$D_{10^{5}}-D_{\rm ref}$
Occupancy	1.0000	$0.6843\!\pm\!0.6832$	$0.1271\!\pm\!0.0211$	$-0.0002\!\pm\!0.0003$
Conditioner only	0.9999	$\mathbf{0.5648\!\pm\!0.5422}$	$0.1416\!\pm\!0.0247$	Escape
Gaussian	0.5149	$1.7013\!\pm\!1.1747$	$\mathbf{0.0708\!\pm\!0.0218}$	Escape
Conditional RealNVP	0.2394	$3.1391\!\pm\!1.5809$	$0.1119\!\pm\!0.0264$	Escape
Latent overshooting	0.3380	$7.0333\!\pm\!1.0566$	$0.1692\!\pm\!0.0212$	Escape
Invariant penalty	0.4399	$4.0644\!\pm\!1.6440$	$0.1120\!\pm\!0.0327$	Escape

Table 1: Conditional accuracy and long-horizon stability across recurrent models. Equal-task means, twelve tasks

\times

three paired seeds; NLL, MSE, and gaps show training-seed SE. NLL is a posterior-initialized particle score under recorded actions, in nats per pixel channel.

D_{\rm ref}

scores decoded Gaussian reference states, a limiting level only for preserving models under their convergence assumptions. “Escape”: no gap estimated after censoring. The last row scales the fitted occupancy model’s output. Bold marks best primary-arm means; the signed gap has no directional optimum.

Figure 4: Absolute marginal discrepancy and static-prediction agreement. (a) Occupancy reference energy statistics against data, using normalized RMS pixel distance. (b) Rollout-minus-reference discrepancies against the same data. Points: three training seeds; squares: task means

\pm

SE. The equal-task gap is

-0.0002

, with SE

0.0003

for independent task–seed cells and

0.0006

when grouping the reused seeds (Appendix C.3 ).

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: All four nonpreserving arms lose chains at

10^{5}

steps. Paired survival contrasts (candidate minus occupancy), with nominal

95\%

paired training-seed percentile-bootstrap intervals over twelve tasks and three seeds (10,000 draws, fixed task suite). The right column gives equal-task survival. Occupancy retains every evaluated chain.

Figure 6: Marginal agreement, conditional fit, and planning are distinct outcomes. (a) The

10^{5}

-step rollout-minus-static energy gap versus the static score for occupancy checkpoints; vertical bars show

\pm 1

paired evaluation SE of that gap; horizontal uncertainty is not shown. (b) Observation NLL versus marginal gap for cells with uncensored estimates; 102 censored cells are excluded from this scatter and retained in survival summaries. (c) Validation-selected gain return contrasts with nominal paired training-seed percentile-bootstrap intervals; failed comparisons remain unestimated.

D_{\rm ref}

is a limiting level for preserving models under their convergence assumptions.

Figure 7: The native DreamerV3 reference is the stronger short-horizon predictor; survival is tied at the shared evaluation horizon. Native

-

occupancy contrasts with nominal 95% paired training-seed percentile-bootstrap intervals (4 priority tasks

\times

3 seeds, horizon 1000): observation NLL, 10-step MSE, and rollout energy all favor the native model with intervals excluding zero, while survival is identical (1.0 vs. 1.0). Bar: native planned return

0.0734\pm 0.0013

with all 12 plan cells completed (a full-grid occupancy return is unavailable because of planner failures, Fig. 8 ). The native

10^{5}

-step marginal is unmeasured, and no Gaussian-reference identity exists for its categorical state.

Figure 8: Planner completion varies across tasks and transition families. Each cell counts failures among three training seeds: the CEM planner has fewer finite candidates than required elites. The margin strip counts $38$ ablation failures. The selected gain shares the same $19$ failed task–seed cells as occupancy, and $12$ planning-sweep jobs fail, giving $172$ failures among $348$ planning evaluations. Primary failure counts range from $14/36$ for Gaussian to $19/36$ for occupancy and conditioner only.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 419 papers

Understanding the limits of latent world models
Oct 27, 2026Rodrigo Gonzalez Laiz, Steffen SchneiderWorld Models Latent World Models
Helmholtz Munich Munich Center for Machine Learning (MCML)
Correcting a learned physical invariant improves world-model rollouts
Aug 24, 2026Richard BaoWorld Models World Model Learning
Independent Researcher
Certified World Models: Predictability Across Configuration, Horizon, and Resolution
Jun 11, 2026Hongbo WangEquivariant Neural Networks Latent World Models
Department of Mathematics, Stony Brook University, Stony Brook, NY 11794, USA

Oct 27, 2026J/K move · Enter open · S save

Understanding the limits of latent world models

Rodrigo Gonzalez Laiz, Steffen Schneider

Helmholtz Munich Munich Center for Machine Learning (MCML)

World Models Latent World Models Latent Dynamics Modeling World Model Evaluation Action-Conditioned World Models

Open paper AskPDF arXiv

World models are widely proposed as core components for enabling agents to reason over long temporal and spatial horizons. Yet current evaluations often emphasize task performance while providing limited insight into whether learned world models acquire predictive latent dynamics that reflect the underlying structure of the environment. Here we introduce the concept of faithful latent world models as an equivariance condition between the environment's transition map and the learned latent dynamics. We propose latent rollout consistency (LRC) as a metric to evaluate action-conditioned latent prediction. Strikingly, we find that across the studied state of the art latent world models, none solves the task by actually faithfully modeling the dynamics of the environment.

Abstract

World Model Learning

Explore similar work

Understanding the limits of latent world models

Correcting a learned physical invariant improves world-model rollouts

Certified World Models: Predictability Across Configuration, Horizon, and Resolution