Driving · 2026-10-08

Sparse Planning in Visual World Models via Cost Gradients

Yingchen Xu, Edward Grefenstette

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

The authors introduce CostGrad, a training-free, goal-conditioned token selector for sparse planning in token-based visual world models. It ranks spatial tokens by the gradient norm of the planning cost w.r.t. each input token, then runs CEM planning using only the selected tokens. On AdaLN-conditioned predictors at 50% sparsity, it matches or exceeds full-token planning on three of four continuous-control benchmarks while giving measured 2.6× wall-clock speedup per planning step.

Why it matters

Token-based world models make latent planning expensive because CEM must evaluate many action sequences across full spatial token grids. CostGrad uses the downstream planning objective (planning-cost gradients) to choose tokens that matter for control, improving the speed–success trade-off and revealing that sparse planning gains depend on the predictor’s action-conditioning architecture.

Method

Compute a one-step planning-cost gradient w.r.t. each token and score tokens by the gradient norm; keep the top-K tokens (goal-conditioned, training-free).
Run CEM for action search using only the selected token subset, keeping the subset fixed throughout each planning step; reselect at the next observation.
Compare against prediction-derived and attention-derived token selection and analyze an AdaLN vs concat action-conditioning failure mode via action-pathway drift diagnostics.

Limitation

The authors report an architecture-dependent failure mode: in a matched AdaLN-vs-concat comparison, pure CostGrad loses its advantage over random selection on concat, which they attribute to action-pathway drift differences when tokens are removed.

Abstract (from arXiv)

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 $\sim 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/

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-07 · Indexed by Pub-AI: 2026-10-08 UTC

preprint