SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices

원문 초록

World-action models (WAMs) jointly predict robot actions and future observations, but their video diffusion backbones are costly for edge deployment. We introduce SmolWAM, a compact latent WAM derived from the 2B-parameter Cosmos Policy using single-frame encoding, cached text features with preserved attention sinks, depth pruning, and token compaction. With 12 of 28 transformer blocks and 1,008 tokens, SmolWAM predicts actions and future latents in one rectified-flow step. After 3.6 A100 GPU-hours of training, it achieves 95.5% success over 400 LIBERO episodes at 57ms, 2.3TFLOP, and 2.4GiB per 16-action chunk on an A100: 11x lower latency and 3.2x lower memory than the five-step teacher. Its world head predicts scene motion; joint training improves success by 6.2 percentage points over an action-only ablation.