Rewards: task success and reconstruction

6.1 Task success

LIBERO returns a done flag when the task's success predicate holds. The task reward is

r^task_t = 1   if success first becomes true during the 10 steps executed at macro-step t
         = 0   otherwise

This is the entire reward in SimpleVLA-RL and RLinf, and it alone produces substantial LIBERO improvements in those papers. It must always be present; the reconstruction term is a shaping signal on top of it, not a replacement.

6.2 The reconstruction reward, adapted to Fast-WAM

WAM-RL's reconstruction reward compares the imagined future with the real future. In Fast-WAM first-frame mode the rollout worker never imagines anything, so the imagination is computed later, in the learner, as follows.

Step 1 — Record the real future. During the rollout, after each macro-step's 10 executed actions, the worker keeps the rendered 224×448 concatenated image at the offsets that line up with the video model's frame rate. The model generates one video frame every 4 actions, so within a 10-step replan window the matching real frames are at offsets +4 and +8. (If you want all 8 imagined frames to have a real counterpart, you have to execute more of the chunk per replan; Chapter 14 lists that as an ablation.) The evaluation code has a skip_get_obs_within_replan option for RoboTwin that disables rendering inside the window; for RL you need those renders.

Step 2 — Imagine the future in the learner. For each stored macro-step, run the Optional-IDM model's imagination path (infer_action with action_infer_mode="idm" returns the video latents; infer_joint additionally decodes them to pixels) from the stored state. Decode the latents with the VAE to get $H$ imagined frames at 224×448.

Step 3 — Score. Following WAM-RL's finding that pixel MSE works best,

$$r^{\mathrm{rec}}_t = -\ \text{mean over matched frames and pixels of } (\hat{x} - x)^2$$

with images in $[0, 1]$. Because the scale of this number is arbitrary, normalise it with a running mean and standard deviation across all macro-steps seen so far (WAM-RL does not say how it normalises; this is the standard choice). The group-normalised advantage in Chapter 7 will also absorb scale, but a running normaliser keeps the mixing weight below interpretable.

Step 4 — Combine.

r_t = r^task_t + β · r̃^rec_t

Choose $\beta$ so that the reconstruction term contributes roughly 0.05–0.2 per macro-step after normalisation. Success must still dominate the return; the reconstruction term exists to break ties between trajectories that both failed and to give a gradient before the first success.

6.3 Which world model produces the imagination

There are two defensible choices, and they are the heart of the experiment in Chapter 14.

• The live world model. The same video expert that is being fine-tuned in Chapter 9. This is what WAM-RL does. The reward co-evolves with the model: as the world model learns to imagine the recovery behaviours that succeed, trajectories that perform those recoveries are rewarded for matching.
• A frozen copy of the pretrained world model. The reward is then stationary, which makes RL easier to reason about, but it can never learn to reward behaviours it has not seen.

If you use LoRA on the video expert (Chapter 10), the frozen copy is free: disable the adapters for the reward forward pass.

6.4 Cost

The imagination is a 5-billion-parameter DiT running 20 denoising steps over 9 latent frames, plus a VAE decode. On the hardware in the README this is on the order of 0.2–0.5 seconds per macro-step, so about 10–20 seconds of learner-side GPU time per 40-macro-step episode. For 256 episodes per iteration that is roughly an hour of single-GPU compute, which parallelises trivially across the learner's GPUs. If that is too slow, drop to 10 video denoising steps; pixel MSE is tolerant of blur.

6.5 Return

The return at macro-step $t$ is the discounted sum of future rewards,

$$R_t = \sum_{t' \ge t} \gamma^{\,t' - t} r_{t'}$$

and it is the quantity the advantage estimator in the next chapter works with.