"""Diffusion-policy baseline — DDPM denoiser conditioned on DINOv3 CLS. Same DINOv3 backbone as :class:`DinoCLSXYZBaseline`. Instead of regressing (T, 3) XYZ directly, the head is a denoising MLP that takes a noisy trajectory + a diffusion-step index + the CLS embedding and predicts the added Gaussian noise. At inference we run K-step ancestral DDPM sampling from N(0, I) to get the clean trajectory. Why CLS-only (no per-pixel features): the user-specified design — the image is fixed per chunk, so the spatial conditioning is the same as the ``DinoCLSXYZBaseline`` (CLS → query). Only the regression head differs. ────────────────────────────────────────────────────────────────────── Forward (training) ────────────────────────────────────────────────────────────────────── rgb : (B, 3, H, W) target_xyz : (B, T, 3) → sample diffusion step t ~ Uniform(1, T_diff) → eps ~ N(0, I_{B,T,3}) → x_t = sqrt(ᾱ_t) * target_xyz + sqrt(1-ᾱ_t) * eps → eps_pred = denoiser(x_t, t, CLS) → loss = MSE(eps_pred, eps) Forward (inference / sampling) ────────────────────────────────────────────────────────────────────── rgb → CLS (one forward through DINOv3 — frozen otherwise) x_K ~ N(0, I) for k = K..1: eps_pred = denoiser(x_k, k, CLS) x_{k-1} = ancestral_step(x_k, eps_pred, k) return x_0 as ``xyz_pred`` The grip / rot heads are kept identical to the deterministic baseline (per-timestep CE on bins) — only the position trajectory is diffused. """ from __future__ import annotations import math import torch import torch.nn as nn import torch.nn.functional as F from .model import AdaLNZeroBlock, DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS, sin_pe def _cosine_beta_schedule(n_steps: int, s: float = 0.008) -> torch.Tensor: """Nichol-Dhariwal cosine schedule. Returns ``betas`` of shape (n_steps,).""" steps = torch.arange(n_steps + 1, dtype=torch.float64) alphas_bar = torch.cos(((steps / n_steps) + s) / (1 + s) * math.pi / 2) ** 2 alphas_bar = alphas_bar / alphas_bar[0] betas = 1.0 - (alphas_bar[1:] / alphas_bar[:-1]) return betas.clamp(0.0001, 0.9999).float() class _TimestepEmbed(nn.Module): """Sinusoidal diffusion-step embedding → linear → silu → linear.""" def __init__(self, dim: int): super().__init__() self.dim = dim self.net = nn.Sequential( nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim * 4), ) self.out_dim = dim * 4 def forward(self, t: torch.Tensor) -> torch.Tensor: # t: (B,) int → (B, dim*4) half = self.dim // 2 freqs = torch.exp( -math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half) args = t.float().unsqueeze(-1) * freqs # (B, half) sinu = torch.cat([torch.sin(args), torch.cos(args)], dim=-1) # (B, dim) return self.net(sinu) class _Denoiser(nn.Module): """Predicts noise given (x_t [B,T,3], diff_step t, cond [B, D_cond]). Uses a small MLP that flattens the trajectory then expands back. For T=32 this is ~96 inputs + cond, well within MLP territory. """ def __init__(self, n_window: int, d_cond: int, d_diff: int = 128, d_hidden: int = 1024): super().__init__() self.n_window = n_window self.t_embed = _TimestepEmbed(d_diff) # → 4*d_diff d_t = self.t_embed.out_dim in_dim = n_window * 3 + d_cond + d_t out_dim = n_window * 3 self.net = nn.Sequential( nn.Linear(in_dim, d_hidden), nn.SiLU(), nn.Linear(d_hidden, d_hidden), nn.SiLU(), nn.Linear(d_hidden, d_hidden), nn.SiLU(), nn.Linear(d_hidden, out_dim), ) def forward(self, x_t: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: # x_t: (B, T, 3); t: (B,); cond: (B, d_cond) B = x_t.size(0) h_t = self.t_embed(t) # (B, 4*d_diff) h = torch.cat([x_t.flatten(1), cond, h_t], dim=-1) # (B, in_dim) return self.net(h).view(B, self.n_window, 3) class DinoCLSXYZBaselineDiffusion(nn.Module): """Diffusion-policy baseline. CLS conditioning + DDPM denoiser → XYZ. Shares the DINOv3 backbone + (grip, rot) head bookkeeping with :class:`DinoCLSXYZBaseline`. The position head is replaced. Hyperparams ----------- n_diffusion_steps_train : int — number of beta steps in the schedule. n_diffusion_steps_infer : int — K ancestral steps at sampling time. K == n_diffusion_steps_train is exact; a smaller K is faster but noisier. """ def __init__( self, n_window: int = 32, d_model: int = 256, d_cond: int = 128, d_sin_t: int = 48, n_blocks: int = 5, n_gripper_bins: int = 32, n_rot_clusters: int = 64, n_diffusion_steps_train: int = 100, n_diffusion_steps_infer: int = 100, d_diff_embed: int = 128, denoiser_hidden: int = 1024, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, ): super().__init__() self.n_window = n_window self.d_model = d_model self.n_gripper_bins = n_gripper_bins self.n_rot_clusters = n_rot_clusters self.n_diff_train = int(n_diffusion_steps_train) self.n_diff_infer = int(n_diffusion_steps_infer) import os self.backbone = torch.hub.load( os.environ.get("DINO_REPO_DIR", dino_repo), "dinov3_vits16plus", source="local", weights=os.environ.get("DINO_WEIGHTS_PATH", dino_weights), ) embed_dim = self.backbone.embed_dim self.embed_dim = embed_dim # Conditioning: project CLS → d_cond. Per-timestep position-embed # (carried over from the deterministic baseline) helps grip/rot # be per-step distinct. self.register_buffer("t_sin", sin_pe(n_window, d_sin_t)) self.t_cond_proj = nn.Linear(d_sin_t, d_cond) self.input_proj = nn.Linear(embed_dim, d_model) self.blocks = nn.ModuleList([ AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks) ]) self.final_norm = nn.LayerNorm(d_model) self.grip_head = nn.Linear(d_model, n_gripper_bins) self.rot_head = nn.Linear(d_model, n_rot_clusters) # Position head = denoiser conditioned directly on the raw CLS # token + diffusion-step embedding. No AdaLN routing here — the # denoiser is a feed-forward MLP over the full (T, 3) trajectory. self.denoiser = _Denoiser( n_window=n_window, d_cond=embed_dim, d_diff=d_diff_embed, d_hidden=denoiser_hidden) # Beta / alpha schedule, registered as buffers. betas = _cosine_beta_schedule(self.n_diff_train) alphas = 1.0 - betas alpha_bars = torch.cumprod(alphas, dim=0) self.register_buffer("betas", betas) self.register_buffer("alphas", alphas) self.register_buffer("alpha_bars", alpha_bars) self.register_buffer("sqrt_alpha_bars", torch.sqrt(alpha_bars)) self.register_buffer("sqrt_one_minus_alpha_bars", torch.sqrt(1.0 - alpha_bars)) def _cls(self, rgb: torch.Tensor) -> torch.Tensor: return self.backbone.forward_features(rgb)["x_norm_clstoken"] def _encode(self, rgb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Run DINOv3 once and return (CLS, patch grid). Used by `forward()` to also publish patch_features for downstream DINO-PCA viz.""" import math out = self.backbone.forward_features(rgb) cls = out["x_norm_clstoken"] toks = out["x_norm_patchtokens"] P = int(round(math.sqrt(toks.shape[1]))) patches = toks.transpose(1, 2).reshape(toks.shape[0], -1, P, P) return cls, patches def _grip_rot(self, cls: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-timestep grip + rot logits from the AdaLN-Zero head.""" B = cls.size(0); T = self.n_window q_in = self.input_proj(cls) q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1) cond_t = self.t_cond_proj(self.t_sin) cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1) h = q_in_bt for blk in self.blocks: h = blk(h, cond_bt) h = self.final_norm(h).view(B, T, self.d_model) return self.grip_head(h), self.rot_head(h) def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor = None, target_xyz: torch.Tensor | None = None) -> dict: """Returns dict. Training (``target_xyz`` provided): - sample t ~ Uniform(1, T_diff), eps ~ N(0, I) - predict eps; loss in ``baseline_diffusion_losses`` is MSE. Inference (``target_xyz=None``): - K-step ancestral DDPM → ``xyz_pred``. Grip + rot logits are computed deterministically in both modes. """ B = rgb.size(0); T = self.n_window device = rgb.device cls, patches = self._encode(rgb) grip_logits, rot_logits = self._grip_rot(cls) if target_xyz is not None and self.training: # Per-sample diffusion step (broadcast over T). t = torch.randint( 0, self.n_diff_train, (B,), device=device, dtype=torch.long) eps = torch.randn_like(target_xyz) ab = self.sqrt_alpha_bars[t].view(B, 1, 1) om = self.sqrt_one_minus_alpha_bars[t].view(B, 1, 1) x_t = ab * target_xyz + om * eps eps_pred = self.denoiser(x_t, t, cls) return { "eps": eps, "eps_pred": eps_pred, "grip_logits": grip_logits, "rot_logits": rot_logits, "cls": cls, "patch_features": patches, } # Sampling. x = torch.randn(B, T, 3, device=device) step_indices = list(reversed(range(self.n_diff_infer))) for k in step_indices: # Map inference step k onto the training schedule (linear). tt = int(round(k * (self.n_diff_train - 1) / max(1, self.n_diff_infer - 1))) t_b = torch.full((B,), tt, device=device, dtype=torch.long) eps_pred = self.denoiser(x, t_b, cls) beta_t = self.betas[tt] alpha_t = self.alphas[tt] alpha_bar_t = self.alpha_bars[tt] coef = beta_t / torch.sqrt(1.0 - alpha_bar_t).clamp(min=1e-8) mean = (x - coef * eps_pred) / torch.sqrt(alpha_t) if k > 0: x = mean + torch.sqrt(beta_t) * torch.randn_like(x) else: x = mean return { "xyz_pred": x, "grip_logits": grip_logits, "rot_logits": rot_logits, "cls": cls, "patch_features": patches, # (B, D, P, P) for DINO PCA } def baseline_diffusion_losses(out: dict, target_xyz: torch.Tensor, target_grip: torch.Tensor, target_rot: torch.Tensor) -> dict: """Loss bundle for :class:`DinoCLSXYZBaselineDiffusion`. ``loss/xyz`` here is the DDPM noise-prediction MSE (not direct XYZ MSE). The grip and rot losses are the same per-step CE used by the deterministic baseline. """ eps_loss = F.mse_loss(out["eps_pred"], out["eps"]) B, T = target_grip.shape grip_loss = F.cross_entropy( out["grip_logits"].reshape(B * T, -1), target_grip.reshape(B * T)) rot_loss = F.cross_entropy( out["rot_logits"].reshape(B * T, -1), target_rot.reshape(B * T)) return { "loss/xyz": eps_loss, "loss/grip": grip_loss, "loss/rot": rot_loss, "loss/total": eps_loss + grip_loss + rot_loss, }