Sparse Planning in Visual World Models via Cost Gradients

AI summary

Overview

Research area: Machine learning — model-based reinforcement learning and visual world models, specifically efficient inference-time planning with token-based (patch-based) latent dynamics models.

Technical level: Intermediate. The paper assumes familiarity with transformers, cross-entropy method (CEM) planning, action conditioning in transformers (AdaLN vs. concatenation), and gradient-based importance scoring, but its central idea can be grasped without deep mathematics.

Scope: The paper proposes and evaluates CostGrad, a training-free token selector that uses planning-cost gradients to choose a sparse subset of spatial tokens for goal-conditioned planning in a frozen visual world model, and it identifies a selector–architecture compatibility failure mode when the predictor's action pathway is changed.

What This Paper Is About

Token-based visual world models represent each observation as a grid of spatial tokens (here, a 16 × 16 grid of 384-dimensional tokens, i.e. P = 256), and plan by running CEM over many candidate action sequences rolled out inside the frozen model. This is expensive: one planning step performs roughly 300 × 30 × 6 × 256 ≈ 1.4 × 10^7 token-level forward computations. Most of those tokens describe background rather than the agent, objects, obstacle, and goal. The paper asks how a fixed, pretrained predictor can select the tokens that actually matter for the current goal — and shows that selecting tokens by the planning objective is different from selecting them by prediction quality, with results that depend on how the predictor is conditioned on actions.

Key Contributions

1. CostGrad, a training-free, goal-conditioned token selector. Given an encoded context and goal, CostGrad runs a single full-token, zero-action probe prediction, computes the gradient of the planning cost with respect to each input token, ranks tokens by gradient norm, keeps the top-K, and then runs CEM action search using only those tokens. Selection is reselected at each new observation; model weights stay frozen. At 50% sparsity (K = 128 of P = 256), it matches or exceeds full-token planning on three of four benchmarks and gives a measured 2.6× wall-clock speedup per environment planning step.

2. A combined sparsity-and-search compute strategy. Combining 50% token retention with a reduced CEM budget (15 iterations instead of 30) yields a ~5× total speedup while still exceeding the full-token baseline, and CostGrad configurations trace the non-dominated speed–success frontier across measured token budgets and CEM settings.

3. A selector–architecture compatibility failure mode. In a matched AdaLN-vs-concat comparison, both predictors achieve comparable full-token success on Wall (84.7% vs. 85.1%), but pure CostGrad beats Random by +25.0 pp on AdaLN-Zero while falling to −0.1 pp versus Random on matched-concat and 5.1 pp below Full.

4. An action-pathway drift diagnostic. The paper measures how token removal redistributes the spatial magnitude of action effects (a KL divergence between normalized action-effect distributions on and off the sparse support), showing CostGrad produces (0.29 ± 0.03)× Random drift on AdaLN-Zero but (1.33 ± 0.17)× on matched-concat — tracking the loss of planning advantage.

Main Findings

• Gains over full-token planning at half the tokens. On AdaLN-Zero at K = 128, CostGrad scores 86.2 ± 6.1 (PointMaze), 92.0 ± 1.8 (Wall), 65.2 ± 6.5 (PushT), and 65.6 ± 4.7 (MetaWorld), averaging 77.3 versus 73.5 for Full. Relative to Full: −2.0 pp on PointMaze (within noise), +7.3 pp on Wall, +2.8 pp on PushT, +6.9 pp on MetaWorld, for an average gain of +3.8 pp using half the tokens.

• Prediction-derived selectors are much weaker. The closest baseline, PredAttn, trails CostGrad by 12.1 pp on average and by more than 20 pp on Wall and MetaWorld. PredGrad averages 59.2 and ToMe 62.3, while Random averages 58.1, which is 19.2 pp below CostGrad. Full averages 73.5. Gains therefore are not a generic sparsification effect.

• Sparsity sweep. CostGrad stays near Full at 25% tokens on Wall, PushT, and MetaWorld and approaches Full at 50% on PointMaze. Below roughly 25%, all sparse methods degrade sharply, and the gap to PredAttn and Random widens as the token budget shrinks.

• Measured speedups. Retaining 50% of tokens at the standard CEM budget gives the 2.6× speedup; the extra forward and backward pass adds 27.5 ms on a single H100, amortized over subsequent CEM iterations. Halving tokens reduces attention flops by 4× and MLP flops by 2×. Combining 50% tokens with 15 CEM iterations reaches 92.0% success at ~17 s per planning step — a 5.2× total speedup over Full (84.7% at 88.6 s). At the same 15-iteration budget, Full achieves 84.4 ± 1.0% versus CostGrad's 92.0 ± 5.1%. CostGrad at K = 192 reaches 92.7% success at 60.4 s per step (+8.0 pp and 1.47× faster than Full), with success entering a plateau around K ≥ 64 and best points near ~92%.

• Released-checkpoint transfer. On released JEPA-WMs checkpoints, CostGrad preserves full-token performance across all four environments without retraining, using half the tokens and outperforming Random by 11.3–46.3 pp (Appendix L, Table 8).

• Selection is task-relevant, not just prediction-hard. Visualizations show kept tokens concentrating on maze corridors, the agent and wall, the pusher and T-block, and the robot arm and object. A quadrant analysis finds PushT uniquely contains many "easy-but-critical" tokens (low prediction error, high planning gradient) at ~2× chance density, while the other three environments are depleted (Appendix J) — prediction-derived selection can downweight exactly the objects that matter for control.

• Attention selection is unstable across seeds; CostGrad is stable. Three Wall AdaLN-Zero predictors with indistinguishable validation loss produce top-128 attention selections overlapping at only IoU 0.44 ± 0.02, with PredAttn planning success varying by 50 pp across seeds. CostGrad selections overlap at IoU > 0.82 with success within evaluation noise.

• Probe-action robustness. CostGrad rankings are stable across 10 tested probe actions (ρ = 0.918 on AdaLN-Zero, 0.958 on matched-concat). On AdaLN-Zero, Gaussian-noise or CEM-mean probes instead of the zero-action probe change Wall success by at most ±5 pp (Appendix E).

• Architecture dependence. Across Wall and PointMaze, CostGrad has no consistent advantage over Random on concat: paired gaps are −0.1 ± 4.9 pp (Wall) and +2.0 ± 4.5 pp (PointMaze). Reversing the ranking (keeping bottom-gradient tokens) collapses planning on both predictors (Appendix I).

• Drift diagnostic and anchor mixing. KL drift ratios versus Random: AdaLN-Zero CostGrad (0.29 ± 0.03)×, AdaLN-Zero with α = 0.8 (80% random anchors) (0.44 ± 0.02)×, matched-concat CostGrad (1.33 ± 0.17)×, matched-concat with α = 0.8 (1.04 ± 0.11)×. Anchor mixing narrows concat drift but planning recovery is checkpoint-dependent, so it is a diagnostic, not a generally better selector.

• Sharper cost landscapes amplify the gain. On Wall, CostGrad's advantage over Full triples when the cost is more degenerate: +22.9 pp under L1 versus +7.3 pp under L2 (Appendix F).

Methodology in Plain English

The researchers start from an existing setup: a frozen DINOv2 ViT-S/14 encoder turns each 224 × 224 image into a 16 × 16 grid of 384-dimensional tokens (P = 256), and a small 6-block transformer predictor (about 17.6M parameters) predicts the next token grid given the current grid and an action. Planning is done with CEM: sample 300 action sequences, roll each out for horizon H = 6, score by distance to the encoded goal, and refit the sampling distribution to the 10 best candidates for 30 iterations. Because every candidate rollout touches all 256 tokens for all 6 steps, the compute adds up fast.

Their selector works in two steps. First, before any CEM search, they run one full-token prediction using a dummy zero action and measure how far the predicted grid lands from the goal grid (a plain squared-error cost). One backward pass tells them how sensitive that cost is to each individual input token — the norm of the gradient at that token. Tokens with large gradient norms are kept; the rest are dropped. This is why the approach is "planning-aware": the score comes from the downstream control objective, not from the model's training loss, so a token can score highly either by being goal-relevant itself or by supplying context for predicting goal-relevant regions. Second, CEM is run as usual, but only over the kept positions, which stay fixed for the whole planning step and are reselected at the next observation. The method adds one forward and backward pass per planning step, needs no retraining, and depends on the current goal.

To test whether this is really about planning rather than prediction, they compare against selectors based on attention received in the last block (PredAttn), prediction-loss gradients (PredGrad), token merging (ToMe), and uniform random selection, all at 50% sparsity with the same predictor and CEM settings. To test architecture dependence, they train a matched "concat" predictor that differs from AdaLN-Zero only in how the action enters: concat appends a 10-dimensional frame-skip action vector to each visual token before the query/key/value projections, whereas AdaLN modulates each token with action-dependent scales, shifts, and gates at every block. Finally, they design a drift metric: measure how the spatial distribution of action effects (the difference between a prediction under action a and under the zero action) changes when only the kept tokens are used, normalized over the same kept positions, and compare CostGrad against random selection at the same budget.

Evaluation uses four continuous-control benchmarks from the DINO-WM and JEPA-WMs suites: PointMaze (U-maze navigation, 2-dim actions, horizon

Authors’ 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 $\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/