SmolWAM: A World-Action Model Small Enough for the Edge

A world-action model that must be served from a datacenter cannot watch over a robot’s motion. This paper asks how small such a model can get while still predicting both actions and the future they create.

Background

An action-conditioned world model does two things at once: it predicts what the robot should do, and what the world will look like afterwards. That second prediction is what makes policy evaluation and planning possible without touching the real robot every time. The strongest world-action models achieve this with video diffusion transformers, but those are far too computationally expensive to run on the resource-constrained edge devices that actually sit on robots. The result is a familiar bind: the models good enough to be useful are too heavy to deploy where the robot is.

Core Innovation

SmolWAM is a latent world-action model distilled from the 2B-parameter Cosmos Policy. Its central move is to leave the pretrained trunk in place and rebuild the inference interface around it, rather than retraining a smaller model from scratch. The rebuilt interface uses single-frame visual latents instead of a video stack, cached text features that preserve the trunk’s attention sinks, only 12 of the 28 transformer blocks, and 1,008 tokens from which both actions and future visual latents are read out in a single rectified-flow step. Because the trunk is preserved, distillation is cheap at 3.6 GPU-hours. And because the world head is kept, the model still produces future predictions rather than collapsing into a plain reactive policy.

Results

• 95.5% success across 400 LIBERO episodes.
• 57 ms per inference and 2.3 TFLOP per 16-action chunk, at 2.4 GiB of GPU memory.
• About 11x faster than its teacher, with 18.5x fewer FLOPs and 3.2x less memory.
• Training cost of 3.6 GPU-hours, and a world head that still forecasts scene motion.

Limitations

The headline numbers come from LIBERO, a simulated manipulation benchmark; the paper does not report a matched on-robot evaluation, so the deployed real-world success rate and latency under real sensor and actuator dynamics remain unestablished. Distilling from a specific teacher ties the result to Cosmos Policy, and the paper does not show how the interface-rebuild recipe transfers to other trunks, or to tasks with longer horizons or heavier physical contact. Keeping 12 of 28 blocks and a single rectified-flow step is an aggressive compression, and the reported 95.5% leaves open how much of the teacher’s broader generality, beyond the evaluated suite, survives.

Industry Implications

The immediate implication is that on-robot world modeling becomes a shippable engineering target rather than a research aspiration: 2.4 GiB and 57 ms per action chunk fit within the memory and latency budgets of realistic edge compute, and a 3.6 GPU-hour training cost makes iteration cheap. For teams evaluating world-action models, the lesson is that the interface around a pretrained trunk can matter as much as the trunk itself; a compression path that keeps the world head is far more valuable than one that merely accelerates a reactive policy. The natural next test for any team is to reproduce the recipe on their own robot and horizon length before betting on it in a product.