test_dino_dynamics.py

#!/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()