# Closing the Train-Test Gap in World Models for Gradient-Based Planning

arXiv: https://arxiv.org/abs/2512.09929 · HTML: https://mlsanity.com/paper/2512.09929/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 52 of its 81 references that are in the map.
- First posted on arXiv 10 Dec 2025.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 30 of 52)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Intriguing properties of neural networks](https://mlsanity.com/paper/1312.6199/) (2013)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Fast is better than free: Revisiting adversarial training](https://mlsanity.com/paper/2001.03994/) (2020)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Dynamics-Aware Unsupervised Discovery of Skills](https://mlsanity.com/paper/1907.01657/) (2019)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Unsupervised Learning for Physical Interaction through Video Prediction](https://mlsanity.com/paper/1605.07157/) (2016)
- [LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics](https://mlsanity.com/paper/2511.08544/) (2025)
- [Deep Dynamics Models for Learning Dexterous Manipulation](https://mlsanity.com/paper/1909.11652/) (2019)
- [Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data](https://mlsanity.com/paper/1605.06432/) (2016)
- [Learning to Poke by Poking: Experiential Learning of Intuitive Physics](https://mlsanity.com/paper/1606.07419/) (2016)
- [Model-Based Imitation Learning for Urban Driving](https://mlsanity.com/paper/2210.07729/) (2022)
- [Nymeria: A Massive Collection of Multimodal Egocentric Daily Motion in the Wild](https://mlsanity.com/paper/2406.09905/) (2024)

## Built on by (most cited first, 11 of 11)

- [AdaJEPA: An Adaptive Latent World Model](https://mlsanity.com/paper/2606.32026/) (2026)
- [Learning Invariant Visual Representations for Planning with Joint-Embedding Predictive World Models](https://mlsanity.com/paper/2602.18639/) (2026)
- [World Model Control by Trajectory Reachability Metrics](https://mlsanity.com/paper/2605.22164/) (2026)
- [Parallel Stochastic Gradient-Based Planning for World Models](https://mlsanity.com/paper/2602.00475/) (2026)
- [Slot-MPC: Goal-Conditioned Model Predictive Control with Object-Centric Representations](https://mlsanity.com/paper/2605.14937/) (2026)
- [Lifting Embodied World Models for Planning and Control](https://mlsanity.com/paper/2604.26182/) (2026)
- [Music-JEPA: Learning a World Model of Sound from Action](https://mlsanity.com/paper/2607.22000/) (2026)
- [Action from Adjacent Set in Physical Space Outperforms the Best Prediction in World Models](https://mlsanity.com/paper/2607.23602/) (2026)
- [WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systems](https://mlsanity.com/paper/2603.14392/) (2026)
- [LePlanner: An Iterative Amortized Controller For World Models](https://mlsanity.com/paper/2609.13845/) (2026)
- [Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models](https://mlsanity.com/paper/2610.07540/) (2026)

## Abstract (arXiv)

World models paired with model predictive control (MPC) can be trained offline on large-scale datasets of expert trajectories and enable generalization to a wide range of planning tasks at inference time. Compared to traditional MPC procedures, which rely on slow search algorithms or on iteratively solving optimization problems exactly, gradient-based planning offers a computationally efficient alternative. However, the performance of gradient-based planning has thus far lagged behind that of other approaches. In this paper, we propose improved methods for training world models that enable efficient gradient-based planning. We begin with the observation that although a world model is trained on a next-state prediction objective, it is used at test-time to instead estimate a sequence of actions. The goal of our work is to close this train-test gap. To that end, we propose train-time data synthesis techniques that enable significantly improved gradient-based planning with existing world models. At test time, our approach outperforms or matches the classical gradient-free cross-entropy method (CEM) across a variety of object manipulation and navigation tasks in 10% of the time budget.