File size: 3,018 Bytes
8acc5ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/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()
|