Closing the Train-Test Gap in World Models for Gradient-Based Planning (2512.09929v1)

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.

• The paper presents two novel methods, Online World Modeling and Adversarial World Modeling, to mitigate compounding errors in gradient-based planning.
• The approach achieves up to 30% improvement in task success rates and over tenfold computational efficiency compared to standard baselines.
• The methods enhance model robustness and gradient reliability, enabling scalable deployment in robotic manipulation and navigation tasks.

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

Introduction

This paper addresses a critical issue in the application of latent world models to gradient-based planning (GBP): the discrepancy between trajectories observed during training and those encountered during planning. While world models are typically optimized with next-state prediction objectives using expert or behavioral policy data, they are deployed at test time for planning, resulting in compounding errors when the model is rolled out on action sequences not represented in the training distribution. The work presents two novel data synthesis and finetuning procedures—Online World Modeling and Adversarial World Modeling—to directly counter this train-test gap, enhancing both the reliability and efficiency of GBP in diverse manipulation and navigation tasks.

Figure 1: Both Online World Modeling and Adversarial World Modeling proactively adapt the world model to out-of-distribution planning trajectories, extending performance beyond expert data.

Background: World Models and Differentiable Planning

Modern world models, leveraging offline datasets and powerful feature encoders (e.g., DINOv2), approximate a latent dynamics function $$f_\theta: \mathcal{Z} \times \mathcal{A} \to \mathcal{Z}$$, enabling predictive rollouts in a compressed, feature-rich latent space. Planning aims to optimize action sequences that minimize the latent distance between a predicted terminal state and a designated goal embedding.

Traditional planning methods, such as Cross-Entropy Method (CEM) and Model Predictive Path Integral Control (MPPI), circumvent differentiability by sampling and evaluating action sequences using model rollouts. Gradient-based planning, in contrast, exploits the smoothness and end-to-end differentiability of world models for direct optimization, presenting substantial computational advantages in high-dimensional scenarios. Despite this, GBP with standard world model training is empirically brittle, suffering from misaligned optimization landscapes and pronounced out-of-distribution errors.

Figure 2: Adversarial World Modeling substantially smooths and broadens the basin of attraction in the planning loss landscape, enabling improved convergence for GBP.

Proposed Methods

Online World Modeling

Online World Modeling (OWM) mitigates compounding errors induced by out-of-distribution planning trajectories by iterative correction with simulator feedback. For each planning instance, actions proposed by GBP are executed in the environment simulator, producing corrected state-action pairs. These corrected rollouts, which naturally sample regions of latent space traversed during active planning, are used to augment the training dataset, and the world model is finetuned on this synthetic data. The procedure is reminiscent of dataset aggregation/DAgger, but is implemented in the latent feature space rather than state space.

Adversarial World Modeling

Adversarial World Modeling (AWM) constructs adversarial perturbations around expert trajectories, designed to maximally increase world model prediction error under bounded latent and action perturbations. Training with these samples smooths the input gradient landscape of the world model, improving the stability of optimization and suppressing the formation of pathological local minima characteristic of non-smooth loss surfaces. Adversarial examples are efficiently generated using the Fast Gradient Sign Method (FGSM) rather than iterative PGD, achieving competitive robustness with significantly reduced computational overhead.

Experimental Results

The work benchmarks the proposed methods versus the DINO-WM architecture on PushT, PointMaze, Wall, Rope, and Granular manipulation and navigation tasks. Planning is evaluated both open-loop and under Model Predictive Control (MPC), comparing stochastic gradient descent (GD), Adam, and CEM optimization procedures.

Key findings:

• Adversarial World Modeling with GBP matches or exceeds CEM baseline performance in all environments, achieving success rates of up to +18%, +20%, and +30% over standard DINO-WM on PushT, PointMaze, and Wall, respectively.
• Adversarial finetuning yields a planning loss surface with a broad basin of attraction, improving the informativeness of gradients and the stability of action sequence optimization.

  Figure 3: GBP with Adversarial World Modeling achieves equivalent solution quality as CEM, but runs over an order of magnitude faster in wall-clock time on PushT.

  

Figure 4: Visualization of experimental setups for PushT, PointMaze, Wall, Rope, and Granular tasks.

Computational efficiency is a decisive advantage: GBP with adversarially trained models produces solutions matching or exceeding CEM in a fraction (10%) of the time budget.

Figure 5: Planning efficiency comparison on PointMaze demonstrates that GBP with OWM/AWM approaches the performance of classical samplers at superior compute efficiency.

Figure 6: On the Wall task, GBP-based planners exhibit marked efficiency increases without sacrificing final score.

Closed-loop MPC with GBP and Adversarial WM demonstrates stable scaling across horizons. On robotic manipulation tasks (Rope, Granular), Chamfer distance metrics confirm that AWM-trained models enable more accurate configurational predictions versus baseline DINO-WM.

Train-Test Gap Analysis

World model prediction error is systematically higher for planning-derived trajectories than expert trajectories when using models trained with standard next-state prediction objectives (see DINO-WM curves). Both OWM and AWM substantially reduce this error disparity, confirming their capacity to close the train-test gap. Empirical ablations on perturbation magnitudes and radii indicate robustness to a range of scaling parameters, with excessively large visual perturbations being detrimental.

Figure 7: Closed-loop MPC planning success rates remain stable for scaling factors $\lambda_z, \lambda_a$ up to 0.2; higher values degrade semantic integrity of visual latent states.

Figure 8: Adversarial World Modeling consistently enables recovery from out-of-distribution latent states, whereas DINO-WM and OWM exhibit failure modes under challenging rollouts.

Figure 9: In certain difficult environments, only Online World Modeling supports successful completion of planning objectives, revealing complementary strengths with Adversarial WM.

Implications and Future Directions

The presented methods establish that targeted train-time data synthesis—either simulator-corrected or adversarially perturbed—can endow latent world models with the robustness required for reliable and efficient gradient-based planning. The implications are substantial for model-based control in robotics, where rollout efficiency and distributional coverage are essential for real-time, long-horizon decision making. These strategies are also directly relevant to multi-timescale and hierarchical planning paradigms, where stability in different abstraction tiers is crucial.

Practically, Adversarial WM enables scalable deployment of world model-based planners in compute-constrained or high-dimensional settings, potentially serving as a critical component for autonomous agents requiring fast, online inference. Theoretically, these results substantiate the value of countering objective mismatch and distributional shift not via indirect regularization, but by explicit expansion of model support during training in the manners most relevant to deployment regimes.

Further research avenues include: evaluating these methods under true real-world sensory noise and environmental stochasticity; extending adversarial/online data synthesis procedures to richer, multi-agent or interactive world models; and integrating robustness objectives with hierarchical memory architectures for continuous, lifelong learning.

Conclusion

Online World Modeling and Adversarial World Modeling effectively address the fundamental train-test gap in latent world models used for gradient-based planning. By adapting model training to the trajectories and loss surfaces encountered during planning, these approaches enable robust, efficient, and scalable model-based planners, matching or exceeding the performance of classical sampling methods while dramatically reducing computational burden. These contributions advance both the practical state-of-the-art and understanding of train-test dynamics in high-dimensional, differentiable planning systems.