SmolWAM speeds up robot action prediction on low power devices
SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices
RoboticsMachine Learning
Summary
Robots need to predict what will happen when they take actions, which usually requires lots of computing power. The authors created SmolWAM, a smaller, faster model that can run on devices with limited resources like edge GPUs. It keeps key parts of a large existing model but uses fewer computations to predict robot actions and future scenes with good accuracy. SmolWAM completes tasks much faster while using less memory than the big model it’s based on.
What this means in practice
- •For robotics engineers: Run efficient robot action and scene prediction models on resource-limited hardware to enable faster manipulation in real-world settings.
- •For embedded systems developers: Deploy predictive models for robot behavior on edge computing devices to reduce latency and power consumption during autonomous tasks.
Authors
Timi Owolabi, Hikmah Olawore, Frances Adelakun, Ayotomiwa Oyewumi
Abstract
World-action models (WAMs) predict a robot's actions together with the observations those actions will produce. However, the strongest models rely on video diffusion transformers that are far too computationally expensive for resource-constrained edge devices. We present SmolWAM, a latent WAM that we derive from the 2B-parameter Cosmos Policy by keeping its pretrained trunk and rebuilding its inference interface: single-frame visual latents, cached text features that preserve the trunk's attention sinks, 12 of 28 transformer blocks, and 1,008 tokens from which actions and future visual latents are read out in one rectified-flow step. Trained in 3.6 GPU-hours, SmolWAM succeeds in 95.5% of 400 LIBERO episodes at 57 ms and 2.3 TFLOP per 16-action chunk with 2.4 GiB of GPU memory, 11x faster than its teacher with 18.5x fewer FLOPs and 3.2x less memory, and its world head still forecasts scene motion.