arXiv:2610.13740 · v2
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.
作者备注：v2: Added action-only ablation, expanded efficiency analysis, and edge GPU benchmarks. Accepted at the NeurIPS 2026 Robot Learning with World Models workshop and the CoRL 2026 Do Robots Need World Models? and The Action Gap workshops. Code: https://github.com/iitimii/SmolWAM