Sparse Imagination for Efficient Visual World Model Planning
Authors: Junha Chun
Youngjoon Jeong
Taesup Kim
Venue: ICLR 2026 Poster (2026)
Keywords: World Model
Planning
Computational Efficiency
Model Predictive Control
Vision Transformer
Primary area: applications to robotics, autonomy, planning
TL;DR: To overcome the high computational cost of world models, our method uses a sparse imagination approach to achieve faster planning for real-time applications while maintaining high task performance.

Abstract: World model based planning has significantly improved decision-making in complex environments by enabling agents to simulate future states and make informed choices.
This computational burden is particularly restrictive in robotics, where resources are severely constrained.
To address this limitation, we propose a Sparse Imagination for Efficient Visual World Model Planning, which enhances computational efficiency by reducing the number of tokens processed during forward prediction. 
Our method leverages a sparsely trained vision-based world model based on transformers with randomized grouped attention strategy, allowing the model to flexibly adjust the number of tokens processed based on the computational resource.
By enabling sparse imagination during latent rollout, our approach significantly accelerates planning while maintaining high control fidelity.
Experimental results demonstrate that sparse imagination preserves task performance while dramatically improving inference efficiency. 
This general technique for visual planning is applicable from simple test-time trajectory optimization to complex real-world tasks with the latest VLAs, enabling the deployment of world models in real-time scenarios.

Published as a conference paper at ICLR 2026
SPARSE IMAGINATION
FOR EFFICIENT VISUAL WORLD MODEL PLANNING
Junha Chun Youngjoon Jeong' Taesup Kim If
Graduate School of Data Science, Scoul National University 2 *Department of Electrical and Computer Engineering, Seoul National University
ABSTRACT
World model based planning has significantly improved decision-making in com- plex environments by enabling agents to simulate future states and make informed choices. This computational burden is particularly restrictive in robotics, where resources are severely constrained. To address this limitation, we propose a Sparse Imagination for Efficient Visual World Model Planning, which enhances compu- tational efficiency by reducing the number of tokens processed during forward prediction. Our method leverages a sparsely trained vision-based world model based on transformers with randomized grouped attention strategy, allowing the model flexibly adjust the number of tokens processed based on the compu- tational resource. By enabling sparse imagination during latent rollout, our ap- proach significantly accelerates planning while maintaining high control fidelity. Experimental results demonstrate that sparse imagination preserves task perfor- mance while dramatically improving inference efficiency. This general technique for visual planning is applicable from simple test-time trajectory optimization complex real-world tasks with the latest VLAs, enabling the deployment of world models in real-time scenarios.
INTRODUCTION
By "imagining" future trajectories a learned world model of the enviromment, agents can perform sophisticated decision making without trial-and-error in the real world (Ha & Schmidhuber) 2018 Hafner et al., 2019a|b) b). Recent world model advancements for visual control tasks, which neces-
sitate reasoning over high-dimensional image observations, increasingly leverage high-level visual features derived from self-supervised vision models (Caron et al. [2011] Oquab et al. 2023 Nair al. 2022; Xiao et al. 2022). For example, DINO-WM (Zhou et al. 2024) introduces a world model that predicts future Vision Transformer (ViT) patch tokens (DINO features) instead of pixels or single vector representation, enabling zero-shot planning for diverse tasks. These approaches re- tain rich spatial tokens, achieving strong generalization on complex control tasks. However, using numerous visual tokens for planning incurs a prohibitive quadratic computational cost, as requires many simulated rollouts. This raises crucial question:
Can we retain the advantages of detailed visual world models while enhancing computational efficiency for planning?
In this paper, address the above question by introducing the concept of sparse imagination, accelerating world model inference by deliberately using only sparse subset of visual tokens during forward prediction. Our key insight is that ViT-based representations are redundant (Raghu al.,
2021 Pan et al.) 2021; Chen et al. 2022; Kim et al. 2024 Simééni et al. 2025): thus, not all patch tokens are necessary for planning.
To achieve this, we propose a random dropout-based token selection mechanism for the world model during inference that dynamically encodes and predicts only subset of patch tokens. Furthermore, we train the world model using randomized grouped attention to accommodate dynamic token se- lection. The resulting system performs planning on the subset of latent patch tokens, reducing
'These authors contributed equally.
1 "Corrsponning author.

---

Published as a conference paper at ICLR 2026
computational cost while maintaining competitive performance. The primary contributions of this work are as follows:
We introduce sparse imagination, a simple yet effective method that enables efficient visual world model planning by leveraging random patch feature dropout during inference.
We demonstrate that our approach is a general technique applicable from simple test-time tra- jectory optimization to complex real-world tasks with VLAs. It achieves substantial planning speedups while maintaining high performance, validated on benchmarks like LIBERO- 10 (Liu et al. 2023a). Meta- World (Yu al. 202I), and on real-wordd wal-world robotic tasks.
Our comparative analysis shows that simple random sampling achieves competitive or superior performance to sophisticated token selection methods due to fundamental "blind spot" prob- lem we identify, where static importance metries fail dynamic planning, making the unbiased coverage of random sampling more robust strategy.
RELATED WORK
World Models for Control. Early world models transitioned from computationally heavy pixel- level models (Finn et al. 2016 Ebert et al., 2018) to latent dynamics models like PlaNet (Hafer jet al. 2019b) and Dreamer (Hafner et al., 2019a 2020 2023), which achieved impressive results by leveraging planning and latent imagination. W While most existing latent world models compress entire images into relatively low-dimensional vectors, these compact representations often lose the fine-grained spatial information necessary for high-precision manipulation tasks (Zhou et al. 2024
Tsagkas et al. 2025).
Visual Representations for Decision Making. A recent trend in decision making leverages pow- erful pre-trained visual encoders for rich state representations (Nair et al. 2022 Xiao et al. 2022) While effective for representation learning, methods using a single image embedding (e.g.. CLS token) from encoders like ResNet VIT often lose crucial fine-grained spatial details, hindering performance on tasks requiring precise spatial understanding, such as manipulation (Tsagkas et al. 2025 Zhou et al., 2024). To address this, approaches using spatial features like ViT patch em- beddings for decision making have been developed, exemplified by DINO-WIY (Zhou et al.) 2024) which uses set of DINO patch tokens as the state. Building on this paradigm, our work utilizes pre-trained ViT patch features for their spatial granularity. Our innovation is an inference-ting token dropping mechanism that significantly enhances computational efficiency during imagination.
Vision Transformers and Token Efficiency. The quadratic scaling of ViT self-attention poses significant computational challenge, particularly for sequential models like world models (Dosovit- skiy et al. 2020). Improving inference efficiency via token reduction thus a key vision research area. Existing methods employ diverse strategies such as learned/attention based selection (Rao ict al. 2021 Luo et al.) 2024]] Yin et al.) 2022, Liang et al., 2202]] Wang et al. 2022. Zhang et al., 2024b), merging (Bolya ya et al. 2022: Feng & Zhang 2023 Haurum et al. (222)) and training- time dropout (Liu et al. 2023b). These demonstrate efficiency gains from ViT representation redun- dancy (Raghu et al. 2021 Pan al.) 2021 Chen et al.] 2022) Kim al.) 2024). In contrast, present a patch-dropping method specifically for inference-time world model planning. Our ap- proach uniquely applies dropout only during the imagination phase, enabling seamless integration with any transformer world model.
METHODOL.OGY
3.1
WORLD MODEL
Our world model is composed of two key modules: (1) pre-trained image encoder do (e.g., DINO) generating visual latent patch tokens from observations, and (2) a transformer-batt based world model fe predicting future latent tokens based historical observations and actions. The image encoder remains fixed during training to ensure stable visual representations, while the transformer based world model learns temporal dynamics. Structured causal transformer decoder, the world model
2

---

Published as a conference paper at ICLR 2026
0t-hit
Sparse Imagination
0c-1+1:41
- hit
gg
21-A+1#+1
DIOO
WM
DIOO
MPC Step
Random Token Drop
Random
Actions
Lmpe
LIND
DIO
g
Og
0g
Figure 1: Sparse Imagination. Sparse imagination accelerates planning by performing model pre- dictive control (MPC) rollouts on random subset of visual tokens. new dropout pattern is dynam- ically sampled at each MPC iteration, and both predictions and optimization for CEM are computed using only the selected patches improve efficiency and robustness.
allows each token to attend only to tokens from prior time steps, enabling efficient sequence model- ing.
Specifically, given an image observation or Hx at time t, the image encoder Is extracts visual features 2, E R=,,,,,>,, where the number of visual tokens is H, W,, The world model Se predicts future features 2++1 using temporal context of visual tokens and actions: Zx411 = [=(===Aze Ast, at-A:t,, with a history length of + 1. Our model predicts dynamics within the pre-trained DINO feature space. Actions and proprioceptive information, when available, linearly projected and concatenated to visual tokens along the feature dimension. The world model training is guided by minimizing the mean squared error (MSH) prediction loss:
N
1
Lwm
Σ N
11 - 2411,11
(1)
where and N represent token index and the total visual tokens per frame, respectively. An optional decoder, detached from gradient updates, can be trained separately for visualization purposes.
To enhance generalization and robustness when only sparse subsets tokens are provided as inputs, introduce randomized grouping strategy during training depicted in Fig. 2 Specifically, visual tokens from each frame are randomly partitioned into two groups. This approach ensures different subset patterns during training, promoting adaptability across different levels of sparsity. To implement this, we apply attention masks within the transformer layers, restricting token interactions exclusively those within the same spatial group while maintaining temporal consistency. By enforcing this structured token separation, our model learns to process arbitrary subsets of visual tokens while optimizing for the same prediction loss, ultimately improving its ability to generalize across varying input conditions.
3.2 SPARSE IMAGINATION FOR MPC
We MPC for planning, and for each step, candidate action sequences {0r.4 H-1 are sampled from distribution (CEM) or pre-trained policy. The trained world model fo then predicts future states by iteratively rolling out these action sequences from the current state. In a goal-conditioned setting, we evaluate candidate sequences using the MSE loss between the predicted and goal state features:
Cape || 2.
When using CEM for test-time trajectory optimization, the distribution candidate actions is sub- sequently updated based the op-performing sequences. This iterative optimization continues for a predefined iteration step M, after which the optimized action sequence executed in the environ- ment.
3

---

Published as a conference paper at ICLR 2026
Transformer
Ot hct
0++1
hct
Afterdion Map
Rt+1
2441
Lum
DINO
ONO
Randomized Grouped Attention Strategy. Our randomized grouping strategy used during training generalize to arbitrary token subsets. Visual tokens randomly partitioned into two groups, and attention is masked to occur only within the same spatial group. This trains the model process sparse inputs effectively, and its necessity shown in ablations.
Figure 2:
MPC is computationally intensive, world model rollout occurs K x M x H per every planning step, and this cost scales quadratically with the number of visual tokens (the number of candidate sequences is K and planning horizon is H), making real-time exccution increasingly challenging. Since attention layers in the world model are the primary source this computational burden, reducing the number of tokens processed is crucial for improving planning efficiency.
We address this problem through sparse imagination. Given that our model is explicitly trained to operate arbitrary token subsets, we implement random token dropout during planning as shown in Fig. T At each planning step, we randomly generate a token dropout mask characterized by dropout fraction p [0. 1) of the user's choice, denoting the proportion of tokens discard, keeping randomly sampled (1 N) tokens. Sparse rollout imagination can effectively reduce redundant computations and improve planning efficiency. Planning a single step might fall to wrong direction when dropping key features but has minimal impact on task performance since it can recover from failures by cemaapllngg at each planning iteration. The drop ratio p serves a controllable parameter balancing computational speed and task accuracy. We empirically evaluate and analyze this trade-off' the following Sections.
EXPERIMENTS
4.1 IMPLEMENTATION DETAILS
Environments and Datasets. We evaluate our method on eight simulated environment and two real-world tasks using LeRobot (Cadene et al. 2024) (See Fig. 3.. Detailed descriptions of the environments, datasets, and model configurations provided in the Appendix A
CrAaied 40tontinn
F.ll Attencionn
00.0%
23.35
Rope
Pushing
Dietonmert
Figure 3: Environments. We assess the world model across eight simulation environments and two real-world tasks, arranged follows tleft-to-right, top-to-bottom): Pointmaze, Wall, PushT, Granular, Rope, Block Pushing, LIBERO, Meta-World and LeRobot tasks (PickPlace & Drawer).
Figure 4: Contribution of grouped attention during pretraining to plan- ning. We compare the average plan- ning success rates achieved by grouped attention and full attention under vari- ous drop ratios, highlighting the benefit of grouped attention.
4

---

Published as a conference paper at ICLR 2026
Table 1: Performance results (Mean Success Rate, across different environments with varying drop ratios.
Pointmaze (60 evals)
PushT (60 evals)
Granular (20 evals)
Rope evals)
Block Pushing (50 evals)
(60 evals)
Avg.
Full CLS
98.3
91.7
75.0
75.0
63.3
22.0
69.8
96.7
85.0
43.3
20.0
36.7
16.0
50.7
Drop (Ratio)
91.7
78.3
80.0
73.3
28.0
74.1
80.0
66.7
16.0
69.6
30%
98.3
61.7
85.0
70.0
18.0
71.0
40%
96.7
56.7
70.0
76.7
18.0
68.6
50%
100.0
95.0
70.0
60.0
73.3
20.0
69.7
60%
83.3
91.7
46.7
70.0
70.0
20.0
63.6
70%
70.0
93.3
28.3
55.0
53.3
20.0
53.3
80%
73.3
90.0
21.7
50.0
66.7
24.0
54.3
90%
71.7
86.7
20.0
40.0
50.0
12.0
46.7
In each environment, the goal is reaching a target observation randomly sampled as a desired goal state, starting from initial positions. For the LIBERO environment, we conducted experiments on LIBERO- 10, suite of Long-horizon tasks, to demonstrate the effectiveness of our planning acceler- ation. For the real world, the task is to pick the blue block and place it in the cup.
Baselines. We evaluate our approach across several visual control domains and compare it against key baselines. The primary baselines include: (1) Full-Patch world model, DINO-WM (Zhou et al. 2024) based on DINO (Caron et al. 2021). which shares the same architecture as ours but utilizes patches during planning (i.e., p 0) and applies positional encoding to each patch- representing high-fidelity yet computationally expensive setting: and (2) CLS-Token world model. which employs only the global CLS token from the DINO encoder as the visual representa- tion, thereby offering single-vector representation bascline. These two baselines are based on the approach presented in DINO-WM.
Planning Setup.
To evaluate our framework across diverse scenarios, we utilize three distinct
planning schemes:
MPC-CEM: For Pointmaze, Wall, PushT, Block Pushing, we use MPC-CEM. This in- volves replanning at each step in receding- horizon manner.
CEM (Open-loop): For computationally intensive deformable object simulations (Granu- lar, Rope), we employ a standard open-loop CEM approach. This optimizes single, fixed sequence of actions without replanning manage the high processing cost.
Policy-Guided Planning: For complex, long- long-horizon tasks (LIBERO, Meta- Woold,,,d,,, LeR- obot), where random action sampling is inefficient, we guide the planner by sampling candidate action sequences from a pre-trained policy such as Vision-LaanuaggeAAtion (VLA) models. Our planner then evaluates these candidates.
Our contribution highlights significant efficiency gains while preserving Full-Patch (DINO-WM) baseline performance, demonstrating our sparse imagination as effective, computationally effi cient replacement, rather than focusing on absolute performance maximization. Therefore, for fair comparison, most of hyperparameters are adopted from the DINO-WM baseline, as detailed in the Appendix A.3 and B.I
4.2
TEST-TIME TRAJECTORY OPTIMIZATION
Our sparse imagination framework is first evaluated on test-time trajectory optimization using MPC- CEM or CEM. Moderate token dropout (10-50%) maintains task performance relative to the Full- Patch baseline while substantially reducing planning time. For example, in the PushT environment, a 50% dropout ratio cuts the planning time per iteration from 173s to 82s (a 52.6% reduction) without sacrificing performance (Table
5

---

Published as a conference paper at ICLR 2026
Table Planning time and change compared Full baseline for different environments at various drop ratios.
Pointaaee
Wall
PushT
Block Pashing
Avg.
Planning Time (vitte)
Chinge (96)
Planning Time (xlter)
Change (%)
Planning Time (slter)
Change (%)
Planning Time (sliter
Change (96
Planning Time (witer)
Change
Full CLS
184
0.0
79
0.0
173
0.0
297
0.0
183.3
0.0
-734
-19.1
32
-$1.5
163
-45.1
71,0
-61.3
Drop (Ratio)
10%
-10.3
-7.6
-1399
278
166.3
-9.3
20%
141
-23.4
69
-12.7
131
-24.3
259
-1288
150.0
-18.1
30%
126
-31.5
65
-17.7
114
-341
240
-19.2
136.3
-25.6
40%
-42.4
-21.5
43.9
-23.9
119.8
-34.7
50%
-49.5
-32.9
-52.6
208
-30.0
109
-40.5
60%
80
-36.5
-36.7
60
461.1
200
-32.7
99.8
-55.6
70%
69
-62.5
46
-4.88
59
45.9
184
-38.0
89.5
-11.2
80%
-09.6
-1.8
-71.7
175
81.5
-55.5
-73.9
-6.88
Table 3: Performance results and planning time across difficulties in Meta- World simulation.
All
Easy (96
Medium
(%)
Very Hard
Planning Time (slepisode)
VLA-only
41.87
60.95
20.61
17.78
10.67
1.95
Full
48.80
66.90
29.09
28.89
14.67
3.63
Drop 50%
47.73
65.24
27.27
33.33
12.00
2.37
This efficiency gain is especially critical complex environments. Although the CLS-Token of- fers faster inference, its dramatic performance drop in these scenarios renders it impractical. For instance, in spatially demanding scenes like Granular and Rope, our approach achieves up to 85.0% and 76.7% success respectively, whereas the CLS-Token fails (20.0%, 36.7%) by losing critical spatial information. Similarly, on tasks like PushT, our method with 50% dropout (70.0%) surpasses CLS-Token (43.3%) baselines. Refer to Appendix B.8 for more detailed and fine-grained analysis.
4.3 TOWARDS COMPLEX AND REAL-WORLD TASKS: VLA-GUIDED PLANNING
To tackle complex and long- tasks, we integrate our planning framework with the latest VIAA model such SmoIVL.A (Shukor et al. (025). We evaluate this approach on the LIBERO-10 and Meta-World simulation and the real-world enviromments using SO-101 robot arm, focusing on two tasks: PickPlace (picking up block and placing into cup) and Drawer (picking up block, placing it into drawer, and then closing the drawer). We find that guiding the search with pre- trained policy crucial in these demanding settings. For the real-world LeRobot experiments PickPlace and Drawer, we used a manually selected goal image and proprioceptive state from the dataset for planning (See Appendix A.8 for details).
Simulation Results. The SmolVLA policy used for simulation experiments is trained to predict action chunks spanning 50 steps and at each decision step we sample 50 candidate action chunks from the policy for LIBERO-10 and 10 for Meta- World to perform planning- Wlasedng-based evaluation. In the original work of SmoIVLA, only a single action is performed in every policy inference. Since we find impractial deploy single-step policies in real world, we execute all actions in the action chunk despite lower absolute performance. In LIBERO-10 10 evaluation. we test our method in the evaluation benchmark setting with total of 100 trials. Our method (50% token drop, * in the Fig- 5) boosted the base VLA's success rate from 29% to 33%. This performance matches the computa- tionally expensive Full-Patch planner but completes episodes nearly twice as fast (29.7s vs. 53.4s episode). Although it is slower than the base policy due to additional planning overhead, this can efficiently handled real-world deployment which will be stated in following section. In Meta- World evaluation, we test on 50 tasks with 15 trials each. As shown in Table 3 method also show comparable performance to the Full-Patch planner with lower computation overhead.
Meta-World
Real-world LeRobot Results. In our real-world LeRobot evaluation on PickPlace and Drawer, tested by fixing the cup's position or drawer pose and placing the block in initial locations unseen during training for both tasks. We sample 10 distinct candidate chunks from the policy at each decision step instead of 50 for real-world deployment. On physical robot, sparse imagination

---

Published as a conference paper at ICLR 2026
Real ReelWorld PictPlace
Real Drewer
UBCRO-1 10
24% fetter
9070
34%
FJR
10%4
59%
25%
50%
32%
500025 (ats Ph)
407%
12%
1556
1300
70%
42%
65%
20%
Vilvnnty SR
559
Rado
Vuhann S4
VUk-anly 5RC
25%
70%
16
12
50
40
20
10
(srepisdde
Gregisode)
Figure 5: Trade-off between Inference Time and Performance in LeRobot and LIBERO-10. We show performance of SmolVLA with sparse imagination planning with a particular drop ratio in Left: Real-World PickPlace; Center: Real- World Drawer; Right: LIBERO-10 where (*) represent the operating point we highlight (50% drop), which provides a practical choice between speed and reliability.
with a 50% drop had an even greater impact, significantly increasing the success rate from 60% (VLA-only) 80% and 70% PickPlace and Drawer, respectively. Again, this matched the Full-Pater planner's success rate while reducing planner latency from 19.1s to 10.4s on PickPlace and from 14.0s to 10.6s on Drawer per episode. By leveraging asynchronous inference, this major performance gain across both tasks is achieved with minimal overhead compared to the VLA-only baselines (8.3s and 8.8s), proving our method is a practical solution for real-time robotic planning scenarios.
4.4 TOKEN REDUCTION METHODS
We compare our random sampling against broad set of token-selection strategies, grouped into four families as summarized in Table 4 (i) Random Sampling (Random, Fixed, LHS), (ii) Learning- Based Pruning (LTRP), (iii) Attention-Based Pruning (Attention-Encode WM), and (iv) Cluster & Merging (ATC). All methods share the same world model and planning hyperparameters, and retain the same number of tokens for a fair comparison.
STAR, Attention-
subset of patch tokens independently every planning iteration, while Fixed draws a single subset once for each rollout and reuses it for all rollouts within the episode. We additionally include LHS (Latin Hypereube Sampling: (McKay ct al. 2000)), which performs stratified sampling over the 2D patch grid so that retained tokens are approximately uniformly distributed across the image.
Random Sampling. The Random baseline samples
LTRP (Learning to Rank Patches: (Luo et al. 2024) trains a ranking network on top of a pre-trained MAE (He al., 2022) to predict patch importance based on recon- struction sensitivity; at test time, the top- K tokens according to this learned score are kept.
Learning-Basse Based Pruning.
Attention-Based Pruning. Attention-Enconn portance from the self-attention maps of a pre-trained DINO encoder, while Attention-WWa WM uses attention maps from our ViT-based world model itself. In all three cases, we rank tokens by their attention-based scores and retain the top-K K as the "important" subset.
and STAR (Zhang et al., 2024b) derive token im-
Cluster & Merging. ATC (Haurum al. 2024) clusters spatially similar patch tokens (using DINO features) and replaces each cluster with its average, effectively merging redundant patches. For goal- conditioned planning, both the current observation and goal features are clustered, and the distance between predicted futures and goal is computed after aligning clusters via Hungarian matching. A full description of each method and its training configuration is provided in Appendix [A.6]
Despite their complexity, empirical results across various datasets consistently show that none of these methods significantly outperform our simple random sampling (Table 4) Fig. visualizes the distinct dropout patterns for each approach.
7

---

Published as a conference paper at ICLR 2026
The average planning success rate (%) achieved by various token dropout and merging methods under different drop ratios. Our results suggest that no method achieves significantly su- perior performance compared to the Random sampling. "Fixed" is excluded for Granular and Rope because they only use CEM, making "Random" and "Fixed" strategies equivalent. For the "Fixed" variant Granular Rope tasks, the Avg. computed by reusing the drop ratios the corre- sponding "Random" setting.
Table 4:
Type
Methods
Pointmaze
Wall
PushT
Granular
Rope
Block Pushing
Avg.
Random Fixed
85.8
93.8
50.4
82.5
67.5
66.7
Random Sampling
87.1
85.0
47.4
16.5
64.3
LHS (McKay et al.
2000
84.2
92.5
49.6
82,5
64.2
21.0
65.7
Learaing-Basee Basnn Pruning
LTRP Luo et al.
2024
79.2
94.6
40.9
75.0
16.5
59.5
Attention-Encoder
79.2
97.1
43.9
70.0
67.5
20.5
63.0
Attention- Based Pruning
STAR (Zhung al. 2024b
82.9
85.8
46.9
72.5
63.3
17.0
61.4
Attention-W WM
80.0
93.8
44.3
82.5
16.5
64.0
Cluster Merging
Haurum
78.3
81.2
15.0
11.5
41.7
(a) Random
(b) LHS
(c) LTRP
(d) Attention-Fin
(c) STAR
(f) Attention- WM
Figure 6: Visualization of token dropout patterns for different token reduction methods. Pat- terns at 50% dropout ratio are shown from left to right for (a) Random, (b) LHS, (c) LTRP, (d) Attention-Encodie (e) STAR, and (f) Attention-WM
5 ABLATION AND ANALYSIS
5.1
ARCHITECTURA. GENERALIZATION
To demonstrate broader applicability, we evaluated our sparse imagination framework with MoCo-
v3 VIT-S (Chen et al. 202I). MAE ViT-B (He et al. 2022), and DINOv3 VIT-S (Siméoni et al. 2025) in the PointMaze environment (60 episodes). With 50% token drop, method matched the MoCo-v3 Full-Patch baseline's 96.7% success and the DINOv3 Full-Patch baseline's 98.3% success, while outperforming the MAE Full-Patch baseline (75.0% 68.3%). These results confirm that sparse imagination generalizes effectively across diverse vision encoders.
5.2 ROBUSTNESS TO TOKEN SPARSITY DURING PLANNING
Training the world model with our randomized grouped attention is crucial for leveraging these sparse token subsets. This strategy significantly reduces prediction errors on token subsets compared to standard full-attention model (c.g., 0.016 0.036 of normalized L2 error for 50% drop). which directly translates to improved planning performance (Fig. 4 We also find that the CEM planner robust enough to handle the small remaining prediction noise, based the experimental results, which consistent with prior studies (De Boer et al.) 2005) Hu & Hu 2009 Lale et al.) (024)).
Detailed experiments evaluating CEM's performance with visual token errors our framework in Appendix B.3
5.3 INFORMATION SUFFICIENCY SPARSE TOKEN SUBSETS
Our analysis confirms that sparse subset of tokens is sufficient for effective planning. We found that even a small, random fraction of tokens retains substantial ground-trus truth state information, often more than the global CLSS token, measured by both fformation-theoreee analysis using normalized
8

---

Published as a conference paper at ICLR 2026
0.9
0.16
0.8
0.14
CLS
Validation Loss
Random Drop
Afti-Far Drop
Normalized
0.08
0.5
0.06
0.4
0.02
0.3
Random Drop
Attn-Ene Enc Drop
0.00
02
Drop 0.1
Drop
Drop
Drop 0.4
Drop D.5
Drop 0.6
Drop
Drop 0.6
Drop
Laave
0.0
0.2
0.4
0.6
0.8
1.0
Drop Rate
Token Selection
Figure 7: nHSIC between visual tokens and environment states, nHSIC scores with respect to dropout rate for random dropout compared to Full, CLS, and Attention- Encoder dropout in Granular dataset. High HSIC scores are maintained even under sub- stantial dropout, indicating minimal loss of critical state information.
Figure 8: Atfentive probing performance. Valida- tion losses for state prediction with attentive probing compared to Attention-Encoode ular dataset. Surprisingly, a single random token shows predictive accuracy comparable (albeit with variance) to the CLS token, indicating state infor- mation distributed across tokens.
bascline in Gran-
Hilbert-Schmidt Independence Criterion (Gretton et al. 2005), (Fig. and an attentive probing module (Bardes et al. 2024) (Fig. 8). Methodological details are deferred to the Appendix A.4 and
5.4 EXPLAINING THE COUNTER-INTUITVE SUCCESS OF RANDOM SAMPLING
Counter-intuitively, sophisticated importance- based sampling methods failed to show a clear advan- tage over simple random sampling, often underperforming it (Table vs. Fig. 6
The "Blind Spot" Problem of Importance-Baze Sampling. We trace this phenomenon fundamental "blind spot" problem. Methods based on static importance metries, by design, learn to ignore regions of the scene deemed unimportant from the static initial and goal images. This proves fatal in dynamic tasks. If an object of interest enters blind spot, the predicted visual features for different candidate actions become indistinguishable to the planner. This causes the candidate se- lection process to collapse, as it cannot identify the best action based on visual progress. The failure starkly demonstrated in a vision-only planning experiment (Wall, 60 episodes, 50% dropout, pro- prioception distance is not used): the Attention-Ennooder method's success rate collapsed to 21.7%, whereas random sampling, which avoids such blind spots, remained robust at 58.3%. A control ex- periment using only proprioception confirmed vision's essential role, achieving just 16. 7% success. We provide more detailed analysis in Appendix B.9 and visualization example of this blind-spot
failure in Figs. and 20
Robustness from Unbiased Coverage and Information Redundancy. The success of random sampling lies its unbiased coverage, which is effective precisely because task-relevant infor- mation is highly redundant and distributed. We quantify its lack of bias using "cluster retention entropy", metric showing random sampling achieves significantly more uniform coverage than EreettoonEEncdde (2.37 vs. 1.91 bits, higher is more uniform). The viability of this unbiased strat- egy is confirmed by a critical finding: planning with only the tokens deemed least important by the iitention-Encodee still yielded a strong 63.7% success rate (compared to 79.2% with the most important tokens; PointMaze, 60 episodes, 50% drop). This proves that a random subset is almost certain to contain sufficient information for planning. This finding of distributed information is fur- ther validated by our information-thhorrtiicc analyses in Section 5.3 More experimental details are presented in Appendix A.7

---

Published as a conference paper at ICLR 2026
(a) Random
(b) Importance- Aattention-nncodee
Figure 9: Visualization Example of the Blind Spot Problem. Both images show 60% token dropout masks in PushT: (a) Random sampling and (b) an importance- method (Attention- Encoder). In the importance-basee case, the blue ball's path toward any direction solve the task almost entirely covered by masked patches, creating a persistent blind spots where the world model cannot observe the object's movement and thus fails plan correctly. Random sampling, by contrast, spreads retained tokens more uniformly, making such task-relevant regions less likely to be systematically ignored.
Our analyses resolve the paradox of random sampling's 's success: its unbiased coverage prevents the fatal "blind spots" that plague static, importance- based methods. This strategy is effective precisely because, as our experiments show, task-relevant information is broadly distributed across the image, ensuring random subset likely contains sufficient information for robust planning.
CONCLUSION
In this paper, we introduced sparse imagination, a novel approach for computationally efficient plan- ning in transformer-bar world models. By processing only subset visual tokens via random dropout, method significantly reduces computational complexity while maintaining high task performance across tasks ranging from simple test-time optimizations to complex real-world appli- cations with the latest VLAs. Our comparative analysis showed that various sophisticated sampling methods failed to outperform simple random sampling, and our investigation into the "blind spot" problem suggests this is due fundamental mismatch between their static importance metrics and the dynamic nature of planning.
While more advanced, dynamic-aware strategies might exist, they must justify their inherent com- plexity and computational overhead. For practical applications under resource constraints, we argue that a simple, near-zero zero-oerheed strategy like random dropout provides powerful and effective baseline, whose practicality may ultimately outweigh the marginal gains of more complex alterna- tives. The efficiency gains from this approach can then be reallocated widen the action search or process longer history of observations. This work establishes random dropout as simple, robust, and highly practical baseline for deploying visual world models real-time robotic applications.
Limitations. Our work has two primary limitations. First, the efficacy of our approach is inher- ently dependent on the quality of the underlying visual representations from the pre-trained encoder. Second, while demonstrate that moderate dropout ratio is robust default in our benchmarks, we do not claim that any particular fixed ratio is universally optimal, and the ideal sparsity level likely task-dependent. Although explore a preliminary uncertainty Appendix [B.6, developing a more sophisticated method to adaptively adjust this ratio remains a promising avenue for future research.
adaptive mechanism
Reproducibility Statement. We strongly encourage the reproducibility of our work. To facili- tate this, provide the maximum possible level of detail regarding experimental setup in the Appendix[AI This includes comprehensive description of the datasets used, the full training proce- dures, model architectures, and detailed hyperparameter settings.
10

---

Published as a conference paper at ICLR 2026
Acknowledgements. This work supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2024-00545889, Research on AI Robustness Against Distribution Shift in Real-Worle Scenarios; and No. RS- RS-2023-00222663, Center for Optimizing Hyperscale Al Models and Platforms).
REFERENCES
Adrien Bardes, Quentin Garrido, Jean Ponce, Michael Rabbat, Yann LeCun, Mahmoud Assran, and Nicolas Ballas. Revisiting feature prediction for learning visual representations from video. AXXv:22004.08471, 2024.
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. arXiv preprint arXiv:2210.0946/1 2022.
Remi Cadene, Simon Alibert, Alexander Soare, Quentin Gallouedec, Adil Zouitine, Steven Palma, Pepijn Kooijmans, Michel Aractingi, Mustafa Shukor, Dana Aubakirova, Martino Russi, Francesco Capuano, Caroline Pascal, Jade Choghari, Jess Moss, and Thomas Wolf. Lerobot: State-of-ther thrt machine learning for real-world robotics in pytorch. hucgingface/ lerobotr 2024.
com/
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self- vision transformers. In Proceedings of the IEFF/CVF international conference on computer vision, pp. 9650-9660, 2021.
Tianlong Chen, Zhenyu Zhang, Yu Cheng, Ahmed Awadallah, and Zhangyang Wang. The principle of diversity: Training stronger vision transformers calls for reducing all levels of redundancy. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp-
12020-12030, 2022.
Xinlci Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp- 9640-9649 2021.
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion, 2024. URL
https: /ttp://ariv...//. org/ abs 2303.04137
Pieter-Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Rubinstein. tutorial on the cross-entropy method. Annals of operations research, 134:19-67, 2005.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 11929, 2020.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual fore- sight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv:1812.00550 2018
Zhanzhou Feng and Shiliang Zhang. Efficient vision transformer via token merger. IEEE Transac- fions on Image Processing, 32:4156-4169, 2023.
Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. Advances in neural information processing systems, 29, 2016.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
A. Gretton, O. Bousquet, AJ. Smola, and B. Schölkopf. Measuring statistical dependence with hilbert-schmidt norms. Technical Report 140, Max Planck Institute for Biological Cybernetics, Tübingen, Germany, June 2005.
11

---

Published as a conference paper at ICLR 2026
David Ha and Jürgen Schmidhuber. World models. 2018. doi: 10.22117ENNOD0.20761.. URL
https:// /ewodozooooooooooo wttp:oddd..gg//eoooooooo ott/reecod//1177333
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2555-2565. PMLR, 09-15 Jun 2019b.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02199, 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 S.XXv:2300100100,, 2023.
Joakim Bruslund Haurum, Sergio Escalera, Graham W Taylor, and Thomas B Moeslund. Agglom- erative token clustering. In European Conference Computer Vision, pp. 200-218 Springer, 2024.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Jiaqiao Hu and Ping Hu. On the performance of the cross-entropy method. In Proceedings of the 2009 winter simulation conference (WSC), pp. 459-468. IEEE, 2009.
Manjin Kim, Paul Hongsuck Seo, Cordelia Schmid, and Minsu Cho. Learning correlation structures for vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 18941-18951, 2024.
Sahin Lale, Peter I Renn, Kamyar Azizzadenesheli. Babak Hassibi, Morteza Gharib, and Anima Anandkumar. Falcon: Fourier adaptive learning and control for disturbance rejection under ex- treme turbulence. npj Robotics, 2(1):6, 2024.
Youwei Liang, G E. Chongjian, Zhan Tong, Yibing Song, Jue Wang. and Pengtao Xie. EViT: Ex- pediting vision transformers via token reorganizations. In International Conference Learning Representations, October 2021.
Bo Liu, Yifeng Zhu, Chongkai Gao, Yihao Feng, Qiang Liu, Yuke Zhu, and Peter Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning. Advances in Neural Information Processing Systems, 36:44776-44791, 2023a.
Yue Liu, Christos Matsoukas, Fredrik Strand, Hossein Azizpour, and Kevin Smith. Patchdropout: Hconomizing vision transformers using patch dropout. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 3953-3962, 2023b.
Yang Luo, Zhineng Chen, Peng Zhou, Zuxuan Wu, Xieping Gao, and Yu-Gang Jiang. Learning to rank patches for unbiased image redundancy reduction. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 22831-22840, 2024.
Michael D McKay, Richard J Beckman, and William Conover. A comparison of three methods for selecting values of input variables in the analysis of output from a computer code. Technometrics, 42(1):55-61, 2000.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A univer- sal visual representation for robot manipulation. arXiv preprint arXiv:2203. 2022.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXivv2304.077993 2023.
12

---

Published as a conference paper at ICLR 2026
Bowen Pan, Rameswar Panda, Yifan Jiang, Zhangyang Wang, Rogerio Feris, and Aude Oliva. la- red2: interpretability-waalt redundancy reduction for vision transformers. Advances in neural information processing systems, 34:24898-22 0224898-249111 2021.
Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in neural information processing systems, 34:12116 121116-122 12128, 2021.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 34: 1493779-1,, 2021.
Mustafa Shukor, Dana Aubakirova, Francesco Capuano, Pepijn Kooijmans, Steven Palma, Adil Zouitine, Michel Aractingi, Caroline Pascal, Martino Russi, Andres Marafioti, al. Smolvla: A vision-lane sssionllnguage-actiion model for affordable and efficient robotics. arXiv preprint arXiv:2506.01 2025.
Oriane Siméoni, Huy Vo, Maximilian Seitzer, Federico Baldassarre, Maxime Oquab, Cijo Jose, Vasil Khalidov, Mare Szafraniec, Seungeun Yi, Michaöl Ramamonjisoa, Francisco Massa, Daniel Haziza, Luca Wehrstedt, Jianyuan Wang. Timothée Darcet, Théo Moutakanni, Leonel Sentana, Claire Roberts, Andrea Vedaldi, Jamie Tolan, John Brandt, Camille Couprie, Julien Mairal, Hervé Jégou, Patrick Labatut, and Piotr Bojanowski. DINOv3, 2025. URL battps://arxisn org/
abs 22t//258...0 10104
Nikolaos Tsagkas, Andreas Sochopoulos, Duolikun Danier, Sethu Vijayakumar, Chris Xiaoxuan Lu, and Oisin Mac Aodha. When pre-trained visual representations fall short: Limitations in visuo-motor robot learning, 2025. URL htRpstt:: http:///axvvvvviii org. /ttp////bb//a abs 2502. 03270
Zhenyu Wang, Hao Luo, Pichao Wang, Feng Ding, Fan Wang, and Hao Li. Vtc-Ife: Vision trans- former compression with low-frequency components. Advances in Neural Information Processing Systems, 39:13974-13988, 2022.
Fen Xia, Tie-Yan Liu, Jue Wang, Wensheng Zhang, and Hang Li. Listwise approach to learning to rank: theory and algorithm. In Proceedings of the 25th international conference Machine
learning, pp. 1192-1199, 2008
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.0617 2022.
Hongxu Yin, Arash Vahdat, Jose Alvarez, Arun Mallya, Jan Kautz, and Pavlo Molchanov. A-vit: Adaptive tokens for efficient vision transformer. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10809-10818, 18, 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Avnish Narayan, Hayden Shively, Adithya Bellathur, Karol Hausman, Chelsea Finn, and Sergey Levine Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning, 2021. URLhttts::/arxiv org/
abbs/1900. 10910.00897 10897
Kaifeng Zhang, Baoyu Li, Kris Hauser, and Yunzhu Li. Adaptigraph: Material-adaptive graph-based neural dynamics for robotic manipulation. arXiv preprint arXiv:2407.07888) 2024a.
Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric, 2018. URLAttp:://axxi..orr
0B8C103922
Yuyao Zhang, Lan Wei, and Nikolaos Freris. Synergistic patch pruning for vision transformer: Uni- fying intra- & inter-layer patch importance. In The Twelfth International Conference on Learning Representations, October 2024b.
Gaoyue Zhou, Hengkai Pan. Yann LeCun, and Lerrel Pinto. DINO-WIY -MO: World models on pre- trained visual features enable zero-shot planning. arXiv preprint 2rXiv:2411.04998 2024. URL
13

---

Published as a conference paper at ICLR 2026
APPENDIX
A IMPLEMENTATION DETAILS
A.1 DATASETS AND ENVIRONMENT DETAILS
We elaborate various datasets and environments where method is experimented on (Table 5 Offline trajectory datasets Pointmaze, Wall, PushT, Granular and Rope as well as some of dataset- specific configurations are adopted from those provided by Zhou al. (2024 Configurations specific to dataset include history trajectory length for the world model and frame skip length, which interval of frames and from the world model, reducing redundancy between consecutive inputs. We also differ training epoch due to difference of dataset size and complexity.
Table 5: Dataset details and dataset-specific configurations. Note that history length, frame skip, and train epochs are for training world model.
Dataset
Size
Frames
History len
Frame skip
Train epoch
Pointmaze
2000
100
5
I0
Wall
1920
50
5
100
PushT
20000
100~300
5
Rope
1000
20
1
100
Granular
1000
20
1
100
Block Pushing
11000
200
5
LIBERO 1693
75~505
5
10
Meta-World
2500
14~500
5
10
LeRobot 100
208~459
5
Pointmaze. A navigation task introduced in D4RL (Fu al. 2020 where agent interacts with given 2D maze environment. There are 2000 trajectories with length of 100 created with random actions. Since visualized each frame.
task-agnostic world model with goal-conditioned planning, goal point is not
Wall A visually simplistic 2D navigation task where agent needs pass door to separated room. There are 1920 trajectories with length of 50 created with random actions.
PushT. A 2D control task introduced in Chi et al. (0044) where an agent needs manipulate T-shaped block to designated location. Dataset is generated by rollouting expert demonstration action trajectory with additional action noise in the environment. This allows the world model to learn various dynamics.
Granular and Rope. Fine-grained manipulation tasks involving deformable objects introduced in Zhang et al. (2024a). Both datasets have 1000 trajectories with length of 20 created with random actions.
Block Pushing. Robotics control task requiring precise spatial reasoning to position blocks cor- rectly. We generate 11000 trajectories following PushT generation process.
LIBERO. The LIBERO benchmark is a multitask robot learning environment featuring diverse set of objects and tasks. For training, we utilized only the fixed-view images from the LIBERO- Spatial, LIBERO- LIBERO-Objctt, LIBERO-Goal, and LIBERO- suites to train both our world model and the SmoIVLA policy. This dataset, comprising 1,693 trajectories with 75 to 505 frames each, is publicly available on Hugging Face (available heref). We conducted our evaluation exclusively on the LIBERO-10 10 benchmark, as its challenging long-horizon tasks provide ideal testbed for validating the efficacy of our planning framework.
Meta- Meta-Worid. The Meta- Meta-World benchmark (Yu et al. 2021 is a comprehensive robotic manipula- tion suite designed for multi-task and meta-learning, We utilize the MT50 multi-task dataset, which consists of demonstrations for 50 diverse manipulation tasks spanning a wide range of horizons and
14

---

Published as a conference paper at ICLR 2026
difficulty levels. The dataset contains 2,500 episodes. each ranging from 14 500 frames, and is publicly available on Hugging Face (available herd This collection of demonstrations allows our world model and planning framework to learn dynamics across broad variety of manipulation behaviors.
LeRobot. We collected real-world datasets for the task "pick the blue block and place it in the cup (PickPlace)" and "put the blue block in the drawer and close it" (Drawer) via teleoperation, using the LeRobot SO-101 leader-follower system with fixed-view camera. Each dataset comprises 100 trajectories, each containing between 208 to 459 frames. During data collection for PickPlace, the cup's position remained fixed while the block was placed at 10 distinct initial positions, with 10 trajectories recorded for each. During data collection for Drawer, the drawer's initial pose remained fixed while the block was placed at 4 distinct initial positions, with 25 trajectories recorded for each. This dataset then used to train the trained agent was tested on 5 novel block positions held out from the training set, with the final success rate calculated over 10 total trials (two attempts for each unseen position). For Drawer task, the trained agent was tested on 2 block positions included in the training set, with the final success rate calculates over 10 total trials (five attempts for each position). The SmolVLA policy was loaded from pretrained VLM weights then only the state projector and the action expert are trained when finetuning. The architectural settings follow SmoIVLA with 0.45B parameter in the original work. We also leverage asynchronous inference pipeline of Smo/VLA. Whenever the remaining action chunk ratio reaches certain threshold, robot client sends the observation from to the policy server and let the remote server compute the next action chunk in advance.
world model and the SmoIVLA policy. For PickPlace task,
Table 6: Hyperparameters for SmoIVLA.
Name
Value
Batch size
64
Training steps
50,000(Lerobot), 100,000 Meta-World)
Optimizer
Adam W
Learning rate
0.0001
Scheduler
Cosine
Warmup steps
1,000
Decay steps
30,000
Action chunk size
50
Async inference thres.
0.5
A.2 MAIN METHOD IMPLEMENTATION DETAILS
Our model is built on pre-trained DINO Caron et al. 2021 VIT backbone which is frozen during training. Specifically, we utilize DINO-VIT VYT- on 224x224 image, which creates a 14x14 spatial tokens per image. Visual tokens from context frames are flattened as a input to the world model. Actions and proprioceptive vectors are linearly projected on a dimension of 10 respectively and concatenated to every visual token along the feature dimension axis. The world model is causal transformer decoder, which predicts the tokens of the next timestep. The transformer has attention layers, 16 attention heads per layer, and 2048 embedding dimensions. The world model is trained with Adam optimizer at a learning rate of 5e-4 without scheduler for 100 epochs, batch size of 32, dropout probability of 0.1. All experiments were conducted using at most 4 Nvidia RTX 3090s.
A.3 PLANNING DETAILS
MPC rollouts utilize the CEM with action distribution defined as a Gaussian distribution. CEM is performed with planning horizon of 5, optimizing by mean and variance of top- 10 action sequences among 100 candidate sequences. CEM optimization is done for 10 iterations and returns the mean of the action distribution as the final action sequence, executing all the planned action sequence for every MPC step. MPC iterates until the task is completed until it reaches maximum limit of 10 iterations.
15

---

Published as a conference paper at ICLR 2026
For the PointMaze, Wall, PushT, and Block Pushing environment, the MPC horizon is set to H 5. For each optimization step, we sample 100 candidate actions and select the top 10 based on the distance computed between the current and goal observations. Note that in sparse imagination scenarios, the distance is computed using only a subset of patches rather than the full patches. Each CEM optimization is performed for 10 steps. For the Wall and Block Pushing environments, the maximum number of MPC iterations is set to 10, while for Pointmaze and PushT, we set the maximum iterations to 15, For the Granular and Rope environments only employ the CEM approach for 10 steps.
A.4 HSIC ANALYSIS DETAILS
Hilbert-Schmidt Independence Criterion (HSIC). We measure statistical dependence between two random variables X and Y using the Hilbert-Schmidt Independence Criterion (HSIC) (Gretton al. 2005). defined via kernels i and E on their respective domains. Let K and L be the n n Gram matrices with entries Kú = k(z,,Tz) Ty) and Lij = E(y:: Vj), and H = In 1 11 T the centering matrix. Then the empirical HSIC is
1 HSIC(X,Y) (n - 1)2
tr tr(KHIHH)
and we normalize it as
HSIC(X, Y)
HHIIC(XYY) HSIC X) HSIC(Y, Y)
so that nHSIC [0, 1), with indicating independence and maximal dependence.
Experimental Setting, We compute nHSIC between visual tokens and corresponding environ- ment state vectors under varying dropout probabilities. A single dropout mask is fixed within each batch of size n = 128. Then, we repeat with 20 independent masks of same dropout rate and re- port the mean nHSIC Visual tokens whose handwidth is set by the median heuristic. We utilized state vectors from the Granular dataset, consisting of the and y coordinates for 12,774 particles.
linear kernel, while state vectors
RBF kernel
A.5 ATTENTIVE PROBING DETAILS
Attentive Probing. Attentive probing is a nonlinear probing method priorly used for downstream tasks since visual tokens do not satisfy linear separable guarantee (Bardes et al.] 2024). We use attentive probing to fairly compare predictive power of varying number of tokens. Instead of linear probing, attentive probing uses cross-attention layer to pool all input features query vector. After query residual connection, output is fed into 2-layer MLP with GeLU and Lay- erNorm, projecting original input dimension. Then, final linear layer predicts the given probing target.
single learnable
Experimental Setting. Using attentive probing predict the environment state with visual tokens under varying dropout probabilities, we present the mean and standard deviation of validation loss when training the model for times with different dropout masks for each dropout rate. We also add leave-1 1 case where extreme dropout is performed to leave only a single random visual token. We use feature dimension of 384 and 4 heads for the cross- cotss-attnntion layer and hidden dimension of times the feature dimension for the MLP layer. We trained the probing module for 500 epochs. The optimization was performed using Adam with a learning rate of 1 x 10-5 and a batch size of 128, incorporating cosine scheduler and omitting warmup period. Parallel to our HSIC analysis, used Granular dataset state information. Rather than attempting to precisely predict 12,774 particle locations (impractical due to the order of particles and dimensionality), our probing module predicts the mean and standard deviation of x and y coordinates for all particles, thereby estimating their approximate position and spread state information.
A.6 TOKEN SELECTION METHODS DETAILS
LHS. Latin Hypercube Sampling (LHS; McKKay et al. 2000)) is used as a simple uniform baseline that enforces spatial coverage. We treat the H W patch grid as 2D domain and partition it into K
16

---

Published as a conference paper at ICLR 2026
strata corresponding to the desired number of tokens. At each step, sample one patch from each stratum, which yields a set of retained tokens that are approximately uniformly distributed across the image.
LTRP. LTRP (Learning to Rank Patches; Luo et al. (2024)) utilizes a pre-trained MAE (Masked Autoencoder; He et al. (0222) to learn a ranking model that assesses the importance of individual patch tokens. With 90% of the patch tokens masked, LTRP compares the reconstruction results (using L1 distance) when each remaining token is individually removed versus when it is retained. Tokens contributing more significantly to the reconstruction error are considered more important. The ranking model is trained using ListMLE (Xia et al. 2008), optimizing the order of the top 20 items within each list. For the ranking model, we employed a ViT-Small architecture, while the MAE utilized ViT-Base checkpoint. The ranking model was trained on ImageNet-IK IK for 50 epochs using four NVIDIA RTX 3090 GPUs. For aspects not explicitly mentioned, training configurations adhered to the defaults outlined in the original paper. The LTRP code is available here and the MAE ViT-Base checkpoint is from h
Attention-Encoded Similar Liang et al. (2021); Wang al. (022); Zhang al. (2024b), we gauge token importance by examining the attention map between the CLS token and other patch
tokens within a pre-trained DINO encoder. Tokens receiving higher attention weights from the CLS token are deemed most significant. We implement token dropping by retaining the top-K K tokens sorted by their attention weights. Our experiments utilized the DINO-ViT-S/L checkpoin f specifically leveraging the attention map from the final transformer block of the DINO encoder.
STAR. STAR (SynergisTic pAtch pRuning; Zhang et al. (2024b)) determines a comprehensive importance score for each layer by combining two distinct components: an intra-layer importance score, derived from the attention weights between the CLS token and individual patch tokens, and an inter-layer importance score, quantified using Layer-wise Relevance Propagation (LRP). The inter-layer scores were pre-computed on the ImageNet- IK dataset, consistent with the methodology outlined in the original publication. Given that the original STAR implementation was based on DeiT, we modified their codebase to ensure compatibility with DINO. The aggregated score for layer I, denoted St, is formulated S Siimma" X S Simter") We set 0.3, which is one of the values proposed in the original paper. Analogous to the Attention-Enconn was realized by sorting tokens according their scores from the final encoder layer in descending order and preserving only the top-K. The official code of STAR can be found here
strategy, token dropping
Attention-WM. Unlike masder- methods like Attention Attenti,n-Encoder, this approach stems from the idea that it might be more efficient to retain tokens deemed important by the world model during dynamics prediction for given action. Thus, we leverage the attention maps from our trained ViT- based world model. Since our world model does not utilize a CLS token, we instead consider the attention maps involving all individual patch tokens. Specifically, for each patch token acting key, we calculate the average attention it receives from all query patch tokens. We then assume that tokens with higher average attention are more important. A drawback of this method the two-pass forward computation during planning. First, the world model must be forwarded with all tokens obtain the attention maps across all layers. These maps are then used as importance scores to select the most significant patch tokens. Subsequently, second forward pass is executed with only the selected tokens, introducing an additional computational cost. In this approach, we averaged attention maps across all layers.
Agglomerative Token Clustering (ATC). Inspired by HHuurum al. (2024), we employ a mod- ified version of Agglomerative Token Clustering (ATC) for our clustering approach. Necessary modifications were made as the original ATC was primarily designed for single-frame applications, such as classification segmentation. Adapting it for goal-conditioned world model-based plan- ning necessitated a modified mechanism to address multiple history frames, and to match clusters between the goal image features and predicted latent features.
httpss//gttt0b.coo/mmmm//e/////
17

---

Published as a conference paper at ICLR 2026
Our clustering methodooogy proceeds follows. First, extract DINO features from the input image and apply agglomerative clustering using either L2 distance or cosine similarity (the choice between these two metrics did not significantly affect results). The number of clusters is dynam- ically determined by the desired drop ratio (e.g.. 137 clusters for a 20% drop). Using the cluster assignments from the initial frame as anchors, subsequent frames undergo K-means clustering. Af- ter clustering, patches within each cluster are aggregated via average pooling, maintaining the same number of patches clusters. To address inconsistencies in clustering results across visually simi- lar scenes, we observed through trial and error that storing and re-using K-means centroids from the previous frame as initial centroids for the subsequent frame ensures better consistency by aligning clusters across consecutive frames.
To integrate this into MPC, we initially cluster the goal image and use as an anchor for clustering the current observation, aiming for inherent cluster alignment. However, this alone did not guarantee perfect alignment, especially with an increased number of clusters. To mitigate this, we introduced an additional matching step: during the loss calculation within the CEM, specifically when selecting the top-K actions, we employ Hungarian matching between the world model's s predictions and the goal clusters to ensure proper alignment. Without this crucial matching step, observed that the results degrade significantly. For instance, in the Wall Environment, using 128 clusters, experiments over 30 episodes showed that the success rate improved from 70.0% without matching to 86. 7.7% with matching applied.
A.7 CLUSTER RETENTION ENTROPY DETAILS
To quantitatively measure the spatial bias of different token sampling methods, we introduce Cluster Retention Entropy. This metric assesses how uniformly a sampling strategy distributes its kept tokens across the spatial layout of an image. The calculation proceeds follows:
1. Spatial Clustering: First, we establish a set of spatial anchors. We take all visual patch tokens from a single, random image (e.g., Block Pushing) and cluster them into K groups based on their spatial coordinates using the K-Means algorithm. This gives us K spatial regions of interest.
2. Distribution of Retained Patches: For single image and a given sampling method, we identify the subset of patches that are retained. We then count how many these retained patches fall into each of the K pre-defined spatial clusters. This count creates discrete probability distribution over the clusters.
3. Entropy Calculation: Finally, we compute the Shannon entropy of this distribution. A higher entropy value indicates that the retained patches are spread more uniformly across the spatial clusters, signifying low spatial bias (i.c., good coverage). Conversely, a low en- tropy value indicates that the sampling method heavily biased, concentrating its attention on only few spatial regions.
In the analysis presented in the main text, we computed the mean entropy for each method on the Block Pushing dataset. To ensure robust comparison, averaged the results across range of settings, varying the number of clusters K E {2,4, 8, 12} and the token drop ratio from 10% to 90%.
A.8 REAL- WORLD LEROBOT PLANNING DETAILS
To perform goal-conditioned planning in the real-world LeRobot environment, we used the goal image shown in Figure 19 The corresponding goal proprioceptive state, which represents the robot's joint information, was set to [23.19, 17.32, 12.97, 79.24, -4.68, 16.58] (PickPlace) and [-31.44, 33,.44,00 0.68, 16.94, 67.90. 26.77. 0.61] (Drawer).
18

---

Published as a conference paper at ICLR 2026
(a) PickPlace
(b) Drawer
Figure 10: The goal images used for planning in the real-world LeRobot environment.
B ADDITIONAL ABLATIONS AND EXPERIMENTS
B.1 HYPERPARAMETER
CEM Optimization Steps. We selected 10 CEM optimization steps for consistency with the DINO-WM baseline, ensuring fair comparison. Our ablation study in the PushT environment (60 episodes, various drop ratios) validated this choice: steps yielded 49.3% success, while 10 steps achieved 51.7%. Increasing to 15 steps (50.7%) offered significant benefit. Thus, 10 optimization steps represent a well-justified trade-off between planning quality and computational cost.
Longer Planning Horizon in MPC. To further assess the robustness of our approach, we evalu- ated its performance under more demanding, longer planning horizon. Specifically, conducted experiments in the PushT environment with the planning horizon extended to H=7. Under this chal- lenging condition, the Full-Patch baseline, which utilizes all visual tokens for planning, achieved a success rate of 48.3%. Our sparse imagination method remained highly competitive, achieving 50.0% success rate with a 20% token drop and 46.7% with a 40% drop. This demonstrates that our framework's 's performance does not degrade relative to the Full-Patch baseline even on longer- horizon tasks, confirming the robustness planning with a sparse subset of tokens.
Frame Skips in World Model Training. We performed ablation the frame skip value in the Pointmaze environment. We found that a frame skip of 1, which processes every frame, resulted in low success rate of 36.7%, likely due redundant observations. In contrast, increasing the frame skip to 5 significantly boosted performance to 85,8%. However, further increasing the frame skip to 10 yielded no additional benefit, with the success rate remaining at 85.5%. These results demonstrate clear point of diminishing returns and justify our final choice of using frame skip of which provides an optimal balance between task performance and computational efficiency.
B.2 GROUPED ATTENTION
This section investigates the influence of randomized grouped attention (applied during training) on world model prediction error during planning, relative to full attention (Fig. [T]. The analysis spans multiple datasets and various dropout ratios. In contrast to the main text's primary focus on findings for the Wall dataset under 50% dropout condition, the present discussion offers a more comprehensive overview of results across these diverse conditions. We quantify this impact by pre- senting the mean relative L2 distance of retained tokens the validation set of each dataset, both with and without dropout, for each attention mechanism. This metric normalizes the L2 distance by the tokens' norm. All reported values represent the mean and standard deviation derived from
19

---

Published as a conference paper at ICLR 2026
five independent experiments. Our analyses consistently demonstrate that grouped attention leads to lower prediction error and reduced variance compared to full attention across the Wall, Granular, Rope, and Block Pushing datasets. We observe apparent correspondence between these improve- ments and notable enhancements in success rates across the evaluated datasets (as detailed in the main text). This observed pattern seems to hold for majority of the dropout ratios investigated.
0.0
0.05
Nean Ratebe Dirtnnc
0.02
0.00
0.03
0.02
Drizou Ratio
(a) Wall
(b) Granular
Grouped imerttnn
Grouped intention
full Aftention
Ful Attertionn
0.00
Meen SAtite Dietarce
0.25
0.02
9.00
0.2
0.5
Drophel Retia
Diazout Rati
(c) Rope
(d) Block Pushing
Figure 11: Contribution of grouped attention during pretraining to planning on various drop ratios across various datasets.
B.3 PLANNING ROBUSTNESS TO NOISES
In this section, we experimentally validate our claim that the remaining prediction error, consistent with prior studies CEM's robustness noise and uncertainties, is sufficient for effective planning. To investigate this, we first examined whether the prediction error, incurred by token dropping in the world model, correlates with the world model's input across different dropout ratios. We reduced the high-dimensional input features to 2D representation using UMAP and generated a scatter plot comparing these UMAP coordinates with the magnitude of the error (Frobenius norm) (Fig. 12)
Visual inspection of the UMAP projection of input features shows no obvious global or local patterns linking UMAP coordinates of model inputs the error norm. Data points, colored by their error magnitude (Frobenius norm), are generally intermixed throughout the low-dimensional embedding. Therefore, this specific UMAP representation of input features, based on visual analysis, does clearly indicate strong, direct relationship between the spatial embedding of these features and their corresponding error magnitude.
Given the findings, we assume that the prediction error inherent in world model when operating on token subsets could be approximated zero-mean Gaussian noise, exhibiting no direct correlation with the input tokens. To empirically test the robustness of MPC-CEM planning, we proceeded to add this synthesized noise to the input tokens during the actual planning.
To simulate the behavior of a dropout-affected world model, begin by generating predictions ing the full world model without applying dropout. Next, we modify these predictions by randomly dropping subset of the predicted tokens and adding zero-mean Gaussian noise to the remaining ones. These modified predictions are then used for MPC-CEM planning. The magnitude of the Gaussian noise scaled relative the token norm and experimented under varying noise scales
20

---

Published as a conference paper at ICLR 2026
Drop Radio 90%
Drop Ratio 80%
Drop Ratio 70%
70
memeioon
Eroo Frbbmni Noms
UMMP D mersion
20
-gnu Enor Proohiuz horm
USAP D.merslon2
Enoor Probemus Nomm
20
20
-20
20
UMAP Dimension
UMAP Dimension
UMKP Dimension 1
Drop
Drop
Drop
Enor Frodzonuz Nomi
UMMP Dimersioo 2
20
Emrr Fibbenlus Nom
20
Frodenuus Nom
Dimensionn
Dimaentoon
15
10
10
-10
20
20
20
UMAP Dimencion 1
UMAP Dimenston
UMAP Dimenslon
Radio
Ratio 20%
Drop Ratio
Dimertionn
Erder Robomnu us Norm
UMAP climersion 2
10.0
Error Roben us Norm
UMAB Sliversiion 2
Eneer Rosenuus Nom.
M
2.5
-10
Dimendion
Dimenston
Figure 12: Scatter plots visualizing the embedding of world model inputs in a 2D UMAP mani- fold, colored by their prediction error magnitude (Frobenius norm), presented across different token drop ratios.
and dropout rates. We then evaluated the robustness of MPC-CEM planning by measuring the suc- cess rates of actual rollouts in the environment using these noisy predictions. This experiment was conducted in the point maze environment, with the agent allowed maximum of 3 MPC planning steps per episode. The presented success rates were averaged over 60 independent evaluations.
The results (Fig. [3) indicate that MPC-CEM planning sustains effective performance at dropout levels with minimal performance decrease, which aligns with our main experiments. While its performance remains stable noise levels increase up to specific threshold, beyond which it begins to decline. This suggests that MPC-CEM robust to prediction errors caused by token dropout. reinforcing claim that manageable for planning. These observations also align with prior works on the resilience of CEM to noise and uncertainties.
B.4 DROPOUT VS. CLUSTERING IN EXECUTION TIMES
Table 7 compares the execution times of single MPC iterations between dropout and clustering methods at various compression rates (drop ratios). The clustering approach requires additional computations, including the calculation of pairwise distances similarities among patches, incur- ring a quadratic complexity relative to the number of clusters. Furthermore, since clustering merges patches from multiple spatial locations, the resulting pooled features lose fixed spatial positions, cessitating additional matching procedures when computing distances between goal image patches and observations during MPC planning. Consequently, the clustering method requires significantly
21

---

Published as a conference paper at ICLR 2026
08-
10
Relative Average Success Rate
Success Rats
9.0
9.8
Dropout
0.4
drop 0.0
drop 6.1
drop 0.3
drap 0.5
0.7
03
677
Dropout Rate
0.003 s note s 0.0%
0.0
1.0
log scaled noiso
2.0
25
10
0.01 noise a 2.05
0.05 < noise a 3.2
0.2 nular
(a) MPC-CEM success rate vs. dropout rate under various noise conditions.
(b) Relative MPC-CEM success rate log noise scale for different dropout ratios.
Robustness MPC-CEM planning prediction dropout and noise. The results show minimal degradation of success rates by increasing dropout rates and performance loss only after certain noise threshold.
Figure 13:
Table 7:
Comparison of planning time and change across different drop ratios (compression rates)
between drop and clustering apporaches.
Pointmaze
Wall
PushT
Block Pushing
Planning Time (sfiter) Drop
Planning Time (s/iter)
Planning Time (s/iter)
Planning Time (s/iter)
Drop
Chunge (%)
Change (%)
Change (%)
Chunge
(Compression Rate)
(%)
Cluster
Drop
Cluster
Drop
Cluster
Drop
Cluster
10%
73
+237.0
328
478
20%
141
297
+110.6
69
235
+240.6
131
295
+125.2
259
444
+71.4
30%
40%
106
245
+131.1
62
207
+233.9
97
237
+144.3
214
376
+75.7
50%
214
+130.1
53
+262.3
213
+159.8
208
60%
193
+1413
50
184
+268.0
69
187
+171.0
200
318
+59.0
70%
172
+149.3
46
+256.5
170
+188.1
295
80%
56
152
+171.4
46
145
+215.2
49
139
+183.7
175
266
+52.0
90%
137
+185.4
42
117
+178.6
115
+202.6
167
248
+48.5
more computation, with runtime increases of up 268.0% compared dropout, as shown in Ta-
ble 7 In contrast, our dropout approach merely samples patch indices, which results in nearly constant computational overhead regardless of dropout ratio. Additionally, since the same patch subsets are consistently applied to both observation and goal images, the spatial positions remain fixed, climinating the need for extra matching processes.
B.5 GENERALIZATION DIFFERENT PLANNING PARADIGMS
To demonstrate that our sparse imagination framework not limited sampling-bast based planning algorithms like the CEM, we evaluate its compatibility with a gradient-based trajectory optimizer (MPC-GD). We replace the CEM planner with gradient descent-based optimizer the PointMaze and Wall environments, while keeping the randomized grouped attention and sparse token selection mechanisms identical to the main experiments.
Table g presents the planning success rates across varying dropout ratios. The results indicate that our method maintains high performance even when utilized with gradient-basee optimization. No- tably, with 50% token drop, the success rates remain comparable to the Full-Patch baseline (e.g., 94.0% vs. 94.0% PointMaze, and 88 68.0% vs. 88.0% in Wall). This suggests that random token sparsification preserves the optimization landscape sufficiently for gradient-based planners
22

---

Published as a conference paper at ICLR 2026
verge to effective solutions, confirming the general applicability of our approach across different planning paradigms.
Table 8: Performance of Sparse Imagination using Gradient- -essentt based Planner (MPC-GD) on PointMaze and Wall environments (Mean Success Rate, %).
Drop Ratio
PointMaze
Wall
Full (0%)
94.0
88.0
10%
94.0
90.0
30%
98.0
92.0
50%
94.0
88.0
70%
90.0
82.0
90%
78.0
78.0
B.6 EXPLORING UNCERTAINTY- ADAPTIVE DROPOUT
While our main experiments characterize performance function of fixed sparsity levels (demon- strating robustness up to approximately 50% dropout), we do not claim that single fixed ratio is universally optimal. In scenarios where the optimal drop ratio is unknown or task complexity fluctu- ates, an online adaptive mechanism becomes desirable. Since our world model is explicitly trained to be robust under random sparsification via randomized grouped attention, the Sparse Imagination framework naturally supports dynamic adjustment of the dropout ratio at test time without retraining.
To explore this capability, we implement a simple uncertainty-awarr adaptive planner. We utilize the variance of latent rollouts across MPC samples as proxy for prediction uncertainty. The logic straightforward: if uncertainty increases between iterations, the planner reduces the drop ratio (retaining more tokens); if uncertainty is low, the drop ratio is increased to save computation.
We evaluate two variants: (1) Ada-Rand, which updates the scalar drop ratio based on the change in variance and resamples new random mask; and (2) Ada-Inc, which updates the mask incre- mentally, preferentially keeping high-varianee tokens. Table summarizes the results. In the Wall environment, Ada-Rand achieves a 95.0% success rate (matching the strong fixed 50% baseline) with an average drop ratio of 26. 1.. In PushT, Ada-Rand reaches 68.3% success (matching the 10-50% fixed range). Crucially, these adaptive methods avoid the sharp performance collapses ob- served with overly aggressive fixed sparsity (e.g.. >70% in PushT) by automatically settling at a safe sparsity level. While they do not necessarily outperform the best-tuned fixed ratio, they demonstrate the potential of our framework to support online optimization when the optimal sparsity is unknown.
Table 9: Feasibility study of Uncertainty-Aware Adaptive Planning variants (Mean Success Rate, %) and the resulting Average Drop Ratio (%).
Wall
PushT
Method
Success Rate
Avg. Drop
Success Rate
Avg. Drop
Ada-Rand
95.0
26.1
68.3
29.0 35.8
Ada-Inc 91.7
29.8
65.0
B.7 WORLD MODEL PREDICTION QUALITY
To assess whether our training mechanism affects the intrinsic quality of the world model's predic- tions, train the world models with an optional decoder and measure the image reconstruction quality the predicted features using the Learned Perceptual Image Patch Similarity (LPIPS) met- ric (Zhang et al. 2018) (lower is better). Table 10 compares the LPIPS scores of the standard Full (full attention) baseline against our model trained with Randomized Grouped Attention ("Drop"), where both models decode from full patches evaluation time.
Across various environments, the LPIPS values for our method (Drop) remain very close to the Full- Patch baseline. Notably, in visually or physically complex tasks such Granular and Rope,
23

---

Published as a conference paper at ICLR 2026
method achieves slightly better reconstruction scores, potentially due to the regularization effect of the grouped attention strategy. These results indicate that our training mechanism does not degrade the world model's 's prediction quality under full- full-patch evaluation while enabling robust performance under sparse imagination at planning time.
Table 10: Comparison of World Model Prediction Quality (LPIPS, lower is better).
Environment
Full
Drop (Ours)
PointMaze
0.00028
0.00031
Wall
0,00074
0,00078
Granular
0.05607
0.04739
Rope
0.03347
0.01906
PushT
0.00561
0.00644
B.8 DETAILED PLANNING TIME ANALYSIS
To investigate the source of computational efficiency, we profile the latency of major components within the planning loop. Table II presents the fine-grained latency breakdown for the PushT envi- ronment.
The analysis confirms that the dominant acceleration stems from the World Model's Self-Attention mechanism. As shown in the tables, Self-Attention latency drops dramatically (e.g., in PushT: 100% + 28.8% at 50% drop), closely following the expected quadratic reduction. In contrast, the Projection and MI.P layers scale linearly, while fixed costs remain stable.
Note: The latency values reported here were measured in separate additional profiling session and may exhibit minor variances compared to the end-to-end timings in Table 2
Table 11: Latency breakdown per planning iteration on PushT. Format: Time in seconds (Reduc- tion % compared Full baseline).
Variable Costs (World Model)
Fixed Costs
Drop (%)
Self-Attn
Proj (QKV)
MLP (FFN)
Encoder
Env Step
Others
0
87.4 (-0.0%)
30.5 (-0.0%)
34.5 (-0.0%)
11.9
2.8
10.8
10
73.3 (-16.0%)
27.5 (-9.9%)
31.1 (-9.7%)
11.7
3.4
10.3
20
57.2 (-34.6%)
(-19.2%)
27.9 (-19.0%)
11.9
4.7
30
47.0 (-46.2%)
21.6 (-29.29%)
24.5 (-28.8%))
11.9
3.4
10.6
40
32.9 (-62.3%)
(-38.7%)
(-38.7%)
3.5
50
25.1 (-71.2%)
(-49.3%)
17.5 (-49.25)
11.8
2.8
60
16.0 (-81.%)
12.5 (-59.2%)
1..9
11.8
4.0
10.5
70
10.6 (-87.9%)
9.5 (-68.85)
10.7 (-69.15)
11.8
2.4
11.1
80
5.0 (-94.3%)
(-79.2%)
(-79.3%)
3.3
10.4
90
2.4 (-97.2%)
3.5 (-88.5%)
(-8899%)
3.4
10.3
B.9 ADDITIONAL ANALYSIS ON BLIND SPOTS
To provide quantitative characterization of the "blind spot" phenomenon visualized in Figures 20] we perform controlled analysis on the PushT enviromment. each MPC planning itera- tion, we track the ground-truth position of the agent (blue ball) for all candidate action sequences. We project the agent's 's position into the image space. We approximate the blue ball as a circle whose diameter equals one patch size and this circle to determine which patch tokens it overlaps. A blind-spot event counted if all tokens overlapping with the agent are dropped by the sampling mask for given timestep.
Table presents the blind-spot occurrence rates evaluated at 50% and 60% dropout ratios (60 episodes). At 50% drop, Random sampling produces blind spots in only 3.10% of states, whereas baseline methods fail to capture the agent in 17-26% cases. At 60% dropout, while failure rates increase due to higher sparsity, Random sampling (12.26% %) still maintains lower blind-spot rate
24

---

Published as a conference paper at ICLR 2026
compared to baselines (22 (22-33%) These results quantitatively demonstrate that Random sampling consistently yields fewer blind-spot events compared to other selection strategies.
Table 12: Quantitative analysis of Blind Spot occurrence rates (%) in the PushIT environment (60
episodes). Lower is better.
Method
50% Drop
60% Drop
Random (Ours)
3.10
12.26
STAR
17.74
22.87
LTRP
18.56
27.05
Inention-Encodee
20.99
33.04
Attention-We WM
26.60
31.17
25

---

Published as a conference paper at ICLR 2026
C VISUAL DEMONSTATIONS
C.1 IMAGINED ROLLOUT BY WORLD MODEL
(a) LIBERO
(b) World
(c) LeRobot PickPlace
(d) LeRobot Drawer
Figure 14: Imagined rollout of the LIBERO, Meta-World, real-world LeRobot Pickplace and Drawer by world model. Top row: Ground truth. Bottom row: Imagined rollout by world model.
26

---

Published as a conference paper at ICLR 2026
C.2
EXECUTED ROLLOUT DURING PLANNING
Os
1s
2s
3s
4s
5s
65
7s
8s ~9s
Drop
50%
VLA-
Only
Full-
Patch
Figure 15: Executed rollout demonstration real-world PickPlace task in LeRobot environ- ment (50% Drop, VLA-only, and Full-Patch). A green border indicates successful episode, while red border indicates incompleteness. Both method and the VLA-only baseline succeed under seconds, while the Full-Patch planner fails to do so in the same timeframe due to its significant planning latency.
0s
2s
4s
6s
85
~10s
Drop
50%
VLA-
Only
Full-
Patch
Figure 16: Executed rollout demonstration on real-world Drawer task in LeRobot environment (50% Drop, VLA-only, and Full-Patch). A green border indicates successful episode, while red border indicates incompleteness. Both our method and the VLA-only baseline succeed in under 10 seconds, while the Full-Patch planner fails to do so in the same timeframe due to its significant planning latency.
27

---

Published as a conference paper at ICLR 2026
Init.
Goal
Full (Success)
CLS
(Fail)
Ours
(Success)
Figure 17: Executed rollout demonstration on Block Pushing environment of Full, CLS, and our methods (60% Drop).
Init.
Goal
Full (Success)
CLS
(Fail)
Ours (Success)
Figure 18: Executed rollout demonstration on Wall environment of Full, CLS, and our methods (50% Drop).
28

---

Published as a conference paper at ICLR 2026
C.3 ADDITIONAL BLIND SPOT EXAMPLES
(a) MoCo-v3
(b) DINO-v3
Figure 19: Visualization Examples of the Blind Spot Problem in Additional Visual Encoders (60% Drop in PushT).
(a) Random
(b) Importance-Based Aatentiio--Encodel
Figure 20: Visualization Example of the Blind Spot Problem in Real-Worle LeRobot Environ- ment.
D LLM USAGE
LLMs were utilized solely for improving the clarity, grammar, and style of the paper and supple- mentary materials, No LLM was used for research ideation, experimental design, data generation,
result analysis. The authors take full responsibility for the content of this submission.
29


Reviews:
[{"review_id": "Fyv0tTT1Ae", "reviewer": "Reviewer_ASzo", "summary": "This paper introduces a method called sparse imagination, which aims to increase the efficiency for test-time planning with patch token-based transformer world models while maintaining planning performance. The method works by randomly dropping half of the patch tokens when calculating attention, ensuring robust predictions on any subset of tokens. During MPC planning, an arbitrary percentage of tokens can be dropped in each iteration to accelerate world model rollouts. Experiments on seven simulated environments and one real-world environment show that the method achieves comparable planning performance to the full-patch world model, with significantly reduced planning time.", "questions": "- In Table 1, why does a nonzero drop ratio sometimes outperform planning with full patches? Why isn’t the success rate strictly decreasing as the drop ratio increases?\n- For the LIBERO and LeRobot tasks, the dataset used to train the world model appears to consist solely of expert trajectories, which provides limited state and action space coverage. How does the world model generate the counterfactual information essential for planning with the policy?\n- How does the token drop ratio during training affect the world model’s prediction quality?\n-  In Table 2, what does the change in time for CLS mean, given that no tokens are dropped?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The idea is well-motivated, as planning time and inference cost are major concerns for patch token-based world models that encode rich information about the environment.\n- The reduction in inference and planning time is significant, while achieving comparable or even better planning success rates across various benchmarks.\n- The analysis of token information in Section 5.3 is quite interesting, providing insights into the information content and redundancy of patch features.", "weaknesses": "- The application of the proposed method seems somewhat limited, as it only applies to world models with patch tokens, a transformer backbone, and MPC as the planning algorithm. However, the idea of randomly dropping tokens during training and inference appears more general. Could this approach be extended to other use cases?\n- At planning time, the method relies on resampling different tokens across MPC iterations to capture the full task information. For open-loop CEM planning, however, information is lost since token sampling occurs only once. How are the CEM-only results in Table 1 affected by this limitation?\n- The paper could be strengthened by evaluating patch token-based world models using features other than DINO-v2. Does the analysis of patch token information in Section 5.3 generalize to other pre-trained patch features?\n- There is no evaluation of the world model’s prediction quality when using full patches while training with random token dropout. Does this training mechanism improve or degrade the world model’s quality?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "snV3vGZLgo", "reviewer": "Reviewer_C7kW", "summary": "This paper introduces a method for accelerating visual world-model planning by processing only a randomly selected subset of vision tokens during latent rollouts. The approach trains a transformer-based world model with randomized grouped attention so that, at test time, the planner can flexibly reduce token count without modifying the architecture. During planning, a fraction of tokens is randomly retained and used for MPC/CEM rollouts, yielding substantial computational savings with minimal degradation in performance. The work compares against learned token-importance and compression baselines, showing that random token selection matches or outperforms these more complex strategies. Experiments across serveral simulated domains and a real-world robotic task demonstrate that dropping 10–50% of tokens preserves task success while significantly reducing planning time. The authors argue that static importance metrics suffer from “blind spots” in dynamic planning settings, while random sampling avoids such failures with negligible overhead. Overall, the paper positions random token dropout as a simple and practical baseline for efficient visual planning.", "questions": "1. How sensitive is the approach to the choice of token-drop ratio across environments? Did you observe cases where higher or lower sparsity meaningfully changed success rates or planning stability?\n\n2. Can the authors formalize or provide intuition for the conditions under which random token subsets preserve planning quality? For example, how do factors like scene observability, task complexity, and visual redundancy influence performance?\n\n3. Do you expect similar efficiency gains on non-manipulation tasks such as navigation, locomotion, or long-horizon household planning? Have you tested (or can you comment on) performance in these settings?\n\n4. Would this method apply to other planning paradigms beyond CEM-based MPC, such as actor-critic agents, or transformer-based learned planners?\n\n5. Can you consider providing some insight into adaptive mechanisms for adjusting the token-drop ratio at test time (e.g., confidence- or uncertainty-based heuristics)? Do you foresee a principled approach for online adaptation?\n\n6. Can the authors provide a breakdown of computation cost across components (e.g., attention layers, latent rollouts, sampling iterations, VLA calls) to clarify where efficiency gains originate?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper is generally well-written and easy to follow, with a clear problem formulation, motivation, and method description. The overall narrative is coherent and the technical contributions are communicated effectively.\n\n2. The core idea of training a transformer-based world model with randomized grouped attention such that it can perform test-time planning by randomly dropping visual patch tokens is conceptually simple, architecture-agnostic, and broadly applicable to visual world-modeling settings.\n\n3. Experiments across multiple simulated control environments (e.g., PointMaze, Wall, PushT, Rope) demonstrate that dropping approximately 50% of visual tokens yields substantial planning-time reductions (40–60%) while preserving control performance. These results show that the approach delivers consistent and meaningful efficiency gains with minimal performance degradation.\n\n4. The method is compared against learned or attention-based token selection or merging techniques (e.g., LTRP, STAR, ATC), and consistently matches or outperforms them. This provides compelling evidence that simple random sampling can avoid failure modes associated with static importance-scoring heuristics and offers a strong, low-overhead baseline for token reduction in visual planning.\n\n5. Real-robot results and integration with a Vision-Language-Action model (SmolVLA) further strengthen the empirical validation. The technique transfers effectively to real-world robotic control and maintains competitive task success while reducing real-world episode execution time by roughly half, demonstrating practical benefit in realistic deployment scenarios.", "weaknesses": "1. The writing, particularly in the Experiments section, needs improvement for readability. Important details and baselines are placed in Appendix B.6 instead of the main text, which makes it difficult to follow comparisons and understand experimental context. Additionally, some methods (e.g., Latin Hypercube Sampling, McKay et al., 2000) are referenced without explanation, making it harder for readers unfamiliar with these techniques to interpret results.\n\n2. The paper argues that training the world model with random patch dropout induces robustness to token sparsification at test time, but this claim is supported primarily by empirical evidence. A simple formal or conceptual analysis clarifying when and why random token subsets preserve planning quality would strengthen the contribution. For instance, discussing conditions related to scene observability, task complexity, visual redundancy, and environmental structure would help clarify the boundaries of the method’s effectiveness. Explicitly connecting these factors to the observed planning success patterns would provide a stronger theoretical grounding and a clearer understanding of the method’s limitations and applicability.\n\n3. The experiments focus on goal-conditioned manipulation tasks, and it is unclear whether the approach extends to broader embodied settings (e.g., DMControl, Meta-World, RLBench, BEHAVIOR-1K, Habitat 3.0, or long-horizon household environments). Likewise, the evaluation uses a narrow set of visual encoders. Given recent advances in visual backbones (e.g., DINOv3, FastVLM, OpenVision-2, SAIL-VL2), testing across a wider set of pretrained representations would strengthen the claim that random patch sparsification generalizes beyond the chosen model family.\n\n4. The proposed approach is evaluated primarily within an MPC/CEM planning framework, with an additional VLA-guided sampling variant for real-robot experiments. However, it remains unclear whether the computational and performance benefits of this method extend to other planning paradigms. In particular, it would be useful to see results under implicit planners used in model-based RL (e.g., actor-critic methods) and learned planning approaches (e.g., transformer-based sequential planners or diffusion-policy rollouts). Because different planners have varying tolerance to approximation errors and representation sparsity, demonstrating consistent benefits would help establish that random token dropout is a generally applicable mechanism rather than one primarily suited to CEM-style sampling-based planning.\n\n5. The method relies on a manually chosen token-drop ratio, tuned per task. While the simplicity of random dropout is appealing, the lack of adaptive control limits its practicality in settings where computation budgets or scene complexity change over time (e.g., dynamic camera views or occlusions). An adaptive mechanism that adjusts sparsity based on model confidence, rollout divergence, uncertainty estimates, or planning consistency could potentially offer better performance–efficiency trade-offs. The paper briefly acknowledges this possibility, but does not explore even lightweight heuristics (e.g., annealing token sparsity or entropy-based token retention). As a result, it is unclear whether token sparsity could be automatically optimized online, which would be important for real-time deployment and applications with fluctuating computational constraints.\n\n6. Although the reported runtime savings are substantial, the compute analysis remains relatively coarse-grained. The paper primarily reports end-to-end planning time, but does not clearly decompose where efficiency gains arise (e.g., savings in vision attention layers, latent rollout computation, sampling loops, and planning iterations). Providing latency per module across token-drop ratios would help understand where the speedups come from and how to tune the method for different hardware and latency budgets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CZ5XaOMr2a", "reviewer": "Reviewer_HGmx", "summary": "This paper proposes Sparse Imagination, a simple yet effective method designed for transformer-based world models. The key idea is to accelerate planning by applying random dropout on visual tokens, significantly improving inference speed while largely preserving predictive performance. Extensive experiments and analyses demonstrate that the random dropout strategy, along with the corresponding attention mechanism used during training, is effective on a variety of robotics tasks.", "questions": "1. From my understanding of this work, actions are sampled either from MPC or from a VLA policy. The world model then imagines the future states resulting from these actions. My question is: during execution, how are the candidate actions selected based on the imagined states, especially in open-ended robotics tasks where no golden state is available?\n\n2. On line 258, the authors mention that the world model can evaluate and refine trajectories. How exactly is the refinement performed? I could not find a detailed description of this process in the paper.\n\n3. On line 141, the authors mention splitting the visual tokens into two groups. What is the motivation for this design? How are the groups specifically divided — for example, first half vs. second half, interleaved, or one group being dropped out? Why does such a partition help strengthen subtoken-level patterns?\n\n4. To address the issue of full-patch visual tokens being overly redundant for world model learning, there is also the concept of latent action [1] in robotic manipulation, which models the state transition a -> s using only a small number of tokens. I would like to ask how the authors consider the relationship between sparse imagination and latent action. Have you considered making a comparison, since both approaches largely fall within the scope of robotics tasks with action annotations?\n\nReferences:\n\n[1] Latent Action Pretraining from Videos.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper introduces Sparse Imagination, a surprisingly simple yet highly effective approach for accelerating world models. Counter-intuitively, such a straightforward method outperforms importance-based sampling strategies, offering both practical value and novel insight that the full patch tokens are redundant.\n\n2. The authors conduct extensive experiments to demonstrate the superiority of random dropout–based Sparse Imagination. Comparisons across multiple dropout ratios and different dropout strategies provide strong empirical support for the central claim.", "weaknesses": "1. The real-world evaluation is limited to only a single task, which raises concerns about the robustness of the reported results. In general, I would suggest using at least three tasks, or at a minimum two different embodiments. This is particularly important because the workspace of the LeRobot arm is quite small, meaning that the effective action chunk workspace is very limited.\n\n2. I have some concerns regarding the comparative experimental settings. In the case of random sampling, does fixed imply that the same subset of tokens is dropped consistently across different frames? Concerning the blind spot analysis, the paper suggests that blind spots may originate from static aspects of the scene during initialization. In this context, are the important tokens primarily those associated with the end-effector (EEF)? Moreover, I believe the analysis would be more intuitive and convincing if the authors could provide additional visualizations and concrete examples of blind spots as part of a quality assessment.\n\n3. A minor suggestion: It would be helpful to report an average score across multiple tasks in the experiments. This would make it more intuitive to see the overall superiority of Sparse Imagination compared to other baselines or ablations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ciIBTOhaws", "reviewer": "Reviewer_gLWK", "summary": "This paper proposes Sparse Imagination, a simple yet effective approach for improving computational efficiency in visual world-model-based planning.The key idea is to perform model predictive control (MPC) rollouts on only a random subset of ViT patch tokens, rather than processing all of them.To make this feasible, the authors train a transformer world model with a randomized grouped-attention strategy, enabling robustness to missing visual tokens at inference time.\nDuring planning, a random dropout mask is applied dynamically at each iteration, which substantially reduces the quadratic computational cost of self-attention.Extensive experiments on seven simulated environments (e.g., Pointmaze, PushT, Rope, LIBERO-10) and a real-world LeRobot setup show that the method achieves comparable task success to the full-token baseline while reducing planning time by more than 50%.", "questions": "1. Can the dropout ratio be made adaptive to uncertainty or task complexity?\n2. Does grouped attention reduce long-range spatial reasoning ability?\n3.  What is the real-world end-to-end latency improvement on hardware?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Clear motivation and simplicity**.\nThe paper tackles an important bottleneck of visual world-model planning (quadratic attention cost) with a conceptually simple and practical solution.\n\n2. **Strong empirical validation**. \nExperiments are extensive, covering both simulation and real robotic control tasks. Results convincingly demonstrate significant efficiency gains with minimal performance loss.", "weaknesses": "1. **Limited algorithmic novelty**.\nWhile effective, the proposed approach mainly combines known components (token dropout + grouped attention) and lacks deeper theoretical or architectural innovation.\n\n2. **Shallow theoretical grounding**. \nThe work frames the idea as “redundancy reduction,” yet offers little formal analysis or justification of when and why sparse imagination preserves sufficient task information.\n\n3. **Empirically chosen design choices**.\nThe dropout ratio and grouping strategy are empirically chosen; no adaptive mechanism or general rule for balancing performance and speed is proposed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)