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

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. The gist: finetuning world models with Online World Modeling and Adversarial World Modeling substantially improves the reliability of gradient-based planning (GBP) by narrowing the train-test gap, allowing GBP to match or outperform sampling-based planners like CEM.

The Train-Test Gap Problem

The core issue addressed is a fundamental mismatch between how world models are trained and how they are used for planning. World models are typically trained using a next-state prediction objective on datasets of expert trajectories, but at test time, they are used to optimize a planning objective over sequences of actions. This mismatch underlies poor empirical performance because:

1. During planning, the intermediate sequence of actions explored by gradient descent drive the world model into states that were not encountered during training, causing model errors compound in these out-of-distribution states.

2. The action-level optimization landscape induced by the world model may be difficult to traverse, containing many poor local minima or flat regions, which hinders effective gradient-based optimization.

Proposed Solutions: Online World Modeling (OWM)

Online World Modeling is proposed to address the distribution shift between offline expert trajectories and predicted trajectories from planning. This method iteratively corrects the trajectories produced by GBP and finetunes the world model on the resulting rollouts. Specifically, it involves:

1. Performing GBP to obtain an action sequence.

2. Using a environment simulator to correct states along a trajectory produced by performing GBP.

3. Adding these corrected trajectories to the dataset, effectively expanding the region of latent states that the world model can reliably predict, thereby mitigating compounding prediction errors during planning.

Proposed Solutions: Adversarial World Modeling (AWM)

Adversarial World Modeling targets regions where the world model is expected to perform poorly by explicitly training it on worst-case perturbations. This method aims to smooth the induced action loss landscape, making it easier to optimize via gradient-based planning. The objective function used is:

)&min E(ot,at,ot+1)∼T max δa∈Ba,δz∈Bz∥fθ(Φµ(ot) + δz, at + δa) − Φµ(ot+1)∥22 where Ba and Bz constrain the magnitude of perturbations. This procedure smooth[s] the loss surface of the planning objective, improving optimization stability. The paper notes that this approach can be used to generate adversarial samples over entire large-scale offline imitation learning datasets efficiently using techniques like FGSM. 23

Key Findings and Performance Gains

The finetuning methods lead to substantial improvements in planning performance. Applying Adversarial World Modeling enables GBP to match or exceed the performance of search-based CEM on a variety of robotic object manipulation and navigation tasks. Crucially, this performance is achieved with a 10× reduction in computation time compared to CEM, underscoring its practicality. Furthermore, the methods demonstrate that they can reverse the train-test gap in world model error, showing lower error during planning compared to training for PointMaze and Wall tasks.

Planning Enhancements

To further improve GBP, the paper introduces several auxiliary techniques:

1. Weighted Goal Loss (WGL): This alternate loss function encourages intermediate latent states to also be close to the goal latent state, using a formulation like LWGL = 1/H Σ X +1 i=2 wi∥zˆi − zgoal∥22, where the weights wi are chosen based on the task (e.g., exponentially upweighting later states for PointMaze).

2. Initialization Network (IN): An initialization network gθ: Z × Z → AT is trained to initialize action sequences for GBP, although results show it only performs comparably to random initialization in some environments.

3. Horizon Extension: Evaluation on longer horizons (e.g., 50 steps instead of 25) shows that Adversarial World Modeling outperforms DINO-WM on PushT and both OWM and AWM outperform DINO-WM on PointMaze, indicating better robustness for long-horizon planning.

Conclusion

The work concludes that Online World Modeling and Adversarial World Modeling are effective techniques for addressing the train-test gap in world models used for gradient-based planning. By narrowing this gap, GBP can be a practical alternative to sampling-based planners, especially where computational efficiency is critical. Future work suggests evaluating these methods on real-world systems and exploring their application in multi-timescale or hierarchical world models. 500 words (approximate).

Here are specific improvements that can be made to current AI systems by implementing the concepts from this scientific paper, along with a description of what these improved systems could achieve:

)System Improvement 1: Implement Online World Modeling (OWM) for Robust Gradient-Based Planning (GBP).

This improvement involves modifying the world model training pipeline to incorporate OWM. Instead of training solely on expert trajectories, the system will be iteratively finetuned using trajectories generated by the planner itself.

• The system will use Gradient-Based Planning (GBP) to generate candidate action sequences.

• Instead of immediately evaluating these actions against the fixed world model, the trajectory is executed in a simulator to obtain ground truth next states.

• These corrected trajectories are then added to the training dataset, and the world model parameters are retrained on this expanded dataset.

• This improved system can achieve:

• Significantly reduced prediction errors when planning in novel or out-of-distribution (OOD) state spaces encountered during execution.

• Improved long-horizon planning performance, as compounding prediction errors are mitigated by continuous data augmentation from the planner's own exploration.

)System Improvement 2: Implement Adversarial World Modeling (AWM) to Smooth Optimization Landscapes.

This improvement involves modifying the world model training objective to include adversarial samples that specifically target regions where the model is expected to perform poorly during action optimization.

• The system will generate adversarial latent states and actions by maximizing the loss function, effectively smoothing the input gradient landscape for GBP.

• This can be achieved using techniques like Fast Gradient Sign Method (FGSM) or Projected Gradient Descent (PGD) applied during training.

• This improved system can achieve:

• More stable and reliable optimization for gradient-based planning, preventing the planner from getting stuck in poor local minima or navigating non-smooth loss surfaces.

• The ability to match or exceed the performance of computationally expensive search-based planners (like CEM) while maintaining superior computational efficiency (up to 10x faster).

)System Improvement 3: Utilize Weighted Goal Loss (WGL) during Planning.

This improvement involves replacing the standard goal loss function in the GBP optimization objective with a Weighted Goal Loss that incorporates intermediate states into the loss calculation.

• The planning objective will be modified to include a term that penalizes deviations from the goal state at intermediate timesteps, weighted by a hyperparameter (e.g., exponentially upweighting later states for navigation tasks).

• This improved system can achieve:

• Improved performance on complex multi-step tasks where reaching intermediate milestones is as critical as reaching the final goal.

• Better generalization across different task types (navigation vs. non-navigation), as it provides a richer gradient signal during the optimization process.

)System Improvement 4: Employ Initialization Networks for Planning Optimization.

This improvement involves training an initialization network that maps initial state and goal states to an action sequence, which is then used to bootstrap the main gradient descent optimization.

• An initialization network will be trained using supervised learning on expert trajectories to predict initial action sequences given start and goal latent states.

• The GBP process will use this predicted sequence as a starting point for its gradient descent optimization instead of random sampling.

• This improved system can achieve:

• Faster convergence during the planning phase, particularly in environments where the optimization landscape is rugged or difficult to traverse (e.g., PushT).

• Potentially better performance compared to purely random initialization in certain scenarios.

)Overall System Capability Summary:

By integrating these four improvements—Online World Modeling for data augmentation, Adversarial World Modeling for landscape smoothing, Weighted Goal Loss for richer gradients, and Initialization Networks for faster convergence—the resulting AI system will be a world model paired with gradient-based planning that is significantly more robust to the train-test gap. This system can perform complex manipulation and navigation tasks with higher success rates than current state-of-the-art methods (matching or exceeding CEM performance) while operating at orders of magnitude faster inference times, making it practical for real-world, high-dimensional robotic control.

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.