Gradient-based Planning for World Models at Longer Horizons

grasp-results-table table { font-size: 0. 875rem; line-height: 1. 35; width: 100%; }. grasp-results-table th,. grasp-results-table td { padding: 0. 35rem 0. 5rem; } /* Consistent whitespace between major sections (this post is long and hr-heavy) */ article. post-content h2 { margin-top: 2. 75rem; margin-bottom: 0. 75rem; } article. post-content h2:first-of-type { margin-top: 2. 25rem; } article. post-content h3 { margin-top: 1. 65rem; margin-bottom: 0. 5rem; } article. post-content hr { margin-top: 2.

5rem; margin-bottom: 2. 5rem; } 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 we avoid brittle “state-input” gradients through high-dimensional vision models.

Large, learned world models are becoming increasingly capable. They can predict long sequences of future observations in high-dimensional visual spaces and generalize across tasks in ways that were difficult to imagine a few years ago. As these models scale, they start to look less like task-specific predictors and more like general-purpose simulators. But having a powerful predictive model is not the same as being able to use it effectively for control/learning/planning.

In practice, long-horizon planning with modern world models remains fragile: optimization becomes ill-conditioned, non-greedy structure creates bad local minima, and high-dimensional latent spaces introduce subtle failure modes. In this blog post, I describe the problems that motivated this project and our approach to address them: why planning with modern world models can be surprisingly fragile, why long horizons are the real stress test, and what we changed to make gradient-based planning much more robust.

This blog post discusses work done with Mike Rabbat, Aditi Krishnapriyan, Yann LeCun, and Amir Bar (* denotes equal advisorship), where we propose GRASP. What is a world model? These days, the term “world model” is quite overloaded, and depending on the context can either mean an explicit dynamics model or some implicit, reliable internal state that a generative model relies on (e. g. when an LLM generates chess moves, whether there is some internal representation of the board).

We give our loose working definition below. 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. Formally, it defines a predictive distribution on a sequence of observed states $s_{t-h:t}$ and current action $a_t$: that approximates the environment’s true conditional $P(s_{t+1} \mid s_{t-h:t},\; a_t)$.

For this blog post, we’ll assume a Markovian model $P(s_{t+1} \mid s_{t-h:t},\; a_t)$ for simplicity (all results here can be extended to the more general case), and when the model is deterministic it reduces to a map over states: In practice the state $s_t$ is often a learned latent representation (e. g. , encoded from pixels), so the model operates in a (theoretically) compact, differentiable space.

The key point is that a world model gives you a differentiable simulator; you can roll it forward under hypothetical action sequences and backpropagate through the predictions.

Planning: choosing actions by optimizing through the model Given a start $s_0$ and a goal $g$, the simplest planner chooses an action sequence $\mathbf{a}=(a_0,\dots,a_{T-1})$ by rolling out the model and minimizing terminal error: Here we use $\mathcal{F}^T$ as shorthand for the full rollout through the world model (dependence on model parameters $\theta$ is implicit): In short horizons and low-dimensional systems, this can work reasonably well. But as horizons grow and…