---
tags: [robotics, world-model, dynamics-model, dinov2, mimicgen]
---

# DINO-WM multi-view dynamics — MimicGen coffee_d0

A [DINO-WM](https://arxiv.org/abs/2411.04983) 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, 16,123 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.4573** | 1.4514 | — |
| PSNR | **23.25** | 18.06 | 24.80 |
| SSIM | 0.9167 | 0.9000 | 0.9218 |
| LPIPS | 0.1185 | 0.0977 | 0.1023 |

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.1023 — 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_coffee_d0_224x224_mtdit_flow_60k_success`
- `chomeed/mimicgen_coffee_d0_224x224_mtdit_flow_60k_failure`
- `chomeed/mimicgen_coffee_d0_224x224`

250,306 train windows, 16,123 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`** (66D = 9 robot + 57 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.