chomeed/mimicgen_threading_d0_224x224_sd3_dynamics

SD3 pixel-space dynamics — MimicGen threading_d0

An action-conditioned world model that predicts future pixels, not features. A full finetune of the Stable Diffusion 3.5-medium MMDiT (2.25B parameters) with Ctrl-World-style conditioning: context frames and the noisy target live in one canvas so the transformer attends over all of them jointly.

{I_t}^agentview,eye_in_hand ,  a[t:t+8] ,  s_t   ->   {I_{t+8}}^both views

Result: it beats copy-the-current-frame

That is the bar that matters, and it is not a given — an earlier design in this codebase that fed frames only through encoder_hidden_states reached 7.9 dB against copy's 19.0 dB, i.e. worse than doing nothing. The canvas conditioning fixes that.

step	PSNR	copy-the-current-frame	gap
1000	14.65	18.12	−3.48
3000	16.37	18.12	−1.76
10000	18.12	18.12	±0.00
25000	18.88	18.12	+0.76
28000	19.24	18.12	+1.12 ← peak
30000 (this checkpoint)	18.96	18.12	+0.83

Checkpoints are saved every 5,000 steps, so the 28k peak is not on disk; this is step 30,000.

It looks better in latent space than in pixels

PSNR is a pixel measure, and pixels are not what this model predicts — the ODE runs entirely on 16×28×28 VAE latents and the decode is presentation only. Scoring in latent space, which is also what a planner would do, on 160 held-out windows (eval_sd3_latent.py):

sampler	latent MSE	vs copy	sharpness
20 steps, mean of 1	0.0768	0.719	0.998
20 steps, mean of 4	0.0489	0.458	0.965
copy-the-current-latent	0.1068	1.000	—

The pixel number and the latent number disagree about how good this model is:

pixel  (from PSNR 18.96 vs 18.12)   MSE ratio 0.824
latent, one sample                  MSE ratio 0.719
latent, mean of 4                   MSE ratio 0.458

The VAE decode masks part of the gain. Judged where it actually operates, the model is well clear of copy, not marginally so — and averaging just 4 samples nearly halves the error while keeping sharpness at 0.965. Use predict(..., return_latent=True) and score there.

sharpness = std(pred)/std(target) over the 784 spatial positions, per channel. 0.998 at N=1 is essentially perfect: unlike an MSE-trained head, a diffusion sampler has no incentive to blur.

Use 2 sampler steps, not 20

Step count trades the same way sample-averaging does, and in the opposite direction — fewer steps land nearer the conditional mean, more steps give a sharper draw (N=1 throughout):

steps	latent MSE	sharpness	passes
1	0.0544	0.953	1
2	0.0535	0.950	2
4	0.0655	0.984	4
8	0.0734	0.994	8
16	0.0768	0.998	16
32	0.0791	0.999	32
50	0.0785	1.000	50

One denoising step from noise lands near E[z | context, a]; more steps integrate toward an actual sample, which is sharper but further from any single ground truth.

The practical consequence is large. 20 steps × N=4 costs 80 passes for latent MSE 0.0489. 2 steps × N=1 gives 0.0535 for 2 passes — within 10% of it at 1/40th the cost. For planning, where you score thousands of candidate action sequences, that ratio decides whether the model is usable at all.

goal	setting	cost
lowest latent MSE / planning	2 steps, N=1	2 passes
best MSE regardless of cost	20 steps, N=4	80 passes
sharpest sample / visualisation	16+ steps (50 buys +0.002 over 16)	16–50

Caveat: 160 windows, one seed. The ranking is clear but the third digit is not. The same tradeoff appears in the DINO flow variant, far more violently — there, averaging drives sharpness below the deterministic baseline's.

Architecture

Input — two 224×224 RGB cameras, VAE-encoded to 28×28×16 latents and tiled into one canvas (rows = timesteps, columns = cameras):

to_canvas: (B, R, V, C, h, w) -> (B, C, R*h, V*w)

           agentview   eye_in_hand
   t      │  28x28   │   28x28    │   context row (clean)
   t+8    │  28x28   │   28x28    │   TARGET row (noised)
          canvas = 16 x 56 x 56  ->  784 MMDiT tokens

Because both views share one sequence they are denoised jointly, and the target row attends to both context frames.

Conditioning — 9 cross-attention tokens:

act   = act_tok(action)         (B, 8, 1536)   one token per action step
state = state_tok(state)        (B, 1, 1536)   eef_pos(3)+quat(4)+gripper(2)
ctx   = cat([act, state])       (B, 9, 1536)   -> encoder_hidden_states
pooled = pool(ctx.mean(1))      (B, 1536)      -> pooled_projections

act_tok, state_tok and pool are zero-initialised, so step 0 is the pretrained MMDiT undisturbed rather than a randomly perturbed one.

Objective — rectified flow matching on the target row only; context rows are never scored.

Usage

import torch
from sd3_dynamics import SD3MultiViewDynamics
from diffusers import FlowMatchEulerDiscreteScheduler

SD3 = "stabilityai/stable-diffusion-3.5-medium"
model = SD3MultiViewDynamics.from_pretrained(SD3, n_views=2, action_dim=7, state_dim=9).cuda()
model.transformer.to(torch.float32)                       # this checkpoint is a full fp32 finetune
model.load_state_dict(torch.load("step_30000.pt"), strict=False)   # vae.* absent, reloaded above
model.eval()

sched = FlowMatchEulerDiscreteScheduler.from_pretrained(SD3, subfolder="scheduler")

pred = model.predict({                   # -> (B, 2, 3, 224, 224) in [0,1]
    "context": ctx,                      # (B, 1, 2, 3, 224, 224)  current frame, both cameras
    "action":  actions,                  # (B, 8, 7)   delta OSC_POSE chunk
    "state":   state,                    # (B, 9)      observation.state[:9]
}, sched, steps=20)

# what a planner should use: the latents the ODE actually produced, no decode (~30% cheaper)
z      = model.predict(batch, sched, steps=20, return_latent=True)   # (B, 2, 16, 28, 28)
z_true = model.encode_grid(batch["future"][:, None])[:, 0]           # same (z-shift)*scale convention

Three things that will bite if missed:

1. history=1. This checkpoint was trained with a single context frame, so context has a time axis of length 1. Passing two frames changes the canvas geometry and the weights will not apply.
2. fp32. The MMDiT loads in bf16 by default; these weights are fp32. Cast before loading.
3. Buffers are absent by design — the checkpoint stores trainable parameters only. That means all vae.* keys (the VAE is frozen and reloaded from the SD3 repo) and transformer.pos_embed.pos_embed, a deterministic sincos table that carries no gradient and is rebuilt identically by from_pretrained. Use strict=False, and check that the missing list contains nothing beyond those two — anything else is a real mismatch.

Running it downloads stabilityai/stable-diffusion-3.5-medium (Stability AI Community License, gated); that licence governs the base weights.

Training

--lora-rank 0            full finetune (not LoRA); the MMDiT is cast to fp32 first, because
                         AdamW on bf16 weights diverges -- Adam's second moment does not survive
                         8 mantissa bits
--lr 1e-5                the script's 1e-4 default is a LoRA learning rate, ~10x too high here
--history 1 --steps 30000 --batch-size 8 --grad-accum 2   (effective batch 16)
--state-dim 9 --image-size 224 224 --action-dropout 0.2

2.25B trainable, 46 GB peak, 1.03 s/step, ~8.6 h on one B200.

Data: chomeed/mimicgen_threading_d0_224x224_mtdit_flow_55k_{success,failure} + chomeed/mimicgen_threading_d0_224x224 — 260,931 train / 16,369 validation windows.

Limitations

• +0.83 dB over copy is a modest margin. It clears the bar, which the previous design did not, but it is not a large improvement over a trivial baseline.
• Single-step only. Trained for one 8-action jump; multi-step rollout compounding is untested.
• No planning evaluation. Good PSNR does not guarantee useful dynamics for CEM/MPC.
• Single seed, one task. No cross-task or cross-embodiment claims.
• Checkpoint granularity is 5,000 steps, so the 19.24 dB peak at 28k is not recoverable from this repo.