Parallel Stochastic Gradient-Based Planning for World Models

TL;DR

This paper introduces GRASP, a parallelizable stochastic gradient-based planner for world models that efficiently handles long-horizon control tasks from visual inputs, outperforming existing methods in success rate and convergence time.

Contribution

The paper presents a novel differentiable planning algorithm, GRASP, that leverages parallel computation and stochasticity to improve long-horizon planning in visual world models.

Findings

Outperforms CEM and GD in success rate.

Faster convergence in long-horizon tasks.

Effective handling of high-dimensional visual inputs.

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.