By Michael Psenka, Mike Rabbat, Aditi Krishnapriyan, Yann LeCun, Amir Bar
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 and observe states (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 and current action :
that approximates the environment’s true conditional . For this blog post, we’ll assume a Markovian model 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 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 and a goal , the simplest planner chooses an action sequence by rolling out the model and minimizing terminal error:
Here we use as shorthand for the full rollout through the world model (dependence on model parameters is implicit):
In short horizons and low-dimensional systems, this can work reasonably well. But as horizons grow and models become larger and more expressive, its weaknesses become amplified.
So why doesn’t this just work at scale?
Why long-horizon planning is hard (even when everything is differentiable)
There are two separate pain points for the more general world model, plus a third that is specific to learned, deep learning-based models.
1) Long-horizon rollouts create deep, ill-conditioned computation graphs
Those familiar with backprop through time (BPTT) may notice that we’re differentiating through a model applied to itself repeatedly, which will lead to the exploding/vanishing gradients problem. Namely, if we take derivatives (note we’re differentiating vector-valued functions, resulting in Jacobians that we denote with ) with respect to earlier actions (e.g. ):
We see that the Jacobian’s conditioning scales exponentially with time :
leading to exploding or vanishing gradients.
2) The landscape is non-greedy and full of traps
At short horizons, the greedy solution, where we move straight toward the goal at every step, is often good enough. If you only need to plan a few steps ahead, the optimal trajectory usually doesn’t deviate much from “head toward ” at each step.
As horizons grow, two things happen. First, longer tasks are more likely to require non-greedy behavior: going around a wall, repositioning before pushing, backing up to take a better path. And as horizons grow, more of these non-greedy steps are typically needed. Second, the optimization space itself scales with horizon: , further expanding the space of local minima for the optimization problem.
Distance to goal along the optimal path is non-monotonic, and the resulting loss landscape can be rough.
A long-horizon fix: lifting the dynamics constraint
Suppose we treat the dynamics constraint as a soft constraint, and we instead optimize the following penalty function over both actions and states :
This is also sometimes called collocation in planning/robotics literature. Note the lifted formulation shares the same global minimizers as the original rollout objective (both are zero exactly when the trajectory is dynamically feasible). But the optimization landscapes are very different, and we get two immediate benefits:
Each world model evaluation depends only on local variables, so all terms can be computed in parallel across time, resulting in a huge speed-up for longer horizons, and
You no longer backpropagate through a single deep -step composition to get a learning signal, since the previous product of Jacobians now splits into a sum, e.g.:
Being able to optimize states directly also helps with exploration, as we can temporarily navigate through unphysical domains to find the optimal plan:
Collocation-based planning allows us to directly perturb states and explore midpoints more effectively.
However, lunch is never free. And indeed, especially for deep learning-based world models, there is a critical issue that makes the above optimization quite difficult in practice.
An issue for deep learning-based world models: sensitivity of state-input gradients
The tl;dr of this section is: directly optimizing states through a deep learning-based is incredibly brittle, à la adversarial robustness. Even if you train your world model in a lower-dimensional state space, the training process for the world model makes unseen state landscapes very sharp, whether it be an unseen state itself or simply a normal/orthogonal direction to the data manifold.
Adversarial robustness and the “dimpled manifold” model
Adversarial robustness originally looked at classification models , and showed that by following the gradient of a particular logit from a base image (not of class ), you did not have to move far along to make classify as (Szegedy et al., 2014; Goodfellow et al., 2015):
Depiction of the classic example from (Goodfellow et al., 2015).
Later work has painted a geometric picture for what’s going on: for data near a low-dimensional manifold , the training process controls behavior in tangential directions, but does not regularize behavior in orthogonal directions, thus leading to sensitive behavior (Stutz et al., 2019). Another way stated: has a reasonable Lipschitz constant when considering only tangential directions to the data manifold , but can have very high Lipschitz constants in normal directions. In fact, it often benefits the model to be sharper in these normal directions, so it can fit more complicated functions more precisely.
As a result, such adversarial examples are incredibly common even for a single given model. Further, this is not just a computer vision phenomenon; adversarial examples also appear in LLMs (Wallace et al., 2019) and in RL (Gleave et al., 2019).
While there are methods to train for more adversarially robust models, there is a known trade-off between model performance and adversarial robustness (Tsipras et al., 2019): especially in the presence of many weakly-correlated variables, the model must be sharper to achieve higher performance. Indeed, most modern training algorithms, whether in computer vision or LLMs, do not train adversarial robustness out. Thus, at least until deep learning sees a major regime change, this is a problem we’re stuck with.
Why is adversarial robustness an issue for world model planning?
Consider a single component of the dynamics loss we’re optimizing in the lifted state approach:
Let’s further focus on just the base state:
Since world models are typically trained on state/action trajectories , the state-data manifold for has dimensionality bounded by the action space:
where is some optional space of augmentations (e.g. translations/rotations). Thus, we can typically expect to be much lower than , and thus: it is very easy to find adversarial examples that hack any state to any other desired state.
As a result, the dynamics optimization
feels incredibly “sticky,” as the base points can easily trick into thinking it’s already made its local goal.1
- This adversarial robustness issue, while particularly bad for lifted-state approaches, is not unique to them. Even for serial optimization methods that optimize through the full rollout map , it is possible to get into unseen states, where it is very easy to have a normal component fed into the sensitive normal components of . The action Jacobian’s chain rule expansion is
See what happens if any stage of the product has any component normal to the data manifold.
Our fix
This is where our new planner GRASP comes in. The main observation: while is untrustworthy and adversarial, the action space is usually low-dimensional and exhaustively trained, so is actually reasonable to optimize through and doesn’t suffer from the adversarial robustness issue!
The action input is usually lower-dimensional and densely trained (the model has seen every action direction), so action gradients are much better behaved.
At its core, GRASP builds a first-order lifted state / collocation-based planner that is only dependent on action Jacobians through the world model. We thus exploit the differentiability of learned world models , while not falling victim to the inherent sensitivity of the state Jacobians .
GRASP: Gradient RelAxed Stochastic Planner
As noted before, we start with the collocation planning objective, where we lift the states and relax dynamics into a penalty:
We then make two key additions.
Ingredient 1: Exploration by noising the state iterates
Even with a smoother objective, planning is nonconvex. We introduce exploration by injecting Gaussian noise into the virtual state updates during optimization.
A simple version:
Actions are still updated by non-stochastic descent:
The state noise helps you “hop” between basins in the lifted space, while the actions remain guided by gradients. We found that specifically noising states here (as opposed to actions) finds a good balance of exploration and the ability to find sharper minima.2
- Because we only noise the states (and not the actions), the corresponding dynamics are not truly Langevin dynamics.
Ingredient 2: Reshape gradients: stop brittle state-input gradients, keep action gradients
As discussed, the fragile pathway is the gradient that flows into the state input of the world model, . The most straightforward way to do this initially is to just stop state gradients into directly:
Let be the same value as , but with gradients stopped.
Define the stop-gradient dynamics loss:
This alone does not work. Notice now states only follow the previous state’s step, without anything forcing the base states to chase the n
[truncated for AI cost control]