GRASP: Gradient-based Planning for World Models at Longer Horizons

BAIR has introduced GRASP (Gradient RelAxed Stochastic Planner), a new gradient-based planner designed to make long-horizon planning practical for learned world models. GRASP solves the fragility of long-term planning by lifting trajectories into virtual states for parallel optimization, injecting stochasticity into state iterates for exploration, and reshaping gradients to rely on stable action Jacobians rather than brittle state-input gradients.

The Challenge of Long-Horizon Planning in World Models

While large world models can act as general-purpose simulators, using them for control and planning is often fragile. Long-horizon planning typically fails due to three primary factors:

• Ill-conditioned Computation Graphs: Differentiating through a model applied to itself repeatedly (Backpropagation Through Time) leads to exploding or vanishing gradients. The Jacobian's conditioning scales exponentially with the time horizon $T$.
• Non-Greedy Landscapes: Long-term tasks often require non-greedy behavior (e.g., moving away from a goal to bypass an obstacle). As horizons grow, the optimization space expands, increasing the likelihood of getting trapped in local minima.
• State-Input Gradient Sensitivity: In deep learning-based world models, gradients with respect to the state input ($D_s F_\theta$) are often untrustworthy. Because world models are trained on a low-dimensional data manifold, they are susceptible to "adversarial" examples where small perturbations in the state can trick the model into predicting a desired outcome, making the optimization landscape "sticky" and unreliable.

GRASP's Technical Approach

GRASP addresses these failures by combining collocation-based planning with specific gradient reshaping and exploration strategies.

1. Lifting Dynamics Constraints (Collocation)

Instead of a serial rollout where $s_{t+1} = F_\theta(s_t, a_t)$, GRASP treats dynamics as a soft constraint. It optimizes a penalty function over both actions and states simultaneously:

$$\min_{\mathbf{s},\mathbf{a}} \mathcal{L}(\mathbf{s}, \mathbf{a}) = \sum_{t=0}^{T-1} \big|F_\theta(s_t,a_t) - s_{t+1}\big|_2^2, \quad \text{with } s_0 \text{ fixed and } s_T=g$$

This "lifted" formulation allows for two major improvements: all $T$ terms can be computed in parallel across time, and the gradient signal no longer requires backpropagating through a deep $T$-step composition.

2. Gradient Reshaping via Stop-Gradients

To avoid the adversarial sensitivity of state Jacobians, GRASP ensures the optimization depends only on action Jacobians ($D_a F_\theta$), which are typically lower-dimensional and more densely trained.

GRASP implements a stop-gradient dynamics loss, where gradients are blocked from flowing into the state input of the world model. To prevent the system from collapsing into trivial minima, GRASP adds a dense goal shaping term:

$$\mathcal{L}(\mathbf{s},\mathbf{a}) = \mathcal{L}{\text{dyn}}^{\text{sg}}(\mathbf{s},\mathbf{a}) + \gamma \mathcal{L}{\text{goal}}^{\text{sg}}(\mathbf{s},\mathbf{a})$$

This objective balances feasible dynamics with a consistent signal toward the goal without relying on brittle state gradients.

3. State-Iterate Stochasticity for Exploration

To navigate the non-convex landscape of long-horizon planning, GRASP injects Gaussian noise directly into the virtual state updates during optimization:

$$s_t \leftarrow s_t - \eta_s \nabla_{s_t}\mathcal{L} + \sigma_{\text{state}} \xi, \qquad \xi\sim\mathcal{N}(0,I)$$

By noising the states rather than the actions, the planner can "hop" between basins in the lifted space while maintaining guided action updates.

4. Periodic Synchronization

Because the lifted stop-gradient objective is an approximation, GRASP periodically performs a "sync" phase every $K_{\text{sync}}$ iterations. It rolls out from $s_0$ using current actions and takes a few small gradient steps on the original serial loss to ensure the plan remains grounded in real trajectories.

Performance and Results

GRASP demonstrates significant improvements in success rates and computation speed over baseline planners like Cross-Entropy Method (CEM) and standard Gradient Descent (GD) in the Push-T environment. As the horizon $H$ increases, GRASP maintains higher success rates and lower median times to success.

Horizon	CEM	GD	LatCo	GRASP
H=40	61.4% / 35.3s	51.0% / 18.0s	15.0% / 598.0s	59.0% / 8.5s
H=50	30.2% / 96.2s	37.6% / 76.3s	4.2% / 1114.7s	43.4% / 15.2s
H=60	7.2% / 83.1s	16.4% / 146.5s	26.2% / 49.1s
H=70	7.8% / 156.1s	12.0% / 103.1s	0.0% / —	16.0% / 79.9s
H=80	2.8% / 132.2s	6.4% / 161.3s	0.0% / —	10.4% / 58.9s

(Success rate % / median time to success)

Future Directions

The researchers suggest several paths for extending GRASP, including applying it to diffusion-based world models, developing more sophisticated noising strategies, and integrating the planner into closed-loop systems or RL policy learning for adaptive long-horizon planning.