Sparse world action models speed up robot video prediction twice as fast
Sparse-WAM: Accelerating World Action Models via Action-Guided Sparse Imagination
RoboticsArtificial Intelligence
Summary
Robot controllers sometimes predict what will happen next in a video to decide their actions. This prediction can be slow because it processes lots of video frames fully. The authors found a way to speed this up by paying attention only to the parts of each video frame that matter most for the robot's future actions. They designed a method to pick these important parts efficiently, reducing the amount of video data processed and making the robot’s predictions nearly twice as fast while keeping performance high.
What this means in practice
- •For robotics engineers: Improve robot decision-making speed by processing only action-relevant video regions during future state prediction.
- •For autonomous vehicle software teams: Speed up visual prediction models that guide driving decisions by focusing computation on critical future regions influenced by planned actions.
Authors
Xinling Xie, Haodong Wang, Jiazhi Mi, Zhiming Liu, Zicong Hong, Xiaoyi Pang, Qianli Liu, Yangjia Hu, Ying Chen, Zhengyang Yan, Song Guo
Abstract
World-action models (WAMs) leverage pretrained video models to improve generalization in robot control by jointly predicting future visual states and actions. This capability comes at a substantial inference cost, as dense future-frame tokens are repeatedly processed during denoising. Prior methods address this by token pruning that prioritizes visual fidelity to reduce denoising costs in video diffusion models. However, these methods do not use action relevance to determine which future-frame tokens to retain during joint denoising in WAMs. In this paper, we propose Sparse-WAM, a training-free framework for action-guided sparse imagination that selectively processes future-frame tokens to accelerate WAM inference. We observe substantial overlap in the spatial distribution of attention from action tokens to future-frame tokens (action-to-future attention) between consecutive denoising steps, despite continued updates to the future representations. Motivated by this, we develop Action-Guided Token Selection to retain frame-specific action-relevant regions together with cross-frame context. However, a naive implementation can incur attention-scoring and token-packing overhead that offsets the computational savings from pruning. We therefore introduce Pilot, an efficient engine that reduces sparse inference overhead through lightweight scoring and cross-step reuse of token selections. On LIBERO with FastWAM-Joint and RoboLab-120 with Cosmos 3 Edge, Sparse-WAM achieves inference speedups of approximately $2.0\times$ and $1.8\times$, respectively, over dense eager inference on an NVIDIA RTX 4090, while largely preserving task performance.