Pitfalls specific to Fast-WAM

about 1 minutes

The video timestep for the first frame is zero. infer_action passes timestep_video = 0 when encoding the first frame. The grad-capable learner path must do the same; if it samples a timestep the way training_loss does, the cache distribution shifts and the feature KL explodes.
Attention masks are built per sequence length. _build_mot_attention_mask and the group_diagonal action mask depend on the number of video tokens and action tokens. Build them in the learner exactly as infer_action does; do not hard-code the 32-token case if you ablate chunk length.
Two gripper sign flips. The dataset flips the gripper sign to match other datasets, and evaluation flips it back. Compute all log-probabilities and all MSEs in the normalised model space, before any of that.
Compiled buffers. Never keep a reference to a tensor returned by the compiled prefill without cloning it; the code already clones for this reason.
The Optional-IDM mixing. When running video SFT on the Optional-IDM model, keep action_idm_prob so that both inference modes stay valid; setting it to zero silently turns the model into the plain first-frame variant.
Group degeneracy. At 98% success, uniform task sampling wastes most rollouts. Implement hard-task oversampling from the very first iteration.
Evaluation must be deterministic and at full step count. Report numbers only from $a = 0, K = 20$ so that they are comparable to the published Fast-WAM table.