Four ways to condition a visuo-tactile world model

Same 30-episode training split, same 100 epochs, same optimizer, same seed — four different conditioning architectures, evaluated on ten identical held-out windows.

Diffusion-Forcing Mixture-of-Transformers over Wan2.2 latents (48 ch, 4 latent frames, 16×16). Every number below comes from last.ckpt at step 117 199 — the split warns that val aliases test, so selecting best-epoch=* would be selecting on the very episodes reported here.

The task

Each window is 16 fps6 frames → 4 latent frames. The first 2 latent frames are given as clean context; the model generates the other 2 from the action alone. Decoded through the Wan2.2 VAE that is 13 pixel frames, of which frames 0–4 come from the context and frames 5–12 are generated. The strips at the bottom outline the generated frames in red.

All four models are the same trunk: 12 blocks, hidden 768, patch (1,2,2), temporal-causal joint attention, cross-first block order, per-frame independent noise levels per modality. They differ only in how the action reaches the tokens.

The four configurations

action_ctx: false
pixflow:      {enable: true, video: true, tactile: true, share_tactile: true}
vector_film:  {enable: true, hidden: 256}
action_image: {enable: false}
action_motion:{enable: false}
video_motion: {enable: false}
pred_flow_loss:{enable: false}

action_ctx: false
pixflow:      {enable: false}      # the only difference
vector_film:  {enable: true, hidden: 256}
action_image: {enable: false}
action_motion:{enable: false}
video_motion: {enable: false}
pred_flow_loss:{enable: false}

action_ctx: true                     # (default)
action_image:  {enable: true, sigma_cells: 0.8}
action_motion: {enable: true, checkpoint: amx_gel5_s2/best.pt,
                freeze: true, inject_to_dit: true}
pred_flow_loss:{enable: true, checkpoint: latent_flow_probe/best.pt,
                weight: 0.05}

action_ctx: true
action_image:  {enable: true, sigma_cells: 0.8}
action_motion: {enable: true, checkpoint: amx_gel5_s2/best.pt, freeze: true}
video_motion:  {enable: true, checkpoint: vmx_E/best.pt,
                freeze: true, inject_to_dit: true, flow_to_action_image: true}
pred_flow_loss:{enable: true, weight: 0.05}

What conditions what

The raw action vector is a real gap the middle two columns share: with action_ctx: false and no rasters, an earlier FiLM-only variant left action[18:36] — the 18-D tactile SE(3) step — reaching the network nowhere at all. vector_film restores it, standardized with the model’s own action statistics (the raw 9-D is metres ~1e-2 next to Rot6D components near 1, which a zero-init MLP head cannot cope with).

Training result — validation loss

Mean ± sd over the last 10 epochs. Total val_loss is not comparable across the families: the two baselines carry ≈0.15 of frozen-probe loss inside it. The two components above are.

Held-out result — decoded pixel PSNR

Paired per-window differences

Because every config saw the same ten windows, the difference can be taken window by window instead of comparing two noisy means.

Reading

• The frozen pixel flow does not earn its place. Against its own control it loses 0.51 dB on tactile-left, with the control ahead on 9 of the 10 windows. On view it is 0.07 dB behind and the two split 5–5. On tactile-right its +0.28 dB mean is one window (+2.98 dB on 005:0); the other nine are inside ±0.3 dB and the median is −0.02. Validation loss agrees: +0.00035 ± 0.00031 on video over the last 50 epochs, with the control ahead in 41 of those 50. Running the control is what makes this readable — on its own pixflow would have looked like a clear win over the baselines.
• The new family really is better than the baselines, by about 2 dB on tactile and 0.65 dB on view, winning 10/10 and 8/10 windows. But that gain belongs to the configuration change as a whole — dropping cross-attention, the analytic ActionImage, the actionflow injection and the probe loss, while restoring the raw action vector through FiLM. It is not attributable to the pixel flow.
• The view stream is nearly saturated across all four (22.6–24.0 dB, and gelforce_vmx matches the new pair). Two clean context frames plus the camera-projected action apparently determine the next two latent frames about as well as any of these mechanisms can.
• Tactile PSNR is dominated by whether contact happens, not by conditioning quality — 38 dB on contact-heavy windows against 52 dB on quiet ones, a 14 dB range inside a single config. Ranking configs on the tactile mean alone would mostly be ranking which windows landed in the sample.
• Both new runs bottom out early (epochs 57 and 63) and drift up slightly afterwards, where gelforce’s minimum is at epoch 95. The lighter conditioning converges faster and then mildly overfits.

Ten held-out windows

Ground truth on top, then the four configs. Frames 0–4 are the clean context (identical for every row); frames 5–12 outlined in red are generated. Per-row PSNR is over the whole 13-frame clip, computed on the raw decoded tensors — the images here are h264, which costs a little fidelity equally in every row.

Each window also gets an error map: |prediction − ground truth| on one shared scale across all four configs. That panel is the one worth reading for the tactile streams — a gel image is nearly uniform, so the side-by-side strip cannot show where the configs actually differ.

005:0 — sample_000

005:34 — sample_001

005:68 — sample_002

005:102 — sample_003

005:136 — sample_004

006:0 — sample_005

006:150 — sample_006

006:300 — sample_007

006:450 — sample_008

006:600 — sample_009

Reproducing

# cache the frozen Stage-1 flows (32 episodes)
python data_preprocessing/build_stage1_pixel_flows.py

# gate the architecture before spending a GPU-day
python -m vm_diffusion.scripts.verify_pixflow

# train
sbatch vm_diffusion/scripts/train_mot.sh mot_df_vt_pixflow \
       vm_diffusion/config/train_mot_df_vt_pixflow.yaml
sbatch vm_diffusion/scripts/train_mot.sh mot_df_vt_vecfilm \
       vm_diffusion/config/train_mot_df_vt_vecfilm.yaml

# evaluate all four on the same ten windows and decode to pixels
sbatch vm_diffusion/run_eval_4way_test10.sh

Held-out episodes motherboard_0510_episode_005 and 006, split episode_holdout_0510_gel (30 train / 2 held out; the 4 pusht_0618 episodes are excluded — they have no force export and therefore no gel-frame action map). Rollouts are teacher-forced on the conditioning streams: the frozen encoders read ground-truth frames, so these numbers are an upper bound on what a closed-loop rollout would give.