GRASP: Gradient-Based Planning for World Models

GRASP is a new gradient-based planner for learned dynamics (a "world model") that makes long-horizon planning practical by (1) lifting the trajectory into virtual states so optimization is parallel across time, (2) adding stochasticity directly to the state iterates for exploration, and (3) reshaping gradients so actions get clean signals while avoiding brittle "state-input" gradients through high-dimensional vision models.

What is a world model?

These days, the term "world model" is quite overloaded. We give our loose working definition: suppose you take actions $a_t \in \mathcal{A}$ and observe states $s_t \in \mathcal{S}$ (images, latent vectors, proprioception). A world model is a learned model that, given the current state and a sequence of future actions, predicts what will happen next:

$$P_\theta(s_{t+1} \mid s_{t-h:t}, a_t)$$

approximating the environment's true conditional $P(s_{t+1} \mid s_{t-h:t}, a_t)$. When deterministic, it reduces to a state map:

$$s_{t+1} = F_\theta(s_t, a_t)$$

The key point: a world model gives you a differentiable simulator—you can roll it forward under hypothetical actions and backpropagate through predictions.

Planning: choosing actions by optimizing through the model

Given start $s_0$ and goal $g$, the simplest planner chooses actions $\mathbf{a}=(a_0,\dots,a_{T-1})$ by rolling out the model and minimizing terminal error:

$$\min_{\mathbf{a}} \| s_T(\mathbf{a}) - g \|_2^2, \quad \text{where } s_T(\mathbf{a}) = \mathcal{F}_{\theta}^{T}(s_0,\mathbf{a})$$

Here $\mathcal{F}^T$ is shorthand for the full rollout:

$$\mathcal{F}_{\theta}^{T}(s_0, \mathbf{a}) = F_\theta(F_\theta(\cdots F_\theta(s_0, a_0), \cdots, a_{T-2}), a_{T-1})$$

In short horizons and low-dimensional systems, this works reasonably well. But as horizons grow and models become larger, its weaknesses become amplified.

Why long-horizon planning is hard (even when everything is differentiable)

1) Long-horizon rollouts create deep, ill-conditioned computation graphs

Differentiating through a model applied to itself repeatedly leads to exploding/vanishing gradients. The Jacobian's conditioning scales exponentially with time $T$:

$$\sigma_{\text{max/min}}(D_{a_0}\mathcal{F}_{\theta}^{T}) \sim \sigma_{\text{max/min}}(D_s F_\theta)^{T-1}$$

2) The landscape is non-greedy and full of traps

At short horizons, the greedy solution (move straight toward the goal) is often good enough. But long horizons require non-greedy behavior, and greedy methods get trapped in local optima.

3) Deep learning models have brittle gradients

State-input gradients $D_s F_\theta$ are extremely fragile (adversarial robustness issues), while action-input gradients $D_a F_\theta$ are relatively smooth.

GRASP's three core ingredients

Ingredient 1: Lift states for parallel-in-time optimization

GRASP converts serial rollouts to collocation optimization—optimizing all states and actions simultaneously with a consistency constraint:

$$\mathcal{L}_{\text{dyn}}(\mathbf{s},\mathbf{a}) = \sum_{t=0}^{T-1} \|F_\theta(s_t, a_t) - s_{t+1}\|_2^2$$

Stochasticity is added for exploration (Langevin-like dynamics): $s_t \leftarrow s_t - \eta \nabla_{s_t} \mathcal{L} + \sqrt{2\eta/\beta} \epsilon$.

Ingredient 2: Reshape gradients—stop brittle state gradients, keep action gradients

Stop-gradient dynamics loss:

$$\mathcal{L}_{\text{dyn}}^{\text{sg}}(\mathbf{s},\mathbf{a}) = \sum_{t=0}^{T-1} \|F_\theta(\bar{s}_t, a_t) - s_{t+1}\|_2^2$$

Dense goal term:

$$\mathcal{L}_{\text{goal}}^{\text{sg}}(\mathbf{s},\mathbf{a}) = \sum_{t=0}^{T-1} \|F_\theta(\bar{s}_t, a_t) - g\|_2^2$$

Final objective:

$$\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})$$

Ingredient 3: Periodic sync

Every $K_{\text{sync}}$ iterations, roll out from $s_0$ with current actions and take a few gradient steps on the original serial loss for refinement.

Results

Push-T results: GRASP achieves higher success rates at long horizons while finding solutions faster.

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	2.0% / 231.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. Bold = best in row.

What's next?

Extension to diffusion-based world models, more sophisticated optimizers, and integrating GRASP into closed-loop systems or RL policy learning are natural next steps. World model planners are at an interesting sweet spot—the background literature is mature, but pure planning optimization over modern large-scale world models is still heavily underexplored. Once we figure out the right ideas, world model planners will likely become as commonplace as RL.

Source: BAIR Blog | Paper: arXiv:2602.00475