SmolWAM

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

Timi Owolabi, Hikmah Olawore, Frances Adelakun, Ayotomiwa Oyewumi · ML Collective · NeurIPS 2026 Workshop on Robot Learning with World Models

Project page · Code · Paper

Cosmos Policy 2B (left) and SmolWAM (right) on the same LIBERO tasks and initial states. Before every 16-action chunk each robot waits for its model's measured A100 inference time (624 ms vs. 57 ms), so the clocks show wall time at A100 speed. Both policies succeed on all four tasks; together they take 33.6 s with SmolWAM and 57.3 s with Cosmos Policy.

SmolWAM is a 0.88B-parameter world-action model for robot manipulation. Given two camera images, the robot state and a language instruction, one forward pass predicts a 16-step action chunk and the latent future of both camera views 16 steps ahead. It is derived from Cosmos Policy by keeping 12 of its 28 pretrained diffusion-transformer blocks and rebuilding the inference interface: single-frame Wan2.1 VAE latents with a small adapter, cached BERT instruction features padded to keep the trunk's attention sinks, 28 state-conditioned action queries, and a one-step rectified-flow readout. On LIBERO it keeps 97% of the teacher's success at 11× lower latency.

	Cosmos Policy 2B	SmolWAM
LIBERO success (400 episodes)	98.2%	95.5%
Trunk parameters	1.96 B	0.88 B
Compute per 16-action chunk	42.6 TFLOP	2.30 TFLOP
Latency per chunk (A100, bf16)	624 ms	57 ms
Peak GPU memory	7.55 GiB	2.36 GiB
Denoising steps	5	1

Usage

pip install "smolwam @ git+https://github.com/iitimii/SmolWAM"

from smolwam import SmolWAM

policy = SmolWAM.from_pretrained("timiiowolabi/SmolWAM")  # CUDA, Apple silicon (MPS) or CPU

obs = {
    "primary_image": primary,  # (H, W, 3) uint8, upright third-person camera image (LIBERO: 256 x 256)
    "wrist_image": wrist,  # (H, W, 3) uint8, upright wrist camera image
    "proprio": proprio,  # (9,) gripper joint positions (2), end-effector position (3), quaternion (4)
}
actions = policy.predict_action(obs, "put the bowl on the plate")  # (16, 7): end-effector deltas (6), gripper (1)

from_pretrained downloads this repository and the Wan2.1 VAE encoder/decoder (Wan-AI/Wan2.1-T2V-1.3B, Wan2.1_VAE.pth). Closed-loop LIBERO evaluation, benchmarking, world-model decoding and training are documented in the GitHub repository.

Files

File	Contents
config.json	Model configuration and the action / proprioception normalisation statistics
model.safetensors	All SmolWAM weights in fp32 (trunk, text projection, vision-to-language layers, state MLP, action queries, latent adapter)
text_cache.safetensors	Cached BERT features of the LIBERO instructions, used for exact reproduction of the paper's evaluation
media/, replay.mp4	The comparison video above (replay.mp4 is a web-encoded copy for the Hub's video preview)

Model details

Architecture	Diffusion transformer, 12 blocks, width 2048, 16 heads, AdaLN-LoRA; 1,008 tokens per query
Parameters	0.88 B (network); plus the Wan2.1 VAE encoder (54 M) and BERT-base (109 M, only for unseen instructions)
Inputs	Third-person and wrist RGB images, 9-D proprioception, a language instruction
Outputs	16 × 7 action chunk (delta end-effector pose and gripper) and the latent future of both cameras 16 steps ahead
Inference	One forward pass at fixed noise (rectified-flow one-step estimate); no denoising loop
Base model	nvidia/Cosmos-Policy-LIBERO-Predict2-2B

Training. Distilled from the base model on the 1,700 successful demonstrations of LIBERO-Cosmos-Policy with L1 losses on the action chunk and future latents, against both the demonstrations and the teacher's one-step predictions. The first 1,000 of 5,000 steps keep the teacher's 1,764-token layout, then training switches to the compact layout. AdamW, learning rate 2×10⁻⁵ for inherited and 10⁻⁴ for new weights, batch 32, bf16, 3.6 hours on one NVIDIA A100-40GB.

Evaluation

LIBERO, 40 tasks × 10 trials, seed 195, success rate (%):

Model	Spatial	Object	Goal	Long	Average
Cosmos Policy 2B (5 steps)	97.0	100.0	99.0	97.0	98.2
SmolWAM (1 step)	99.0	100.0	94.0	89.0	95.5

The teacher's per-suite results come from an earlier evaluation pipeline; Cosmos Policy reports 98.5%. On 96 held-out contexts, SmolWAM's world head lowers the latent L1 error of the 16-step forecast by 24% relative to copying the current frame, and the optical-flow end-point error from 1.63 to 1.43 pixels. Per-episode records are in the repository's results/.

Intended use and limitations

SmolWAM is a research model for non-commercial study of efficient world-action models. It was trained and evaluated only on LIBERO in simulation (Franka Panda, two cameras) and is not expected to transfer to real robots, other embodiments or other camera setups without further training. Long-horizon tasks degrade the most under pruning (LIBERO-Long 89%). All costs were measured on one A100, and results come from a single training seed.

License

The weights are a derivative of nvidia/Cosmos-Policy-LIBERO-Predict2-2B and are released under the NVIDIA One-Way Noncommercial License (NSCLv1): non-commercial research and educational use only. The code is Apache-2.0.

Citation

@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}
}