Parallel Stochastic Gradient-Based Planning for World Models

TL;DR

We address planning with learned vision-based world models, where long horizons and high-dimensional state spaces create difficult optimization landscapes. GRASP (Gradient RelAxed Stochastic Planner) lifts intermediate states to enable time-parallel optimization, uses Langevin-style state perturbations for exploration, and employs stop-gradient on state inputs to stabilize gradients, with periodic full-gradient rollouts to refine solutions. The method outperforms zero-order methods like CEM and vanilla gradient-based planning on long-horizon tasks, while remaining competitive on short-horizon problems. This approach offers practical improvements for robust, scalable planning in visual world models and suggests broader potential for parallelized, gradient-based planning in learned dynamics.

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.

Figures (7)

• Figure 1: Difficulty of the planning problem. Subfigure (a) shows the distance to the goal in $L^2$ norm throughout a successful trajectory. This illustrates the difficulty of planning optimization away from a minimizer: successful trajectories often have to first move away from the goal to successfully plan towards it later, resulting in greedy strategies failing. Subfigures (b)-(c) depict the loss landscape at convergence of standard rollout-based planners vs. our planner. The example given is in the Push-T environment at horizon length 50. The axes plotted over are with respect to two random, orthogonal, unit-norm directions in the full action space $\mathbb{R}^{50 \times 2}$. Our planner loss is taken as in Eq. (\ref{['eq:loss_full']}), and for GD the loss is taken as in Eq. (\ref{['eq:planning_loss_main']}).
• Figure 2: Graphical depiction of (a) a standard serial-based setup for optimization-based planning, where states are rolled out using the actions and the loss is evaluated on the goal state, (b) our setup, which parallelizes the world model evaluations by optimizing "virtual states" directly and only supervising pairwise dynamics satisfaction. The crossed lines and skipped connections for our method's depiction (b) are detailed in Section \ref{['sec:no_state_grads']}, which keeps the full planning graph connected while not requiring state gradients of the dynamics $F_\theta$. For our planner, we find it helpful to alternate between (a) and (b) throughout the planning optimization.
• Figure 3: Sensitivity of state gradient structure. Examples of three states far away from the goal on the right (either in-distribution or out-of-distribution), such that taking a small step along the gradient $s' = s - \epsilon \nabla_s \mathcal{L}(s), \mathcal{L}(s) = \| F_\theta(s, a=\mathbf 0) - g\|_2^2$, leads to a nearby state $s'$ that solves the planning problem in a single step: $F_\theta(s', \mathbf 0) = g$. Thus, optimizing states directly through the world model $F_\theta$ can be quite challenging.
• Figure 4: Virtual states learned through planning. All examples are instantiations of our planner at horizon 50 in the Point-Maze, Wall-Single, and Push-T environments. Regardless of the dynamics constraint relaxation and state noising, directly optimized states find realistic, non-greedy paths towards the goal.
• Figure 5: Success rate over time at a fixed horizon. Success rate over fixed set of open-loop planning tasks for CEM, GD, LatCo rybkin2021model, and our planner for a fixed horizon of 50. Curves summarize how quickly each planner makes progress under the learned world model setting when evaluated at a fixed planning horizon. Shaded regions are Wald 95% confidence intervals.

Theorems & Definitions (18)

• Theorem 1: informal
• Theorem 2: Shooting: Exploding Smoothness
• proof
• Theorem 3: Lifted: Stable Smoothness
• proof
• Theorem 4: Gaussian smoothing contracts gradients and yields scale control
• proof
• Corollary 1: Noise induces smoothed gradients in expectation
• Theorem 5: Nonexistence of exact dynamics-enforcing losses with Jacobian-free state gradients
• proof