Sparse Planning in Visual World Models
via Cost Gradients

Yingchen XuEdward Grefenstette

NeurIPS 2026

Paper & code forthcoming

01 / IN ACTION

Keep what matters for the goal.

PointMaze trajectory with retained visual tokens highlighted and discarded tokens dimmed — PointMazeNAVIGATION

Wall navigation trajectory showing the changing COSTGRAD token selection — WallNAVIGATION

PushT trajectory showing selected tokens around the pusher and T-shaped object — PushTMANIPULATION

MetaWorld reach trajectory showing selected tokens around the robot arm and workspace — MetaWorldREACH

Bright: retained tokens Dim: discarded tokens

Recorded rollouts · 50% tokens · retrained predictors

128/256

spatial tokens per sparse rollout

2.6×

measured planning-step speedup on Wall

gradient probe per planning step

02 / THE IDEA

Ask the planning objective.

Prediction and control need
different things from an image.

COSTGRAD scores each visual token by how strongly it affects the planning cost. A single gradient probe chooses the tokens; the planner then searches over actions using that smaller representation.

Probe the goal

Encode the observation and goal. Predict one step with a zero-action probe and measure the latent distance to the goal.

L_probe = ‖p(z, 0) − z^goal‖²

Score the tokens

Backpropagate the cost to the input tokens. Retain the K positions with the largest gradient norms.

s_j = ‖∂L_probe / ∂z_j‖₂

Plan on the subset

Run the full CEM action search on those positions. Keep the subset fixed within a planning step; select again when replanning.

S = TopK(s) → sparse rollouts

03 / PERFORMANCE

Less rollout compute.
Strong planning performance.

50% TOKEN RETENTION

On the evaluated retrained AdaLN predictors, COSTGRAD achieves higher mean success than Full on three of four environments.

Retrained AdaLN checkpoints: planning success (%), mean ± standard deviation over three evaluation seeds.
Environment	Full 256 tokens	COSTGRAD 128 tokens	Random 128 tokens	CG − Full percentage points
PointMaze	88.2 ± 8.7	86.2 ± 6.1	77.0 ± 9.4	−2.0
Wall	84.7 ± 2.6	92.0 ± 1.8	67.0 ± 3.6	+7.3
PushT	62.4 ± 4.4	65.2 ± 6.5	41.3 ± 5.9	+2.8
MetaWorld reach	58.7 ± 1.8	65.6 ± 4.7	46.9 ± 3.1	+6.9

96 episodes per evaluation seed on a fixed predictor per environment. Positive differences are observed mean improvements, not a claim of statistical significance.

Wall benchmark: COSTGRAD configurations form the measured speed–success frontier. At 50% tokens, fewer CEM iterations further reduce planning time. — **Token selection and search budget are complementary.** The Wall experiments measure a 2.6× planning-step speedup at 50% tokens, including selection overhead. Combining token selection with fewer CEM iterations yields approximately 5× speedup at the measured operating point.

04 / WHAT MAKES IT WORK?

Good tokens need
a compatible predictor.

Choosing relevant tokens is only half the problem. The predictor must still support useful action-conditioned rollouts when those tokens are the only ones present.

In matched AdaLN–concat experiments, COSTGRAD’s advantage over Random depends on how actions enter the predictor. Selected-token removal reduces measured action-pathway drift for AdaLN, but increases it for concat.

Read the scope and limitations

Action-pathway drift

Selected / random KL ratio · lower is less drift

AdaLN0.29 ± 0.03

Matched concat1.33 ± 0.17

Three training seeds per architecture. The ratio compares COSTGRAD-selected removal with random removal on the same predictor; it is a diagnostic, not a success guarantee.

Where the evidence stops

The empirical conditioning comparison covers AdaLN and channel-wise concat, with concat replication on Wall and PointMaze. FiLM and cross-attention remain untested. Action-pathway drift is evidence for a compatibility hypothesis, not a proven necessary or sufficient condition.

All experiments use visual token grids, CEM planning, and simulated continuous-control tasks. Public MetaWorld results show that inference-time applicability does not guarantee preserved performance on every checkpoint. The displayed MetaWorld demo uses the retrained reach predictor.

05 / REFERENCE

Cite this work.

@inproceedings{xu2026sparseplanning,
  title     = {Sparse Planning in Visual World Models
               via Cost Gradients},
  author    = {Xu, Yingchen and Grefenstette, Edward},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}

Sparse Planning in Visual World Modelsvia Cost Gradients