chomeed/mimicgen_threading_d0_224x224_dinowm_flow

DINO-WM with a rectified-flow predictor — MimicGen threading_d0

The published DINO-WM baseline predicts future DINOv2 features by minimising MSE. That objective is minimised by the conditional mean, so the model is trained to hedge across every future consistent with the context. These checkpoints replace that head with rectified flow matching, which models the full conditional distribution and has no such incentive. Everything else is identical — same frozen facebook/dinov2-small, same 13.2M predictor, same 1-action-token conditioning, same data, same 60k steps, same LR — so the objective is the only variable.

Two checkpoints:

file	models
dino_step_60000.pt	z_{t+8} — absolute features	the direct substitution
dino_step_60000_residual.pt	z_{t+8} − z_t — the change	better; start here

The hypothesis did not survive measurement. Flow does produce sharper predictions, but the margin is small and it costs 1.7–2.5× the feature MSE. This repo documents that, and ships the scripts that show it.

Result

640 held-out windows, all arms on identical batches, 20 Euler steps, via eval_flow.py:

	feature MSE	vs copy	sharpness
residual, mean of 1	1.1280	0.852	0.968
residual, mean of 8	0.7807	0.590	0.897
residual, mean of 32	0.7438	0.562	0.888
absolute, mean of 1	1.2518	0.946	0.965
absolute, mean of 8	0.7710	0.583	0.863
absolute, mean of 32	0.7199	0.544	0.851
copy-the-current-frame	1.3232	1.000	—
MSE head (published recipe)	0.4479	0.339	0.925

Sharpness = std(prediction) / std(target) across the 196 patch tokens, per feature dim, then averaged. 1.0 means the prediction varies over space exactly as much as the true next frame.

The MSE head wins on MSE by 1.7× even against the best flow configuration, at 1 forward pass against 640. If feature MSE is your objective, use the baseline. The flow models earn their place only where you need samples — a plausible future, or an uncertainty estimate — which the MSE head cannot produce at all.

Residual vs absolute

Rectified flow transports N(0,1) noise onto the target distribution, so the target's geometry is the whole game:

absolute   z_{t+8}        per-dim mean [−4.00, +13.31]   std [0.96, 8.07]
residual   z_{t+8} − z_t  per-dim mean [−0.07,  +0.10]   std [0.65, 2.34]

The residual target is already near zero-mean and unit-scale, so the flow has far less work to do. It won at all 30 matched validation checkpoints from step 2k to 60k, and the absolute run never reached the residual run's final value at any point in its training:

step	residual	absolute
2000	2.1775	3.3462
10000	1.6073	1.9392
30000	1.4107	1.5304
60000	1.2883	1.4116

(single-sample val on 128 windows — noisy point-to-point; the 30-of-30 sign test is what makes it credible, not any single value.)

On the held-out eval, residual is −9.9% MSE at N=1 and holds sharpness far better under averaging (0.888 vs 0.851 at N=32). It is marginally worse on MSE at N=8/32 — the gain is concentrated in the single-sample regime.

Choosing the sampler: the two knobs oppose each other

Euler steps and sample count both cost forward passes, and they pull in opposite directions. Residual checkpoint, 640 windows:

Steps buy sharpness and cost MSE (N=1):

steps	MSE	sharpness	passes
1	0.9630	0.915	1
2	0.9757	0.930	2
4	1.0375	0.948	4
8	1.0915	0.961	8
16	1.1236	0.967	16
32	1.1449	0.969	32
64	1.1519	0.970	64

This is not a bug. One Euler step from noise lands near E[z | context, a] — the conditional mean. More steps integrate toward an actual draw, which is sharper but further from any single ground truth. The step count is a second control on the same mean-versus-sample tradeoff that averaging governs, and it runs the opposite way on cost.

At a fixed budget, spend it on samples, not steps (steps × N = passes):

passes	best config	MSE	sharpness
1	1 step × 1	0.9630	0.915
8	1 step × 8	0.8273	0.886
32	1 step × 32	0.8128	0.883
32	4 steps × 8	0.8197	0.904
64	2 steps × 32	0.8001	0.893
128	4 steps × 32	0.7955	0.899
640	20 steps × 32	0.7438	0.888

Diminishing hard: 640 passes buys 6% MSE over 128. Recommended defaults:

goal	setting	cost
lowest MSE	use the MSE head instead — 0.4479	1 pass
cheap conditional-mean estimate from this model	1 step, N=8	8
best MSE from this model at sane cost	2 steps, N=32	64
sharpest single sample	16–32 steps, N=1 (64 is wasted: +0.001)	16–32

Implementation

z_t   = (1-t)·data + t·noise            t ~ sigmoid(N(0,1))   logit-normal, as SD3 uses
target velocity = noise - data
sampling: x ← x − v(x,t)·dt, Euler from t=1 (noise) to t=0 (data)
residual: data = z_{t+8} − z_t, and z_t is added back after integration

The predictor is reused verbatim as the velocity field. Two additions:

• Time conditioning — t_mlp, a 2-layer SiLU MLP over a sinusoidal embedding of t, added to the query tokens.
• Feature normalisation (feat_mean, feat_std buffers, from 501,760 patch tokens at the start of training) — see the geometry table above for why. The buffers ship in the checkpoint and eval_flow.py asserts they loaded; defaulting silently to (0, 1) does not error, it just samples garbage.

Usage

import torch
from transformers import AutoModel
from dino_dynamics import DinoDynamics

dino  = AutoModel.from_pretrained("facebook/dinov2-small")
model = DinoDynamics(dino, vae=None, n_views=2, action_dim=7, state_dim=9,
                     history=2, flow=True, residual=True).cuda()   # BOTH flags required
model.load_state_dict(torch.load("dino_step_60000_residual.pt"), strict=False)
model.eval()

batch = {"context": ctx,       # (B, 2, 2, 3, 224, 224)  frames t-8 and t, both cameras
         "action":  actions,   # (B, 8, 7)   delta OSC_POSE chunk
         "state":   state}     # (B, 9)      observation.state[:9]

pred = model(batch, flow_steps=20)                  # (B, 2, 196, 384) -- one sampled future

with torch.no_grad():                               # uncertainty across samples
    z = model.encode_views(batch["context"])
    s = torch.stack([model.sample(z, batch["action"], batch["state"], 2) for _ in range(32)])
mean, uncertainty = s.mean(0), s.std(0)

Five things that will bite if missed:

1. residual=True must match the checkpoint — the single easiest way to misuse these weights. The two files have identical tensor sets, so a mismatched flag produces no missing key, no unexpected key, and no error; it loads cleanly and silently returns wrong predictions. Symptom seen in the wild: inferred-noise std ≈ 3.7 instead of ≈ 1.0, and a round-trip error that fails its gate while the absolute checkpoint passes comfortably.

   Don't pass the flag by hand — infer it:

   from dino_dynamics import DinoDynamics, flags_from_checkpoint
   sd    = torch.load(path, weights_only=True)
   model = DinoDynamics(dino, vae=None, ..., **flags_from_checkpoint(sd))
   model.load_state_dict(sd, strict=False)
   

   The signal is feat_mean, which is fitted on whatever the model emits — a delta is centred, absolute DINOv2 features are not:

   checkpoint	feat_mean.abs().max()	feat_std.mean()
   dino_step_60000.pt	13.3128	1.7999
   dino_step_60000_residual.pt	0.0990	1.0703

   A 134× separation, so the 1.0 threshold is not delicate. Use feat_mean, not feat_std — the latter separates by only 1.7× and is not a safe discriminator. eval_flow.py now asserts the constructed flags match the checkpoint's buffers and fails loudly on a mismatch.

2. flow=True at construction. Without it the model builds no t_mlp, load reports it missing, and forward takes the deterministic path. Load with strict=False but check the missing keys all start with dino./vae..

3. model(batch) already integrates the ODE. In flow mode forward calls sample. Calling predict_from_features directly gives you a velocity, not a prediction.

4. It is stochastic. Two calls give two different futures. Seed if you need determinism, and use common random numbers (one fixed noise draw, reused) when comparing candidate actions — otherwise sampling noise swamps the action effect.

5. flow_steps and N multiply. See the budget table above.

eval_flow.py reproduces every table here; --steps and --means both accept lists and sweep with the model loaded once.

Training

--flow --flow-steps 20 --flow-stat-batches 40 [--residual]
--steps 60000 --batch-size 32 --lr 5e-4 --decoder-lr 3e-4
--history 2 --history-stride 8 --horizon 8 --action-tokens 1 --state-dim 9 --val-episodes 20

13.2M predictor + t_mlp, 12.7 GB peak, ~0.21 s/step, ~3.5 h on one B200. Data: the same three-dataset mixture as the baseline, 251,811 train / 15,889 val windows.

The val/feat_mse logged during training is single-sample at 20 steps and is not comparable to an MSE run's curve.

A correction about the motivation

This work was launched on an internal measurement that the MSE head's sharpness was 0.866, the figure quoted in the header comments of dino_dynamics.py and run_flow.sh. The evaluation script published here measures 0.925 on the same checkpoint. The 0.866 came from a different, unversioned measurement that this script does not reproduce, so treat 0.925 — with the definition above, from code you can run — as the real one. The over-smoothing the experiment set out to fix was less severe than believed, which is much of why the payoff is small.

Limitations

• Worse than the MSE head at every setting, and at N≥8 worse on sharpness too. If you want one number from this repo, it is that the substitution did not pay off.
• Sharpness is not sample quality. std(pred)/std(tgt) ≈ 1 says the prediction is not blurred; it does not say the detail is correct. Nothing here verifies that sampled futures are individually plausible — that needs a decoder study or a planning evaluation, neither of which was run. The +0.043 at N=1 may not be a usable gain.
• No planning evaluation. Whether sharper-but-higher-MSE features help or hurt CEM/MPC is open.
• Single-step, single seed, one task, inheriting every baseline limitation.
• The step/budget sweeps are on the residual checkpoint only; the absolute one was swept at 20 steps.
• Normalisation was added after a first attempt without it; only the normalised runs are published.