Fast-WAM as it exists today

Before changing a model it helps to know precisely what it is. This chapter describes the Fast-WAM architecture and the parts of the code you will be touching. All file paths are relative to the FastWAM repository root.

2.1 The two experts and the MoT

Fast-WAM is built on the Wan2.2 TI2V-5B video diffusion model. Internally it is a mixture of transformers (MoT): two transformer stacks that run side by side and attend to each other's tokens.

The video expert is the Wan2.2 DiT itself, about 5 billion parameters, 30 layers of hidden size 3072. It operates on video latents produced by Wan's VAE (the VAE compresses a 224×448 RGB frame into a much smaller latent grid and is frozen throughout). This is the world model.

The action expert (ActionDiT, in src/fastwam/models/wan22/action_dit.py) is a smaller transformer, 30 layers of hidden size 1024 with a 4096-wide feed-forward block, that was initialised by interpolating the Wan DiT's weights down to the smaller width. It operates on a sequence of action tokens, one per timestep in the action chunk. This is the actor.

The MoT (mot.py) runs both stacks layer by layer. At each layer, action tokens can attend to video tokens. The attention mask that governs this is built in _build_mot_attention_mask in fastwam.py; the important rule is that action tokens attend only to the video tokens of the first frame. That is the whole reason Fast-WAM can skip imagination: the actor only ever needs the encoding of the current observation.

Two small extras: a proprio_encoder (a single linear layer) maps the robot's 8-dimensional state into the text-embedding space and appends it to the text context, and the text itself is encoded by a frozen T5 encoder (or loaded from a precomputed cache).

2.2 Inputs, outputs, and normalisation

For LIBERO (configs/data/libero_2cam.yaml) the model sees:

• An image of shape [3, 224, 448]: the third-person "agentview" camera and the wrist camera, each resized to 224×224 and concatenated horizontally.
• A proprioceptive state of 8 numbers (end-effector pose as 6 numbers plus two gripper values), normalised with min/max statistics stored in a dataset_stats.json file.
• A T5 embedding of the task prompt, shape [128, 4096] with a padding mask.

It produces an action chunk of shape [32, 7]: 32 future timesteps, each a 6-dimensional end-effector delta plus a gripper command, all normalised to roughly $[-1, 1]$. During training the 32 actions line up with 9 video frames (one video frame every 4 actions, action_video_freq_ratio: 4, num_frames: 33).

At evaluation time (experiments/libero/eval_libero_single.py) the model is queried every 10 environment steps (replan_steps: 10), the first 10 of the 32 predicted actions are executed, and the rest are discarded. The gripper value is denormalised, sign-flipped back to the simulator's convention, and binarised. Remember these deterministic post-processing steps; Chapter 5 explains why they do not interfere with the RL maths.

2.3 How actions are sampled: the flow-matching ODE

The action expert is a flow-matching model. Rather than predicting the action directly, it predicts a velocity field that transports Gaussian noise to the action. The scheduler (schedulers/scheduler_continuous.py) defines the forward noising process as

$$x_\sigma = (1 - \sigma) \cdot x_0 + \sigma \cdot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I), \quad \sigma \in [0, 1]$$

where $x_0$ is the clean action chunk and $\sigma = 1$ is pure noise. The network is trained to predict the velocity $v = \varepsilon - x_0$ (this is training_target in the scheduler), and the training timesteps are drawn from a "shifted" distribution $\sigma = \varphi(u) = \mathrm{shift} \cdot u \,/\, (1 + (\mathrm{shift}-1) \cdot u)$. With the default action shift of 1.0, $\varphi$ is the identity and the schedule is uniform.

To sample, infer_action starts from $x \sim \mathcal{N}(0, I)$ at $\sigma = 1$ and takes $K = 20$ Euler steps of the ordinary differential equation

$$x_{k+1} = x_k + v_\varphi(x_k, \sigma_k) \cdot \Delta_k, \qquad \Delta_k = \sigma_{k+1} - \sigma_k \quad \text{(negative)}$$

until $\sigma = 0$. This is a deterministic map from the initial noise to the final action. Given the same noise, it always produces the same chunk. That matters enormously for RL and is the subject of Chapter 5.

Before the denoising loop, infer_action encodes the first frame once through the video expert and caches the per-layer key/value tensors (mot.prefill_video_cache_tensor). Each of the 20 denoising steps then runs only the action expert, attending into that cache (_denoise_action_with_video_cache). With compile_action_infer=true both pieces are compiled with torch.compile(mode="reduce-overhead"), which uses CUDA graphs. This is the path that gives the 110–210 ms inference time. It is also a path you cannot backpropagate through, which Chapter 10 comes back to.

2.4 The three model variants

configs/model/ contains three variants that share the architecture but differ in training objective:

• fastwam ("uncond"): trained jointly on video and action denoising, acts from the first frame only. This is the headline Fast-WAM.
• fastwam_idm: the action expert reads the imagined future video (inverse-dynamics style). This is the "WAM with imagination" configuration, closest to what WAM-RL trained.
• fastwam_optional_idm: one set of weights trained with a 50/50 mix (action_idm_prob: 0.5) of both behaviours, so that at inference you can pick action_infer_mode=idm or action_infer_mode=first_frame without retraining. The released checkpoint scores 98.55% (IDM) and 97.75% (first-frame) on LIBERO.

This document recommends building on the Optional-IDM checkpoint. The reason is practical: it is the only variant in which a single model can both act cheaply without imagination (the mode you will roll out in the simulator) and generate the future conditioned on the same features (which you need for the reconstruction reward and for video fine-tuning). Using it lets you ask the WAM-RL question cleanly: does improving the world model help an actor that does not read the world model's imagination at test time?

2.5 What the existing trainer already does

src/fastwam/trainer.py implements ordinary supervised training. Its training_loss (in fastwam.py) takes a batch with a 33-frame video, a 32-step action chunk, the proprio state, and the text context; it noises the video latents and the actions independently, runs both experts jointly, and returns a weighted sum of the video loss and the action loss. _apply_dit_only_train_mode freezes everything except the MoT (both experts) and the proprio encoder. Checkpoints are saved with save_checkpoint and include the optimizer state.

You will reuse training_loss almost unchanged for the world-model update in Chapter 9, and you will reuse the data-loading and normalisation code to turn simulator rollouts into training samples in exactly the format the model expects.