cs.LGOct 7, 2026
SaveSparse Planning in Visual World Models via Cost Gradients
Organizations: University College London
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 sparsity, COSTGRAD matches or exceeds full-token planning on three of four continuous-control benchmarks, while giving a measured wall-clock speedup per environment planning step. Combining token sparsity with reduced CEM search increases this to a 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/
Figures & tables
Figure 1 : CostGrad : select once, then plan on the selected tokens. (a) A full-token, zero-action probe supplies planning-cost gradients to select top- positions for context and goal. (b) CEM searches for actions using only , then executes the first action. The subset stays fixed throughout action search and is reselected at the next observation; model weights remain frozen. Token grids are schematic; proprioception is omitted.
| Method ( ) | PointMaze (MZ) | Wall | PushT (PT) | MetaWorld (MW) | Avg |
|---|---|---|---|---|---|
| Full ( ) | |||||
| CostGrad | |||||
| PredAttn | |||||
| PredGrad | |||||
| ToMe | |||||
| Random |
Table 1: Main results at tokens. Planning success rate ( ) across four continuous-control environments. Mean std over 3 evaluation seeds (96 episodes each). Full uses all tokens; all other methods keep tokens. Bold : best non-Full method per environment. CostGrad preserves or exceeds Full on three of four environments while giving a measured wall-clock speedup per environment planning step.
Figure 2 : Sparsity sweeps on four environments. Planning success ( ) vs. keep ratio for CostGrad , PredAttn , and Random; horizontal dashed line is Full ( ). Shaded bands: std across evaluation seeds at keep ratios / / ( episodes per cell per seed); the point uses Table 1 ’s 3-eval-seed mean. CostGrad approaches or exceeds Full at tokens on three of four environments (Wall, PushT, MetaWorld); the gap to prediction-derived baselines ( PredAttn , Random) widens as sparsity increases. All sparse methods degrade sharply below .
Figure 3 : Wall speed–success trade-off on AdaLN-Zero. (a) Five canonical operating points at the standard sparse setting. CostGrad -50% with CEM iterations achieves the same mean success as standard CostGrad -50% while reducing planning time to s, a speedup over Full. (b) Pareto frontier across methods at CEM iterations, plus the CostGrad , -iteration point. Among measured points, the non-dominated frontier is traced by CostGrad configurations; Full, PredAttn , and Random are dominated. Timing details and the full marker convention are in Appendix K and Appendix H .
Figure 4 : CostGrad selection across environments. Per-row: selection overlay at seven timesteps through one episode (bright = kept, dim = dropped, kept throughout). CostGrad adapts to each environment’s task structure: maze corridors on PointMaze, agent and wall on Wall, T-block and pusher on PushT, arm and object on MetaWorld. Decoded one-step predictions from the same sparse token sets are in Appendix C .
| Predictor | Full | CostGrad | Random | CG Rand | CG Full |
|---|---|---|---|---|---|
| AdaLN-Zero | |||||
| Concat (matched) |
Table 2: Full and sparse planning on Wall. AdaLN results are from Table 1 ; concat reports mean std across three training checkpoints, with episodes per method per checkpoint. CG Full : sparsification is profitable; : net harmful.
| Predictor | Selection | KL random | KL selected | Sel. / random |
|---|---|---|---|---|
| AdaLN-Zero | CostGrad | |||
| AdaLN-Zero | ||||
| Concat (matched) | CostGrad | |||
| Concat (matched) |
Table 3: Multi-seed action-pathway diagnostic on Wall. Mean sample standard deviation across three training seeds per architecture; , observations action probes per checkpoint. Selected/Random ratios below one indicate less drift than Random. uses random anchors.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5 : Decoded one-step predictions from CostGrad -50% sparse tokens. Per-row: the predictor’s decoded output at seven timesteps through one episode, when only the top- CostGrad -scored tokens are kept (corresponding to the selection overlays in Figure 4 ).
| Cost | Full | CostGrad ( ) | |
|---|---|---|---|
| L 2 | |||
| Cosine | |||
| L 1 |
Table 4: Wall planning success vs. choice of CEM cost function. The selector uses an L 2 probe throughout; its largest gain is under the L 1 CEM cost, where full-token planning performs worst.
Figure 6 : CostGrad operating curve on Wall AdaLN-Zero. Every measured CostGrad cell. Color encodes cluster (sequential viridis: purple , teal , green , yellow ); marker shape encodes iter count ( , open , , filled ). Thin lines connect points within each cluster (sorted by iter count). labels at the iters marker of each cluster. The headline operating point ( with iters , 17 s, ) gives a wall-clock speedup over Full at higher mean success.
| Environment | Seed 1 | Seed 2 | Seed 3 | Mean std |
|---|---|---|---|---|
| Wall | ||||
| PointMaze |
Table 5: Concat planning across environments and training seeds. Paired CostGrad –Random success-rate gaps (pp), with of tokens and episodes per method per checkpoint. Mean std is across the three training seeds.
| Environment | Easy + critical (count) | Multiple of expected |
|---|---|---|
| PushT (PT) | ||
| PointMaze (MZ) | ||
| Wall | ||
| MetaWorld (MW) |
Table 6: Easy-but-critical token counts per environment ( obs each, ): bottom-quartile prediction error intersected with top-quartile CostGrad score. Counts are mean std; the independence reference is . PushT is the only environment with enrichment above this reference; the other three are depleted.
| Component | K=256 (Full) | K=128 ( CostGrad ) |
|---|---|---|
| Encode context + goal (ms) | ||
| CostGrad forward + backward (ms) | — | |
| Predictor forward at batch (ms) | ||
| Per CEM iter ( , ms) | ||
| Total predictor over CEM step ( iters, s) | ||
| Planning call total, excl. shared encoder (s) |
Table 7: Timing breakdown for one CEM planning step ( candidates iterations ). CostGrad ’s extra forward+backward is a one-time cost amortized over CEM iterations, while the per-iteration cost drops by from the predictor-compute reduction. The reported total includes predictor rollout time and CostGrad overhead; the shared ms encoder cost is shown separately and is negligible at this scale. Net wall-clock speedup at vs. is .
| Environment | Full | CostGrad ( ) | Random ( ) |
|---|---|---|---|
| PointMaze | |||
| Wall | |||
| PushT | |||
| MetaWorld (Reach) |
Table 8: Planning success (%) on released checkpoints without retraining. Mean standard deviation over three evaluation seeds, with episodes per seed. Full uses all tokens; CostGrad and Random keep ( ).