Sparse Imagination for Efficient Visual World Model Planning

TL;DR

A Sparse Imagination for Efficient Visual World Model Planning is proposed, which enhances computational efficiency by reducing the number of tokens processed during forward prediction by enabling sparse imagination during latent rollout.

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.

Figures (20)

• Figure 1: Sparse Imagination. Sparse imagination accelerates planning by performing model predictive control (MPC) rollouts on a random subset of visual tokens. A new dropout pattern is dynamically sampled at each MPC iteration, and both predictions and optimization for CEM are computed using only the selected patches to improve efficiency and robustness.
• Figure 2: Randomized Grouped Attention Strategy. Our randomized grouping strategy used during training to generalize to arbitrary token subsets. Visual tokens are randomly partitioned into two groups, and attention is masked to occur only within the same spatial group. This trains the model to process sparse inputs effectively, and its necessity is shown in ablations.
• Figure 3: Environments. We assess the world model across eight simulation environments and two real-world tasks, arranged as follows (left-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 planning. We compare the average planning success rates achieved by grouped attention and full attention under various drop ratios, highlighting the benefit of grouped attention.
• 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 ($\bigstar$) represent the operating point we highlight (50% drop), which provides a practical choice between speed and reliability.