Fast-WAM: Do World Action Models Need Test-time Future Imagination?

TL;DR

Fast-WAM (Tsinghua IIIS + AIR) is a controlled ablation study that asks whether World Action Models (WAMs) actually need to imagine future video at inference time or whether the benefit comes from the video-prediction training objective alone. It builds a Mixture-of-Transformers WAM on Wan2.2-5B (video DiT + 1B action expert, ~6B total) that keeps video co-training during training but drops future-video denoising at inference, and mirrors two representative imagine-then-execute paradigms (joint denoising à la Motus/UVA; video-then-action à la LingBot-VA/Vidar) as internal baselines under a shared framework. Across LIBERO and RoboTwin 2.0 and a real-world towel-folding task, Fast-WAM matches or nearly matches the imagine-then-execute variants (97.6 / 91.8 avg without embodied pretraining) while removing video co-training causes a much larger drop (93.5 / 83.8) — and Fast-WAM runs at 190 ms latency, >4× faster than the video-then-action variant. The paper’s headline is a design-space claim: the load-bearing role of video prediction in WAMs is representation-shaping during training, not future generation at test time.

Key claims

Fast-WAM matches state-of-the-art WAMs without embodied pretraining: 97.6% avg on LIBERO (vs LingBot-VA 98.5, Motus 97.7, both with pretraining) and 91.8% on RoboTwin 2.0 (vs LingBot-VA 92.2, Motus 87.8, both with pretraining) [Tables 1, 2].
Under matched training (same backbone, tokenizer, and recipe), removing video co-training is the dominant failure mode: Fast-WAM w/o video co-train drops to 83.8% RoboTwin (−8 pp) and 93.5% LIBERO (−4.1 pp), with the largest degradations on Spatial (89.2%) and Long (90.0%) subsets [Tables 1, 2].
Under the same matched training, imagine-then-execute variants (Fast-WAM-Joint, Fast-WAM-IDM) are only marginally different from Fast-WAM: LIBERO 98.5 / 98.0 vs 97.6; RoboTwin 90.6 / 91.3 vs 91.8 — the video-co-training gap is much larger than the imagine-vs-no-imagine gap [Tables 1, 2].
Real-world towel folding: all Fast-WAM variants with video co-training substantially outperform π₀ without embodied pretraining; Fast-WAM-IDM has highest success rate, Fast-WAM has best completion time. Removing video co-training collapses to 10% success rate and slowest completion — the largest single-factor gap in the study [Fig. 4].
Inference latency: Fast-WAM 190 ms vs Fast-WAM-Joint (joint denoising with attention across video+action) and Fast-WAM-IDM 810 ms — >4× speedup from dropping test-time future imagination, on a single RTX 5090D V2 [§4.3.3].
Architecture is a shared-attention Mixture-of-Transformers on Wan2.2-5B: clean first-frame latent tokens act as a shared visual anchor consumed by both branches; during training, future noisy video tokens attend within the video branch and to clean anchors, action tokens attend within the action branch and to clean anchors but never to future video tokens (structured mask prevents future-info leakage into actions); at inference the future video branch is removed entirely [§3.2, Fig. 2].
Training uses joint flow-matching with $\lambda=1$ balance between action loss and video-latent loss on Wan2.2 VAE features; single-step encoding at inference plus 10-step denoising for actions with CFG=1.0 [Eqs. 5–9, §4.1].
The controlled variants are the paper’s methodological contribution as much as Fast-WAM itself: they isolate imagine-then-execute paradigms (Joint / IDM) and no-video-co-training under a shared framework, making the “which factor matters” comparison meaningful rather than confounded by backbone or recipe differences [§3.3, Fig. 1].

Method

Fast-WAM is a Mixture-of-Transformers WAM: a video DiT (Wan2.2-5B backbone, frozen T5 text encoder and video VAE reused) plus a shape-matched 1B action expert with reduced hidden dimension. Every block shares attention across the two branches. Input tokens are grouped into (a) clean latent tokens of the first observation frame — the shared visual anchor, (b) noisy latent tokens of future frames — used only during training, and (c) action tokens for the action chunk. A structured attention mask lets action tokens attend bidirectionally within themselves and to the clean anchor, but not to the future video tokens, so the video branch supervises representations without leaking future info into the action head. At inference the future video branch is dropped and the video DiT runs one forward pass over the clean anchor to produce latent world features that the action expert denoises actions against with 10 flow-matching steps.

Training is a weighted sum of two flow-matching losses on the same interpolation schedule ( $x^\tau_1 = (1-\tau)\epsilon + \tau x_1$ ): one supervises the action expert on action chunks, the other supervises the video branch on future-frame VAE latents. Both share the logit-normal noise schedule from Wan2.2.

The controlled comparison is instantiated by two variants under the same framework: Fast-WAM-Joint allows attention between action tokens and (noisy) future video tokens, so joint denoising of video+action drives inference; Fast-WAM-IDM first fully denoises future video from the clean anchor, then predicts actions conditioned on the generated future (with noise-augmentation with $p=0.3$ following LingBot-VA to reduce train-test distribution mismatch). A third variant simply drops the video co-training loss — same architecture, same inference — to isolate the role of the objective.

Simulation evaluation runs 20k steps on LIBERO (500 demos × 4 suites, 2000 trials × 40 tasks) and 30k steps on RoboTwin 2.0 (2,500 clean + 25,000 randomized demos across 50+ bimanual tasks, 100 trials/task under clean and randomized). Real-world is a Galaxea R1 Lite towel-folding task, 60 hours of teleop demos, 30k training steps.

Results

LIBERO 4-suite average [Table 2]: Fast-WAM 97.6 (Spatial 98.2 / Object 100.0 / Goal 97.0 / Long 95.2), no embodied pretraining. Fast-WAM-Joint 98.5, Fast-WAM-IDM 98.0, Fast-WAM w/o video co-train 93.5 (Spatial drops to 89.2, Long 90.0). LingBot-VA (with pretraining) 98.5, Motus (with pretraining) 97.7, π₀.₅ (with pretraining) 96.9.
RoboTwin 2.0 [Table 1]: Fast-WAM 91.8 (Clean 91.88 / Rand 91.78), no pretraining. Motus without pretraining 77.3, LingBot-VA without pretraining 80.6; Motus with pretraining 87.8, LingBot-VA with pretraining 92.2. Fast-WAM w/o video co-train 83.8 (−8 pp).
Real-world towel folding [Fig. 4]: pretrained π₀ leads on both success rate and completion time. Among Fast-WAM variants (no pretraining), Fast-WAM-IDM tops success rate, Fast-WAM tops completion time. Removing video co-training collapses to 10% success rate and slowest completion — largest single-factor gap in the study. All Fast-WAM variants with video co-training beat π₀ without pretraining.
Inference latency on a single RTX 5090D V2 [§4.3.3]: Fast-WAM 190 ms, Fast-WAM-Joint substantially slower, Fast-WAM-IDM 810 ms (>4× Fast-WAM).

Why it’s interesting

Fast-WAM is the sharpest controlled disentanglement filed so far of “does the WAM benefit come from the video-prediction training objective or from test-time future imagination?” — and its answer is the training objective. That directly complements the load-bearing empirical finding in Dyna-2: A 1-Million-Hour Scaling Law for World-Action Models — Dyna-2’s action-only supervision loses on 39/39 tasks at every scale, and joint denoising + video co-training on unlabeled human video is the specific lever that makes cross-embodiment transfer scale — but arrives at it via a smaller-scale, ablation-first study on a shared 6B backbone rather than via a >1M-hour scaling law. The Slack pointer “Fast-WAM (what Dyna-2 is essentially)” captures the resemblance: both are MoT WAMs on video-diffusion backbones with clean-frame anchoring, both keep video co-training as the load-bearing objective, and both find joint denoising outperforms action-only supervision. Fast-WAM has the sharper isolation but ~1000× less pretraining data.

On the VLA Models recipe-lever board, Fast-WAM adds an architectural answer distinct from the “WAM-as-policy-internal-substrate” recipe of Causal World Modeling for Robot Control (LingBot-VA) (unified causal AR flow-matching) and the “frozen 3D-trace WM + action expert” recipe of μ₀: A Scalable 3D Interaction-Trace World Model: keep the WAM structure of MoT + video co-training but explicitly drop the imagined-future branch at inference. It thereby stakes out a middle position — video-modeling representation without video-generation latency — that has been implicit in FLUX 3 x mimic: The Next Generation of Video-Action Models‘s FLUX-mimic (frozen video WM features → small action head) and ImageWAM: Do World Action Models Really Need Video Generation, or Just Image Editing?‘s ImageWAM (KV-cache injection instead of decode) but never quite framed as an explicit imagine-vs-no-imagine ablation on matched hardware.

The paper’s methodological contribution — a shared-framework instantiation of the two dominant imagine-then-execute paradigms plus a no-video-co-training control — is the clean counterpart to Dyna-2’s scaling ablation, and together they provide the strongest evidence to date for the “video co-training is the lever, future imagination at test time is optional” position on the WAM design space.