Accepted at the NeurIPS 2026 Workshop on Robot Learning with World Models
SmolWAM
A Latent World-Action Model for Robotic Manipulation on Edge Devices
All against the teacher, Cosmos Policy (98.2% over 400 episodes), on one A100-40GB with batch 1, bf16 and text features cached for every model.
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, 2.3 TFLOP, and 2.4 GiB of GPU memory per 16-action chunk, 11× faster and 3.2× leaner than its teacher, and its world head still forecasts scene motion.
Get started
From pip install to actions
Install the package, load the weights from Hugging Face, and predict actions from an observation and an instruction. Setup, LIBERO evaluation and training are documented in the repository.
pip install git+https://github.com/iitimii/SmolWAM
from smolwam import SmolWAM model = SmolWAM.from_pretrained("timiiowolabi/SmolWAM") actions = model.predict_action(obs, instruction)
Method
Keep the trunk, rebuild the interface
Cosmos Policy turns a video model into a WAM by treating actions, future observations and values as latent frames of one denoising sequence. Its inference stack was built for video generation: a 4.8 B T5-XXL text encoder, a 33-frame clip encode per query, and five denoising passes over 1,764 tokens. Our premise is that much of that cost comes from the interface, not from what control requires.
Swipe sideways to see the whole diagram.
Five changes, each backed by a diagnostic
Most of these changes break the pretrained trunk if applied zero-shot, so each one comes with a test of what it costs, and the trunk is retrained with distillation.
| Change | What SmolWAM does | Diagnostic |
|---|---|---|
| Single-frame latents | Encodes each camera image once, as a single frame, with the Wan2.1 encoder (54 M). A 0.5 M-parameter adapter maps it to the clip-mode latents the trunk expects. | Single-frame latents raise the teacher's action error from 0.0089 to 0.038; the adapter brings it back to 0.0092. |
| Cached text that keeps the attention sinks | Replaces T5-XXL with cached BERT-base features, zero-padded to 512 positions. Six zero-initialised vision-to-language layers let the instruction attend to the images. | The cross-attention key and value projections have no bias, so T5's ~500 zero-padded positions act as an attention sink. Truncating the context to 32 or 128 positions drops the teacher to 0/5 successes; restoring the removed softmax mass analytically returns it to 5/5. |
| Boundary-preserving pruning | Keeps blocks 0–2 and 19–27, 12 of 28 (0.85 B parameters). | Blocks 3–22 are individually close to the identity, but removals don't compose: skipping block 7, 14 or 20 alone fails no episode, skipping blocks 12–15 together fails all of them, and skipping block 1, 26 or 27 fails every episode. |
| State-conditioned action queries | A zero-initialised MLP adds the robot state to 28 queries that replace the 196-token action frame at its original positions. Only the blank and image frames stay: 1,008 tokens. | Zero-shot, dropping the future-state and value frames barely moves the action error (0.009 to ≤ 0.019), but dropping the blank (0.30) or state (0.058) frame, or compacting the vector frames to 28 tokens (0.21), does. The layout has to be learned. |
| One-pass readout | Evaluates the trunk once, at the maximum noise level with fixed noise, and reads actions and future latents from the closed-form clean estimate x̂0 = xτ − τvθ. | The teacher itself solves all 60 screening episodes in one step. |
1,764 tokens down to 1,008
Swipe sideways to see the whole diagram.
Distillation with a layout curriculum
Pruning and compaction are separate distribution shifts, so the student recovers from one before facing the other. It is trained on the demonstrations and on the teacher's own one-step predictions at once:
ℒ = ‖â − a‖₁ + ‖â − aT‖₁ + ‖ẑ − z‖₁ + ‖ẑ − zT‖₁
Results
97% of the teacher's success at 11× lower latency
SmolWAM succeeds in 382 of 400 LIBERO episodes (95.5%; 95% Wilson interval 93.0–97.1), against the teacher's 393 (98.2%; 96.4–99.1). There is no detectable difference on Spatial and Object. The deficit lies in the multi-step Goal and Long suites, where failed rollouts usually complete the first subgoal and stall on the second.
| Model | Params | TFLOP | ms | GiB | Screen | Spatial | Object | Goal | Long | Full |
|---|---|---|---|---|---|---|---|---|---|---|
| Cosmos Policy (5 steps) | 1.96 B | 42.6 | 624 | 7.55 | 96.7 | 97.0 | 100.0 | 99.0 | 97.0 | 98.2† |
| Cosmos Policy (1 step, 1 frame) | 1.96 B | 7.52 | 126 | 4.39 | 100.0 | – | – | – | – | – |
| Pruned + LoRA distillationearlier student | 1.40 B | – | – | – | – | 83.0 | 100.0 | 82.0 | 85.0 | 87.5† |
| SmolWAM | 0.88 B | 2.30 | 57 | 2.36 | 93.3 | 99.0 | 100.0 | 94.0 | 89.0 | 95.5 |
| + 7.5k training steps | 0.88 B | 2.30 | 60 | 2.36 | 93.3 | 94.0 | 100.0 | 94.0 | 89.0 | 94.2 |
| INT8 weights§ | 0.88 B | 2.30 | 65‡ | 2.36 | 91.7 | – | – | – | – | – |
| 8 blocks | 0.60 B | 1.76 | 134‡ | 1.84 | 83.3 | – | – | – | – | – |
| SmolWAM-LA (4 steps) | 0.92 B | 4.28 | 239 | 2.43 | 53.3 | – | – | – | – | – |
Screen: 60-episode screening set, which includes the tasks an earlier pruned student failed most, so it scores below the full protocol. Suites and Full: 400 episodes. Cost per 16-action chunk on one A100 (batch 1, bf16); Params excludes the VAE and text encoder, TFLOP counts the network and VAE encoder, GiB is peak GPU memory. † Earlier evaluation pipeline (Cosmos Policy reports 98.5%). ‡ Shared GPU. § Simulated: INT8 weights, bf16 kernels.
Where the 11× comes from
Running the unmodified teacher for one step on single-frame inputs already cuts latency 5× and compute 5.7×, to 126 ms and 7.5 TFLOP, with no loss on the screening set. SmolWAM's trunk redesign adds 2.2× in latency, 3.3× in compute and 1.9× in memory.
At 57 ms per chunk (6.8 ms preprocessing, 12.9 ms VAE, 36.9 ms trunk) it issues 17.7 chunks, or 282 actions, per second. Chunks execute open loop, so the chunk rate bounds how often the policy reacts to new observations.
Depth, data and weights
An earlier student that removed only 8 blocks and recovered with LoRA distillation on 414 teacher observations reached 87.5%. SmolWAM differs in distillation coverage, full fine-tuning and the learned layout; we have not separated those effects.
Training 7,500 more steps does not help (94.2%), and an 8-block variant falls to 83.3% on the screening set. INT8 quantisation of the 172 largest linear layers shrinks the stored network from 1.64 to 0.85 GiB for one screening episode.
One task per suite, closed loop
pick up the black bowl in the top drawer of the wooden cabinet and place it on the plate
pick up the alphabet soup and place it in the basket
put the wine bottle on the rack
put the white mug on the left plate and put the yellow and white mug on the right plate
SmolWAM's four rollouts from the video above, one task from each LIBERO suite, played at simulation speed (20 Hz control) without inference time. Recorded on a MacBook GPU.
Student and teacher, same initial state
One screening episode per suite from the same initial state, first trial, one frame per executed action chunk. On the Goal task SmolWAM opens the drawer, then stalls before placing the bowl: the typical failure, which completes the first subgoal and not the second.
World model
It still forecasts the future
The pass that produces the actions also forecasts both camera views 16 steps ahead. These future latents are SmolWAM's world head: supervised in training, never decoded at deployment. Below they are decoded with the frozen Wan VAE for display only.
Held-out contexts, primary camera. Forecasts keep the scene and place the arm near its future pose, but blur its appearance.
| 96 held-out contexts | Latent L1 ↓ | LPIPS ↓ | PSNR ↑ | Flow EPE ↓ |
|---|---|---|---|---|
| Copy the current frame | 0.133 | 0.115 | 24.3 | 1.63 |
| SmolWAM | 0.101 | 0.128 | 24.7 | 1.43 |
| SmolWAM-LA | 0.102 | 0.130 | 24.6 | 1.47 |
| SmolWAM-LA + reward fine-tuningrange over 5 rewards | 0.104–0.105 | 0.167–0.177 | 23.6–23.8 | 1.66–1.72 |
SmolWAM lowers latent L1 by 24% relative to copying the current frame, and the end-point error between RAFT optical flow to the predicted and to the real future from 1.63 to 1.43 pixels. LPIPS stays at the copy-frame level: the world head captures motion rather than texture.
Negative results
Richer WAM components hurt at this scale
A latent-action cascade
SmolWAM-LA ports ABot-M0.5 onto SmolWAM (39 M more parameters): tokens from a frozen 32-dimensional latent-action model behind an asymmetric attention mask, mixture-of-transformer towers for the latent-action, arm and gripper streams, four-step conditional flow matching, and dream forcing, which trains the action head on predicted latents.
| Model and action sampler | Screen |
|---|---|
| SmolWAM, one-step readout | 93.3 |
| SmolWAM-LA, 4 flow steps (as trained) | 53.3 |
| SmolWAM-LA, 2 flow steps | 70.0 |
| SmolWAM-LA, 1 flow step | 63.3 |
Its one-step action estimate matches SmolWAM's offline (validation L1 0.089 vs. 0.091), but four-step sampling raises the error to 0.155. Its latent-action predictions are worse than predicting the mean code (MSE 1.15 against a variance of 0.64), and its world head is no better.
Rewarding the world head
We fine-tuned SmolWAM-LA's world head with reward-weighted denoising regression under five rewards built from optical flow, depth and LPIPS.
- Every reward worsened every image metric (LPIPS from 0.130 to 0.167–0.177).
- Closed-loop success collapsed to 5.0–16.7% within 300 steps.
- The likely cause: the objective has no action term while the shared trunk trains.
A pretrained video WAM tolerates large interface and depth changes when each is distilled rather than applied zero-shot, but not auxiliary objectives that train its shared trunk without an action anchor.
Limitations
What's next
- Benchmarks. SmolWAM is evaluated only on LIBERO. RoboCasa and, ultimately, real robots come next.
- Long horizons. Reducing transformer depth hurts long-horizon tasks disproportionately, most of all on LIBERO-Long. Pruning that preserves long-horizon control remains open.
- Hardware. Costs are measured on one A100, and INT8 is simulated with bf16 kernels; on-device measurements are future work.
Citation
BibTeX
@inproceedings{owolabi2026smolwam,
title = {SmolWAM: A Latent World-Action Model for Robotic
Manipulation on Edge Devices},
author = {Owolabi, Timi and Olawore, Hikmah and Adelakun, Frances
and Oyewumi, Ayotomiwa},
booktitle = {NeurIPS 2026 Workshop on Robot Learning with World Models: Capabilities, Frontiers, and Challenges},
year = {2026},
url = {https://smolwam.vercel.app}
}