Parallel Stochastic Gradient-Based Planning for World Models

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 approach 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.

World models are becoming a popular form of AI used for robots, as they directly interface and interact with the real world. In order to use these world models, we typically want to generate sequences of actions for a robot to execute in order to achieve a specific goal; this is what we call planning. Planning at longer horizons remains a central challenge for world model planning. Furthermore, planning with modern world models that operate at a higher scale has many technical challenge that existing planning algorithms don't adequately address. We present a new world model planner that addresses these technical challenges, both by utilizing classical ideas from optimal control literature and incorporating new ideas specifically tailored to modern world models.