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

1. Quick overview of the paper

Difficulty rating: ★★★★☆. Need to understand WAM/VLA, video diffusion transformer, flow matching, action chunk diffusion, attention mask to prevent future leakage, and simulation/real robot evaluation.

Keywords: World Action Model, video co-training, test-time future imagination, flow matching, Mixture-of-Transformer, Wan2.2-5B, LIBERO, RoboTwin, towel folding.

Figure 1. Three types of WAM paradigms: (A) future video/action joint denoising; (B) first generate future video, and then use inverse dynamics/action predictor; (C) Fast-WAM retains video co-training during training, and directly generates actions from the current latent world representation during inference.

2. Motivation

2.1 The attraction and cost of WAM

Standard VLA directly maps visual observations and language to actions, mainly inheriting the semantic priors of web-scale visual language pre-training. But robot control also requires an understanding of how the physical world evolves under action. The appeal of WAM is to put future visual prediction and action modeling into the same framework, allowing the policy to explicitly contact the task-relevant temporal structure.

The problem is that most WAMs iterate over denoise future videos during inference and then generate actions based on this imagined future. Iterative sampling of video diffusion is very expensive, and every few hundred milliseconds in a real robot closed loop may slow down and blunt the strategy.

2.2 The real question the paper asks

This question used to be difficult to answer because many WAMs tied two factors together: the same model both learned video predictions during training and generated future videos during testing. The contribution of Fast-WAM is to separate these two things and use controlled variables to experimentally determine which is more important.

3. Related work context

4. Detailed explanation of method

4.1 Formalization of the problem

Suppose the current observation is $o$, the task language is $l$, and the action chunk is $a_{1: H}$. Standard visuomotor policy learning:

A typical imagine-then-execute WAM introduces future vision $v_{1: T}$, written as:

Fast-WAM is changed to a direct policy interface:

4.2 Architecture: Wan2.2 video DiT + action expert DiT

Figure 2a. Fast-WAM architecture: Wan2.2-5B video DiT serves as world modeling backbone, T5 coding language, video VAE encodes multi-camera images latent; action expert DiT generates action chunks.

Model input tokens are divided into three groups:

• clean first-frame latent tokens: The clean latent of the current observation is the shared visual anchor.
• noisy future video latent tokens: Only exists during training and is used for future video flow matching.
• action tokens: Handled by action experts for action chunk generation.

All tokens receive language embedding via cross-attention. A Mixture-of-Transformer structure is used between the video and action branches, and a structured attention mask is used to control the information flow.

4.3 Attention mask: share context and prevent future leakage

Figure 2b. Training/inference mask: action tokens can see the current clean first-frame tokens, but not future noisy video tokens; first-frame tokens do not look at other tokens to prevent future information from contaminating the current anchor point.

During training, future video tokens have bidirectional attention within the video branch and can access first-frame tokens; action tokens have bidirectional attention within the action branch and can also access first-frame tokens. The most critical constraint is that action tokens cannot attend to future video tokens. In this way, video modeling and action prediction both rely on the same current visual context, but the action does not peek into the ground-truth future.

During inference, the future video branch is completely deleted: only the first-frame latent tokens are retained, the video backbone generates latent world features in a single forward direction, and then performs action denoising for the action expert.

4.4 Training goal: action loss + video co-training loss

Fast-WAM uses the same flow matching form for both actions and videos. For any target variable $y$, which can be action chunk $a_{1: H}$ or future video latents $z_{1: T}$, sample noise $\epsilon\sim\mathcal{N}(0, I)$ and time step $t\in(0, 1)$:

4.5 Three controlled variants

5. Experiment

5.1 Implementation details

• Backbone: Wan2.2-5B, including video DiT, T5 text encoder and video VAE.
• Action expert: Isomorphic to the video branch but the hidden dimension is reduced to $d_a=1024$; the action expert is about 1B, and the total model has about 6B parameters.
• Action horizon: $h=32$.
• Video chunk: The video frame time is downsampled to $4\times$, each chunk is 9 frames; multi-camera images are first assembled into one picture and then sent to VAE.
• Flow matching: Training/inference uses logit-normal $t$ schedule; inference action denoising 10 steps, CFG scale 1.0.
• Optimization: AdamW, learning rate $1\times10^{-4}$, weight decay 0.01, cosine annealing, mixed precision, gradient clipping 1.0.
• Latency measurement: Single NVIDIA RTX 5090D V2 32GB.

5.2 Benchmark settings

Figure 3. Realistic towel folding task. The author emphasizes that completion time is equally important as success rate, because slow trial-and-error success does not represent a good strategy.

5.3 RoboTwin 2.0 Main Results

Fast-WAM does not use embodied pretraining, but reaches 91.8%, which is significantly higher than Motus from WAN2.2 (77.3) and LingBot-VA from WAN2.2 (80.6), which also have no embodied pretraining, and is close to LingBot-VA with embodied pretraining (92.2). Appendix Table 3 gives RoboTwin clean/rand details for each task; overall, Fast-WAM competes with the strongest baseline on many tasks, but the average value of no-video-co-train is significantly lower. Appendix Table 3.

5.4 LIBERO main results

The average Fast-WAM on LIBERO is 97.6%, exceeding $\pi_{0.5}$'s 96.9, and close to Motus/LingBot-VA. It does not have embodied pretraining, which is a data efficiency point emphasized by the author.

5.5 Control variables: future imagination vs video co-training

This is the key evidence of the paper: the gap between Fast-WAM and the two imagine-then-execute variants is small; but the drop is more obvious after removing video co-training. RoboTwin dropped from 91.8 to 83.8, LIBERO dropped from 97.6 to 93.5, and LIBERO Spatial/Long dropped particularly significantly. Based on this, the authors believe that the main value of WAM is more likely to come from the video prediction goal during training, rather than actually generating future videos during inference.

5.6 Real Towel Folding: Performance and Latency

Figure 4. Real towel folding results: The left picture is closer to the upper left, the better; the right picture shows the inference delay. Fast-WAM is 190 ms, Fast-WAM-IDM is 810 ms, and Fast-WAM-Joint is 580 ms. The figure also shows that no-video-co-train has a success rate of about 10% and has the longest completion time.

In real tasks, pre-training $\pi_{0.5}$ is still the strongest method, with the highest success rate and the shortest completion time. The performance between Fast-WAM families is similar: Fast-WAM-IDM has the highest success rate, and Fast-WAM completion time is better. More importantly, all Fast-WAM variants with video co-training are significantly stronger than $\pi_{0.5}$ without pretraining, while no-video-co-train collapses to 10% success. This again supports video co-training as the main reason.

In terms of latency, Fast-WAM is 190 ms, which is of the same order as no-video-co-train's 190 ms; Fast-WAM-Joint is 580 ms, and Fast-WAM-IDM is 810 ms. Fast-WAM therefore becomes a better compromise for deployment: retain most of WAM performance but avoid explicit future video sampling overhead.

6. Reproducible auditing

6.1 Components required to reproduce

6.2 Minimum recurrence route

1. First implement Fast-WAM w.o. video co-train on LIBERO: only use the current first-frame latent + language + action DiT to run through action flow matching.
2. Add future video latent branch and $\mathcal{L}_{vid}$, but keep action tokens and cannot see future video tokens to verify whether Fast-WAM is improved.
3. Implement Fast-WAM-Joint: Release the mutual attention of action/video tokens and test whether it is close to Fast-WAM.
4. Implement Fast-WAM-IDM: first generate future video representation, then condition action; note the use of $p=0.5$ ground-truth video token noise augmentation.
5. Reproduce the LIBERO table and then migrate to RoboTwin multi-task training; finally consider real towel folding.
6. Delay evaluation must be done separately: Fast-WAM has no future branch, but still has 10 steps of action denoising; the additional overhead of IDM/Joint comes from future video generation/joint sampling.

6.3 Recurrence risk points

7. Analysis, Limitations and Boundaries

7.1 The most valuable part of this paper

The most valuable thing is the problem solving itself. Many WAM papers default to "generating future videos" as a necessary step, but Fast-WAM splits it into two factors: training goal and inference mechanism, and verifies it with variants in the same framework. This experimental design is more instructive than simply proposing a new model: it tells us that the value of the world model may be mainly reflected in the training representation, rather than the explicit imagination at deployment time.

The second value point is deployment orientation. The gap of 190 ms vs 580/810 ms is realistic for real robots. Fast-WAM brings WAM close to the inference interface of VLA while retaining the WAM training signal, which is a very practical compromise.

7.2 Why the results hold up

• Control variables are clear: The backbone and training recipes of Fast-WAM, Joint, IDM, and no-co-train should be aligned as much as possible.
• Consistent across benchmarks: RoboTwin, LIBERO, and real towel folding all show that the video co-training drop is larger than the future imagination variant gap.
• The efficiency indicators are clear: Real tasks report not only success, but also completion time and latency.
• Appendix fine-grained support: The RoboTwin per-task table shows that the improvement does not only come from a few tasks, but the average drop in no-co-train is obvious in clean/rand.

7.3 Limitations and points that need to be questioned

7.4 Questions that can be asked in group meetings

1. Will test-time future imagination become important again if the task requires explicit intermediate subgoals, such as complex assembly or navigation?
2. What exactly does the representation learned by Fast-WAM's video co-training encode? Can probing/attention/feature prediction be used to prove that it captures physical dynamics?
3. Action tokens cannot see future video tokens, but the video branch and action branch share the first-frame anchor; is this mask optimal, or is there a thinner causal mask?
4. Wan2.2 is a general-purpose video generation backbone. After switching to the robot video pre-training backbone, will the gap between Fast-WAM and Joint/IDM become larger or smaller?
5. Fast-WAM has better completion time than IDM in real tasks, but its success rate is not necessarily the highest. How should the actual system choose between delay, success rate and action stability?

Attachment: This report covers inspections