Beyond Myopic World Models: Long-Horizon End-to-End Training for Direct Future Prediction
2026-08-07 • Machine Learning
Machine Learning
AI summaryⓘ
The authors found that training world models step-by-step makes it hard to predict far into the future because small errors build up over time. They propose a new model called DPWM that looks at the whole sequence of actions at once and predicts only the final outcome without repeating predictions over each step. This approach improves accuracy for long-term predictions, especially when looking far ahead. They also show that changing the training goal to focus on the final result helps other models too, highlighting that the training method matters more than the model type.
world modelslong-horizon predictionend-to-end trainingrecursive rolloutendpoint predictionembeddingautoregressive modelsgradient propagationcontinuous controlpixel-based benchmarks
Authors
Xinyi Li, Zaishuo Xia, Chenjie Hao, Yubei Chen
Abstract
World models are expected to support imagination over extended temporal horizons, yet most are still trained through local few-step prediction objectives and deployed by recursively rolling out their own predictions. This creates a fundamental mismatch: few-step losses optimize local transition fidelity, while long-horizon prediction depends on how errors and gradients propagate through the entire trajectory. As a result, transitions with different downstream influence on the endpoint are treated uniformly during training, and small local errors are amplified through recursive inference. We argue that long-horizon accuracy is better achieved by optimizing directly, through an end-to-end endpoint prediction objective. To instantiate this paradigm, we introduce the Direct Prediction World Model (DPWM), a non-recursive architecture that compresses an action sequence of arbitrary length into a single embedding and predicts the endpoint observation in a single forward pass. This design avoids recurrent rollout in both prediction and gradient propagation, making long-horizon end-to-end training practical at horizons where unrolled autoregressive training becomes unstable. Empirically, DPWM substantially improves long-horizon endpoint prediction over recursive world-model baselines on continuous-control and pixel-based benchmarks, with larger gains as the prediction horizon increases. We further show that recurrent baselines benefit similarly when retrained with the same long-horizon endpoint objective, supporting our central claim that the training objective, rather than the particular backbone choice, is the main driver of long-horizon prediction accuracy. Our results suggest that world models can benefit from being trained and evaluated at the temporal scales where they are ultimately used, shifting the focus from local transition modeling toward long-horizon predictive accuracy.