Visuo-tactile world model · Mixture-of-Transformers · Diffusion Forcing

Action-conditioned
video and touch,
trained in three stages.

latest config train_mot_df_vt_pixflow_gelz*.yaml (2026-08)
large-trunk variant train_mot_df_vt_wan5b_vmx.yaml
latents 48ch · 16×16 · Wan2.2 VAE · 6 fps · window 8+8 frames
action 36-D = camera proj (18) | gel-frame SE(3) (18)

vision stream tactile stream (left / right gel) action contact force frozen after its stage

Each stage trains one thing, freezes it, and hands it to the next. Stage 1 turns the robot action into a dense motion field for each sensor; Stage 2 trains small probes that read motion and force back out of latents and act as judges; Stage 3 is the generative model. Every stage holds out the same two episodes (motherboard_0510_005/006).

34-second walkthrough of the three stage diagrams. Stills: overview · stage 1 · stage 2 · stage 3

INPUTS

What the model sees

One middle RGB camera and two GelSight sensors, 6 fps, encoded by the Wan2.2 VAE into 48-channel 16×16 latents. The action is the two sensors' motion, expressed twice: projected into the camera, and as an SE(3) step in each gel's own frame.

Raw episode motherboard_0510_005 (held-out). Green / red rings: projected left / right gel centres with a 12-frame trail. Right: the two GelSight streams the tactile expert predicts.

The ActionImage: the view half of the action rasterised on the latent grid — 6 position Gaussians (2 sensors × {pos, normal, up}) plus 6 endpoint-motion channels (du, dv, weight per sensor).

Calibrated sensor pose projected into all three cameras — the chain that produces the 18 camera-projection action dims.

Projection validated against AllTracker point tracks (<1 px / frame drift-free EPE).

STAGE 1

Motion priors

Turn the action into a dense flow field — once for touch, once for vision. Now in raw pixel space, on AllTracker's own 24×32 grid; the VAE is not involved.

task: per adjacent frame pair: (action raster, current frame) → optical flow on the 24×32 grid; target = AllTracker flow, native px
notation: n_gel = gel normal (unit vector, from calibration); F_n = normal force in newtons; gel-frame axes (t1, t2, n_gel), so press replaces the old dz.
gel frame: world poses are composed with the gel calibration before the SE(3) converter, so its third axis is the gel normal n_gel (≈ −y in the sensor rigid frame): press = displacement along n_gel, shear_u/v = in-plane, and the normal force F_n rides on the same axis. Median flow cosine of the rigid-body prior rises 0.44 → 0.53 (left) and 0.32 → 0.52 (right) from this alone
force: react-calib normal force F_n (along n_gel; a scalar, no direction stored), k = 2 N/mm (assumed), taken at the transition-start frame; channels 3–4 = log1p(F_n) + binary contact. 42 % of frames are in contact
optim: lr 3e-4 · batch 128 pairs · 8 epochs · best ckpt by active-pixel EPE · minutes per run
why pixels: the Wan latent carried ~45 % per-window encode variance and sits 4 frames apart, not the 5 the flow chunks assumed; the frame carries neither problem and gives 5× more samples
predecessors: latent-space amx_gel5_s2 (EPE 1.129 vs zero-flow 1.61) and vmx_E (EPE 2.42 vs 7.07) — still used by the Wan-5B variant below

STAGE 2

Latent probes

Frozen judges that read motion and force back out of tactile latents.

task: from clean tactile latents, regress the measured flow between consecutive latents, and the contact state / force of one latent
optim: lr 3e-4 · batch 64 · 6–8 epochs · must beat the zero-flow / all-zero baseline on held-out episodes before use
role: frozen; used only to score the world model's predicted latents, never at sampling time. The force probe's held-out accuracy is the ceiling for anything it says about a rollout

STAGE 3

World model

Mixture-of-Transformers over video + touch, trained from scratch with Diffusion Forcing. Config pixflow_gelz (+F force head, +jitv x-prediction).

objective: Diffusion Forcing: independent per-frame noise level per modality · sigmoid schedule, 1000 steps · fused-SNR weighting (clip 5, decay 0.96). gelz: pred_v. jitv: net emits clean x₀ but is scored in v-space (loss_space: v) — JiT's x-prediction with the pred_v weighting exactly restored, so the only variable is the prediction target
conditioning: three zero-init paths, all identity at init: vector FiLM per branch (18-D view / 11-D per gel), frozen Stage-1 flow as a token bias (mathematically the same as concatenating 2 channels to the latent without widening the diffusion variable), and — in gelzF — a force head that shapes the tactile hidden state h_G the diffusion head reads. No action cross-attention
alignment: force enters at the transition-start frame, never at the frame being generated (a measurement of the outcome would leak the answer)
trainable: everything: 240 M. Frozen: pixc5_D, pixview_D, probes
optim: AdamW lr 1e-4 (β .9/.95, wd .01) · 500-step warm-up · batch 16 × 2 GPU · bf16 · grad-clip 1 · 100 epochs · split manifest drops the 4 pushT episodes (no force export)
sampling: pyramid schedule, 100 DDIM steps, 2 context frames, chunk 1 · both streams decoded by the Wan2.2 VAE · force head and probes are never called

Large-trunk variant · `wan5b_vmx`

Same MoT layout on the frozen Wan2.2-Fun-5B DiT (30 layers, dim 3072, 24 heads, one trunk shared by both experts) with per-modality LoRA r = 32. Differences from pixflow_gelz:

motion priors: latent-space amx_E (gel map + z_t + physics prior) and vmx_E (ActionImage + z_view), both frozen; features injected as token bias; vmx flow appended to the ActionImage
action path: ActionEncoder 18→256 per branch with cross-attention before joint attention; ActionImage fed through Wan's pretrained dense-control patch-embed via a zero-init 12→48 adapter
tactile target: residual: d₀ = z₀, d_j = (z_j − warp(z_{j−1}, f_j)) / 0.3303; sampling rebuilds absolutes closed-loop. Later measured: warping a Wan latent by this flow is indistinguishable from a 3×3 blur or a shuffled flow, which is why the pixflow family keeps absolute targets
aux loss: Stage-2 flow probe on predicted latents, weight 0.05
optim: batch 4 × 2 GPU × accum 4 · early stop patience 10 · gradient checkpointing

HELD-OUT

Where the family stands

Decoded pixel PSNR (dB) on held-out windows; view = middle camera, tl / tr = GelSight left / right. Tactile PSNR swings ~14 dB between contact-heavy and quiet windows, so rankings use the window-matched 908-window set, 3 seeds.

run	trunk	motion prior	action path	view	tl	tr	note
wan5b_vmx	Wan2.2-5B + LoRA	amx_E + vmx_E (latent)	x-attn + ActionImage	24.36	—	—	best held-out short view
pixflow_gelz	240 M scratch	pixc5_D + pixview_D	vector FiLM (gel + force)	24.03	—	—	best view of the pixflow family
pixflow_gelzF	240 M scratch	pixc5_D + pixview_D	vector FiLM + force head	24.00	44.4	44.6	force head 0.02
pixflow_gelnoF	240 M scratch	pixc5_D + pixview_D	vector FiLM, no force	—	44.66	44.74	best tactile of the family
aimotion	240 M scratch	amx_E (latent)	x-attn + ActionImage-Motion	—	—	—	best all-round 240 M without long-horizon collapse

A dense, perfect flow explains only ~12.5 % of frame-to-frame change on a gel image — a vision-based tactile sensor images surface normals, and a deepening dimple changes pixels in place. That is the ceiling any flow prior has on the tactile side, and why the force channel was added.

STAGE 3 OUTPUT

Rollouts

Ground truth (left) against a 16 s autoregressive rollout (right) from the MoT-DF ActionImage run, held-out episode. The model saw 2 context frames and the action sequence; everything after is generated.