What WAM-RL did, and what changes for Fast-WAM

3.1 The WAM-RL recipe

WAM-RL (Chapter references are in Appendix C) takes the Genie-Envisioner-ACT model, a DiT video generator with an actor that reads the DiT's intermediate features, and trains it online in LIBERO and RLBench. Three mechanisms are combined.

A policy gradient on the actor. The actor is a flow-based model, so it is made stochastic with the Flow-SDE formulation from the πRL paper, and its log-probability is the sum of Gaussian transition log-densities over the denoising steps. The update is the standard $\nabla \log \pi \cdot A$ policy gradient with a clipped ratio.

A reconstruction reward. Instead of relying only on a sparse success signal, each action chunk gets a dense reward

r_t = sim( x̂_{t+1:t+H},  x_{t+1:t+H} )

where $\hat{x}$ is the future the world model imagined before acting and $x$ is the future that actually happened. The paper tried pixel MSE, optical-flow MSE, DINOv2 features, and V-JEPA2 features. Pixel MSE worked best even though it separated successes from failures least sharply, and the paper's explanation is that pixel MSE is the same objective the world model was trained with, so the reward is "speaking the model's language".

Online video self-supervised fine-tuning with a KL feature regulariser. The world model is fine-tuned on the video of successful rollouts using its ordinary video-prediction loss. Because the actor depends on the world model's features, the paper adds a KL term that keeps the updated model's latent features close to the frozen pretrained model's, modelling each as a Gaussian with a diagonal covariance estimated by an exponential moving average:

$$L_{\mathrm{WM}} = L_{\mathrm{video}} + \lambda_{\mathrm{KL}} \cdot \mathbb{E}_t \big[\, \mathrm{KL}\big( \mathcal{N}(z_t, \Sigma_\theta) \,\|\, \mathcal{N}(z_t^{\mathrm{old}}, \Sigma_{\mathrm{old}}) \big) \,\big]$$

The paper's headline numbers are on a weak base model: LIBERO-Object goes from 68% (pretrained) to 78% (actor-only RL) to 82% (full WAM-RL). The ablation it emphasises is that without the world-model update, the policy cannot recover from an initial mistake, while with it, the world model starts imagining recovery behaviours such as re-grasping.

3.2 What transfers directly

Almost all of the actor side transfers without modification. Fast-WAM's action head is a flow-matching model with exactly the structure πRL's Flow-SDE assumes. The MoT attention cache is precisely the "latent feature the actor depends on" that WAM-RL's KL term protects. The existing training_loss is the video-SFT loss. LIBERO is the benchmark, and the evaluation harness already exists.

3.3 What does not transfer: the missing imagination

In WAM-RL's base model, the actor reads the imagined future. If the imagination says the gripper will close on the cup and the real video shows it missing, then either the actor failed to follow its own plan or the plan was wrong; either way the mismatch is informative about the action.

In Fast-WAM's first-frame mode there is no imagination at inference. The actor and the video expert are both conditioned on the same first-frame features, but the actor never sees a generated future, and the generated future (if you bother to compute one) is not conditioned on the chunk the actor produced (the video expert is action_conditioned: false in every released config). So "imagined versus executed" in Fast-WAM measures something different: it measures whether the world evolved the way the world model would have predicted from this state, regardless of what the actor did. That is closer to a progress or plausibility signal than to a plan-adherence signal.

This means three things for the design.

First, to compute a reconstruction reward at all, the imagination has to be generated somewhere. It should not be generated in the rollout worker (that would throw away the entire speed advantage of Fast-WAM); it should be generated in the learner, after the fact, where it costs a few hundred milliseconds per chunk and is off the critical path.

Second, the reconstruction reward's value is an open question, not a given. Chapter 14 sets up the experiment that answers it: success-only reward versus success plus reconstruction, with the imagination coming either from the live (fine-tuned) world model or from a frozen copy.

Third, the world-model update (video SFT plus KL) might help the actor even without a reconstruction reward, simply because the actor reads the first-frame features and those features get better at representing states reached during RL. That is a second, independent thing to test.

3.4 The other thing to keep in mind: a nearly saturated benchmark

WAM-RL improved a 68% model. Fast-WAM starts at about 98%. At that level the majority of rollouts succeed, group-based advantages (Chapter 7) are often zero because every rollout in a group had the same outcome, and the measurable headroom is a couple of points, mostly in LIBERO-10. Chapter 7 describes oversampling of failing tasks, and Chapter 14 suggests harder evaluation settings (unseen instructions, RoboTwin, where the base model is weaker) so that you have room to see an effect.

---