WAM-Cache reduces robot video processing for faster action execution
WAM-Cache: Staleness-Bounded KV Reuse for Efficient World Action Models
RoboticsComputer Vision and Pattern RecognitionMachine Learning
Summary
Robots use video models to understand their surroundings and decide what to do next, but processing all video data each time takes a lot of computing power. The authors introduced WAM-Cache, a method that stores and selectively updates parts of the video information instead of redoing everything repeatedly. They discovered that updating areas where the robot focuses its attention is more important than just refreshing parts of the video that changed visually. This approach saves significant computing effort while keeping the robot's actions nearly as accurate as when processing all video data.
What this means in practice
- •For robotics engineers: Build robot control systems that run faster by caching and selectively updating video model computations without retraining.
- •For machine learning developers: Integrate efficient key-value reuse mechanisms into video-based policy models to reduce processing loads during real-time decision making.
Authors
Kai Ding, Yang He, Ruijie Quan, Yi Yang
Abstract
World Action Models (WAMs) enable generalist robot manipulation by conditioning an action expert on representations from a pretrained video Diffusion Transformer (DiT). In closed-loop control, the video DiT runs at every chunk to encode the current observation into layerwise key-value (KV) pairs that the action expert queries. This prefill dominates the per-chunk computational cost, yet existing training-free accelerations leave it fully dense. We present WAM-Cache, a training-free framework that retains layerwise key-value representations across chunks and recomputes only a sparse refresh set of tokens. Crucially, we find that the intuitive heuristic of refreshing visually drifted tokens plateaus far below the dense baseline, even with an oracle predicting ground-truth KV drift. Downstream action accuracy is instead governed by where the action expert attends, not by what moved. WAM-Cache therefore selects the refresh set by uniting the action expert's cross-attention with visual latent surprise, complemented by a strict age bound that suppresses compounding error. On Fast-WAM, WAM-Cache cuts video DiT prefill FLOPs by 32-42% across RoboTwin 2.0, LIBERO, and real-world experiments, while staying within 0.7-1.8 percentage points of the dense policy in simulation and 2.5 points on a real robot.