chomeed/mimicgen_threading_d0_224x224_dinowm

DINO-WM multi-view dynamics — MimicGen threading_d0

A DINO-WM style world model for two-camera robot manipulation. It predicts future DINOv2 patch features, not pixels:

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

One model step advances 8 environment actions (0.4 s at 20 fps), as a single deterministic forward pass — no diffusion, no sampling loop.

Results

Held-out episodes (20 per source dataset, 15,889 windows). The bar is copy-the-current-frame, which is strong because t+8 is only 0.4 s ahead.

	this model	copy current frame	oracle (decoder ceiling)
feature MSE	0.4987	1.5545	—
PSNR	24.12	20.17	25.54
SSIM	0.9213	0.9156	0.9263
LPIPS	0.1169	0.0771	0.1011

Latent error is 68% below the baseline.

The LPIPS column is decoder-bound, not a dynamics failure. Decoding ground-truth features through the same head gives LPIPS 0.1011 — already worse than copying. No improvement in prediction can win that column through this 0.2M-parameter head. Judge this model on feature MSE.

All four tasks

Same architecture, same hyperparameters, 60k steps each.

task	feat MSE	copy	PSNR	copy	SSIM	copy	LPIPS	copy
coffee_d0	0.457	1.451	23.25	18.06	0.9167	0.9000	0.1185	0.0977
threading_d0	0.499	1.554	24.12	20.17	0.9213	0.9156	0.1169	0.0771
square_d1	0.670	1.428	19.82	16.99	0.8294	0.8577	0.2424	0.1329
hammer_cleanup_d1	0.655	1.886	20.38	15.77	0.8248	0.8233	0.2828	0.1730

Relative latent error (feat/copy) is 0.31 coffee, 0.32 threading, 0.47 square, 0.35 hammer. Square is the hardest, and the only task where SSIM lands below copy — its oracle is below copy too, so both the decoder and the prediction contribute there.

Architecture

• Encoder: facebook/dinov2-small (ViT-S/14, 384-d), frozen. Images 224→196, giving 14×14 = 196 patch tokens per view per frame; CLS dropped.
• Predictor: 6-layer pre-norm transformer, 6 heads, MLP 2048, over [2 frames × 2 views × 196 patches] + [action, state] + [2 × 196 query tokens] = 1178 tokens, full bidirectional attention. Identity is additive: pos_emb + view_emb + frame_emb.
• Loss: plain MSE in feature space.
• Decoder (visualisation only, trained on detached features): features → SD3 VAE latent → frozen SD3.5 VAE → RGB.

Trainable 13.4M. 12.5 GB peak at batch 32, ~96 samples/s on one RTX 5090, ~5.5 h per task.

Differences from the reference implementation

• Multi-view: both cameras share one sequence with learned view embeddings and are predicted jointly. The reference is single-view.
• Explicit query tokens: the reference feeds frames 0..N−1 with full attention and scores against frames 1..N, so every position but the last has its target visible in the input.
• Proprio is input-only: the reference predicts it and uses it in the planning objective (loss_visual + alpha * loss_proprio). This checkpoint is not directly usable with that objective without adding a state head.

Ablations

Measured on coffee_d0, 10k steps, all arms scored on one common window set via eval_common.py (per-run validation sets differ because history length changes window validity):

arm	feature MSE	Δ
history 3	0.5252	−0.6%
history 2 (this config)	0.5282	—
no state conditioning	0.5386	+2.0%
history 1	0.5419	+2.5%
no action conditioning	0.5748	+8.8%

Removing the action chunk hurts most — the model is genuinely action-conditioned, not extrapolating visual motion. History saturates by 2 frames.

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. 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 inside observation.state (37D = 9 robot + 28 object); only the leading 9 robot dims are read, and observation.object / observation.sim_state are never touched. assert_object_free in sd3_dynamics.py enforces this at every training start.

Usage

The checkpoint holds only the trainable parts (13.4M params, fp32); DINOv2 and the SD3 VAE are downloaded separately and stay frozen.

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).cuda()
model.load_state_dict(torch.load("dino_step_60000.pt"), strict=False)  # dino./vae. keys absent
model.eval()

pred = model({"context": ctx,      # (B, 2, 2, 3, 224, 224), frames t-8 and t, both cameras
              "action":  actions,  # (B, 8, 7)
              "state":   state})   # (B, 9)
img = model.features_to_image(pred)   # inspection only

predict_from_features skips re-encoding fixed context, for planning loops.

Limitations

• Single-step only. Trained for one 8-action jump; multi-step rollout compounding is untested.
• No planning evaluation. The paper's headline metric is task success under CEM/MPC in the simulator. Good latent MSE does not guarantee good planning.
• Reconstructions are soft. DINOv2 features were never trained to be invertible.
• Single seed, one task per checkpoint. No cross-embodiment claims.
• Action-gradient cost: ~123 ms per batch of 32 (fp32, unoptimised), OOM at batch 256 on 32 GB. A 3×256 MLP state-space model computes the same gradient ~20,000× faster per candidate — if your planning objective only needs state, this is the wrong tool.

Dependencies and licensing

No third-party weights are included, but running it downloads facebook/dinov2-small (CC-BY-NC 4.0) and the stabilityai/stable-diffusion-3.5-medium VAE (Stability AI Community License, gated). Those licences govern those components.