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

**Authors:** Arjun Parthasarathy, Nimit Kalra, Rohun Agrawal, Yann LeCun, Oumayma Bounou, Pavel Izmailov, Micah Goldblum  
**arXiv:** [2512.09929](https://arxiv.org/abs/2512.09929)  
**Date:** 2025-12-10  
**Summary page:** https://www.summarizepaper.com/en/paper/2512.09929v1/

## Key points

- Understanding how an agent's actions impact the environment is crucial for prediction and planning in robotic tasks.
- Traditional methods use analytical models derived from known principles, while learning-based approaches directly infer models from data.
- World models predict the next state based on the current state and an action, enhancing planning with model predictive control (MPC).
- Gradient-based planning offers a computationally efficient alternative to traditional MPC methods but historically lags behind other approaches in performance.
- Enhanced training methods for world models improve gradient-based planning efficiency by bridging the train-test gap through data synthesis techniques.
- The proposed approach outperforms or matches classical gradient-free cross-entropy method (CEM) in various object manipulation and navigation tasks within a 10% time budget.
- By optimizing training strategies and closing the train-test gap, improved performance is demonstrated in gradient-based planning with world models, advancing robotic task anticipation and planning efficiency.

## Summary

In the field of robotic tasks, understanding how an agent's actions impact the environment is crucial for prediction and planning. Traditional methods rely on analytical models derived from known principles, while learning-based approaches use data to directly infer models. This allows for a better capture of complex dynamics and increased robustness to uncertainty. World models have emerged as a powerful tool in this regard, predicting the next state based on the current state and an action. This paper focuses on improving gradient-based planning with world models paired with model predictive control (MPC). While traditional MPC methods can be slow due to search algorithms or iterative optimization, gradient-based planning offers a more computationally efficient alternative. However, the performance of gradient-based planning has historically lagged behind other approaches. The authors propose enhanced training methods for world models to enable more efficient gradient-based planning. They highlight that while world models are typically trained on next-state prediction objectives, they are used at test-time to estimate sequences of actions. The goal is to bridge this train-test gap by introducing train-time data synthesis techniques that significantly enhance gradient-based planning using existing world models. The results show that their approach outperforms or matches the classical gradient-free cross-entropy method (CEM) across various object manipulation and navigation tasks within a 10% time budget. By closing the train-test gap and optimizing training strategies, the authors demonstrate improved performance in gradient-based planning with world models. This showcases advancements in robotic task anticipation and planning efficiency.