One tactile world model, three conditioning recipes

Three experiments on the same Diffusion-Forcing U-DiT backbone predicting GelSight-style tactile futures in Wan2.2 VAE latent space: ① tactile-only with an 18-D action (9-D relative SE(3) step per sensor), ② visual-tactile with a 36-D action (18 absolute view dims + the 18 tactile dims), and ③ flow-supervised tactile — AllTracker optical flow enters both as a frozen motion-encoder conditioning signal and as an auxiliary loss through a frozen latent-flow probe.

shared backbone: U-DiT · patch 2 · dim 384 · depth 12 · heads 6 · RoPE · spatial-bidirectional / temporal-causal · per-frame noise (Diffusion Forcing, pred-v)

① tactile-only 18-D

44.15 dB

short-window test PSNR ± 1.94 (n = 16)

② visual-tactile 36-D

42.62 dB

tactile avg ± 2.03 (n = 6) · view stream only 24.1 dB

③ flow-supervised

38.28 dB

short-window test PSNR ± 1.88 (n = 5)

Stage 3 · staged flow (final)

42.98 dB

short test ± 1.98 (n = 5) · long test ≥ baseline

Headline (one-shot experiments): neither the visual stream nor one-shot flow conditioning + supervision beat the plain 18-D tactile-only baseline on short-horizon tactile PSNR — the one-shot flow variant was clearly behind (38.3 vs 44.2 dB). Long-horizon behavior is the same everywhere: error jumps at the first predicted frame, then plateaus (~31–32 dB). Sample sizes differ (n = 16 / 6 / 5), so treat gaps < ~2 dB as within noise.

Update — the staged recipe recovers the loss. Re-doing flow conditioning in three stages (probe ablation → frozen-Gφ conditioning → joint finetune, results below) brings short-test from 38.3 back to 43.0 dB (within noise of the 44.2 baseline) and puts every long-horizon test cell at-or-above the baseline (e.g. 16 s: 32.08 vs 31.94) with tighter sd. The finetuned Gφ still predicts AllTracker flow better than its pretrained self — the 0.05·L_action-flow anchor worked.

Architecture: what actually differs

Blue = frozen, orange = trainable, dashed = data. All three share the identical U-DiT denoiser; the differences are entirely in the token groups, the action vector, and (for ③) an extra spatial conditioning stream.

① tactile-only, 18-D action

run tac_only_wan22_fps6act · best ep 54, val 0.0354

tactile latents tl | tr

2 × (48, 4, 16, 16) Wan2.2

↓ patchify per sensor → 2 token groups

U-DiT ×12

spatial bidir → temporal causal, per-frame adaLN(noise, action)

↓

denoised tl | tr latents

action 18-D

[left 9d | right 9d] = rel SE(3): Δp(3) + Rot6D(3×2) per sensor

② visual-tactile, 36-D action

run view_tactile_wan22_action36 · best ep 99, val 0.0502

view + tl + tr latents

3 × (48, 4, 16, 16)

↓ patchify per stream → 3 token groups

U-DiT ×12

same weights budget; joint attention over 3 groups

↓

denoised view | tl | tr latents

action 36-D

[view 18 | tactile 18] — view half = 3 tracked points × (u,v,z) per sensor, absolute image-space position

③ flow-supervised tactile

run tac_only_wan22_fps6act_actionflow · best ep 35, val 0.2080*

tactile latents tl | tr + action 18-D

same as ①

ActionMotionEncoder G_φ ❄

action-map + z_t + physics-flow prior → motion feature (128, 16, 16)

↓ motion_projector (conv, zero-init) → additive per-token bias, frame t+1

U-DiT ×12

= ① plus the motion token-bias

↓ predicted clean latents

LatentFlowProbe ❄

(ẑ_t, ẑ_t+1) → flow → aux loss vs AllTracker (w 0.05)

*③'s val_loss includes the 0.05-weighted flow auxiliary — it is not comparable to ①/②'s plain denoising val_loss. Checkpoints: 612 MB (①/②) vs 618 MB (③; +frozen aux modules + projector).

Supervision: what each run optimizes

experiment	loss	flow's role	inactive-but-configured
① tactile-only	L_DF (pred-v, per-frame noise)	—	—
② visual-tactile	L_DF over 3 streams	—	—
③ flow-supervised	L_DF + 0.05 · L_pred_flow	input side: frozen G_φ token-bias · loss side: frozen probe readout	motion_loss (gated off — G_φ frozen), flow_head (never enabled)

How pixel-level AllTracker flow supervises a latent-space model

AllTracker measures motion in pixels; the world model predicts Wan2.2 VAE latents (48 ch, 16×16 grid, 4× temporal compression). The two never meet in pixel space — there is no VAE decode anywhere in the training loop. Instead, three alignment steps bring the supervision to the latents:

Target rail — pixel motion, aligned to the latent grid

tactile video

192 × 256 px @ 6 fps

↓ AllTracker (offline)

adjacent-frame flow

(N−1, 2, 192, 256) — dx, dy in px

↓ step 1 · spatial align: adaptive avg-pool to 16×16, dy × 256⁄192

flow on the latent grid

(N−1, 2, 16, 16) — still pixel units, one per frame-pair

↓ step 2 · temporal align: sum 5 frame-pairs per latent transition

target F

(T_lat−1 = 3, 2, 16, 16) per window — composed displacement, px

Prediction rail — reading motion out of latents

U-DiT (training)

predicted clean latents ẑ_t, ẑ_t+1 — (48, 16, 16) each

↓ concat pair (96 ch)

step 3 · LatentFlowProbe ❄

conv 96→128→128→2 · pretrained on GT latent pairs → pooled AllTracker flow, validated vs zero-flow baseline on held-out episodes, then frozen

↓

predicted flow F̂

(2, 16, 16) per transition — same units as target

↘ ↙

L_pred_flow = mean( w · |F̂ − F| )

w = 1 + 2·|F| (clamped at 10) — moving regions weigh more · total weight 0.05

↑ gradients flow back through the frozen probe into the DiT

frozen ❄ trainable data / precomputed

The probe is the bridge. It is the only latent↔pixel translator in the system: a 3-layer conv net that learned, from ground-truth latent pairs, how pixel motion is encoded in Wan2.2 latents. Frozen, it becomes a differentiable "flow meter" — the loss asks "does the predicted latent pair contain the right motion?" without ever decoding to pixels. Supervision stays in pixel-motion units end to end; only the readout lives in latent space.

The companion input path: how flow-derived signal enters the DiT

Distinct from the loss above — this is conditioning, computed from given (not predicted) latents at both train and inference time.

action-map (3 ch, metres ×1000) + current tactile latent z_t (48 ch) + physics-flow prior (2 ch, px)

↓ per (sensor, latent-transition), 16×16 grid

ActionMotionEncoder G_φ ❄ (amx_E)

pred_flow = gain · physics_flow + residual · motion feature (128, 16, 16)

↓ motion_projector: conv patch-2 → dim 384 (zero-init = starts as a no-op)

additive bias on that sensor's token group, frame t+1 → U-DiT

Metrics: all six cells, three experiments

① tactile-only 18-D ② visual-tactile 36-D (tactile avg) ③ flow-supervised (one-shot) Stage 2 (frozen Gφ) Stage 3 (joint, final)

Tactile pixel PSNR (tl/tr average), error bars ±1 sd across samples. ②'s view stream (not shown) sits far lower: 24.1 dB short-test, ~21 dB long — the visual stream is much harder than tactile at this model size.

cell	① PSNR (n)	② tac PSNR (n)	③ PSNR (n)	Stage 2 (n)	Stage 3 (n)

Comparability caveats. Sample counts differ per run (① n=16, ② n=6, ③ n=5 per cell) and the long-rollout sample rows are not identical across runs (③'s long-train covers 5 distinct episodes incl. one 0519 outlier at ~21–26 dB). ②'s number is the tl/tr average of a jointly-generated 3-stream rollout. ③'s training val_loss (0.208) includes the 0.05-weighted flow auxiliary; per-run wandb: ① val 0.0360 · ② val 0.0502 · ③ val_pf 3.32 at epoch 99. Long-rollout GT is stitched from different VAE encode windows, deflating all long PSNRs equally.

Qualitative: ground truth | prediction

Held-out (test) samples. Left half = GT, right half = model rollout.

The full training paradigm — what each step is for

Every stage initializes from the previous one's checkpoint (weights + conditioning path carry through) and changes exactly one variable, so attribution stays clean. Vision is never used in this lineage — tactile latents (tl+tr) and 18-D sensor actions only.

Stage 1 — probe done

"Is (tactile state, action) actually predictive of gel motion?"

Gφ(z_ctx, a) → F_AllTracker

World model untouched; train only the tiny encoder. The 3-way ablation answered it: a→F 1.36 < z→F 1.17 < (z,a)→F 1.13 (zero-flow 1.61) — tactile says where, action says which way; the physics prior adds nothing.

Stage 2 — frozen conditioning done

"Does the flow representation help generation?"

freeze Gφ · train DiT · plain L_DF

No auxiliary losses at all — clean attribution. val_loss 0.0354 exactly matches the no-flow baseline; every long-horizon test cell at-or-above it. (The one-shot version of this idea had lost 6 dB; staging recovered all of it.)

Stage 3 — joint finetune done

"Let the DiT adapt Gφ without destroying it."

unfreeze · L_DF + 0.05·L_action-flow

The AllTracker anchor keeps Gφ interpretable: after finetuning its weights moved substantially (‖ΔW‖≈33) yet it predicts flow better than its pretrained self (EPE 1.123 vs 1.130). Short test 42.98 dB.

Stage 4 — rollout robustness done

"Stop autoregressive drift from compounding."

+ scheduled sampling · snr_clip 5→20

Scheduled sampling trains on the model's OWN drifted context (targets exposure bias); relaxing snr_clip restores gradient on the low-noise denoising steps where fine detail forms. Contact-rich 16 s texture energy 0.66→0.72, hardest segment 0.57→0.77, +0.9 dB.

Stage 5 — high-frequency supervision done

"Make blur stop being loss-optimal."

+ 0.1 · ‖∇x_pred − ∇x_gt‖²

A latent gradient-domain loss directly supervises the spatial high-frequency structure of predicted latents — attacking the root cause: for unpredictable texture phase, plain MSE rewards averaging (= blur). Outcome: texture mean 0.72→0.84 (gate 0.85, ceiling 0.86); hard segments 0.89–0.96; short test 43.08 dB — the staged line's best.

Escalation — context anchoring if needed

"Carry the imprint instead of regenerating it."

context texture → token bias (+ optional vision cond.)

Tactile context anchors imprints that already exist; a frozen-encoder vision conditioning could announce imprints about to form (the scene camera is the only source of "what is about to touch"). Triggered only if Stage 5 misses the 0.85 texture gate (VAE ceiling 0.86).

Inference-side, orthogonal to training: best-of-N seed selection with the texture-energy criterion (rescues short-horizon collapse), and the texture evaluation protocol itself (gradient-energy ratio / gradient correlation — PSNR is blind to this failure mode).

Staged results: the Stage-1 ablation

amx12_* runs, 12 epochs, held-out episodes 005/006. Active-pixel EPE (target motion > 0.5 px), lower is better.

variant	inputs	active-EPE ↓	full EPE	dir cos ↑
zero-flow baseline	—	1.612	1.155	—
physics prior	—	1.606	1.308	—
A: physics-only	prior	1.381	1.016	0.557
B: a → F	action map	1.362	1.005	0.553
C: z → F	tactile state	1.168	0.852	0.694
D: (z, a) → F — winner	both	1.130	0.837	0.726
E: all	both + prior	1.135	0.843	0.720

What Gφ learns — action vs tactile-state contributions (held-out, high-motion transitions)

Rows = held-out transitions; columns = raw tactile frame, the action motion-map input (affine over the gel — direction only, no localization), predictions from the three ablations, and AllTracker ground truth. The tactile state answers where motion happens (violet localizes to the imprint), the action answers which way (orange has direction but smears it everywhere); (z, a) combines both. The sharp high-magnitude bands in the GT that no variant captures are the contact-region-shape structure discussed above.

Stage-1 prediction comparison — Click to open full size. Per-panel EPE annotated (px on the 16×16 latent grid).

Gφ drift check (Stage 3). After joint finetuning the encoder moved substantially in weight space (‖ΔW‖ ≈ 33) yet predicts AllTracker flow slightly better than its pretrained self on held-out episodes (active-EPE 1.123 vs 1.130, direction cosine 0.736 vs 0.726) — the 0.05·L_action-flow anchor kept it an interpretable flow module rather than an arbitrary hidden feature.

Roadmap: a staged path for flow conditioning

③ wired flow in with pretrained-elsewhere, frozen modules and underperformed the plain baseline. The staged plan below isolates why, one question at a time: G_φ is the (tiny) motion encoder, F is AllTracker flow, and the expensive world model only trains once the flow representation has proven itself.

Stage 1 — probe the signal done

"Is action + current tactile actually predictive of AllTracker flow?"

tactile world model ❄ (not involved)

train G_φ only — tiny, cheap

a → F vs z_ctx → F vs (z_ctx, a) → F

The critical ablation: if action alone predicts flow as well as (z, a), the tactile state adds nothing; if neither beats a zero-flow baseline, flow conditioning can't help downstream. Existing analogues: runs/action_motion/amx_A…E + latent_flow_probe — rerun as this explicit 3-way comparison.

Stage 2 — frozen conditioning done

"Does a meaningful flow representation improve tactile generation?"

G_φ ❄ (loaded from Stage 1)

↓ F̂ = G_φ(z_ctx, a) conditions the DiT

train world model only

Clean attribution: any gain comes from the conditioning signal, not from moving G_φ. Experiment ③ is essentially this stage (frozen amx_E + token-bias injection) — and its −6 dB result says the Stage-1 question needs a sharper answer first.

Stage 3 — joint finetune done · final model

"Keep G_φ interpretable while the DiT adapts it."

G_φ unfrozen + world model, jointly

L = L_DF + 0.05 · L_action-flow

The AllTracker loss anchors G_φ to real motion instead of letting the DiT repurpose it as an arbitrary hidden feature. This run's motion_loss block (currently gated off by freeze: true) is exactly the L_action-flow needed — unfreeze and go.

Contact-initiated test-set samples — Stage-3 final model

Five windows from tactile_test_set whose first frame already has firm GelSight contact (raw contact area > 0.25, ranked by first-frame area, one window per episode). 2 latent frames given, 2 predicted. The sensor carrying the active contact is the hard one (31–42 dB); the quiet sensor exceeds 50 dB. Aggregate over the 5 contact windows: tl 37.98 dB · tr 44.54 dB. All six windows are from training episodes — the top contact-area windows in this test set happen not to be from the held-out episodes. *Sample 005 is the opposite regime, added on request: a push-T window starting almost contact-free (area 0.077) but with the highest tactile activity in the whole test set (0.92) — the model tracks it well on both sensors.

sample	episode @ fps6 start	first-frame area	tl PSNR	tr PSNR
000	0511_ep017 @ 2423	0.544	42.33	31.53
001	0510_ep009 @ 159	0.461	39.04	51.91
002	0511_ep005 @ 2365	0.435	37.92	52.05
003	0519_ep001 @ 69	0.410	34.47	35.50
004	0511_ep008 @ 993	0.400	36.16	51.69
005*	pusht_ep001 @ 900	0.077	45.89	47.66

Sampling matters: texture collapse and best-of-N seed selection

Rollout imprint texture (e.g. connector ribs) often collapses to blur. Diagnosis: it is seed-dependent diffusion mode selection, not a model ceiling — some seeds sample sharp texture, more DDIM steps (100→250) change nothing, and PSNR cannot see the difference (sharp-but-phase-shifted texture scores the same as blur). Fix: sample N seeds and pick, with a GT-free criterion — match the predicted frames' Sobel texture energy to the context imprint's energy (penalizes both blur and hallucination). On 5 contact windows the picker lands within 0.01–0.02 gradient-correlation of the GT-oracle on 4/5.

Short horizon: best-of-6 rescues the texture

best-of-seed short-window comparison — Worst window (connector ribs, mid-motion frame; left=GT, right=prediction): default seed loses the imprint entirely; the no-GT picker selects a seed with the texture intact — visually better than even the gradient-correlation oracle. Click for full size.

short window · default seed

GT | prediction — imprint collapses

short window · best-of-6 picked

GT | prediction — imprint preserved

Long horizon: an honest null — per-step best-of-4 does not fix compounding drift

long-rollout best-of comparison — t≈15 s of 16 s rollouts on contact-rich segments (left=GT, mid=default, right=per-step best-of-4): both arms lose the imprint. Once the fed-back context is blurred, all N candidates start from blur — late-third texture-energy ratios are mixed (0.5–0.76 both arms), PSNR +0.3 dB. Long-horizon texture loss is compounding drift, not per-step mode selection; it needs context-texture anchoring or training-time scheduled sampling, not more sampling.

16 s rollout · default

0511_017@1332 · left sensor · GT | prediction

16 s rollout · per-step best-of-4

same segment · GT | prediction

Training-side fixes: scheduled sampling + low-noise training signal (in progress)

Ceiling first: the Wan2.2 VAE round-trip itself keeps 86% ± 6% of raw texture energy — so most of the long-horizon gap is the world model, not the representation. First training fix (tac_s4_ss): enable scheduled sampling (the model trains on its own drifted context — directly targets AR exposure bias; previously never enabled) + relax snr_clip 5→20 (restores training weight on the low-noise denoising steps where fine detail forms). 50-epoch finetune from Stage 3.

contact-rich 16 s segment	s3 default	s3 best-of-4/step	s4 + sched. sampling
late-third texture-energy ratio (avg tl/tr over 3 segments)	0.66	0.67	0.72
pixel PSNR (3 segments)	29.41	29.76	30.32

Stage 5: + latent gradient-domain loss — the texture gate is met

contact-rich 16 s segment (tl/tr)	Stage 3	+ sched. sampling	+ gradient loss (final)
ep005@24 (held-out)	0.72/0.76	0.70/0.75	0.71/0.76
017@1332	0.57/0.60	0.59/0.64	0.89/0.96
017@1512	0.75/0.57	0.84/0.77	0.86/0.89
mean (gate 0.85 · VAE ceiling 0.86)	0.66	0.71	0.84
contact-segment PSNR	29.41	30.32	31.12

Gate met. Adding 0.1·‖∇x_pred − ∇x_gt‖² on the predicted latents lifts the texture-energy mean to 0.84 — at the 0.85 gate and essentially at the VAE round-trip ceiling (0.86) on the hard segments (0.89–0.96). Short-window quality improved too (test 43.08 dB, the best of the staged line). The one flat row (ep005@24) is flat across ALL arms — its late texture is a newly-forming contact the model cannot know about, which is exactly the future vision-conditioning case.

four-way texture comparison — t≈15 s, segment 017@1332, right sensor: GT shows a chip plus rows of pin dots. Stage 3 is nearly blank (0.60); scheduled sampling recovers the chip edges (0.64); the gradient-domain loss recovers the pin-dot rows too (0.96). Click for full size.

16 s · Stage 3

017@1332 · right sensor · GT | prediction

16 s · final (+SS +grad loss)

same segment · GT | prediction

scheduled sampling texture comparison — t≈15 s of a 16 s rollout on the hardest contact segment (right sensor): GT shows a chip with pin dots; the Stage-3 default hallucinates a wrong wedge; the scheduled-sampling model recovers the chip's corner structure. Short-window quality is unaffected (42.92 vs 42.98 dB test). Verdict: real progress (0.66 → 0.72) but below the 0.85 gate (ceiling 0.86) — a latent gradient-domain loss (E2b) is training next, with context-texture anchoring (E3) as the following escalation.

16 s · Stage-3 default

017@1512 · right sensor · GT | prediction

16 s · + scheduled sampling

same segment · GT | prediction

Per-experiment Spaces

① tactile-world-model-rollouts · ② view-tactile-world-model-rollouts · ③ tactile-wm-actionflow-rollouts — each hosts every rollout clip, per-cell metrics, and drift curves.

Data: 36 episodes (motherboard + push-T), Wan2.2 VAE latents @ 6 fps, 16-frame windows (8 history + 8 future → 4 latent frames, 2 given). Eval: short = one window; long = autoregressive keep-first-future-frame, 8/16 s. Decode: Wan2.2 VAE + the encode-time LoRA adapter. AllTracker flow pooled to the 16×16 latent grid at build time (build_flow_latent_targets.py).