SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices

Timi Owolabi; Hikmah Olawore; Frances Adelakun; Ayotomiwa Oyewumi

cs.ROOct 9, 2026

SmolWAM: A Latent World-Action Model for Robotic Manipulation on Edge Devices

Authors: Timi Owolabi, Hikmah Olawore, Frances Adelakun, Ayotomiwa Oyewumi

Organizations: ML Collective

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.

Figures & tables

Panel of Figure 1 — Figure 1: From Cosmos Policy to SmolWAM. Latent frames (width $\propto$ tokens), kept DiT blocks, denoising steps, and cost per 16-action chunk on an A100. SmolWAM drops the state and value frames (dashed), replaces the action frame with 28 state-conditioned queries ( $a{+}s$ ), and reads actions and future camera latents (its world head) from one step.

	Inference cost				Success rate (%)
Model	Params	TFLOP	ms	GiB	Screen	Spatial	Object	Goal	Long	Full
Cosmos Policy (5 steps)	1.96B	42.6	624	7.55	96.7	97.0	100.0	99.0	97.0	98.2 †
Cosmos Policy (1 step, 1 frame) ∗	1.96B	7.52	126	4.39	100.0	–	–	–	–	–
Pruned + LoRA distillation	1.40B	–	–	–	–	83.0	100.0	82.0	85.0	87.5 †
SmolWAM	0.88B	2.30	57	2.36	93.3	99.0	100.0	94.0	89.0	95.5
+ 7.5k training steps	0.88B	2.30	60	2.36	93.3	94.0	100.0	94.0	89.0	94.2

Table 1: LIBERO success and inference cost. Screen: 60 episodes; suites and Full: 400. Cost per 16-action chunk on one A100 (batch 1, bf16); Params excludes the VAE and text encoder; TFLOP counts network and VAE encoder; GiB: peak memory. ∗ Unmodified weights, single-frame latents via the 0.5M-parameter adapter. † Suites and Full from an earlier evaluation pipeline, on an L4 (Cosmos Policy reports 98.5% [ 15 ] ). ‡ Shared GPU. § Simulated: INT8 weights, bf16 kernels.

Figure 2: World-head forecast , decoded only for display.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: LIBERO success against cost for every row of Table 1 with an A100 cost. Filled markers are the 60-episode screen and hollow markers the 400-episode protocol (the teacher’s as in Table 1 ); the 8-block latency was measured on a shared GPU.

Figure 4: Distillation of SmolWAM (batch 32): normalized L1 of the one-step action chunk and future latents against the demonstrations, on training batches and held-out validation rows. Vertical lines mark the switch from the teacher’s 1,764-token layout to the compact one, and the start of 7,500 further steps.

Configuration (cumulative)	s/chunk	deviation	closed loop
Cosmos Policy 2B, 5 denoising steps (as released)	1.83	–	5/5, 3/3 Spatial
1 denoising step	1.22	0.005	5/5
+ no future-image decode	0.54	0.005	5/5
+ single-frame VAE encode + exact text sink	0.23	0.044	5/5, 3/3 Spatial, 3/3 Goal

Appendix

Table 2: Zero-training speedups on the released weights (L4, uncontended single-query benchmark). Deviation: mean absolute difference from the five-step model’s actions (raw units) on the benchmark’s single observation; over a 161-observation offline set, one step deviates by 0.003. Closed-loop scores are single-trial screens, so they screen rather than resolve small differences.

Figure 5: Single-flag ablations of the SmolWAM recipe, each trained for 2,500 steps (half the main run) and scored on the 60-episode screen. The shaded band spans five episodes either side of the reference.

Figure 6: Rollouts of SmolWAM (left) and Cosmos Policy (right) from the same initial state, one task per suite: first, middle, and last frames. Both succeed on all four, in 131, 147, 154, and 234 steps (SmolWAM) and 143, 147, 138, and 247 steps (Cosmos Policy). Recorded on an Apple M2 Pro.

Explore similar work

Light-WAM: Efficient World Action Models with State-Fusion Action Decoding
Jun 6, 2026Ziang Li, Dongzhou Cheng, Yibin Wang +5Robot Policy Learning World Models for Robotics
Wuhan University · Shanghai Innovation Institute · Southeast University +2
Faster-WAM: Do World Action Models Need Deep Action Modules?
Aug 3, 2026Liheng Ma, Rui Heng Yang, Amin Abyaneh +8World Model Learning Efficient Inference for World Action Models
Huawei Noah’s Ark Lab Montreal · McGill University & Mila-Quebec AI Institute · University of Toronto +2
LiLa-WAM: Lightweight Latent Reasoning World-Action Model for Robotic Manipulation
Aug 4, 2026Fan Yang, Yuting Su, Xiaobo Wang +7World Models for Robotics Latent Visual Reasoning
1Tianjin University · 2Shenzhen University of Advanced Technology · 5Sangfor Technologies Inc. +2

Jun 6, 2026J/K move · Enter open · S save

Light-WAM: Efficient World Action Models with State-Fusion Action Decoding

Ziang Li, Dongzhou Cheng, Yibin Wang, Shiyue Wang+4

Wuhan University · Shanghai Innovation Institute · Southeast University+2

Robot Policy Learning World Models for Robotics Visuomotor Policy Learning Efficient Inference for World Action Models World Action Models

Open paper AskPDF arXiv

World Action Models (WAMs) extend robot policy learning by incorporating future prediction as an additional training objective, encouraging the policy to encode task-relevant temporal structure in its representations. Current WAMs often rely on large-scale generative architectures that incur high training costs and inference latency, making them difficult to deploy as efficient closed-loop policies. We propose Light-WAM, a lightweight World Action Model for efficient robot manipulation. Specifically, it is built with a compact video backbone and performs future-video supervision in a downsampled latent space, reducing the cost of video co-training while retaining its benefits for representation learning. For action prediction, Light-WAM introduces the StateFusionActionExpert, which reads adapted states from multiple backbone layers, fuses them through learned-query pooling, and directly predicts action chunks in a single forward pass. This design provides an efficient interface between video backbone representations and robot actions, avoiding the need for heavy generative action experts. Experiments demonstrate that Light-WAM maintains strong performance on LIBERO and achieves usable multi-task performance on RoboTwin 2.0, while using only 0.44B trainable parameters. It also achieves 72.03ms inference latency with 4.1GiB peak GPU memory and improved training throughput.

Abstract

Explore similar work

Light-WAM: Efficient World Action Models with State-Fusion Action Decoding

Faster-WAM: Do World Action Models Need Deep Action Modules?

LiLa-WAM: Lightweight Latent Reasoning World-Action Model for Robotic Manipulation