mlsanityData to 7 Oct 2026

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

arXiv 2512.09929 ↗First posted 10 Dec 2025cs.LG cs.ROcited by 10

On the map: World models → Latent predictive world models

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.

Key facts. Cited by 10 papers in the map (data through 7 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.

Builds on

52 of 81 references.

Closest ancestors · 3

Same direction, newest first

  1. 2024Revisiting Feature Prediction for Learning Visual Representations from Video15 Feb
  2. 2023Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture19 Jan
  3. 2018Learning Latent Dynamics for Planning from Pixels12 Nov

Other references · 49

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  3. 2014Explaining and Harnessing Adversarial Examplescited by 7,287
  4. 2017Towards Deep Learning Models Resistant to Adversarial Attackscited by 6,326
  5. 2013Intriguing properties of neural networkscited by 5,114
  6. 2017Neural Discrete Representation Learningcited by 4,650
  7. 2023Diffusion Policy: Visuomotor Policy Learning via Action Diffusioncited by 2,730
  8. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,598
  9. 2018World Modelscited by 1,444
  10. 2020D4RL: Datasets for Deep Data-Driven Reinforcement Learningcited by 1,406
  11. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
  12. 2023Mastering Diverse Domains through World Modelscited by 1,193
37 more
  1. 2018DeepMind Control Suitecited by 935
  2. 2020Fast is better than free: Revisiting adversarial trainingcited by 723
  3. 2019Model-Based Reinforcement Learning for Ataricited by 495
  4. 2021Implicit Behavioral Cloningcited by 442
  5. 2022Transformers are Sample-Efficient World Modelscited by 317
  6. 2019Dynamics-Aware Unsupervised Discovery of Skillscited by 308
  7. 2017Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuningcited by 288
  8. 2016Deep Visual Foresight for Planning Robot Motioncited by 276
  9. 2016Unsupervised Learning for Physical Interaction through Video Predictioncited by 256
  10. 2025LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristicscited by 227
  11. 2019Deep Dynamics Models for Learning Dexterous Manipulationcited by 200
  12. 2016Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Datacited by 154
  13. 2016Learning to Poke by Poking: Experiential Learning of Intuitive Physicscited by 152
  14. 2022Model-Based Imitation Learning for Urban Drivingcited by 149
  15. 2024Nymeria: A Massive Collection of Multimodal Egocentric Daily Motion in the Wildcited by 109
  16. 2018Imitating Latent Policies from Observationcited by 101
  17. 2022Accelerated Policy Learning with Parallel Differentiable Simulationcited by 92
  18. 2020Objective Mismatch in Model-based Reinforcement Learningcited by 87
  19. 2023Diffusion Models for Reinforcement Learning: A Surveycited by 84
  20. 2024World Models for Autonomous Driving: An Initial Surveycited by 79
  21. 2022RoboCraft: Learning to See, Simulate, and Shape Elasto-Plastic Objects with Graph Networkscited by 69
  22. 2024AdaptiGraph: Material-Adaptive Graph-Based Neural Dynamics for Robotic Manipulationcited by 52
  23. 2022StretchBEV: Stretching Future Instance Prediction Spatially and Temporallycited by 47
  24. 2024DINO-Foresight: Looking into the Future with DINOcited by 44
  25. 2024Policy-Guided Diffusioncited by 40
  26. 2018Augmenting Physical Simulators with Stochastic Neural Networks: Case Study of Planar Pushing and Bouncingcited by 36
  27. 2020Model-Predictive Control via Cross-Entropy and Gradient-Based Optimizationcited by 32
  28. 2021Model-Based Offline Planning with Trajectory Pruningcited by 30
  29. 2022DaXBench: Benchmarking Deformable Object Manipulation with Differentiable Physicscited by 22
  30. 2023SoftZoo: A Soft Robot Co-design Benchmark For Locomotion In Diverse Environmentscited by 18
  31. 2023Gradient-based Planning with World Modelscited by 18
  32. 2019Learning Dynamics Model in Reinforcement Learning by Incorporating the Long Term Futurecited by 17
  33. 2024Video Representation Learning with Joint-Embedding Predictive Architecturescited by 14
  34. 2019Robust or Private? Adversarial Training Makes Models More Vulnerable to Privacy Attackscited by 10
  35. 2024Exploring the limits of Hierarchical World Models in Reinforcement Learningcited by 9
  36. 2021Deformable Linear Object Prediction Using Locally Linear Latent Dynamicscited by 7
  37. 2025State-Aware Perturbation Optimization for Robust Deep Reinforcement Learningcited by 7

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 10.

10 papers.

  1. 2026AdaJEPA: An Adaptive Latent World Modelcited by 16
  2. 2026Learning Invariant Visual Representations for Planning with Joint-Embedding Predictive World Modelscited by 11
  3. 2026World Model Control by Trajectory Reachability Metricscited by 7
  4. 2026Parallel Stochastic Gradient-Based Planning for World Modelscited by 6
  5. 2026Slot-MPC: Goal-Conditioned Model Predictive Control with Object-Centric Representationscited by 3
  6. 2026WestWorld: A Knowledge-Encoded Scalable Trajectory World Model for Diverse Robotic Systemscited by 1
  7. 2026Lifting Embodied World Models for Planning and Controlcited by 1
  8. 2026Music-JEPA: Learning a World Model of Sound from Actioncited by 1
  9. 2026Action from Adjacent Set in Physical Space Outperforms the Best Prediction in World Modelscited by 1
  10. 2026LePlanner: An Iterative Amortized Controller For World Modelscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026Unifying Object-Centric World Models and Diffusion Policy: A Hierarchical Framework for Multi-Stage Robotic Tasks
  2. 2025Enter the Void - Planning to Seek Entropy When Reward is Scarce
  3. 2026Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning
  4. 2020World Model as a Graph: Learning Latent Landmarks for Planning
  5. 2026Parallel Predictive World Models for Accurate and Efficient Long-Horizon Planning
  6. 2026Sparse Planning in Visual World Models via Cost Gradients
  7. 2026The Planning Limits of Latent World Models
  8. 2026IMWM: Intuition Models Complement World Models for Latent Planning