World models keep long-run predictions true to real data distribution
World Models with Predictable Long-Horizon Marginals
Artificial Intelligence
Summary
Predicting one step into the future doesn’t mean a model will stay accurate when it forecasts many steps ahead. The authors created a method that makes long-term predictions better by ensuring the overall distribution of predicted data stays close to reality. They do this by using a special approach that keeps the model’s predictions anchored to a known reference distribution, even as it simulates many actions over time. Their method was tested on many tasks and showed very stable performance for very long predictions.
What this means in practice
- •For robotics engineers: Build control systems that maintain realistic state distributions during long simulated action sequences for safer robot planning.
- •For game ai developers: Create game AI simulations that stay consistent with real gameplay state distributions over extended play sessions.
Tested on simulated data.
Authors
Yuhao Du, Shunian Chen
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.