A world-action model must predict actions and their consequences
A world-action model (WAM) does two jobs at once: it predicts the robot's next action chunk and the observations those actions should produce. Cosmos Policy builds this from a video diffusion transformer, treating actions, future camera latents and values as latent frames in one denoising sequence, and reaches 98.5% success on LIBERO. That design is also why it is expensive: its inference stack was built for video generation, not control.
The paper's premise is that much of the teacher's cost sits in that interface rather than in what manipulation requires. Cosmos Policy spends 624 ms, 42.6 TFLOP and 7.55 GiB per 16-action chunk on an A100, even with cached text features, because it runs a 4.8B-parameter T5-XXL encoder, encodes a 33-frame clip per query, and takes five denoising steps over 1,764 tokens. SmolWAM asks whether the same pretrained trunk can be kept while that interface is rebuilt.
Four interface changes shrink the trunk without retraining it
SmolWAM keeps the teacher's weights, token positions and rectified-flow readout, then changes four things. It encodes each camera image once as a single frame through the Wan2.1 encoder, with a 0.5M-parameter adapter that maps those latents back to clip-mode latents; this restores the teacher's one-step action error from 0.038 to 0.0092, close to the 0.0089 baseline. It replaces T5-XXL with cached BERT-base features zero-padded to 512 positions, because the trunk's bias-free key and value projections turn T5's padded positions into an attention sink: truncating the context to 32 or 128 positions drops the teacher to 0/5 closed-loop successes, while restoring the removed softmax mass returns it to 5/5.
It also keeps 12 of 28 transformer blocks, preserving blocks 0–2 and 19–27, and replaces the 196-token action frame with 28 state-conditioned queries. The resulting sequence is 1,008 tokens, and actions plus future latents are read out in one rectified-flow step. The paper reports that most of these changes break the trunk zero-shot, so a distillation curriculum first trains in the teacher's 1,764-token layout for 1,000 of 5,000 steps, then switches to the compact layout. Training takes 3.6 GPU-hours on one A100.
Success falls modestly while cost falls sharply
On the full 400-episode LIBERO protocol, SmolWAM succeeds in 382 episodes (95.5%, 95% Wilson interval 93.0–97.1), against the teacher's 393 of 400 (98.2%, 96.4–99.1). The paper detects no difference on the Spatial and Object suites; the deficit sits in the multi-step Goal suite (94 vs. 99) and Long suite (89 vs. 97), where failed rollouts usually complete the first subgoal and stall on the second.
The cost comparison is the paper's central result. Per 16-action chunk on one A100, SmolWAM uses 0.88B parameters, 2.30 TFLOP, 57 ms and 2.36 GiB, versus the teacher's 1.96B parameters, 42.6 TFLOP, 624 ms and 7.55 GiB. That is 11x lower latency, 18.5x fewer FLOPs and 3.2x less memory. Of the 57 ms, the transformer takes 37 ms, the VAE encode 13 ms and preprocessing 7 ms. At that rate the model issues 17.7 chunks, or 282 actions, per second.
Most of the speedup comes from the interface, not the pruning
The paper separates the sources of savings. Running the teacher's unmodified weights for one denoising step on single-frame inputs cuts latency 5x and compute 5.7x, to 126 ms and 7.5 TFLOP, with no loss on the 60-episode screening set. With no adapter and the text truncated under the exact sink correction, the released model runs 8x faster on an L4, from 1.83 to 0.23 s, with no failures in an 11-episode screen.
The trunk redesign itself adds 2.2x in latency, 3.3x in compute and 1.9x in memory relative to that one-step, single-frame teacher. In other words, the largest single win is removing video-generation work the controller never reads, such as decoding future images; the compact layout and pruning then build on that. The paper also notes that cutting chunks short does not help: re-querying the teacher after 8 of its 16 actions solves no screening episode at one or five denoising steps, so faster reaction would need training-time changes.
The compressed world head keeps motion but loses texture
SmolWAM still predicts the future. On 96 held-out contexts, its forecast of the view 16 steps ahead lowers latent L1 by 24% relative to copying the current frame (0.101 vs. 0.133), and lowers the end-point error between RAFT optical flow on the predicted and real future from 1.63 to 1.43 pixels. Decoded forecasts keep the scene and place the arm near its future pose, but blur its appearance; LPIPS stays at the copy-frame level (0.128 vs. 0.115).
The paper's interpretation is that the world head captures motion rather than texture. That distinction matters for readers comparing WAMs: a model can retain useful dynamics prediction while its decoded images look poor, so image-quality metrics and control success answer different questions here.
Adding richer components hurts at this scale
The paper tests whether more elaborate WAM machinery helps. Porting ABot-M0.5 onto SmolWAM adds 39M parameters, latent-action tokens, mixture-of-transformer towers, four-step flow matching and dream forcing. The resulting SmolWAM-LA matches SmolWAM's offline one-step action estimate (validation L1 0.089 vs. 0.091), but four-step sampling raises the error to 0.155, and one- or two-step sampling recovers only 63.3% and 70.0% of screening episodes. It succeeds in 53.3% of screening episodes at 239 ms.
Reward-weighted fine-tuning of the world head was worse: under five rewards built from optical flow, depth and LPIPS, every image metric worsened (LPIPS from 0.130 to 0.167–0.177) and success collapsed to 5.0–16.7% within 300 steps. The paper attributes this to an objective with no action term training the shared trunk. The authors' conclusion is that a pretrained video WAM tolerates large interface and depth changes when each is distilled, but not auxiliary objectives that train its shared trunk without an action anchor.
The evidence supports LIBERO control, not general robot deployment
Two boundaries shape how far these numbers travel. All evaluation is on LIBERO, with the teacher's observation pipeline, 16-action chunks executed open loop, seed 195 and 10 trials per task; the paper names RoboCasa and real-robot settings as future work. And reducing transformer depth disproportionately affects long-horizon tasks, with the largest degradation on LIBERO-Long, so pruning strategies that preserve long-horizon control remain open.
Offline error is also a poor guide to closed-loop behavior: validation action L1 settles near 0.09, ten times the teacher's own one-step error, while SmolWAM still succeeds in 95.5% of episodes. An extra 7,500 training steps lowered training loss but neither validation loss nor success. The practical priority the paper establishes is therefore the interface rebuild: one-step, single-frame inference with a preserved attention sink delivers most of the 11x speedup, while deeper pruning trades away long-horizon success, as the 8-block variant's drop to 83.3% on the screen shows.
Sources used in this answer
SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices
SmolWAM retains 95.5% LIBERO success over 400 episodes at 57 ms, 2.3 TFLOP and 2.4 GiB per 16-action chunk, 11x faster than its Cosmos Policy teacher with 18.5x fewer FLOPs, while its world head still forecasts scene motion.