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

• URL: http://arxiv.org/abs/2512.09929v1
• Date: Wed, 10 Dec 2025 18:59:45 GMT
• Title: Closing the Train-Test Gap in World Models for Gradient-Based Planning
• Abstract summary: We propose improved methods for training world models that enable efficient gradient-based planning.<n>At test time, our approach outperforms or matches the classical gradient-free cross-entropy method.
• Score: 64.36544881136405
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Dream-MPC: Gradient-Based Model Predictive Control with Latent Imagination [17.4778210092639]
  State-of-the-art model-based Reinforcement Learning (RL) approaches either use gradient-free, population-based methods for planning, learned policy networks, or a combination of policy networks and planning.<n>We propose Dream-MPC, a novel approach that generates few candidate trajectories from a rolled-out policy and optimize each trajectory by gradient ascent using a learned world model.<n>Our results on 24 continuous control tasks show that Dream-MPC can significantly improve the performance of the underlying policy and can outperform gradient-free MPC and state-of-the-art baselines.
  arXiv  Detail & Related papers  (2026-05-06T07:13:11Z)
• Evolution of Optimization Methods: Algorithms, Scenarios, and Evaluations [98.44542103979735]
  First-order gradient descent methods, epitomized by gradient descent (SGD) and Adam, serve as the cornerstone of modern training pipelines.<n>Large-scale model training, stringent privacy requirements, and distributed learning paradigms expose critical limitations in these conventional approaches regarding privacy protection and memory efficiency.<n>We retrospectively analyze the evolutionary trajectory of deep learning optimization algorithms and present a comprehensive empirical evaluation of mainstreams across diverse model architectures and training scenarios.<n>We distill key emerging trends and fundamental design trade-offs, pinpointing promising directions for future research.
  arXiv  Detail & Related papers  (2026-04-14T17:01:36Z)
• Parallel Stochastic Gradient-Based Planning for World Models [39.699893143984916]
  We propose a robust and highly parallelizable planner that leverages the differentiability of the learned world model for efficient optimization.<n>Our method treats states as optimization variables ("virtual states") with soft dynamics constraints, enabling parallel and easier optimization.<n>Our planner, which we call GRASP (GradAxed Planner), can be viewed as a valid version of a non-condensed or collocation-based optimal controller.
  arXiv  Detail & Related papers  (2026-01-31T02:57:47Z)
• Bounding Distributional Shifts in World Modeling through Novelty Detection [15.354352209595973]
  We use a variational autoencoder as a novelty detector to ensure that proposed action trajectories during planning do not cause the learned model to deviate from the training data distribution.<n>The proposed method improves over state-of-the-art solutions in terms of data efficiency.
  arXiv  Detail & Related papers  (2025-08-08T07:42:14Z)
• Optimizing ML Training with Metagradient Descent [69.89631748402377]
  We introduce an algorithm for efficiently calculating metagradients -- gradients through model training -- at scale.<n>We then introduce a "smooth model training" framework that enables effective optimization using metagradients.
  arXiv  Detail & Related papers  (2025-03-17T22:18:24Z)
• Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models [79.2162092822111]
  We systematically evaluate reinforcement learning (RL) and control-based methods on a suite of navigation tasks.<n>We employ a latent dynamics model using the Joint Embedding Predictive Architecture (JEPA) and employ it for planning.<n>Our results show that model-free RL benefits most from large amounts of high-quality data, whereas model-based planning generalizes better to unseen layouts.
  arXiv  Detail & Related papers  (2025-02-20T18:39:41Z)
• Unlearning as multi-task optimization: A normalized gradient difference approach with an adaptive learning rate [105.86576388991713]
  We introduce a normalized gradient difference (NGDiff) algorithm, enabling us to have better control over the trade-off between the objectives.<n>We provide a theoretical analysis and empirically demonstrate the superior performance of NGDiff among state-of-the-art unlearning methods on the TOFU and MUSE datasets.
  arXiv  Detail & Related papers  (2024-10-29T14:41:44Z)
• Simple Hierarchical Planning with Diffusion [54.48129192534653]
  Diffusion-based generative methods have proven effective in modeling trajectories with offline datasets. We introduce the Hierarchical diffuser, a fast, yet surprisingly effective planning method combining the advantages of hierarchical and diffusion-based planning. Our model adopts a "jumpy" planning strategy at the higher level, which allows it to have a larger receptive field but at a lower computational cost.
  arXiv  Detail & Related papers  (2024-01-05T05:28:40Z)
• Gradient-based Planning with World Models [21.9392160209565]
  We present an exploration of a gradient-based alternative that fully leverages the differentiability of the world model. In a sample-efficient setting, our method achieves on par or superior performance compared to the alternative approaches in most tasks.
  arXiv  Detail & Related papers  (2023-12-28T18:54:21Z)
• Locally Optimal Descent for Dynamic Stepsize Scheduling [45.6809308002043]
  We introduce a novel dynamic learning scheduling scheme grounded in theory with the goal of simplifying the manual and time-consuming tuning of schedules in phases. Our approach is based on estimating a locally-optimal practice tuning-rate in the direction of a smooth gradient. Our findings indicate that our method needs minimal tuning when compared to existing approaches.
  arXiv  Detail & Related papers  (2023-11-23T09:57:35Z)
• Temporal Difference Learning for Model Predictive Control [29.217382374051347]
  Data-driven model predictive control has two key advantages over model-free methods. TD-MPC achieves superior sample efficiency and performance over prior work on both state and image-based continuous control tasks.
  arXiv  Detail & Related papers  (2022-03-09T18:58:28Z)
• Extrapolation for Large-batch Training in Deep Learning [72.61259487233214]
  We show that a host of variations can be covered in a unified framework that we propose. We prove the convergence of this novel scheme and rigorously evaluate its empirical performance on ResNet, LSTM, and Transformer.
  arXiv  Detail & Related papers  (2020-06-10T08:22:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.