View a PDF of the paper titled SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices, by Timi Owolabi and 3 other authors
Abstract: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.
Submission history
From: Timi Owolabi [view email]
[v1]
Fri, 9 Oct 2026 15:34:33 UTC (2,069 KB)
[v2]
Sat, 7 Nov 2026 20:42:45 UTC (2,080 KB)
