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× wall-clock speedup per environment planning step. Combining token sparsity with reduced CEM search increases this to a ∼5× 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- K positions S for context and goal. (b) CEM searches for actions using only S , 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 ( 50% )
PointMaze (MZ)
Wall
PushT (PT)
MetaWorld (MW)
Avg
Full ( 100% )
88.2±8.7
84.7±2.6
62.4±4.4
58.7±1.8
73.5
CostGrad
86.2±6.1
92.0±1.8
65.2±6.5
65.6±4.7
77.3
PredAttn
86.7±9.7
67.3±1.7
62.0±8.0
44.8±6.1
65.2
PredGrad
76.8±9.2
67.2±3.3
41.3±5.9
51.4±3.5
59.2
ToMe
74.4±9.1
84.4±2.6
51.4±5.9
39.1±4.1
62.3
Random
77.0±9.4
67.0±3.6
41.3±5.9
46.9±3.1
58.1
Table 1: Main results at 50% tokens. Planning success rate ( % ) across four continuous-control environments. Mean ± std over 3 evaluation seeds (96 episodes each). Full uses all 256 tokens; all other methods keep 128 tokens. Bold : best non-Full method per environment. CostGrad preserves or exceeds Full on three of four environments while giving a measured 2.6× 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 ( 100% ). Shaded bands: ±1 std across 3 evaluation seeds at keep ratios 10 / 25 / 75 ( 96 episodes per cell per seed); the 50% point uses Table 1 ’s 3-eval-seed mean. CostGrad approaches or exceeds Full at 25% 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 ∼25% .
Figure 3 : Wall speed–success trade-off on AdaLN-Zero. (a) Five canonical operating points at the standard K=128 sparse setting. CostGrad -50% with 15 CEM iterations achieves the same mean success as standard CostGrad -50% while reducing planning time to ∼17 s, a 5.2× speedup over Full. (b) Pareto frontier across methods at 30 CEM iterations, plus the CostGrad K=128 , 15 -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 (K,iters) 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, 50% 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
84.7±2.6
92.0±1.8
67.0±3.6
+25.0
+7.3
Concat (matched)
85.1±5.2
80.0±2.9
80.2±3.8
−0.1
−5.1
Table 2: Full and sparse planning on Wall. AdaLN results are from Table 1 ; concat reports mean ± std across three training checkpoints, with 96 episodes per method per checkpoint. CG − Full >0 : sparsification is profitable; <0 : net harmful.
Predictor
Selection
KL random
KL selected
Sel. / random
AdaLN-Zero
CostGrad
0.265±0.005
0.077±0.009
(0.29±0.03)×
AdaLN-Zero
α=0.8
0.265±0.005
0.116±0.009
(0.44±0.02)×
Concat (matched)
CostGrad
0.287±0.017
0.380±0.030
(1.33±0.17)×
Concat (matched)
α=0.8
0.287±0.017
0.298±0.016
(1.04±0.11)×
Table 3: Multi-seed action-pathway diagnostic on Wall. Mean ± sample standard deviation across three training seeds per architecture; K=128 , 20 observations ×5 action probes per checkpoint. Selected/Random ratios below one indicate less drift than Random. α=0.8 uses 80% 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 t+1 output at seven timesteps through one episode, when only the top- 128 CostGrad -scored tokens are kept (corresponding to the selection overlays in Figure 4 ).
Cost
Full
CostGrad ( 50% )
Δ
L 2
84.7
92.0
+7.3
Cosine
82.3
87.5
+5.2
L 1
47.9
70.8
+22.9
Appendix
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 (K,iters) CostGrad cell. Color encodes K cluster (sequential viridis: purple K=26 , teal K=64 , green K=128 , yellow K=192 ); marker shape encodes iter count ( ◊=10 , open ◯=15 , △=20 , filled ∙=30 ). Thin lines connect points within each K cluster (sorted by iter count). K labels at the iters =30 marker of each cluster. The headline operating point ( K=128 with iters =15 , ∼ 17 s, 92.0±5.1% ) gives a 5.2× wall-clock speedup over Full at higher mean success.
Environment
Seed 1
Seed 2
Seed 3
Mean ± std
Wall
+2.1
+3.2
−5.7
−0.1±4.9
PointMaze
+4.1
−3.2
+5.2
+2.0±4.5
Appendix
Table 5: Concat planning across environments and training seeds. Paired CostGrad –Random success-rate gaps (pp), with K=128 of 256 tokens and 96 episodes per method per checkpoint. Mean ± std is across the three training seeds.
Environment
Easy + critical (count)
Multiple of expected
PushT (PT)
31.9±4.3
1.99×
PointMaze (MZ)
10.7±2.8
0.67×
Wall
8.5±4.3
0.53×
MetaWorld (MW)
3.2±2.2
0.20×
Appendix
Table 6: Easy-but-critical token counts per environment ( 30 obs each, P=256 ): bottom-quartile prediction error intersected with top-quartile CostGrad score. Counts are mean ± std; the independence reference is 16 . 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)
7.64
7.64
CostGrad forward + backward (ms)
—
27.53
Predictor forward at batch 300 (ms)
492.25
188.90
Per CEM iter ( ×H=6 , ms)
2953.51
1133.42
Total predictor over CEM step ( ×30 iters, s)
88.61
34.00
Planning call total, excl. shared encoder (s)
88.61
34.04
Appendix
Table 7: Timing breakdown for one CEM planning step ( 300 candidates ×30 iterations ×H=6 ). CostGrad ’s extra forward+backward is a one-time cost amortized over 30 CEM iterations, while the per-iteration cost drops by 2.61× from the predictor-compute reduction. The reported total includes predictor rollout time and CostGrad overhead; the shared 7.64 ms encoder cost is shown separately and is negligible at this scale. Net wall-clock speedup at K=128 vs. K=256 is 2.60× .
Environment
Full
CostGrad ( 50% )
Random ( 50% )
PointMaze
87.7±7.7
86.3±11.9
70.7±15.1
Wall
87.0±2.8
87.0±2.2
75.7±4.6
PushT
66.3±2.6
67.3±3.4
21.0±9.0
MetaWorld (Reach)
58.3±10.0
63.2±5.2
20.8±9.0
Appendix
Table 8: Planning success (%) on released checkpoints without retraining. Mean ± standard deviation over three evaluation seeds, with 96 episodes per seed. Full uses all 256 tokens; CostGrad and Random keep 128 ( 50% ).
Visual world models enable goal-directed control through decision-time action search, but their deployment efficiency is often limited by conservatively large planning budgets. We show that competitive task performance can be achieved without agreement with the Full-budget action, that sufficient budgets vary across model--task pairs, and that iterative planners repeatedly encode solve-invariant context. To address these inefficiencies, we propose {SufficientPlan}, a simple deployment framework that requires no modification to pretrained world models or planner updates. Its {Paired Sequential Budget Certification (PSBC)} component uses paired closed-loop evidence to search for and certify a reduced model--task-specific budget within a predefined Full-performance tolerance. Its {Static-Context Reuse (SCR)} component caches observation and goal representations across search iterations while preserving candidate-dependent planning and selected actions. Experiments across multiple world-model backbones and visual-control tasks show that SufficientPlan substantially reduces search budgets and planning latency while maintaining competitive control performance.
Changbai Li, Sirui Li, Yichen Yang +4
Beihang University · Chengdu University of Technology · Nanyang Technological University
Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.
Planning with a learned latent world model is a promising route to control from raw pixels, but a strong world model alone is not enough. We show this experimentally: even with a perfect world model (operationalized by replacing the learned forward predictor with an idealized rollout of the true environment dynamics), a finite-budget sample-based planner still fails on some tasks, indicating that the bottleneck can lie in search rather than in world-model accuracy. Motivated by this gap, we propose IMWM (Intuition Model + World Model), which pairs the world model with an intuition model trained from demonstrations to recognize promising actions. The two models collaborate through three lightweight components: (i) Retrieval Initialization, which initializes the planner's action proposal from a retrieved demonstration; (ii) Hybrid Cost, which combines the intuition score with the world-model rollout cost; and (iii) a Reliability Gate, which adjusts how much the planner trusts intuition in each setting. Across four pixel-based goal-reaching tasks (Two-Room, Reacher, Push-T, and OGBench-Cube), IMWM has higher mean success than the world-model-only planner on all four, with the largest gains on Two-Room (99.2%, +11.5 percentage points) and OGBench-Cube (94.7%, +28.5 percentage points).
Baoqi Gao, Ruize Han, Miao Wang +1
Beihang University · Shenzhen University of Advanced Technology