mimicgen_threading_d0_224x224_dinowm_heavy

DINO-WM multi-view dynamics — MimicGen threading_d0, scaled predictor

A larger-predictor variant of chomeed/mimicgen_threading_d0_224x224_dinowm. Same task, same data, same frozen facebook/dinov2-small encoder, same 8-action / 8-frame jump — the predictor is scaled from 13.2M → 35.6M and trained with warmup + cosine decay and gradient clipping.

This is the step-55,000 checkpoint, not the end of the run. That is deliberate and is the most important thing to know about this model — see Pick the checkpoint carefully.

Result

Scored against the published baseline on one common set of 1,600 held-out windows, with a paired per-window bootstrap (the same window scores both models, so window difficulty cancels):

−5.29% feature MSE, 95% CI on the paired difference [+0.0201, +0.0259] — excludes zero. Better on 67.6% of individual windows.

That is the whole gain: 2.7× the predictor parameters buys ~5%. Judge this model on feature MSE, as with the baseline; the pixel columns are decoder-bound and cannot separate the two.

Pick the checkpoint carefully

Every checkpoint of this run, same 1,600-window paired eval:

* measured on a 1,920-window sweep; the 1,600-window figure quoted above is 0.4126 / −5.29%. Both put 55k at the top; the difference is the window count.

Two consequences:

1. The run overfits after ~55–75k. By step 100,000 — the end of training — it is only +0.85% over baseline, i.e. essentially back to where it started. Training longer destroyed the gain. Treat 55k–75k as a plateau of ~+4.3% and 55k as its best point; the 55k-vs-70k gap (0.0044) is only ~1.5× the paired CI half-width, so do not over-read the exact peak.

2. The in-training val/feat_mse cannot select the checkpoint. It reported step 70k as best. It sees only --val-batches 4 = 128 windows out of 15,889, and its step-to-step wobble is ±0.02–0.03 with no learning in between — the same size as the differences between checkpoints. Its copy baseline reads 1.3961 against the published run's 1.5545 for the same reason. Any checkpoint chosen from that curve is chosen from noise.

Architecture

Identical to the baseline except the predictor:

• Encoder: facebook/dinov2-small (ViT-S/14, 384-d), frozen. 224→196, 14×14 = 196 patch tokens per view per frame, CLS dropped.
• Predictor: 12-layer pre-norm transformer, 12 heads, MLP 3072 (baseline: 6/6/2048) over [2 frames × 2 views × 196 patches] + [action, state] + [2 × 196 query tokens] = 1178 tokens. The 8-step action chunk is still one token (Linear(8×7 → 384)), as in the baseline.
• Loss: plain MSE in feature space.
• Decoder: 0.2M feature→SD3-latent head, trained on detached features, visualisation only.

35.6M trainable predictor + 0.2M decoder head. 23.4 GB peak at batch 32.

Training

--depth 12 --heads 12 --mlp-dim 3072
--steps 100000 --batch-size 32 --lr 5e-4 --decoder-lr 3e-4
--warmup 2000 --cosine --lr-min-ratio 0.05 --grad-clip 1.0
--history 2 --history-stride 8 --state-dim 9 --val-episodes 20

Note this bundles four changes against the published recipe — capacity, 100k vs 60k steps, cosine schedule, and gradient clipping — so the −5.29% is not attributable to capacity alone. A capacity-only ablation at the published recipe is a separate run.

Training data

• chomeed/mimicgen_threading_d0_224x224_mtdit_flow_55k_success
• chomeed/mimicgen_threading_d0_224x224_mtdit_flow_55k_failure
• chomeed/mimicgen_threading_d0_224x224

251,811 train windows, 15,889 validation windows — identical to the baseline. Cameras agentview and eye_in_hand at 224×224. Robot state is observation.state[:9] = eef_pos(3) + eef_quat(4) + gripper_qpos(2); MimicGen packs object state into the tail of observation.state (37D = 9 robot + 28 object) and only the leading 9 robot dims are read.

Usage

import torch
from diffusers import AutoencoderKL
from transformers import AutoModel
from dino_dynamics import DinoDynamics

dino = AutoModel.from_pretrained("facebook/dinov2-small")
vae  = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-3.5-medium",
                                     subfolder="vae", torch_dtype=torch.bfloat16)
model = DinoDynamics(dino, vae, n_views=2, action_dim=7, state_dim=9, history=2,
                     depth=12, heads=12, mlp_dim=3072).cuda()          # <- not the defaults
model.load_state_dict(torch.load("dino_step_55000.pt"), strict=False)  # dino./vae. keys absent
model.eval()

pred = model({"context": ctx, "action": actions, "state": state})

depth=12, heads=12, mlp_dim=3072 must be passed explicitly — the constructor defaults are the baseline's 6/6/2048 and will fail to load these weights.

eval_heavy_vs_baseline.py reproduces the comparison table above.

Limitations

• ~5% is the whole story. The gain is real and statistically clean, but small for 2.7× the parameters, and the run's peak-then-decay shape suggests this is data-limited rather than capacity-limited at 251k windows.
• Not a clean capacity ablation — four changes at once (see Training).
• Inherits every limitation of the baseline: single-step only (one 8-action jump, multi-step rollout compounding untested), no planning evaluation, soft reconstructions, single seed, one task.
• Reproduction note: an independent re-run of the published recipe under a different seed landed 0.9–1.4% worse than the published checkpoint (CI excludes zero), so treat ~1% as the seed noise floor on this task and disregard any claimed difference smaller than that.