Updating the world model: online video SFT and the feature KL

9.1 Why update the world model at all

The actor attends to the video expert's first-frame features. If those features represent the states that RL reaches poorly (because those states are off the demonstration distribution), the actor is acting on a poor representation. WAM-RL's argument is that fine-tuning the world model on video from successful rollouts improves exactly those representations, and its ablation is that without the update the policy cannot recover from early mistakes. For Fast-WAM the argument is weaker but still plausible: the actor does not read imagined futures, but it does read first-frame features, and those features were trained by the joint video-action objective and should improve if the video objective is continued on new data.

9.2 Turning rollouts into training samples

The pretraining data format (RobotVideoDataset) is a 33-frame video window at stride 1, with the 32 actions between those frames, the proprio at the first frame, and the text context. Successful episodes from the rollout buffer can be sliced into exactly this format:

1. Keep only episodes with $r^{\mathrm{task}} = 1$.
2. For each, choose up to 8 random start indices and cut a 33-frame window of rendered frames (the worker must therefore render every environment step of successful episodes, not just the two frames used by the reconstruction reward; the simplest approach is to render everything and keep only the successful episodes' full video).
3. Resize each camera to 224×224 and concatenate horizontally, using the same Resize transform as the dataset; verify pixel-for-pixel against _obs_to_model_input on one frame.
4. Take the 32 executed normalised actions between the frames as action, the first frame's normalised state as proprio, and the task's cached T5 context.
5. Produce the same dictionary training_loss expects: video [B, 3, 33, 224, 448], action [B, 32, 7], proprio [B, 1, 8], context, context_mask.

Then call model.training_loss(sample).

9.3 Video-only or video-plus-action

training_loss returns $\lambda_{\mathrm{video}} \cdot L_{\mathrm{video}} + \lambda_{\mathrm{action}} \cdot L_{\mathrm{action}}$. WAM-RL's online SFT is video only, corresponding to $\lambda _{\mathrm{action}} = 0$. Setting $\lambda _{\mathrm{action}} = 1$ additionally does behaviour cloning on the actions of successful rollouts, which is self-imitation learning (Oh et al., 2018) applied to a flow policy. It is a strong, stable regulariser for the actor and is worth an ablation of its own, but note that it mixes an imitation signal into the actor that is separate from the policy gradient, so if you are trying to isolate WAM-RL's world-model effect, run the $\lambda _{\mathrm{action}} = 0$ version first.

The Optional-IDM variant's training loss includes the action_idm_prob mixing (the action expert sometimes sees the noised future video), which you should keep so that both inference modes remain usable.

9.4 The feature KL regulariser

The feature the actor consumes is the per-layer K/V cache of the first frame. Denote by $z^{\theta}$ the cache produced by the current video expert and by $z^{\mathrm{old}}$ the cache produced by the frozen pretrained video expert on the same input. WAM-RL models each as a Gaussian with diagonal covariance and penalises the KL between them. Per layer $\ell$ and channel $c$:

$$\begin{aligned} \Sigma_\theta[\ell, c] &\leftarrow \text{EMA over training batches of } \mathrm{Var}\big( z^{\theta}[\ell, :, c] \big), \quad \text{momentum } 0.99 \\ \Sigma_{\mathrm{old}}[\ell, c] &= \mathrm{Var}\big( z^{\mathrm{old}}[\ell, :, c] \big) \quad \text{computed once on a batch of pretraining data and frozen} \end{aligned}$$ $$L_{\mathrm{KL}} = \mathrm{mean}_{\ell, c}\ \tfrac{1}{2} \cdot \big[ \log\big( \Sigma_{\mathrm{old}} / \Sigma_\theta \big) + \big( \Sigma_\theta + (z^{\theta} - z^{\mathrm{old}})^2 \big) / \Sigma_{\mathrm{old}} - 1 \big]$$

This is the closed-form KL between two diagonal Gaussians. In practice it behaves like a per-channel-weighted mean-squared distance between the new and old features, with an extra term that stops the new features from collapsing in variance. The total world-model objective is

L_WM = L_video + λ_KL · L_KL

WAM-RL does not report $\lambda_{\mathrm{KL}}$; start at 0.1 and adjust so that $\lambda_{\mathrm{KL}} \cdot L_{\mathrm{KL}}$ is about a tenth of L_video early in training.

To compute $z^{\mathrm{old}}$ you need a forward pass of the frozen video expert, which is another 5 billion parameters. With LoRA adapters on the video expert (Chapter 10) the frozen model is obtained by disabling the adapters, so there is no second copy. With full fine-tuning you would have to hold the reference on another GPU.

9.5 Scheduling the two updates

The actor update uses every macro-step from every rollout; the world-model update uses windows from successful episodes only. They use different batches and different parameters, so run them as two separate phases within each iteration rather than as one summed loss: first the actor phase, then the world-model phase. WAM-RL describes a "mixed" setting without giving a schedule; alternating once per iteration is the simplest faithful interpretation.