Parallel Stochastic Gradient-Based Planning for World Models

2026-01-31

AI Summary

This work addresses the challenge of planning in vision-based world models, where high-dimensional, unstructured search spaces often lead to optimization difficulties and convergence to local optima. To overcome this, the authors propose GRASP, a novel planner that uniquely integrates stochastic exploration with relaxed gradient-based optimization. By treating states as optimizable variables subject to soft dynamical constraints and modifying the gradient structure—such as through action-input gradient clipping—the method enables efficient parallel optimization. Evaluated on long-horizon visual control tasks, GRASP significantly outperforms both the Cross-Entropy Method (CEM) and standard gradient-based approaches, achieving notable improvements in both success rate and convergence speed. This establishes a new planning framework that is inherently parallelizable and robust against local optima.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Learning for Planning and Scheduling
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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.

Problem

• world models
• planning
• long-horizon control
• visual input
• search space

Innovation

• world models
• gradient-based planning
• stochastic optimization
• parallel planning
• differentiable simulation