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

Oct 27, 2026

Understanding the limits of latent world models

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.

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

Ask PDF

Aug 24, 2026

Correcting a learned physical invariant improves world-model rollouts

World models can predict video without learning dynamics that they reliably preserve. We test whether a frozen DreamerV3 trained only on pendulum video learns a scalar that its own latent transition treats as approximately conserved. A label-free search recovers the same energy-like invariant across independently trained conservative models, while the same procedure finds no comparable invariant in matched damped models. During autonomous rollouts, this quantity drifts. Projecting the latent state back toward its initial level set reduces rollout error in all three conservative models, whereas matched random constraints usually increase it. These results distinguish a dynamically meaningful invariant from a merely decodable correlate and reveal a concrete failure mode: a world model can learn a physical constraint from pixels yet violate that constraint when it imagines forward.

Richard Bao

Independent Researcher

World Models World Model Learning Latent World Models World Model Evaluation Scientific Discovery

Ask PDF

Jun 11, 2026

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

Scale buys interpolation; structure buys certifiable transfer. A world model's average error does not say whether a particular rollout can be trusted, or for how long. For equivariant latent world models we give a predictability certificate: a computable region spanning configuration, horizon, and resolution. Under exact equivariance, rollout error is invariant over the monoid generated by k primitive symmetries and is certified from the k generators (Theorem A); universal orbit-flatness over equivariant targets characterizes equivariance at the function level (Lemma 2), so an unconstrained architecture cannot certify the property by construction. Approximate orbit-transfer defects propagate by the finite-time Lyapunov spectrum (Theorem B): expanding channels give a logarithmic horizon

T_j(ε)\sim\log(1/ε)/λ_j

, neutral channels accumulate recurrent defect linearly, and contracting channels accumulate a bounded nonzero floor. Exact conserved charge values are certified to all horizons only at zero defect; with one-step defect

η

, charge-value error grows at most as

Tη

. Empirically, on a 40-dimensional learned model a

\mathbb{Z}_N

-equivariant network recovers the full Lyapunov spectrum (

R^2=0.98

0.99

) where dense and recurrent baselines fail. A cone/adapted-metric certificate reads an a-priori horizon off the model's own Jacobian, tight on uniformly hyperbolic dynamics and self-abstaining elsewhere; the resulting horizon improves a budgeted re-observation decision. For public non-equivariant world models the tangent spectrum gives a training-free candidate horizon, paired with a held-out divergence cross-check that abstains or corrects when the learned loop over-promises.

Hongbo Wang

Department of Mathematics, Stony Brook University, Stony Brook, NY 11794, USA

Equivariant Neural Networks Latent World Models World Model Evaluation

Ask PDF

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