Action-image conditioned MoT world model: ground truth vs. rollout

The two-expert Mixture-of-Transformers diffusion-forcing world model with dense action-image conditioning on the video expert (mot_df_vt_actionimage). On top of the actionflow baseline's per-modality action cross-attention, the camera-projected action is also rasterized per latent frame into a 6-channel Gaussian action image on the 16×16 latent grid (2 sensors × {position, normal, up}, σ = 0.8 cells — arXiv 2604.06168) and injected as a zero-init token bias, so the action lands on exactly the tokens it should move. It predicts 2 future latent frames of the camera view and both tactile sensors from 2 latent frames of history plus the 36-D action.

240.8M params (+9k projector) Wan2.2 VAE latents, C=48, 16×16 checkpoint: epoch 100 (final) frozen actionflow prior

Where prediction starts: every rollout row is an (episode, start) pair on the 6 fps frame grid. The model is given fps6 frames [start, start+8) — 2 latent frames — as clean context; prediction begins at fps6 frame start+8 (source video frame (start+8)×5 at 30 fps, i.e. (start+8)/6 seconds into the episode). Short cells predict the next 2 latent frames (~1.3 s); long cells continue autoregressively to 16 s.

Architecture in one paragraph

Both experts run in lockstep: each cross-attends its own action context, then their tokens are concatenated (256 video + 512 tactile = 768) for a single joint self-attention, then a per-expert FFN. The joint mask is temporally causal across modalities — camera frame t attends tactile frames ≤ t and vice versa. Each modality carries its own noise scheduler with independently drawn per-frame noise levels, so the two branches sit at unrelated points of their diffusion schedules. On the tactile side a frozen action-motion encoder reads the inverse-action displacement map on the gel plane plus a calibrated physics flow prior and injects a zero-initialised motion feature into the future tactile tokens; a frozen latent-flow probe scores the predicted clean latents against measured AllTracker optical flow.

New in this run: the video expert additionally receives the action as an image. The 18-D camera-projected half is rasterized per latent frame (causally: frame j uses transition ≤ j) into 6 Gaussian channels on the latent grid and projected by a zero-initialised conv into an additive bias on the video tokens — the exact mirror of the tactile side's motion injection. Zero-init means the model starts bit-identical to the baseline and learns to use the spatial signal only if it helps. The projection itself was validated against AllTracker optical flow (drift-free local EPE < 1 px/frame in all three camera views) before training.

The action, projected into pixel space

This is the half of the action the video expert cross-attends — action[0:18]. Both sensors' gel triads are projected into the middle camera and carried as absolute image coordinates: 2 sensors × 3 points × (u, v, z), with u,v normalised to [0,1] and z in metres. The tactile expert never sees this; it gets the body-frame SE(3) step instead.

Rendered at the 6 fps action grid over the 16 s long-rollout span; the short-window rows are the first 16 frames of these. The triads landing on the physical sensors in every frame is also the calibration check — a forward-project/triangulate round trip would return ~0 residual even with wrong extrinsics, so only the image can falsify it.

Metrics

Decoded-pixel PSNR (mean ± sem over n = 6 rollouts per cell), GT vs. rollout, through the Wan2.2 VAE — tactile streams via the tactile LoRA, the camera view via the base VAE. tac is the mean of the two sensors.

The contact-rich set row is the 80 windows of tactile_test_set/, pre-selected by temporal tactile activity across 20 episodes — the frames where the sensors are actually touching something, i.e. where tactile prediction is hardest and most informative. Filenames carry raw 30 fps starts; they map to this model's grid as start_fps6 = raw/5. Only the 0510_005 windows are held out; the other 19 episodes were in training, so treat this cell mainly as contact-window fit, evaluated at a pinned epoch-61 snapshot (eval/ckpt_tacset.ckpt; the four cells above use the final epoch-100 checkpoint).

Textured-start windows: GT vs prediction, and per-pixel motion

17 additional held-out windows selected so the gel visibly carries imprint texture in the very first frame (Laplacian-variance of the gel image at fps6 frame start, ≥1.45× the episode's no-contact baseline, visually verified) with real gel motion in the future frames. For each: a frame strip — top row ground truth, bottom row prediction, red line = where prediction starts (the 2 context latent frames decode to pixel frames f0–f4, so the boundary is decoded frame 5) — and the per-pixel motion comparison (AllTracker adjacent flow on GT vs predicted videos, predicted transitions only: [GT frame | GT flow | predicted flow | EPE heatmap] per stream).

Aggregate over the 17 windows (moving pixels = GT flow > 0.5 px/frame): tactile-left EPE 0.94 px/frame, cosine 0.34, magnitude ratio 0.41; tactile-right EPE 1.22, cosine 0.22, magnitude ratio 0.29; view EPE 2.32, cosine 0.59, magnitude ratio 0.72. The model keeps texture appearance but under-predicts texture motion — it moves the gel imprint at roughly a third of the true speed, and in the worst windows predicts a near-static gel where GT shows a large slide. The camera view's motion is much closer to GT, consistent with the video expert getting the explicit action-image cue.

Against the vector-only actionflow MoT (the A/B baseline)

mot_df_vt_actionflow is this exact architecture, data, seed and budget without the dense action image — its video expert gets the projected action only through cross-attention. The rollout rows below are the identical (episode, start) pairs, and both columns are final epoch-100 checkpoints — this is the definitive A/B.

Δ is action-image − baseline in dB; positive means the action-image run reconstructs better.

Three-way comparison: how should the action be injected?

A third run, mot_df_vt_actionfilm, conditions BOTH experts by the action rendered in their own image space via per-token FiLM (tokens·(1+scale)+shift, zero-init) — video from the same rasterized action image, tactile from the gel-plane inverse-action maps — and removes the action cross-attention entirely (182M params vs 241M, ~17% faster to train). All three runs: same data, seed, budget, final epoch-100 checkpoints, identical rollout rows. PSNR in dB; best per row in bold.

Read: the dense action image on top of cross-attention (actionimage, this Space's clips) is the best long-horizon model — +0.2–0.3 dB on both modalities over the vector-only baseline in the 16 s rollouts, where action grounding matters most. The FiLM-only run matches the baseline at short horizon (and takes held-out tactile, +0.69 dB) with 59M fewer parameters and no action cross-attention at all, but gives back ground on long autoregressive rollouts. Six rollouts per cell — treat sub-0.3 dB gaps as ties.

Motion-prior variants: raw ActionImage-Motion & tactile motion+residual

Two follow-up runs attack the measured texture-motion deficit (the model keeps gel appearance but moves it at ~⅓ speed). Both keep everything else identical to this Space's model. aimotion: the ActionImages carry the projected consequence of the action — 6 sparse camera-frame motion channels (per-sensor endpoint du, dv, weight) added to the video ActionImage, plus a direct 5-ch gel ActionImage bias on the tactile branch. tacresidual: the tactile diffusion target is reparametrized — the previous latent frame is transported by the frozen action→flow model and diffusion generates only the residual correction (d₀ absolute anchor; closed-loop sequential reconstruction at sampling). PSNR dB, final checkpoints, same 12 rows:

No-diffusion-forcing ablation: history-conditioned prediction

mot_hist_vt_actionimage keeps everything from this Space's model (action image, cross-attention, actionflow prior, data, seed, budget) but replaces diffusion forcing with standard history-conditioned diffusion: at training the 2 history latent frames are exactly clean, all future frames share one noise level, and the loss covers the future frames only; at sampling all future frames are denoised jointly (full_sequence). PSNR (dB), all final checkpoints:

Diffusion forcing wins every cell decisively — including the short-horizon cells, which are the exact task the history model trains on. The mixed-noise curriculum apparently functions as strong data augmentation/regularization (every frame supervised under every corruption pattern), not just as long-rollout robustness training. Same 12 rows as all other tables; per-sample frame strips below.

Drift: how error grows over the 16 s rollout

Long rollouts are autoregressive — the model's own prediction becomes the next step's context, so error compounds. Per-latent-frame MSE, averaged over the 6 rollouts in each split, shaded ±1 sd. The first 2 frames (shaded) are ground-truth context, not predictions. Frame 0 is exactly zero; frame 1 is small but non-zero because the continuous GT reference is stitched across per-window VAE encodes while the rollout is seeded from a single encode — the two are both ground truth and differ slightly past frame 0. That offset is a property of the reference, not rollout error.

Short window — 2.7 s, 2 latent frames given, 2 predicted

Left half of each clip is ground truth, right half is the model's rollout.

Long rollout — 16 s autoregressive, predictions fed back

Browse every rollout

How it was run

RUN=.../runs_mot/mot_df_vt_actionimage
CKPT=$RUN/checkpoints/last.ckpt        # epoch 100, fixed step count

# short window (one 16-frame window; 2 latent frames of context)
python -m vm_diffusion.scripts.infer_windows --ckpt $CKPT --config $RUN/config.yaml \
    --windows motherboard_0510_episode_005:0 ... --out_dir $RUN/eval/short_test

# long rollout (sliding window, keep-1, driven by the episode's real action sequence)
python -m vm_diffusion.scripts.long_rollout_mot --config $RUN/config.yaml --ckpt $CKPT \
    --episode motherboard_0510_episode_005.pt --start 0 --seconds 16 --fps 6

# decode both through the Wan2.2 VAE (tactile LoRA for tl/tr, base VAE for view)
python decode_rollout_wan22.py --rollout-dir <cell dir> --fps 6

Caveats

• Six rollouts per cell (two episodes × three start offsets). Treat the sem as indicative, not as a significance test.
• This is a held-out-episode split, not a generalization split. The test episodes are the same board and the same session as many training episodes, so these numbers measure "unseen trajectory, seen conditions".
• PSNR rewards blur. A model that predicts a smooth mean-ish gel image scores well. Watch the clips, not only the table — the tactile texture is where the difference shows.
• The four main cells use the final epoch-100 last.ckpt (fixed step count, the fair comparison — val_loss is ~56% the oscillating frozen-flow-probe term, so selecting on best-* would pick for the wrong reason). The contact-rich cell was evaluated earlier at a pinned epoch-61 snapshot and not re-run.
• Long-rollout GT reference is stitched across per-window VAE encodes, so it differs slightly from any single encode past frame 0. The rollout seed itself is exact.
• The projected action is not always in frame. Across these rows the gel centre is in-frame 100% of the time for most, but as low as 88.5% (train row 0). When a sensor leaves the frame its action image is blank, so the video expert receives no positional signal for that step — while the tactile expert's body-frame SE(3) half is unaffected. The two halves of the action degrade independently.

Diffusion forcing (per-frame independent noise levels), pyramid sampling with 100 denoise steps, 36-D causally-masked action conditioning, 3D axial RoPE. Latents are Wan2.2 VAE, 48 channels on a 16×16 grid at 6 fps.