# Parallel Stochastic Gradient-Based Planning for World Models

arXiv: https://arxiv.org/abs/2602.00475 · HTML: https://mlsanity.com/paper/2602.00475/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 24 of its 42 references that are in the map.
- First posted on arXiv 31 Jan 2026.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 24 of 24)

- [Intriguing properties of neural networks](https://mlsanity.com/paper/1312.6199/) (2013)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning](https://mlsanity.com/paper/2411.04983/) (2024)
- [Entropy-SGD: Biasing Gradient Descent Into Wide Valleys](https://mlsanity.com/paper/1611.01838/) (2016)
- [Diffusion Models Are Real-Time Game Engines](https://mlsanity.com/paper/2408.14837/) (2024)
- [Navigation World Models](https://mlsanity.com/paper/2412.03572/) (2024)
- [Ctrl-World: A Controllable Generative World Model for Robot Manipulation](https://mlsanity.com/paper/2510.10125/) (2025)
- [Understanding World or Predicting Future? A Comprehensive Survey of World Models](https://mlsanity.com/paper/2411.14499/) (2024)
- [Global Convergence of Langevin Dynamics Based Algorithms for Nonconvex Optimization](https://mlsanity.com/paper/1707.06618/) (2017)
- [Sample-efficient Cross-Entropy Method for Real-time Planning](https://mlsanity.com/paper/2008.06389/) (2020)
- [The Dimpled Manifold Model of Adversarial Examples in Machine Learning](https://mlsanity.com/paper/2106.10151/) (2021)
- [Model-Predictive Control via Cross-Entropy and Gradient-Based Optimization](https://mlsanity.com/paper/2004.08763/) (2020)
- [Model-Based Reinforcement Learning via Latent-Space Collocation](https://mlsanity.com/paper/2106.13229/) (2021)
- [Gradient-based Planning with World Models](https://mlsanity.com/paper/2312.17227/) (2023)
- [Parallelised Diffeomorphic Sampling-based Motion Planning](https://mlsanity.com/paper/2108.11775/) (2021)
- [Closing the Train-Test Gap in World Models for Gradient-Based Planning](https://mlsanity.com/paper/2512.09929/) (2025)
- [Surgical Vision World Model](https://mlsanity.com/paper/2503.02904/) (2025)
- [CEM-GD: Cross-Entropy Method with Gradient Descent Planner for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2112.07746/) (2021)
- [Langevin algorithms for very deep Neural Networks with application to image classification](https://mlsanity.com/paper/2212.14718/) (2022)
- [Direct Collocation Methods for Trajectory Optimization in Constrained Robotic Systems](https://mlsanity.com/paper/2304.12908/) (2023)
- [Reliable Solution to Dynamic Optimization Problems using Integrated Residual Regularized Direct Collocation](https://mlsanity.com/paper/2503.09123/) (2025)

## Built on by (most cited first, 7 of 7)

- [stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation](https://mlsanity.com/paper/2605.21800/) (2026)
- [Slot-MPC: Goal-Conditioned Model Predictive Control with Object-Centric Representations](https://mlsanity.com/paper/2605.14937/) (2026)
- [Lifting Embodied World Models for Planning and Control](https://mlsanity.com/paper/2604.26182/) (2026)
- [Unifying Object-Centric World Models and Diffusion Policy: A Hierarchical Framework for Multi-Stage Robotic Tasks](https://mlsanity.com/paper/2606.08775/) (2026)
- [StressDream: Steering Video World Models for Robust Policy Evaluation and Improvement](https://mlsanity.com/paper/2606.00267/) (2026)
- [Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization](https://mlsanity.com/paper/2605.26282/) (2026)
- [Parallel Predictive World Models for Accurate and Efficient Long-Horizon Planning](https://mlsanity.com/paper/2610.08627/) (2026)

## Abstract (arXiv)

World models simulate environment dynamics from raw sensory inputs like video. However, using them for planning can be challenging due to the vast and unstructured search space. We propose a robust and highly parallelizable planner that leverages the differentiability of the learned world model for efficient optimization, solving long-horizon control tasks from visual input. Our method treats states as optimization variables ("virtual states") with soft dynamics constraints, enabling parallel computation and easier optimization. To facilitate exploration and avoid local optima, we introduce stochasticity into the states. To mitigate sensitive gradients through high-dimensional vision-based world models, we modify the gradient structure to descend towards valid plans while only requiring action-input gradients. Our planner, which we call GRASP (Gradient RelAxed Stochastic Planner), can be viewed as a stochastic version of a non-condensed or collocation-based optimal controller. We provide theoretical justification and experiments on video-based world models, where our resulting planner outperforms existing planning algorithms like the cross-entropy method (CEM) and vanilla gradient-based optimization (GD) on long-horizon experiments, both in success rate and time to convergence.