World Models with Predictable Long-Horizon Marginals

2026-09-26

AI Summary

This study addresses the limitation of world models, which, despite accurate single-step predictions, tend to deviate from the data distribution during long-horizon rollouts. To overcome this, the authors propose anchoring a fixed Gaussian reference distribution to constrain mean behavioral shifts, thereby ensuring marginal distribution stability over extended time horizons. Furthermore, they design a joint state-action rotation mechanism coupled with parallel noise injection to precisely preserve transition dynamics. These techniques are integrated with a Gaussian decoder, conditional action graphs, and absolute convergence bound theory to facilitate effective modeling and control. Evaluated across twelve control tasks, all rollout chains maintain distributional consistency with zero divergence over 10⁵ steps. This work effectively resolves the persistent challenge of long-horizon distribution shift in world models.

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.

Problem

• World Models
• Long-Horizon Rollouts
• Stationary Distribution
• Data Distribution
• Occupancy Preservation

Innovation

• World Models
• Long-Horizon Marginals
• Stationary Distribution
• Behavior Occupancy
• Gaussian Reference