Sparse Planning in Visual World Models via Cost Gradients

arXiv:2610.10274 · cs.LG · Submitted 2026-10-07 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Sparse Planning in Visual World Models via Cost Gradients".

Jane: The gist The COSTGRAD method introduces a training-free, goal-conditioned selector that ranks spatial tokens by the gradient norm of the planning cost with respect to each input token,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at this paper called "Sparse Planning in Visual World Models via Cost Gradients," and it seems they're trying to make world models faster by only looking at a small part of the visual information when planning.

Jane: That's right, and the core idea here is introducing COSTGRAD, which is a way to pick which spatial tokens are actually important for planning instead of just using every single token in the model.

Tom: Exactly, it’s training-free because you don't need to retrain anything; you just calculate the gradient norm of the planning cost with respect to each input token.

Lu: What they're doing is deriving importance from that downstream control objective, which means they aren't picking tokens based on how well the model predicts something, but based on what matters for actually reaching the goal >

Meng: But what does this mean in practice? How does this help us when we need to run these models at scale?

Jane: Well, according to their summary of "Sparse Planning in Visual World Models via Cost Gradients," they show that COSTGRAD can match or even exceed full-token planning success on three out of four continuous-control benchmarks.

Tom: That's a solid number. And they claim a measured two point six times wall-clock speedup per environment planning step, which is significant when you’re running these things in real time.

Lalam: From my perspective, the method works by computing the gradient of the planning cost with respect to each input token and ranking tokens by gradient norm, retaining those with the largest scores to define a sparse token set S >

Lu: The paper also points out an architecture-dependent failure mode, saying that pure COSTGRAD succeeds on AdaLN but loses its advantage on matched-concat, even when full-token planning is comparable >

Jane: That suggests that what makes this selector work depends heavily on how the actions enter the predictor; it's not a one size fits all selector >

Tom: It’s a design consideration then, suggesting that sparse planning depends on both choosing relevant tokens and having a predictor architecture that stays stable when you restrict it to those tokens >

Meng: If we want to deploy this, we need to be careful about which predictor architecture we choose because the stability of the action pathways changes based on whether you use AdaLN or concatenation >

Jane: Exactly. Now, let's look at what they actually propose as improvements in "Sparse Planning in Visual World Models via Cost Gradients."

Tom: The main improvement is that COSTGRAD ranks spatial tokens by "the gradient norm of the planning cost with respect to each input token," which allows the AI system to target tokens that matter for planning rather than merely for prediction >

Lu: That's a key point, because it means they are deriving importance from the downstream control objective, addressing some of that objective mismatch you see in model-based reinforcement learning >

Jane: It also shows speedup in inference; COSTGRAD provides a measured two point six times wall-clock speedup per environment planning step by reducing token processing from full-grid computations to only the selected subset >

Tom: And if you combine that with reduced CEM search, it increases that to about five times total speedup while still exceeding the full-token baseline >

Meng: That five times multiplier is what makes me think about deployment. If we can get a five times speedup just by using this selector, that's valuable for real-time decision making on complex visual tasks >

Jane: And they also note that the kept tokens concentrate on task-relevant regions like maze corridors and the agent and wall, while dropping much of the background information >

Lalam: It’s interesting to see how prediction difficulty isn't task relevance, because prediction error selects visually hard-to-predict tokens whereas planning needs tokens whose state affects the goal-reaching cost >

Tom: So, in summary for this section, it's about using gradients from the planning objective to select tokens and getting a speedup that’s around two point six times per step >

Jane: And this leads us into what they conclude about the overall impact of "Sparse Planning in Visual World Models via Cost Gradients."

Tom: The paper concludes that token selection and CEM budget are complementary compute axes, meaning the planner is actually part of the world model’s operating distribution >

Lu: They emphasize that for world models intended for sparse planning, you need to evaluate them under the token subsets and action searches that you'd use at deployment because full-token prediction quality alone doesn't establish suitability for sparse control >

Meng: It also highlights that the planner is part of a world model’s operating distribution, which means we have to think about how these models operate when they are restricted to those token subsets >

Lalam: I think what this means for us is that just having a good visual model isn't enough; you need the planning process itself to be designed around selecting relevant tokens and using the right architecture for that selection >

Jane: And we wrap up with "Sparse Planning in Visual World Models via Cost Gradients," showing how it uses cost gradients to select tokens and achieving a two point six times speedup per environment planning step >

The paper's summary: Tom: So, to wrap up what we just covered, COSTGRAD is this training-free method that uses the gradient of the planning cost to pick which visual tokens are actually useful for planning instead of just looking at everything.

Jane: Right, and those results show it can match or even beat full-token success on three out of four continuous-control benchmarks.

Tom: And a measured two point six times wall-clock speedup per environment planning step, which is pretty significant when you think about running these things in real time on actual hardware.

Lu: What’s interesting is how they derive this importance from the downstream control objective, so it’s not just guessing what looks important to the model, but what actually helps you reach your goal.

Meng: That focus on the control objective is key because when we deploy things at scale, we need to know exactly what's happening in our environment step by step.

Jane: Exactly. And they found that these kept tokens really concentrate on task-relevant stuff like the maze corridors or the agent and wall, while dropping a lot of that background clutter.

Tom: It’s not just about speed, though. They showed this selector is quite robust; it preserves performance even when you change how the actions are fed into the predictor architecture.

Jane: That's a big deal because it suggests that when designing these world models for sparse planning, you have to think about both picking relevant tokens and making sure your predictor architecture stays stable.

Lu: And they pointed out an architectural detail where pure COSTGRAD works on one type of action injection method but struggles with another, which tracks how the action pathways drift between those different methods.

Meng: So for us in production, that means we need to consider not just the prediction quality but also how we structure our model's inputs so that the planner doesn't break when it starts selecting a sparse set of tokens.

Tom: It really boils down to token selection and the CEM search budget being two complementary compute axes, meaning the planner is actually part of how a world model operates.

Jane: And this leads us to think about what it means for deployment—that just having a good visual model isn't enough; you need the planning process itself to be designed around selecting relevant tokens and using the right architecture for that selection.

Tom: Next, we’re going to look at how these results relate to the other papers we’ve seen today.

The paper's improvements: Tom: So, we’re looking at how they suggest improving this COSTGRAD method beyond just the basic gradient selection idea.

Jane: Right, they aren't stopping at just picking tokens based on planning cost; they are looking at how that selection interacts with other parts of the system.

Tom: One big thing is that they identified an architecture-dependent failure mode, meaning how actions get into the predictor really matters for whether this selector works well.

Jane: That’s important because it tells us that sparse planning isn't just about token picking; it depends on choosing a predictor architecture that stays stable when you restrict it to those selected tokens.

Lu: They found that gradient selection produces less drift than random removal when using AdaLN, but more drift with a concatenation method.

Meng: That’s a practical thing for us, because if we use the concatenation approach and rely on this selector, we have to be worried about action pathway instability.

Tom: Another improvement they highlight is that COSTGRAD derives its importance from the downstream control objective instead of just prediction error.

Jane: That means it's not picking tokens based on how well the model predicts something, but what actually helps you reach your goal, which addresses some of that objective mismatch we see in model-based reinforcement learning.

Lu: And they showed this method is fairly robust to changes in the action input because it works with probe actions and proprioceptive inputs consistently.

Meng: If it’s robust to that, that’s good for deployment because we don't have to constantly re-evaluate the selection mechanism every time we slightly tweak the control signal.

Tom: So, these improvements show that sparse planning needs a tighter link between what's selected and how the model is structured.

Jane: It suggests that for any AI system aiming for sparse planning, you have to consider both choosing relevant tokens and making sure your predictor architecture is stable when restricted to those tokens.

Lu: And they also pointed out that the strength of COSTGRAD varies depending on the environment; it performs well on PointMaze and PushT but has a bigger gap on Wall and MetaWorld.

Meng: That’s a caveat for us—we can't just assume this method will work equally well across all environments; we have to test its performance where the gap is largest.

Tom: So, what this means for us is that token selection and the CEM search budget are two complementary compute axes, and we need to check how that interplay changes depending on the environment and our model setup.

Conclusion: Tom: So we’re wrapping up on "Sparse Planning in Visual World Models via Cost Gradients."

Jane: To recap, this paper shows that using the gradient of planning cost to select tokens can match full-token performance and give us a two point six times speedup per planning step.

Lu: The big idea here is that the planner itself is part of the world model’s operating distribution, meaning we have to think about how it works when we only use a subset of tokens.

Meng: From an engineering standpoint, this means if you want to deploy a sparse planner, you need to evaluate it under those exact token subsets and action searches you'd actually use in production because just having good prediction quality isn't enough.

Lalam: I’m pretty sure this is solid; the paper really emphasizes that we have to test these models using the specific token subsets and action searches that we deploy, because full-token prediction quality alone doesn't prove suitability for sparse control.

Tom: And they also pointed out that token selection and the CEM search budget are two complementary compute axes, which is a useful way to think about how these systems work.

Jane: It’s a neat idea for world models intended for sparse planning because it gives us a concrete way to measure if the model is actually good at the task we care about, not just at making sense of every pixel.

Lu: They also mentioned that this method's strength changes depending on the environment; it works well in some places but struggles in others, which is a fair limitation to acknowledge.

Meng: Yeah, so we can’t just take this one solution and apply it everywhere without checking the specific environment constraints first.

Tom: So, "Sparse Planning in Visual World Models via Cost Gradients" gives us a clear mechanism for selecting tokens based on planning cost gradients to achieve better performance with less compute.

Jane: It’s definitely a step forward in making these visual world models more efficient for real-time control tasks.

Lu: What we need to keep watching is how this selector interacts with different predictor architectures, because that seems like a big design axis for future work in sparse planning.

Meng: Next up, we’re going to look at some other stuff we’ve seen recently, like how memorization shows up in diffusion models.

Lalam: That sounds interesting; I want to see what those geometric features are telling us about model behavior.

Yingchen Xu, Edward Grefenstette

University College London

cs.LG

Submitted: 2026-10-07

Updated: 2026-10-07

Comments: Accepted at NeurIPS 2026. 20 pages, 6 figures, 8 tables. Project page and demos: https://ycxuyingchen.github.io/costgrad/

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 81/100

The gist: The gist The COSTGRAD method introduces a training-free, goal-conditioned selector that ranks spatial tokens by the gradient norm of the planning cost with respect to each input token, matching or

Key concepts

COSTGRAD: This is a training-free selector that ranks spatial tokens based on the gradient norm of the planning cost with respect to each input token. It identifies which visual tokens are most critical for achieving a successful plan, effectively creating a sparse set of relevant tokens.
Planning Cost: The planning cost measures how well a proposed action sequence reaches the goal in the world model. COSTGRAD uses this cost to calculate gradients, determining which input tokens have the strongest influence on minimizing this cost.
Token Sparsity: This refers to selecting only a subset of visual tokens deemed important for planning, rather than using all of them. By keeping only the most relevant tokens identified by COSTGRAD, the planning process can be made much faster and more efficient.
Action-Pathway Drift: This measures how much the spatial distribution of action-effect magnitudes changes when moving from one set of selected tokens to another. The paper finds that gradient-selected token removal causes less drift than random removal, indicating a better way to maintain planning performance.

Terminology

World Model — 1× in this paper · explained in 7 episodes

Summary

The gist The COSTGRAD method introduces a training-free, goal-conditioned selector that ranks spatial tokens by the gradient norm of the planning cost with respect to each input token, matching or exceeding full-token planning success on three out of four continuous-control benchmarks while providing a measured 2.6× wallclock speedup per environment planning step<ref:2610.10274#pg6>.

How it works

COSTGRAD scores tokens using the planning objective rather than the predictor’s training objective, by computing the gradient of the planning cost with respect to each input token and ranking tokens by gradient norm<ref:2610.10274#pg4>. Given an encoded observation and goal, it runs a one-step prediction on all P tokens under a probe action a0 (zero by default) and computes Lplan(a0; z, zg) = predictor(z, a0) − zg<ref:2610.10274#pg4>. A single backward pass then yields the per-token score sj = ∇z jLplan, which is the gradient norm of the planning cost with respect to each input token j<ref:2610.10274#pg4>. The retained tokens are those with the largest scores, defining a sparse token set S<ref:2610.10274#pg4>.

Planning and Speedup

Subsequent action search uses only the selected tokens, and COSTGRAD requires one extra forward and backward pass per planning step, operating on a fixed pretrained predictor without retraining<ref:2610.10274#pg3>. On AdaLN-conditioned predictors at 50% sparsity, COSTGRAD matches or exceeds full-token planning success on three of four environments and gives a measured 2.6× wall-clock speedup per environment planning step<ref:2610.10274#pg6>. Combining token sparsity with reduced CEM search increases this to a ∼5× total speedup while still exceeding the full-token baseline<ref:2610.10274#pg6>.

Performance and Relevance

COSTGRAD preserves or exceeds full-token performance on PointMaze (−2.0 pp, within noise) and exceeds Full on Wall (+7.3 pp), PushT (+2.8 pp), and MetaWorld (+6.9 pp), giving an average gain of +3.8 pp while using half the tokens<ref:2610.10274#pg6>. The kept tokens concentrate on task-relevant regions — maze corridors, the agent and wall, the pusher and block, the robot arm and object — while dropping much of the background<ref:2610.10274#pg7>. Prediction difficulty is not task relevance, as prediction error selects visually hard-to-predict tokens whereas planning needs tokens whose state affects the goal-reaching cost<ref:2610.10274#pg9>.

Selector–Architecture Compatibility

The more surprising finding is that COSTGRAD depends on how actions enter the predictor, comparing AdaLN [Peebles and Xie, 2023] which injects actions by modulating each visual token with action-dependent parameters, versus concat which appends the action to each token before self-attention<ref:2610.10274#pg6>. Pure COSTGRAD succeeds on AdaLN but loses its advantage on matched-concat, despite comparable full-token planning<ref:2610.10274#pg6>. This difference tracks action-pathway drift, as gradient-selected token removal produces less drift than random removal on AdaLN but more on concat<ref:2610.10274#pg6>.

Design Implications

The results highlight selector–architecture compatibility as a design axis for sparse world-model planning, suggesting that the planner is part of the model’s operating distribution and token selection changes how the planner uses the model<ref:2610.10274#pg6>. Sparse planning therefore depends on both choosing relevant tokens and on a predictor architecture that remains stable when restricted to those tokens<ref:2610.10274#pg6>. The strength of COSTGRAD is also environment-dependent, as attention-based selection is competitive on PointMaze and PushT while the gap is much larger on Wall and MetaWorld<ref:2610.10274#pg6>.

Limitations

The evaluation covers four simulated environments, one encoder family (DINOv2 ViT-S/14) and one planner (CEM), but physical-system evaluation with perception noise, calibration errors, and model mismatch remains future work<ref:2610.10274#pg14>. The selector–architecture finding is supported for the tested AdaLN and channel-wise concat predictors; its generality beyond this comparison remains untested<ref:2610.10274#pg10>. Future work should evaluate FiLM, cross-attention, and prefixstyle conditioning to test whether action-pathway drift predicts sparse-planning performance across conditioning mechanisms<ref:2610.10274#pg10>. The strength of COSTGRAD is also environment-dependent<ref:2610.10274#pg6>.

Methodology Details

The planner uses CEM [Rubinstein, 1999] in latent space, where at each environment step, CEM samples 300 action sequences over horizon H=6, ranks them by cost to the encoded goal, and refits the sampling distribution to the top elites for 30 iterations<ref:2610.10274#pg3>. The final-step CEM cost is defined as JS (A) = 1/K j Σ j zˆH,j (A) − zg,j squared + α dp pˆH(A) − pg squared<ref:2610.10274#pg15>. The probe-based selectors use zero action and proprioceptive inputs in these coordinates, and CEM rollout uses the observed proprioceptive input<ref:2610.10274#pg4>. The computational cost is one full-token forward and backward pass per planning step, which is amortized over all CEM rollouts<ref:2610.10274#pg4>.

Key Diagnostics

The diagnostic measures action-pathway drift DKL(p∥q) = P i∈S pi log(pi/qi), which quantifies changes in the spatial distribution of squared action-effect magnitudes<ref:2610.10274#pg9>. The easy-but-critical intersection measures disagreement between prediction error and planning sensitivity, identifying tokens with low prediction error but high planning-cost gradient<ref:2610.10274#pg10>. The reverse intersection is below the independence reference in every environment, so it does not show the same enrichment<ref:2610.10274#pg10>. The inverted-COSTGRAD sanity check shows that keeping the bottom-K tokens by gradient norm yields 0/60 successful episodes on AdaLN-Zero and 1/96 (1.04%) on matched-concat, demonstrating that reversing the ranking collapses planning on both predictors<ref:2610.10274#pg18>. The total predictor over CEM step planning call, excluding shared encoder cost, is 88.61 s for COSTGRAD-50% with 30 iterations<ref:2610.10274#pg10>.

Conclusion

Token selection and CEM budget are complementary compute axes, and the planner is part of a world model’s operating distribution<ref:2610.10274#pg6>. World models intended for sparse planning should be evaluated under the token subsets and action searches used at deployment because full-token prediction quality alone does not establish their suitability for sparse control<ref:2610.10274#pg14>. The planner is part of a world model’s operating distribution<ref:2610.10274#pg6>. The planner is part of a world model’s operating distribution<ref:2610.10274#pg6>. The planner is part of a world model’s operating distribution<ref:2610.10274#pg6>. The planner is part of a world model’s operating distribution<ref:2610.10274#pg6>.

Summary

COSTGRAD introduces a training-free, goal-conditioned selector that ranks spatial tokens by the gradient norm of the planning cost with respect to each input token, matching or exceeding full-token planning success on three out of four continuous-control benchmarks while providing a measured 2.6× wallclock speedup per environment planning step<ref:2610.10274#pg6>. The method works by computing the gradient of the planning cost with respect to each input token and ranking tokens by gradient norm, retaining those with the largest scores to define a sparse token set S<ref:2610.10274#pg4>.

Improvements for AI systems

Bold selection of task-relevant visual information during planning: COSTGRAD ranks spatial tokens by the gradient norm of the planning cost with respect to each input token, allowing planners to target tokens that matter for planning rather than merely for prediction. This enables the AI system to focus computational resources on regions critical for goal achievement, as evidenced by its ability to preserve or exceed full-token performance on challenging environments.
Speedup in inference: COSTGRAD provides a measured 2.6× wall-clock speedup per environment planning step by reducing token processing from full-grid computations to only the selected subset, which is further amplified when combined with reduced CEM search for a ∼5× total speedup. This allows for faster real-time decision-making in complex visual tasks.
Architecture-aware design: The system can be tuned based on the predictor architecture; the paper identifies an architecture-dependent failure mode, noting that gradient selection produces less drift than random removal on AdaLN, but more on concat, suggesting a design axis for sparse world-model planning. This guides engineers to select or modify predictors (like AdaLN vs. concat) based on desired stability in action pathways.
Improved control objective alignment: Unlike prediction-based selectors, COSTGRAD derives importance from the downstream control objective, addressing the objective mismatch documented in model-based reinforcement learning [Lambert et al., 2020]. This ensures that selected tokens are relevant to the goal-reaching objective, as seen on PushT where tokens critical for control receive high scores even if their prediction error is low.
Enhanced robustness to architectural variations: The system demonstrates probe-action robustness, with COSTGRAD rankings remaining stable across tested probe actions, indicating that the selection mechanism is less sensitive to minor changes in the action conditioning pathway compared to pure prediction-based methods.

Abstract

Token-based world models enable fine-grained latent planning, but repeatedly processing large spatial token grids makes action search expensive. We introduce COSTGRAD, a training-free, goal-conditioned selector that ranks spatial tokens by the gradient norm of the planning cost with respect to each input token. By deriving importance from the downstream control objective, COSTGRAD targets tokens that matter for planning rather than merely for prediction. On AdaLN-conditioned predictors at 50% sparsity, COSTGRAD matches or exceeds full-token planning on three of four continuous-control benchmarks, while giving a measured 2.6 times wall-clock speedup per environment planning step. Combining token sparsity with reduced CEM search increases this to a about 5 times total speedup while still exceeding the full-token baseline. We also identify an architecture-dependent failure mode: in a matched AdaLN-vs-concat comparison, concat maintains comparable full-token performance but pure COSTGRAD loses its advantage over random selection. This difference tracks action-pathway drift: gradient-selected removal produces less drift than random removal on AdaLN, but more on concat. These results highlight selector-architecture compatibility as a design axis for sparse world-model planning. Project page and demos: https://ycxuyingchen.github.io/costgrad/