DINO-WM multi-view dynamics, MimicGen coffee_d0, 60k steps

README.md

@@ -0,0 +1,142 @@
+---
+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). The forecast is a single
+deterministic forward pass — no diffusion, no sampling loop.
+
+## Results
+
+Held-out episodes (20 per source dataset, 16k windows). The bar to beat is **copy-the-current-frame**,
+which is strong here because t+8 is only 0.4 s ahead.
+
+| | this model | copy current frame |
+|---|---|---|
+| feature MSE | **0.457** | 1.451 |
+| PSNR | **23.25** | 18.06 |
+| SSIM | **0.9167** | 0.9014 |
+| LPIPS | 0.1185 | **0.0977** |
+
+Latent error is 68% below the baseline. Converged by ~step 51k of 60k; no overfitting
+(train 0.409 vs val 0.457).
+
+**The LPIPS loss is the decoder, not the dynamics.** Decoding *ground-truth* features through the
+same head gives LPIPS 0.1023 — already worse than copying. The prediction sits within 0.016 of that
+ceiling, so no dynamics improvement can win LPIPS through this decoder. Same story on PSNR
+(23.25 vs a 24.80 ceiling) and SSIM (0.9167 vs 0.9218).
+
+For context, a Stable Diffusion 3.5 canvas-conditioned model trained on the same data reached
+PSNR 19.01 / SSIM 0.9061 / LPIPS 0.0685 in 7 h without converging. This model reaches PSNR 23.25 in
+5.6 h at half the memory. The split — this model wins L2 metrics, the diffusion model wins LPIPS —
+is the expected regression-vs-diffusion trade: an MSE-trained predictor outputs the conditional
+mean (accurate, blurry), diffusion samples a mode (sharp, less L2-accurate).
+
+## Ablations
+
+10k steps each, all scored on one common window set (see `eval_common.py` — per-run validation sets
+differ because history length changes window validity, which makes raw numbers incomparable).
+
+| arm | feature MSE | Δ vs baseline |
+|---|---|---|
+| history 3 | 0.5252 | −0.6% |
+| **history 2 (this model)** | **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 costs the most, which is the check that matters: the model is genuinely
+action-conditioned rather than extrapolating visual motion. History helps slightly and saturates by
+2 frames. History 1 is 28% faster and 2.3 GB lighter for 2.5% more error — a good trade when iterating.
+
+## Architecture
+
+- **Encoder**: `facebook/dinov2-small` (ViT-S/14, 384-d), **frozen**. Images resized 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 token, state token] + [2 × 196 query tokens]`
+  = 784 context + 2 conditioning + 392 query. Camera identity is a learned view embedding; the
+  target is read from explicit query tokens.
+- **Loss**: plain MSE in feature space.
+- **Decoder** (visualisation/metrics only, trained on *detached* features): predicted features →
+  SD3 VAE latent → frozen SD3.5 VAE → RGB.
+
+Trainable: **13.4M**. Peak VRAM 12.5 GB at batch 32, 96 samples/s on one RTX 5090.
+
+### Differences from the reference implementation
+
+- **Multi-view.** Patch tokens of both cameras share one sequence with learned view embeddings, and
+  both views are predicted jointly. The reference is single-view.
+- **Explicit query tokens.** The reference feeds frames 0..N−1 with full (non-causal) 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 proprio as part of `z` and uses it in the
+  planning objective (`loss_visual + alpha * loss_proprio`). Here state conditions the prediction
+  but is not a prediction target — so this checkpoint is **not** directly usable with that objective.
+
+## Training data
+
+`chomeed/mimicgen_coffee_d0_224x224_{success,failure}` and `chomeed/mimicgen_coffee_d0_224x224`
+(1,200 episodes total, 285k frames, 20 fps). Cameras `agentview` and `eye_in_hand` at 224×224.
+Robot state is `observation.state[:9]` = eef_pos(3) + eef_quat(4) + gripper_qpos(2).
+`observation.object` and `observation.sim_state` are never read.
+
+## Usage
+
+The checkpoint holds only the trainable parts (13.4M params, fp32) — the DINOv2 encoder and 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)
+# pred: (B, 2, 196, 384) DINOv2 features at t+8
+img = model.features_to_image(pred)   # (B, 2, 3, 224, 224), for inspection only
+```
+
+For planning, `predict_from_features` skips re-encoding fixed context.
+
+## Limitations
+
+- **Single-step only.** Trained for one 8-action jump; multi-step rollout error compounding is
+  untested.
+- **No planning evaluation.** The paper's headline metric is task success under CEM/MPC in the
+  simulator. This model is validated on prediction error only — good latent MSE does not
+  guarantee it plans well.
+- **Reconstructions are soft.** DINOv2 features were never trained to be invertible; the decoder is
+  a 0.2M-parameter head. Use images for sanity checks, not as output.
+- **Single seed**, one task (coffee_d0). No cross-task or cross-embodiment claims.
+- **Action-gradient cost**: ~123 ms for a batch of 32 (fp32, unoptimised), and OOM at batch 256 on
+  32 GB. A 3×256 MLP state-space model does the same gradient ~20,000× faster per candidate. If your
+  planning objective only needs state, this model is the wrong tool.
+
+## Dependencies and licensing
+
+The checkpoint contains no third-party weights, 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 — you must accept its terms). Their licences govern those
+components, not this checkpoint.

bench_action_grad.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+"""Cost of d(loss)/d(action) -- the inner loop of gradient-based planning.
+
+Three regimes:
+  dino-wm full    encode context frames, predict, backprop to the action chunk
+  dino-wm cached  context features precomputed once (what a planner actually does across
+                  CEM/GD iterations, since the context is fixed), predict + backprop only
+  mlp 3x256       a 3-layer 256-wide state-space dynamics model, (state, action) -> next state
+
+    policy_learning/lerobot/.venv/bin/python dynamics_model/bench_action_grad.py
+"""
+
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel
+
+from dino_dynamics import DINO_ID, DinoDynamics
+from sd3_dynamics import HORIZON
+
+V, ACT, ST, HIST, SZ = 2, 7, 9, 2, 224
+BATCHES = (1, 32, 256, 1024)
+REPS = 10
+
+
+def timeit(fn, reps=REPS, warmup=3):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for _ in range(reps):
+        fn()
+    torch.cuda.synchronize()
+    return (time.time() - t0) / reps * 1000  # ms
+
+
+def main():
+    dev = "cuda"
+    torch.manual_seed(0)
+    dino = AutoModel.from_pretrained(DINO_ID, dtype=torch.float32)
+    model = DinoDynamics(dino, None, V, ACT, ST, history=HIST).to(dev).eval()
+    for p in model.parameters():
+        p.requires_grad_(False)  # planning optimises actions, not weights
+
+    mlp = nn.Sequential(
+        nn.Linear(ST + HORIZON * ACT, 256), nn.ReLU(),
+        nn.Linear(256, 256), nn.ReLU(),
+        nn.Linear(256, 256), nn.ReLU(),
+        nn.Linear(256, ST),
+    ).to(dev).eval()
+    for p in mlp.parameters():
+        p.requires_grad_(False)
+    print(f"dino-wm predictor params: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M "
+          f"(incl. frozen DINOv2)   mlp params: {sum(p.numel() for p in mlp.parameters()) / 1e3:.0f}k")
+
+    print(f"\n{'batch':>6} {'dino full':>12} {'dino cached':>13} {'mlp 3x256':>11} "
+          f"{'cached/mlp':>11} {'full/cached':>12}")
+    for b in BATCHES:
+        ctx = torch.rand(b, HIST, V, 3, SZ, SZ, device=dev)
+        state = torch.randn(b, ST, device=dev)
+        goal = torch.randn(b, V, model.n_patches, model.dim, device=dev)
+        goal_s = torch.randn(b, ST, device=dev)
+
+        def dino_full():
+            act = torch.randn(b, HORIZON, ACT, device=dev, requires_grad=True)
+            loss = F.mse_loss(model({"context": ctx, "action": act, "state": state}), goal)
+            return torch.autograd.grad(loss, act)[0]
+
+        # cached: encode once, reuse the features across every planner iteration
+        with torch.no_grad():
+            z = model.encode_views(ctx)
+
+        def dino_cached():
+            act = torch.randn(b, HORIZON, ACT, device=dev, requires_grad=True)
+            loss = F.mse_loss(model.predict_from_features(z, act, state), goal)
+            return torch.autograd.grad(loss, act)[0]
+
+        def mlp_grad():
+            act = torch.randn(b, HORIZON, ACT, device=dev, requires_grad=True)
+            loss = F.mse_loss(mlp(torch.cat([state, act.flatten(1)], 1)), goal_s)
+            return torch.autograd.grad(loss, act)[0]
+
+        def safe(fn):
+            try:
+                return timeit(fn)
+            except torch.OutOfMemoryError:
+                torch.cuda.empty_cache()
+                return float("nan")
+
+        t_full, t_cached, t_mlp = safe(dino_full), safe(dino_cached), safe(mlp_grad)
+        fmt = lambda x, w, p: ("OOM" if x != x else f"{x:.{p}f}ms").rjust(w)
+        ratio = lambda x, y: "-" if (x != x or y != y) else f"{x / y:.0f}x"
+        print(f"{b:>6} {fmt(t_full, 12, 2)} {fmt(t_cached, 13, 2)} {fmt(t_mlp, 11, 3)} "
+              f"{ratio(t_cached, t_mlp):>11} {ratio(t_full, t_cached):>12}")
+        del ctx, z
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

dino_dynamics.py

@@ -0,0 +1,340 @@
+#!/usr/bin/env python
+"""DINO-WM style multi-view dynamics: predict future DINOv2 patch features, not pixels.
+
+    {I_{t-8}, I_t}^1..V, a[t:t+8], s_t  ->  DINOv2 features of {I_{t+8}}^1..V
+
+Follows DINO-WM (arXiv:2411.04983): a frozen DINOv2 encoder, a small ViT predictor trained with
+plain MSE in feature space, and a decoder trained separately on detached features for visualisation
+only. No diffusion, no sampling -- the forecast is deterministic.
+
+Differences from the reference implementation, both deliberate:
+  * multi-view: patch tokens of every camera share one sequence, each with a learned view embedding.
+  * the target is read from explicit query tokens. The reference feeds frames 0..N-1 with full
+    (non-causal) attention and scores them against frames 1..N, so every position except the last
+    has its own target visible in the input; only the final position is a true forecast.
+
+Pixel metrics need a decoder: predicted DINO features -> SD3 VAE latent -> frozen SD3 VAE -> image.
+Reusing the SD3 VAE avoids training a pixel decoder from scratch.
+
+    P=policy_learning/lerobot/.venv/bin/python
+    PYTHONUNBUFFERED=1 $P dino_dynamics.py --data <success> <failure> --steps 20000
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sd3_dynamics import (
+    HORIZON,
+    build_datasets,
+    dataset_features,
+    object_state_keys,
+    per_camera_metrics,
+    reconstruction_table,
+)
+
+DINO_ID = "facebook/dinov2-small"  # ViT-S/14, 384-dim, as in the reference config
+DINO_MEAN = (0.485, 0.456, 0.406)
+DINO_STD = (0.229, 0.224, 0.225)
+
+
+class DinoDynamics(nn.Module):
+    """Frozen DINOv2 encoder + ViT predictor over [history x views] patch tokens."""
+
+    def __init__(
+        self,
+        dino,
+        vae,
+        n_views,
+        action_dim,
+        state_dim,
+        history=2,
+        depth=6,
+        heads=6,
+        mlp_dim=2048,
+        horizon=HORIZON,
+        dino_size=196,
+        image_size=224,
+        use_action=True,
+        use_state=True,
+    ):
+        super().__init__()
+        self.dino = dino.eval().requires_grad_(False)
+        self.vae = vae.eval().requires_grad_(False) if vae is not None else None
+        d = dino.config.hidden_size
+        self.dim, self.n_views, self.history = d, n_views, history
+        # ablations zero the embedding rather than drop the token, so token count and
+        # model capacity are unchanged and the only variable is the information itself
+        self.use_action, self.use_state = use_action, use_state
+        self.patch = dino.config.patch_size
+        self.dino_size = dino_size  # multiple of patch size; 196 -> 14x14 = 196 tokens
+        self.side = dino_size // self.patch
+        self.n_patches = self.side**2
+
+        self.view_emb = nn.Parameter(torch.zeros(n_views, d))
+        self.frame_emb = nn.Parameter(torch.zeros(history + 1, d))  # last entry = target frame
+        self.pos_emb = nn.Parameter(torch.zeros(self.n_patches, d))
+        self.query = nn.Parameter(torch.randn(self.n_patches, d) * 0.02)
+        self.act_tok = nn.Linear(horizon * action_dim, d)
+        self.state_tok = nn.Linear(state_dim, d)
+        layer = nn.TransformerEncoderLayer(
+            d, heads, mlp_dim, dropout=0.0, batch_first=True, norm_first=True, activation="gelu"
+        )
+        self.predictor = nn.TransformerEncoder(layer, depth)
+        self.norm = nn.LayerNorm(d)
+
+        # feature -> SD3 latent head, for pixel metrics only; trained on detached features
+        if vae is not None:
+            c = vae.config.latent_channels
+            # the head upsamples the patch grid 2x, so it must land exactly on the VAE latent grid
+            lat_side = image_size // 2 ** (len(vae.config.block_out_channels) - 1)
+            if self.side * 2 != lat_side:
+                raise ValueError(
+                    f"patch grid {self.side}x2={self.side * 2} != VAE latent grid {lat_side} "
+                    f"(image_size={image_size}, dino_size={dino_size}); pick dino_size = "
+                    f"{lat_side // 2 * self.patch}"
+                )
+            self.to_latent = nn.Sequential(
+                nn.Conv2d(d, 4 * c, 3, padding=1), nn.GELU(), nn.PixelShuffle(2), nn.Conv2d(c, c, 3, padding=1)
+            )
+
+    # -- features --
+    @torch.no_grad()
+    def encode(self, imgs):
+        """(N,3,H,W) in [0,1] -> (N, n_patches, d) DINOv2 patch tokens."""
+        x = F.interpolate(imgs, self.dino_size, mode="bilinear", antialias=True, align_corners=False)
+        mean = torch.tensor(DINO_MEAN, device=x.device).view(1, 3, 1, 1)
+        std = torch.tensor(DINO_STD, device=x.device).view(1, 3, 1, 1)
+        out = self.dino(pixel_values=((x - mean) / std).to(self.dino.dtype)).last_hidden_state
+        return out[:, 1:].float()  # drop CLS
+
+    def encode_views(self, imgs):
+        """(B,R,V,3,H,W) -> (B,R,V,P,d)."""
+        b, r, v = imgs.shape[:3]
+        z = self.encode(imgs.flatten(0, 2))
+        return z.reshape(b, r, v, self.n_patches, self.dim)
+
+    # -- prediction --
+    def forward(self, batch):
+        """-> predicted target features (B,V,P,d)."""
+        return self.predict_from_features(
+            self.encode_views(batch["context"]), batch["action"], batch["state"]
+        )
+
+    def predict_from_features(self, z, action, state):
+        """Context features (B,R,V,P,d) -> predicted target features (B,V,P,d).
+
+        Split out from forward so a planner can encode the (fixed) context once and reuse it
+        across every CEM/gradient iteration; only this part is on the action-gradient path.
+        """
+        b, r, v = z.shape[:3]
+        z = z + self.view_emb[None, None, :, None] + self.pos_emb[None, None, None]
+        z = z + self.frame_emb[:r][None, :, None, None]
+        tokens = z.reshape(b, r * v * self.n_patches, self.dim)
+
+        q = self.query[None, None] + self.view_emb[None, :, None] + self.pos_emb[None, None]
+        q = q + self.frame_emb[-1][None, None, None]
+        q = q.expand(b, v, self.n_patches, self.dim).reshape(b, v * self.n_patches, self.dim)
+
+        act = self.act_tok(action.flatten(1)) * float(self.use_action)
+        st = self.state_tok(state) * float(self.use_state)
+        cond = torch.stack([act, st], 1)  # (B,2,d)
+        out = self.predictor(torch.cat([tokens, cond, q], 1))
+        return self.norm(out[:, -v * self.n_patches :]).reshape(b, v, self.n_patches, self.dim)
+
+    def loss(self, batch):
+        pred = self(batch)
+        with torch.no_grad():
+            tgt = self.encode_views(batch["future"][:, None])[:, 0]  # (B,V,P,d)
+        return F.mse_loss(pred, tgt), tgt
+
+    # -- pixels, for metrics and W&B only --
+    def features_to_image(self, feats):
+        """(B,V,P,d) -> (B,V,3,H,W) in [0,1], through the frozen SD3 VAE."""
+        b, v = feats.shape[:2]
+        x = feats.reshape(b * v, self.side, self.side, self.dim).permute(0, 3, 1, 2)
+        lat = self.to_latent(x)
+        z = lat.to(self.vae.dtype) / self.vae.config.scaling_factor + self.vae.config.shift_factor
+        img = (self.vae.decode(z).sample / 2 + 0.5).clamp(0, 1).float()
+        return img.reshape(b, v, *img.shape[1:])
+
+    def decoder_loss(self, batch, tgt_feats):
+        """Train the feature->latent head against the true frames. Detached from the predictor."""
+        with torch.no_grad():
+            imgs = batch["future"].flatten(0, 1)
+            z = self.vae.encode(imgs.to(self.vae.dtype) * 2 - 1).latent_dist.mode()
+            z = ((z - self.vae.config.shift_factor) * self.vae.config.scaling_factor).float()
+        b, v = batch["future"].shape[:2]
+        x = tgt_feats.detach().reshape(b * v, self.side, self.side, self.dim).permute(0, 3, 1, 2)
+        return F.mse_loss(self.to_latent(x), z)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--data", nargs="+", required=True)
+    ap.add_argument("--dino", default=DINO_ID)
+    ap.add_argument("--sd3", default="stabilityai/stable-diffusion-3.5-medium", help="VAE for decoding")
+    ap.add_argument("--state-key", default="observation.state")
+    ap.add_argument("--state-dim", type=int, default=9)
+    ap.add_argument("--image-size", type=int, nargs=2, default=(224, 224))
+    ap.add_argument("--history", type=int, default=2, help="context frames incl. current")
+    ap.add_argument("--history-stride", type=int, default=HORIZON, help="default = horizon: evenly spaced")
+    ap.add_argument("--no-action", action="store_true", help="ablation: zero the action conditioning")
+    ap.add_argument("--no-state", action="store_true", help="ablation: zero the state conditioning")
+    ap.add_argument("--depth", type=int, default=6)
+    ap.add_argument("--heads", type=int, default=6)
+    ap.add_argument("--batch-size", type=int, default=32)
+    ap.add_argument("--steps", type=int, default=20000)
+    ap.add_argument("--lr", type=float, default=5e-4, help="reference predictor_lr")
+    ap.add_argument("--decoder-lr", type=float, default=3e-4)
+    ap.add_argument("--log-every", type=int, default=100)
+    ap.add_argument("--val-episodes", type=int, default=20)
+    ap.add_argument("--val-every", type=int, default=500)
+    ap.add_argument("--val-batches", type=int, default=4)
+    ap.add_argument("--log-images", type=int, default=4)
+    ap.add_argument("--workers", type=int, default=8)
+    ap.add_argument("--out", default=str(Path(__file__).parent / "outputs"))
+    ap.add_argument("--wandb-project", default="sd3-dynamics")
+    ap.add_argument("--run-name", default=None)
+    args = ap.parse_args()
+
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(0)
+    image_size = tuple(args.image_size)
+    train_ds, val_ds, cameras = build_datasets(
+        args.data, image_size, args.state_key, args.state_dim, args.val_episodes,
+        args.history, args.history_stride,
+    )
+    feats = dataset_features(args.data[0])
+    assert args.state_key not in object_state_keys(feats), "object state must not reach the model"
+    action_dim = feats["action"]["shape"][0]
+    print(f"cameras={cameras} action_dim={action_dim} state={args.state_key}[:{args.state_dim}]")
+    print(f"excluded from the model: {object_state_keys(feats)}")
+    print(f"train windows={len(train_ds)} val windows={len(val_ds) if val_ds else 0}")
+
+    from diffusers import AutoencoderKL
+    from transformers import AutoModel
+
+    dino = AutoModel.from_pretrained(args.dino, torch_dtype=torch.float32)
+    vae = AutoencoderKL.from_pretrained(args.sd3, subfolder="vae", torch_dtype=torch.bfloat16)
+    model = DinoDynamics(
+        dino, vae, len(cameras), action_dim, args.state_dim, args.history, args.depth, args.heads,
+        image_size=image_size[0], use_action=not args.no_action, use_state=not args.no_state,
+    ).to(dev)
+    pred_params = [p for n, p in model.named_parameters() if p.requires_grad and not n.startswith("to_latent")]
+    dec_params = list(model.to_latent.parameters())
+    print(f"predictor params: {sum(p.numel() for p in pred_params) / 1e6:.1f}M  "
+          f"decoder head: {sum(p.numel() for p in dec_params) / 1e6:.1f}M  "
+          f"tokens: {args.history * len(cameras) * model.n_patches} ctx + {len(cameras) * model.n_patches} query")
+    opt = torch.optim.AdamW(
+        [{"params": pred_params, "lr": args.lr}, {"params": dec_params, "lr": args.decoder_lr}]
+    )
+
+    lpips = None
+    if val_ds is not None:
+        from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+        lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).to(dev)
+
+    run = None
+    if args.wandb_project:
+        import wandb
+
+        run = wandb.init(
+            project=args.wandb_project, name=args.run_name, config=vars(args) | {"cameras": cameras}
+        )
+
+    def loader(ds, shuffle):
+        return torch.utils.data.DataLoader(
+            ds, batch_size=args.batch_size, shuffle=shuffle, num_workers=args.workers,
+            drop_last=True, pin_memory=True,
+        )
+
+    def to_dev(b):
+        return {k: v.to(dev, non_blocking=True) for k, v in b.items()}
+
+    def log(d, step):
+        print(f"step {step}: " + " ".join(f"{k}={v:.4f}" for k, v in d.items() if isinstance(v, float)))
+        if run:
+            run.log(d, step=step)
+
+    @torch.no_grad()
+    def validate(step):
+        model.eval()
+        sums, n = {}, 0
+        for i, batch in enumerate(loader(val_ds, False)):
+            if i >= args.val_batches:
+                break
+            batch = to_dev(batch)
+            pred = model(batch)
+            tgt = model.encode_views(batch["future"][:, None])[:, 0]
+            cur = model.encode_views(batch["context"][:, -1:])[:, 0]
+            sums["feat_mse"] = sums.get("feat_mse", 0.0) + F.mse_loss(pred, tgt).item()
+            sums["copy_feat_mse"] = sums.get("copy_feat_mse", 0.0) + F.mse_loss(cur, tgt).item()
+            img = model.features_to_image(pred)
+            current = batch["context"][:, -1]
+            for k, v in per_camera_metrics(img, batch["future"], cameras, lpips).items():
+                sums[k] = sums.get(k, 0.0) + v
+            for k, v in per_camera_metrics(current, batch["future"], cameras, lpips).items():
+                sums[f"copy_{k}"] = sums.get(f"copy_{k}", 0.0) + v
+            # decoder ceiling: what the feature->pixel head gives on GROUND-TRUTH features
+            for k, v in per_camera_metrics(
+                model.features_to_image(tgt), batch["future"], cameras, lpips
+            ).items():
+                sums[f"oracle_{k}"] = sums.get(f"oracle_{k}", 0.0) + v
+            n += 1
+            if i == 0 and run:
+                k = min(args.log_images, img.shape[0])
+                run.log(
+                    {"val/reconstructions": reconstruction_table(
+                        current[:k], batch["future"][:k], img[:k], cameras)},
+                    step=step,
+                )
+        model.train()
+        log({f"val/{k}": v / max(n, 1) for k, v in sums.items()}, step)
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    step, running, rdec, t0 = 0, 0.0, 0.0, time.time()
+    model.train()
+    while step < args.steps:
+        for batch in loader(train_ds, True):
+            batch = to_dev(batch)
+            feat_loss, tgt = model.loss(batch)
+            dec_loss = model.decoder_loss(batch, tgt)
+            (feat_loss + dec_loss).backward()
+            opt.step()
+            opt.zero_grad(set_to_none=True)
+            running += feat_loss.item()
+            rdec += dec_loss.item()
+            step += 1
+            if step % args.log_every == 0:
+                dt = (time.time() - t0) / args.log_every
+                log(
+                    {
+                        "train/feat_mse": running / args.log_every,
+                        "train/dec_mse": rdec / args.log_every,
+                        "train/peak_gb": torch.cuda.max_memory_allocated() / 1e9 if dev == "cuda" else 0.0,
+                        "train/s_per_step": dt,
+                        "train/samples_per_s": args.batch_size / dt,
+                    },
+                    step,
+                )
+                running, rdec, t0 = 0.0, 0.0, time.time()
+            if val_ds is not None and step % args.val_every == 0:
+                validate(step)
+            if step % 5000 == 0 or step == args.steps:
+                torch.save(
+                    {k: v for k, v in model.state_dict().items() if not k.startswith(("dino.", "vae."))},
+                    out / f"dino_step_{step}.pt",
+                )
+            if step >= args.steps:
+                break
+
+
+if __name__ == "__main__":
+    main()

dino_step_60000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfaf37eada0a586159d373b61faa03c874cae5c76b8b5d0a1bb427e34d203f7d
+size 53675521

eval_common.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python
+"""Score every ablation on ONE common validation set.
+
+The training runs each used their own valid-window set (requiring more history invalidates more
+windows near episode starts), so their feat_mse values are not comparable. This evaluates every
+checkpoint on the windows valid for the largest history, feeding each model only the context
+frames it was trained with.
+
+    policy_learning/lerobot/.venv/bin/python dynamics_model/eval_common.py
+"""
+
+import torch
+import torch.nn.functional as F
+from diffusers import AutoencoderKL
+from transformers import AutoModel
+
+from dino_dynamics import DINO_ID, DinoDynamics
+from sd3_dynamics import build_datasets, dataset_features, per_camera_metrics
+
+DATA = [
+    "chomeed/mimicgen_coffee_d0_224x224_mtdit_flow_60k_success",
+    "chomeed/mimicgen_coffee_d0_224x224_mtdit_flow_60k_failure",
+    "chomeed/mimicgen_coffee_d0_224x224",
+]
+SD3 = "stabilityai/stable-diffusion-3.5-medium"
+COMMON_HISTORY = 3  # most restrictive -> its windows are a subset of every arm's
+BATCHES, BS = 8, 32
+
+ARMS = [  # name, checkpoint, history it was trained with, use_action, use_state
+    ("h1", "outputs/abl_h1/dino_step_10000.pt", 1, True, True),
+    ("h2 (baseline)", "outputs/dino_step_10000.pt", 2, True, True),
+    ("h3", "outputs/abl_h3/dino_step_10000.pt", 3, True, True),
+    ("noaction", "outputs/abl_noaction/dino_step_10000.pt", 2, False, True),
+    ("nostate", "outputs/abl_nostate/dino_step_10000.pt", 2, True, False),
+]
+
+
+def main():
+    dev = "cuda"
+    torch.manual_seed(0)
+    _, val, cameras = build_datasets(DATA, (224, 224), "observation.state", 9, 20, COMMON_HISTORY, 8)
+    print(f"common val windows: {len(val)}  cameras: {cameras}")
+    action_dim = dataset_features(DATA[0])["action"]["shape"][0]
+
+    dino = AutoModel.from_pretrained(DINO_ID, torch_dtype=torch.float32)
+    vae = AutoencoderKL.from_pretrained(SD3, subfolder="vae", torch_dtype=torch.bfloat16)
+    from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+    lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).to(dev)
+
+    # fixed batches, identical for every arm
+    dl = torch.utils.data.DataLoader(val, batch_size=BS, shuffle=False, num_workers=8, drop_last=True)
+    batches = []
+    for i, b in enumerate(dl):
+        if i >= BATCHES:
+            break
+        batches.append({k: v.to(dev) for k, v in b.items()})
+
+    print(f"\n{'arm':<15}{'feat_mse':>10}{'copy':>9}{'ratio':>8}{'psnr':>8}{'ssim':>8}{'lpips':>8}")
+    for name, ckpt, hist, use_act, use_st in ARMS:
+        model = DinoDynamics(
+            dino, vae, len(cameras), action_dim, 9, history=hist,
+            use_action=use_act, use_state=use_st,
+        ).to(dev)
+        sd = torch.load(ckpt, map_location=dev)
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+        assert not unexpected, f"{name}: unexpected keys {unexpected[:3]}"
+        assert all(k.startswith(("dino.", "vae.")) for k in missing), f"{name}: missing {missing[:3]}"
+        model.eval()
+
+        sums = {}
+        with torch.no_grad():
+            for b in batches:
+                b = dict(b, context=b["context"][:, -hist:])  # only the frames this arm was trained on
+                pred = model(b)
+                tgt = model.encode_views(b["future"][:, None])[:, 0]
+                cur = model.encode_views(b["context"][:, -1:])[:, 0]
+                sums["feat"] = sums.get("feat", 0.0) + F.mse_loss(pred, tgt).item()
+                sums["copy"] = sums.get("copy", 0.0) + F.mse_loss(cur, tgt).item()
+                for k, v in per_camera_metrics(
+                    model.features_to_image(pred), b["future"], cameras, lpips
+                ).items():
+                    if "mean" in k:
+                        sums[k] = sums.get(k, 0.0) + v
+        n = len(batches)
+        f, c = sums["feat"] / n, sums["copy"] / n
+        print(f"{name:<15}{f:>10.4f}{c:>9.4f}{f / c:>8.3f}"
+              f"{sums['mean/psnr'] / n:>8.2f}{sums['mean/ssim'] / n:>8.4f}{sums['mean/lpips'] / n:>8.4f}")
+        del model
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

sd3_dynamics.py

@@ -0,0 +1,595 @@
+#!/usr/bin/env python
+"""Multi-view SD3 dynamics model.
+
+    {I_{t-km..t}^1..V}, a[t:t+8], s_t  ->  {I_{t+8}^1..V}
+
+Conditioning follows Ctrl-World (arXiv:2510.10125): context frames and the noisy target live in
+ONE token sequence so they interact through full self-attention, instead of the target attending to
+frame features through the text/cross-attention pathway. Here that is done by tiling latents into a
+single canvas -- rows are timesteps (history..current, then the target), columns are camera views --
+and letting the unmodified MMDiT attend over it. Loss and sampling touch the target row only.
+
+An earlier design fed frames in as `encoder_hidden_states` only; it trained to a validation PSNR of
+7.9 dB against 19.0 dB for simply copying the current frame, i.e. worse than doing nothing.
+
+Env: policy_learning/lerobot/.venv (lerobot + diffusers + torchmetrics + wandb + peft).
+
+    P=policy_learning/lerobot/.venv/bin/python
+    $P sd3_dynamics.py --data <success> <failure> --inspect
+    PYTHONUNBUFFERED=1 $P sd3_dynamics.py --data <success> <failure> --steps 20000
+"""
+
+import argparse
+import json
+import math
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+HORIZON = 8  # action chunk length AND image lead time, per spec
+
+
+# ── data ──────────────────────────────────────────────────────────────────────
+
+
+def dataset_info(spec) -> tuple[str, Path | None, dict]:
+    """Accepts a local dataset root or an HF repo id -> (repo_id, root|None, info.json)."""
+    root = Path(spec)
+    if (root / "meta" / "info.json").exists():
+        return f"local/{root.name}", root, json.loads((root / "meta" / "info.json").read_text())
+    from huggingface_hub import hf_hub_download
+
+    return spec, None, json.loads(Path(hf_hub_download(spec, "meta/info.json", repo_type="dataset")).read_text())
+
+
+def dataset_features(spec) -> dict:
+    return dataset_info(spec)[2]["features"]
+
+
+def discover_cameras(features: dict) -> list[str]:
+    """Every RGB camera feature, deterministically ordered. No hard-coded names."""
+    cams = sorted(k for k, v in features.items() if v.get("dtype") in ("video", "image"))
+    if not cams:
+        raise ValueError(f"no image/video features found among {sorted(features)}")
+    return cams
+
+
+def object_state_keys(features: dict) -> list[str]:
+    """Simulator/object-state features. Present in MimicGen data; must never reach the model."""
+    return sorted(k for k in features if "object" in k.lower() or "sim_state" in k.lower())
+
+
+class MultiViewWindows(torch.utils.data.Dataset):
+    """(history+current frames, future frames, 8-action chunk, state) windows from one dataset."""
+
+    def __init__(
+        self,
+        spec,
+        cameras,
+        state_key,
+        state_dim,
+        image_size,
+        episodes=None,
+        horizon=HORIZON,
+        history=1,
+        stride=4,
+    ):
+        self.cameras, self.state_key, self.state_dim = cameras, state_key, state_dim
+        self.image_size, self.horizon = image_size, horizon
+        self.history, self.stride = history, stride
+        repo_id, root, info = dataset_info(spec)
+        fps = info["fps"]
+        # oldest history frame first, then the current frame, then the target
+        self.offsets = [-(history - 1 - i) * stride for i in range(history)] + [horizon]
+        dt = {"action": [i / fps for i in range(horizon)]}
+        dt |= {c: [o / fps for o in self.offsets] for c in cameras}
+        self.ds = LeRobotDataset(repo_id, root=root, delta_timestamps=dt)
+
+        # episode boundaries: every offset, past and future, must land in the same episode
+        ep = np.asarray(self.ds.hf_dataset.select_columns("episode_index")["episode_index"])
+        back, fwd = (history - 1) * stride, horizon
+        idx = np.arange(back, len(ep) - fwd)
+        valid = idx[(ep[idx - back] == ep[idx]) & (ep[idx + fwd] == ep[idx])]
+        if episodes is not None:
+            valid = valid[np.isin(ep[valid], list(episodes))]
+        self.indices = valid
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, i):
+        item = self.ds[int(self.indices[i])]
+        frames = [
+            # antialiased resize overshoots [0,1] slightly; LPIPS and the VAE both want it clamped
+            F.interpolate(item[c], self.image_size, mode="bilinear", antialias=True, align_corners=False)
+            .clamp(0, 1)
+            for c in self.cameras
+        ]  # each (history+1, 3, h, w)
+        stacked = torch.stack(frames, 1)  # (history+1, V, 3, h, w)
+        return {
+            "context": stacked[:-1],  # (history, V, 3, h, w), oldest first
+            "future": stacked[-1],  # (V, 3, h, w)
+            "action": item["action"],
+            # first state_dim entries only: eef_pos(3) + eef_quat(4) + gripper_qpos(2) on MimicGen.
+            # Object/sim state lives in separate features and is never read.
+            "state": item[self.state_key][: self.state_dim],
+        }
+
+
+def build_datasets(specs, image_size, state_key, state_dim, val_episodes, history=1, stride=4):
+    """Concat all datasets; hold out the last `val_episodes` episodes of each for validation."""
+    cameras = None
+    train, val = [], []
+    for spec in specs:
+        _, _, info = dataset_info(spec)
+        cams = discover_cameras(info["features"])
+        if cameras is None:
+            cameras = cams
+        elif cams != cameras:
+            raise ValueError(f"camera mismatch: {spec} has {cams}, expected {cameras}")
+        n_ep = info["total_episodes"]
+        holdout = set(range(max(0, n_ep - val_episodes), n_ep))
+        args = (cams, state_key, state_dim, image_size)
+        kw = dict(history=history, stride=stride)
+        train.append(MultiViewWindows(spec, *args, set(range(n_ep)) - holdout, **kw))
+        if holdout:
+            val.append(MultiViewWindows(spec, *args, holdout, **kw))
+    cat = torch.utils.data.ConcatDataset
+    return cat(train), (cat(val) if val else None), cameras
+
+
+# ── model ─────────────────────────────────────────────────────────────────────
+
+
+class SD3MultiViewDynamics(nn.Module):
+    """SD3 MMDiT over a canvas of [history..current | target] frames × camera views.
+
+    Camera identity is positional: view v always occupies column v of the canvas, in the sorted
+    camera order, so it is consistent between training and evaluation. Because every view sits in
+    one sequence, the predicted views also attend to each other -- they are denoised jointly, not
+    in independent batch rows.
+    """
+
+    def __init__(
+        self,
+        vae,
+        transformer,
+        n_views,
+        action_dim,
+        state_dim,
+        ctx_noise=0.0,
+        aux_weight=0.0,
+        aux_layer=None,
+        action_dropout=0.2,
+        horizon=HORIZON,
+    ):
+        super().__init__()
+        self.vae = vae.eval().requires_grad_(False)
+        self.transformer = transformer
+        self.n_views, self.ctx_noise = n_views, ctx_noise
+        self.action_dim, self.horizon = action_dim, horizon
+        self.aux_weight, self.action_dropout = aux_weight, action_dropout
+        d = transformer.config.joint_attention_dim
+        p = transformer.config.pooled_projection_dim
+        # actions and state are not spatial, so they stay on the cross-attention pathway
+        self.act_tok = nn.Linear(action_dim, d)
+        self.state_tok = nn.Linear(state_dim, d)
+        self.pool = nn.Linear(d, p)
+        # zero-init: step 0 is the pretrained MMDiT on empty context, not on random noise
+        for m in (self.act_tok, self.state_tok, self.pool):
+            nn.init.zeros_(m.weight)
+            nn.init.zeros_(m.bias)
+
+        self._aux_feat = None
+        if aux_weight:
+            inner = transformer.config.num_attention_heads * transformer.config.attention_head_dim
+            self.aux_head = nn.Sequential(
+                nn.LayerNorm(inner), nn.Linear(inner, inner), nn.GELU(), nn.Linear(inner, horizon * action_dim)
+            )
+            blocks = transformer.transformer_blocks
+            self.aux_layer = len(blocks) // 2 if aux_layer is None else aux_layer
+            # JointTransformerBlock returns (encoder_hidden_states, hidden_states); grab the image stream
+            blocks[self.aux_layer].register_forward_hook(lambda m, i, o: setattr(self, "_aux_feat", o[-1]))
+
+    @classmethod
+    def from_pretrained(cls, sd3_id, n_views, action_dim, state_dim, dtype=torch.bfloat16, **kw):
+        vae = AutoencoderKL.from_pretrained(sd3_id, subfolder="vae", torch_dtype=dtype)
+        tr = SD3Transformer2DModel.from_pretrained(sd3_id, subfolder="transformer", torch_dtype=dtype)
+        return cls(vae, tr, n_views, action_dim, state_dim, **kw)
+
+    # -- latents --
+    def encode(self, imgs, sample=False):
+        """(N,3,H,W) in [0,1] -> scaled latents (N,C,h,w)."""
+        dist = self.vae.encode(imgs.to(self.vae.dtype) * 2 - 1).latent_dist
+        z = dist.sample() if sample else dist.mode()
+        return (z - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+    def decode(self, z):
+        z = z.to(self.vae.dtype) / self.vae.config.scaling_factor + self.vae.config.shift_factor
+        return (self.vae.decode(z).sample / 2 + 0.5).clamp(0, 1)
+
+    def encode_grid(self, imgs, **kw):
+        """(B,R,V,3,H,W) -> (B,R,V,C,h,w)."""
+        b, r, v = imgs.shape[:3]
+        z = self.encode(imgs.flatten(0, 2), **kw)
+        return z.reshape(b, r, v, *z.shape[1:])
+
+    # -- canvas: rows are timesteps, columns are views --
+    @staticmethod
+    def to_canvas(z):
+        """(B,R,V,C,h,w) -> (B,C,R*h,V*w)."""
+        b, r, v, c, h, w = z.shape
+        return z.permute(0, 3, 1, 4, 2, 5).reshape(b, c, r * h, v * w)
+
+    @staticmethod
+    def from_canvas(canvas, v):
+        """(B,C,h,V*w) single row -> (B,V,C,h,w)."""
+        b, c, h, vw = canvas.shape
+        return canvas.reshape(b, c, h, v, vw // v).permute(0, 3, 1, 2, 4)
+
+    def conditioning(self, action, state, drop=None):
+        """Action chunk and robot state as cross-attention tokens -> (ctx, pooled).
+
+        `drop` is a (B,) bool mask zeroing the action tokens, so the auxiliary head has to recover
+        the action from the frames instead of reading it straight off the conditioning.
+        """
+        dt = self.act_tok.weight.dtype
+        act = self.act_tok(action.to(dt))
+        if drop is not None:
+            act = act * (~drop).to(dt).view(-1, 1, 1)
+        ctx = torch.cat([act, self.state_tok(state.to(dt))[:, None]], 1)
+        return ctx, self.pool(ctx.mean(1))
+
+    def denoise(self, canvas, timestep, ctx, pooled):
+        dt = self.transformer.dtype  # conditioning heads stay fp32; cast at the MMDiT boundary
+        return self.transformer(
+            hidden_states=canvas.to(dt),
+            encoder_hidden_states=ctx.to(dt),
+            pooled_projections=pooled.to(dt),
+            timestep=timestep,
+            return_dict=False,
+        )[0]
+
+    def _context_rows(self, context):
+        """Encoded history+current rows, optionally noise-perturbed for robustness (Ctrl-World §3)."""
+        z = self.encode_grid(context)
+        if self.ctx_noise:
+            z = z + self.ctx_noise * torch.randn_like(z)
+        return z
+
+    # -- training / sampling --
+    def loss(self, batch):
+        """Rectified-flow matching on the target row of the canvas."""
+        z_ctx = self._context_rows(batch["context"])
+        z_tgt = self.encode_grid(batch["future"][:, None], sample=True)  # (B,1,V,C,h,w)
+        b, _, v, c, h, w = z_tgt.shape
+        t = torch.sigmoid(torch.randn(b, device=z_tgt.device, dtype=z_tgt.dtype))  # logit-normal
+        noise = torch.randn_like(z_tgt)
+        tt = t.view(-1, 1, 1, 1, 1, 1)  # (B,R,V,C,h,w) is 6-D; a 5-D view broadcasts into the row axis
+        z_t = (1 - tt) * z_tgt + tt * noise
+        canvas = self.to_canvas(torch.cat([z_ctx, z_t], 1))
+        drop = torch.rand(b, device=canvas.device) < self.action_dropout if self.aux_weight else None
+        ctx, pooled = self.conditioning(batch["action"], batch["state"], drop)
+        pred = self.denoise(canvas, t * 1000, ctx, pooled)[:, :, -h:, :]  # target row only
+        flow = F.mse_loss(pred.float(), self.to_canvas(noise - z_tgt).float())
+        if not self.aux_weight or not drop.any():
+            return flow
+        # inverse dynamics on the action-dropped rows only: with the action still conditioned in,
+        # the head would just read it back off the context tokens and learn nothing.
+        feat = self._aux_feat[drop].float().mean(1)
+        a_pred = self.aux_head(feat).view(-1, self.horizon, self.action_dim)
+        return flow + self.aux_weight * F.mse_loss(a_pred, batch["action"][drop].float())
+
+    @torch.no_grad()
+    def predict(self, batch, scheduler, steps=20):
+        """-> predicted future frames (B,V,3,H,W) in [0,1]."""
+        z_ctx = self._context_rows(batch["context"])
+        b, _, v, c, h, w = z_ctx.shape
+        ctx, pooled = self.conditioning(batch["action"], batch["state"])
+        z = torch.randn((b, 1, v, c, h, w), device=z_ctx.device, dtype=z_ctx.dtype)
+        scheduler.set_timesteps(steps, device=z_ctx.device)
+        for t in scheduler.timesteps:
+            canvas = self.to_canvas(torch.cat([z_ctx, z], 1))
+            vel = self.denoise(canvas, t.expand(b), ctx, pooled)[:, :, -h:, :]
+            row = scheduler.step(vel.float(), t, self.to_canvas(z).float(), return_dict=False)[0]
+            z = self.from_canvas(row.to(z_ctx.dtype), v)[:, None]
+        imgs = self.decode(z.flatten(0, 2))
+        return imgs.float().reshape(b, v, *imgs.shape[1:])
+
+
+# ── evaluation ────────────────────────────────────────────────────────────────
+
+
+def per_camera_metrics(pred, gt, cameras, lpips=None):
+    """pred/gt: (B,V,3,H,W) in [0,1]. -> {'<cam>/mse': ..., 'mean/mse': ...}"""
+    from torchmetrics.functional import structural_similarity_index_measure as ssim_fn
+
+    out = {}
+    for i, cam in enumerate(cameras):
+        p, g = pred[:, i], gt[:, i]
+        mse = F.mse_loss(p, g).item()
+        out[f"{cam}/mse"] = mse
+        out[f"{cam}/psnr"] = 10 * math.log10(1.0 / max(mse, 1e-12))  # data range 1.0
+        out[f"{cam}/ssim"] = ssim_fn(p, g, data_range=1.0).item()
+        if lpips is not None:
+            out[f"{cam}/lpips"] = lpips(p, g).item()
+    for m in ("mse", "psnr", "ssim", "lpips"):
+        vals = [v for k, v in out.items() if k.endswith(f"/{m}")]
+        if vals:
+            out[f"mean/{m}"] = sum(vals) / len(vals)
+    return out
+
+
+def to_wandb_image(img):
+    import wandb
+
+    return wandb.Image((img.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
+
+
+def reconstruction_table(current, gt, pred, cameras, start_id=0):
+    """One row per example, one Current/GT/Pred column group per camera."""
+    import wandb
+
+    cols = ["sample_id"] + [f"{c}/{k}" for c in cameras for k in ("current", "gt", "pred")]
+    table = wandb.Table(columns=cols)
+    for b in range(current.shape[0]):
+        row = [start_id + b]
+        for i in range(len(cameras)):
+            row += [to_wandb_image(x[b, i]) for x in (current, gt, pred)]
+        table.add_data(*row)
+    return table
+
+
+# ── dataset safety checks (spec §15 / §22) ────────────────────────────────────
+
+
+def inspect(specs, state_key, state_dim, image_size, history=1, stride=4, horizon=HORIZON):
+    ref_cams = None
+    for spec in specs:
+        _, _, info = dataset_info(spec)
+        feats = info["features"]
+        cams = discover_cameras(feats)
+        print(f"\n=== {spec}")
+        print(f"  fps={info['fps']} episodes={info['total_episodes']} frames={info['total_frames']}")
+        print(f"  splits={info.get('splits')}")
+        print(f"  all keys: {sorted(feats)}")
+        print(f"  cameras (deterministic order): {cams}")
+        for c in cams:
+            f = feats[c]
+            print(f"    {c}: shape={f['shape']} names={f.get('names')} dtype={f['dtype']}")
+        print(f"  action dim: {feats['action']['shape']}")
+        full = feats[state_key]["shape"][0]
+        print(f"  state: {state_key}[:{state_dim}] of {full}D")
+        assert full >= state_dim, f"{state_key} is only {full}D, need {state_dim}"
+        objs = object_state_keys(feats)
+        print(f"  object/sim-state keys present but NOT read by the model: {objs}")
+        assert state_key not in objs, f"{state_key} is an object-state feature"
+        if ref_cams is None:
+            ref_cams = cams
+        assert cams == ref_cams, f"camera mismatch vs first dataset: {cams} != {ref_cams}"
+
+        ds = MultiViewWindows(spec, cams, state_key, state_dim, image_size, history=history, stride=stride)
+        item = ds[0]
+        print(f"  frame offsets (history..current, target): {ds.offsets}")
+        print(f"  valid windows: {len(ds)} / {info['total_frames']} frames")
+        for c, ctx, fut in zip(cams, item["context"].transpose(0, 1), item["future"]):
+            print(f"    {c}: context {tuple(ctx.shape)}  future {tuple(fut.shape)}")
+        print(f"    action {tuple(item['action'].shape)}  state {tuple(item['state'].shape)}")
+
+        i0 = int(ds.indices[0])
+        raw = ds.ds[i0]
+        fi = int(raw["frame_index"])
+        nxt = ds.ds.hf_dataset[i0 + horizon]
+        print(f"    frame_index {fi} -> {int(nxt['frame_index'])} (must differ by exactly {horizon})")
+        assert int(nxt["frame_index"]) - fi == horizon
+        assert int(nxt["episode_index"]) == int(raw["episode_index"])
+        for c in cams:
+            assert not raw[f"{c}_is_pad"].any(), f"{c} padded at a supposedly valid window"
+        assert not raw["action_is_pad"].any()
+    print(f"\nOK: {len(specs)} dataset(s), cameras={ref_cams}")
+
+
+# ── training ──────────────────────────────────────────────────────────────────
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument(
+        "--data", nargs="+", required=True, help="local dataset roots or HF repo ids (success, failure, ...)"
+    )
+    ap.add_argument("--inspect", action="store_true", help="run dataset safety checks and exit")
+    ap.add_argument("--sd3", default="stabilityai/stable-diffusion-3.5-medium")
+    ap.add_argument("--state-key", default="observation.state")
+    ap.add_argument("--state-dim", type=int, default=9, help="use the first N dims of --state-key (spec: 9)")
+    ap.add_argument("--image-size", type=int, nargs=2, default=(224, 224))
+    ap.add_argument("--history", type=int, default=1, help="context frames incl. current (1 = current only)")
+    ap.add_argument("--history-stride", type=int, default=4, help="frames between history frames")
+    ap.add_argument("--ctx-noise", type=float, default=0.0, help="noise on context latents (Ctrl-World)")
+    ap.add_argument("--aux-weight", type=float, default=0.0, help="auxiliary action-reconstruction loss")
+    ap.add_argument("--aux-layer", type=int, default=None, help="MMDiT block to read (default: middle)")
+    ap.add_argument("--action-dropout", type=float, default=0.2, help="action-drop rate the aux loss uses")
+    ap.add_argument("--batch-size", type=int, default=16)
+    ap.add_argument("--grad-accum", type=int, default=1)
+    ap.add_argument("--steps", type=int, default=20000)
+    ap.add_argument("--lr", type=float, default=1e-4)
+    ap.add_argument("--lora-rank", type=int, default=32, help="0 = full finetune of the MMDiT")
+    ap.add_argument("--log-every", type=int, default=50)
+    ap.add_argument(
+        "--no-grad-ckpt",
+        action="store_true",
+        help="28%% faster per sample but far more activation memory; needs a small batch",
+    )
+    ap.add_argument("--val-episodes", type=int, default=20, help="held-out episodes per dataset")
+    ap.add_argument("--val-every", type=int, default=1000)
+    ap.add_argument("--val-batches", type=int, default=8)
+    ap.add_argument("--log-images", type=int, default=4, help="examples logged to the W&B table")
+    ap.add_argument("--sample-steps", type=int, default=20)
+    ap.add_argument("--workers", type=int, default=8)
+    ap.add_argument("--out", default=str(Path(__file__).parent / "outputs"))
+    ap.add_argument("--wandb-project", default="sd3-dynamics", help="empty string disables W&B")
+    ap.add_argument("--run-name", default=None, help="W&B run name (default: auto-generated)")
+    args = ap.parse_args()
+
+    image_size = tuple(args.image_size)
+    if args.inspect:
+        inspect(args.data, args.state_key, args.state_dim, image_size, args.history, args.history_stride)
+        return
+
+    torch.manual_seed(0)
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    train_ds, val_ds, cameras = build_datasets(
+        args.data, image_size, args.state_key, args.state_dim, args.val_episodes,
+        args.history, args.history_stride,
+    )
+    feats = dataset_features(args.data[0])
+    state_dim = args.state_dim
+    full = feats[args.state_key]["shape"][0]
+    if full < state_dim:
+        raise SystemExit(f"{args.state_key} is only {full}D, cannot take the first {state_dim}D")
+    assert args.state_key not in object_state_keys(feats), "object state must not reach the model"
+    action_dim = feats["action"]["shape"][0]
+    print(f"cameras={cameras} action_dim={action_dim} state={args.state_key}[:{state_dim}] of {full}D")
+    print(f"excluded from the model: {object_state_keys(feats)}")
+    print(f"canvas: {args.history + 1} rows x {len(cameras)} views, history stride {args.history_stride}")
+    print(f"train windows={len(train_ds)} val windows={len(val_ds) if val_ds else 0}")
+
+    model = SD3MultiViewDynamics.from_pretrained(
+        args.sd3,
+        len(cameras),
+        action_dim,
+        state_dim,
+        ctx_noise=args.ctx_noise,
+        aux_weight=args.aux_weight,
+        aux_layer=args.aux_layer,
+        action_dropout=args.action_dropout,
+    ).to(dev)
+    if not args.no_grad_ckpt:
+        model.transformer.enable_gradient_checkpointing()
+    if args.lora_rank:
+        from diffusers.training_utils import cast_training_params
+        from peft import LoraConfig
+
+        model.transformer.requires_grad_(False)
+        model.transformer.add_adapter(
+            LoraConfig(
+                r=args.lora_rank,
+                lora_alpha=args.lora_rank,
+                init_lora_weights="gaussian",
+                target_modules=["to_q", "to_k", "to_v", "to_out.0"],
+            )
+        )
+        cast_training_params(model.transformer, dtype=torch.float32)  # bf16 Adam states diverge
+    for m in (model.act_tok, model.state_tok, model.pool):
+        m.to(torch.float32)
+    if args.aux_weight:
+        model.aux_head.to(torch.float32)
+        print(f"aux action head on block {model.aux_layer}/{len(model.transformer.transformer_blocks)}")
+    trainable = {n for n, p in model.named_parameters() if p.requires_grad}
+    params = [p for p in model.parameters() if p.requires_grad]
+    print(f"trainable params: {sum(p.numel() for p in params) / 1e6:.1f}M")
+    opt = torch.optim.AdamW(params, lr=args.lr)
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.sd3, subfolder="scheduler")
+
+    lpips = None
+    if val_ds is not None:
+        from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+
+        lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).to(dev)
+
+    run = None
+    if args.wandb_project:
+        import wandb
+
+        run = wandb.init(
+            project=args.wandb_project, name=args.run_name, config=vars(args) | {"cameras": cameras}
+        )
+
+    def loader(ds, bs, shuffle):
+        return torch.utils.data.DataLoader(
+            ds, batch_size=bs, shuffle=shuffle, num_workers=args.workers, drop_last=True, pin_memory=True
+        )
+
+    def to_dev(batch):
+        return {k: v.to(dev, non_blocking=True) for k, v in batch.items()}
+
+    def log(d, step):
+        print(f"step {step}: " + " ".join(f"{k}={v:.4f}" for k, v in d.items() if isinstance(v, float)))
+        if run:
+            run.log(d, step=step)
+
+    @torch.no_grad()
+    def validate(step):
+        model.eval()
+        sums, n = {}, 0
+        for i, batch in enumerate(loader(val_ds, args.batch_size, False)):
+            if i >= args.val_batches:
+                break
+            batch = to_dev(batch)
+            pred = model.predict(batch, scheduler, args.sample_steps)
+            current = batch["context"][:, -1]  # newest context frame
+            for k, v in per_camera_metrics(pred, batch["future"], cameras, lpips).items():
+                sums[k] = sums.get(k, 0.0) + v
+            # copy-the-current-frame baseline on the same batches: the number to beat
+            for k, v in per_camera_metrics(current, batch["future"], cameras, lpips).items():
+                sums[f"copy_{k}"] = sums.get(f"copy_{k}", 0.0) + v
+            n += 1
+            if i == 0 and run:
+                k = min(args.log_images, pred.shape[0])
+                run.log(
+                    {
+                        "val/reconstructions": reconstruction_table(
+                            current[:k], batch["future"][:k], pred[:k], cameras
+                        )
+                    },
+                    step=step,
+                )
+        model.train()
+        log({f"val/{k}": v / max(n, 1) for k, v in sums.items()}, step)
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    step, running, t0 = 0, 0.0, time.time()
+    model.train()
+    while step < args.steps:
+        for batch in loader(train_ds, args.batch_size, True):
+            loss = model.loss(to_dev(batch)) / args.grad_accum
+            loss.backward()
+            running += loss.item() * args.grad_accum  # undo the accumulation scaling for reporting
+            if (step + 1) % args.grad_accum == 0:
+                torch.nn.utils.clip_grad_norm_(params, 1.0)
+                opt.step()
+                opt.zero_grad(set_to_none=True)
+            step += 1
+            if step % args.log_every == 0:
+                peak = torch.cuda.max_memory_allocated() / 1e9 if dev == "cuda" else 0.0
+                dt = (time.time() - t0) / args.log_every
+                log(
+                    {
+                        "train/loss": running / args.log_every,
+                        "train/peak_gb": peak,
+                        "train/s_per_step": dt,
+                        "train/samples_per_s": args.batch_size / dt,
+                    },
+                    step,
+                )
+                running, t0 = 0.0, time.time()
+            if val_ds is not None and step % args.val_every == 0:
+                validate(step)
+            if step % 5000 == 0 or step == args.steps:
+                # trainable tensors only: ~93MB of LoRA + conditioning heads, not 5GB of frozen MMDiT
+                torch.save(
+                    {k: v for k, v in model.state_dict().items() if k in trainable},
+                    out / f"step_{step}.pt",
+                )
+            if step >= args.steps:
+                break
+
+
+if __name__ == "__main__":
+    main()

test_dino_dynamics.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python
+"""Self-check for dino_dynamics with a tiny stand-in DINO/VAE. No downloads.
+
+    policy_learning/lerobot/.venv/bin/python dynamics_model/test_dino_dynamics.py
+"""
+
+import torch
+import torch.nn as nn
+from diffusers import AutoencoderKL
+from transformers import Dinov2Config, Dinov2Model
+
+from dino_dynamics import DinoDynamics
+from sd3_dynamics import HORIZON
+
+V, B, SZ, DSZ, ACT, ST, HIST = 2, 2, 112, 196, 7, 9, 2  # DSZ/14=14 patches -> x2 = 28 = SZ/4
+
+
+def tiny():
+    dino = Dinov2Model(
+        Dinov2Config(hidden_size=32, num_hidden_layers=1, num_attention_heads=2,
+                     intermediate_size=64, patch_size=14, image_size=DSZ)
+    )
+    vae = AutoencoderKL(
+        in_channels=3, out_channels=3,
+        down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"),
+        up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"),
+        block_out_channels=(8, 8, 8), layers_per_block=1, latent_channels=16, norm_num_groups=8,
+    )
+    vae.register_to_config(shift_factor=0.0609, scaling_factor=1.5305)
+    return DinoDynamics(dino, vae, V, ACT, ST, history=HIST, depth=1, heads=2, mlp_dim=64, dino_size=DSZ, image_size=SZ)
+
+
+def batch():
+    return {
+        "context": torch.rand(B, HIST, V, 3, SZ, SZ),
+        "future": torch.rand(B, V, 3, SZ, SZ),
+        "action": torch.randn(B, HORIZON, ACT),
+        "state": torch.randn(B, ST),
+    }
+
+
+def main():
+    m = tiny()
+    b = batch()
+    P = m.n_patches
+    assert P == (DSZ // 14) ** 2 == 196, P
+
+    z = m.encode_views(b["context"])
+    assert z.shape == (B, HIST, V, P, m.dim), z.shape
+
+    pred = m(b)
+    assert pred.shape == (B, V, P, m.dim), pred.shape
+
+    # the encoder is frozen: no gradients, and its weights carry none after backward
+    loss, tgt = m.loss(b)
+    dec = m.decoder_loss(b, tgt)
+    (loss + dec).backward()
+    assert all(p.grad is None for p in m.dino.parameters()), "DINOv2 must stay frozen"
+    assert all(p.grad is None for p in m.vae.parameters()), "VAE must stay frozen"
+    assert m.query.grad is not None and m.act_tok.weight.grad is not None
+    # the decoder head must not backprop into the predictor (DINO-WM trains it on detached feats)
+    assert m.to_latent[0].weight.grad is not None
+
+    # conditioning actually matters: different actions -> different prediction
+    torch.manual_seed(0)
+    a = m(b)
+    torch.manual_seed(0)
+    c = m(dict(b, action=torch.randn_like(b["action"])))
+    assert not torch.allclose(a, c), "prediction ignores the action chunk"
+
+    # and so do the context frames
+    torch.manual_seed(0)
+    d = m(dict(b, context=torch.rand_like(b["context"])))
+    assert not torch.allclose(a, d), "prediction ignores the context frames"
+
+    img = m.features_to_image(pred)
+    assert img.shape == (B, V, 3, SZ, SZ), img.shape
+    assert (img >= 0).all() and (img <= 1).all()
+
+    print("ok", {"patches": P, "dim": m.dim, "pred": tuple(pred.shape), "img": tuple(img.shape)})
+
+
+if __name__ == "__main__":
+    main()