You can deploy this model and start earning money today!

chomeed/
mimicgen_threading_d0_224x224_dinowm

Name: mimicgen_threading_d0_224x224_dinowm
Availability: InStock

Views

No views yet

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 five 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
square_d0	0.661	1.806	22.65	18.13	0.9019	0.8900	0.1558	0.1144
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

Every task beats copy-the-current-frame on feature MSE and PSNR. LPIPS loses on every task and is decoder-bound in each (the oracle row is already worse than copy). square_d1 is the weakest and the only task whose 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.

python

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.

chomeed/mimicgen_threading_d0_224x224_dinowm