lerobot_plugin/patchpol/diffusion.py

"""Hand-rolled DDPM action head + the full Patch Policy (trunk + head).

The idea in three sentences:
  FORWARD (fixed, no learning): take a clean action chunk a0, mix in Gaussian
  noise at strength t: a_t = sqrt(ab_t)*a0 + sqrt(1-ab_t)*eps. At t=T the chunk
  is pure noise. REVERSE (learned): a network looks at (a_t, t, obs) and
  predicts the eps that was mixed in; subtracting it step by step walks pure
  noise back into a valid action chunk *for this observation*.

Training is one line: MSE(predicted eps, true eps) at a random t.
Sampling runs the T-step walk. Quick self-test:  uv run python -m patchpol.diffusion
"""

import math

import torch
import torch.nn as nn

from patchpol.model import Block, PatchTrunk


# ---------------------------------------------------------------- schedule --
class DDPMSchedule:
    """Squared-cosine beta schedule (Nichol & Dhariwal) — the one Diffusion
    Policy uses. All tensors are precomputed lookup tables indexed by t."""

    def __init__(self, n_steps: int = 100, device: torch.device = "cpu"):
        self.n_steps = n_steps
        s = 0.008  # small offset so ab(0) ~ 1 without a singularity
        ts = torch.linspace(0, 1, n_steps + 1, device=device)
        f = torch.cos((ts + s) / (1 + s) * math.pi / 2) ** 2
        alpha_bar = f / f[0]                              # ab(0)=1, decays to ~0
        betas = (1 - alpha_bar[1:] / alpha_bar[:-1]).clamp(max=0.999)
        self.betas = betas                                # (T,)
        self.alphas = 1.0 - betas
        self.alpha_bar = torch.cumprod(self.alphas, 0)    # (T,)
        prev = torch.cat([torch.ones(1, device=device), self.alpha_bar[:-1]])
        # posterior variance beta~_t = (1-ab_{t-1})/(1-ab_t) * beta_t
        self.post_var = betas * (1 - prev) / (1 - self.alpha_bar)

    def add_noise(self, x0: torch.Tensor, eps: torch.Tensor, t: torch.Tensor):
        """q(x_t | x_0): one jump to any noise level. t: (B,) long."""
        ab = self.alpha_bar[t].view(-1, 1, 1)             # broadcast over (B,5,2)
        return ab.sqrt() * x0 + (1 - ab).sqrt() * eps

    @torch.no_grad()
    def sample(self, denoiser, cond: torch.Tensor, shape) -> torch.Tensor:
        """Ancestral sampling: pure noise -> action chunk, conditioned on obs."""
        B = cond.shape[0]
        x = torch.randn(B, *shape, device=cond.device)
        for t in reversed(range(self.n_steps)):
            tt = torch.full((B,), t, device=cond.device, dtype=torch.long)
            eps = denoiser(x, tt, cond)
            # mean of p(x_{t-1} | x_t): remove the predicted noise
            x = (x - self.betas[t] / (1 - self.alpha_bar[t]).sqrt() * eps) \
                / self.alphas[t].sqrt()
            if t > 0:  # last step is deterministic
                x = x + self.post_var[t].sqrt() * torch.randn_like(x)
        return x


# ---------------------------------------------------------------- denoiser --
def sinusoidal_emb(t: torch.Tensor, dim: int) -> torch.Tensor:
    """Classic transformer position encoding, here for the *noise level* t."""
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    ang = t[:, None].float() * freqs[None, :]
    return torch.cat([ang.sin(), ang.cos()], dim=-1)  # (B, dim)


class DiffusionHead(nn.Module):
    """Transformer denoiser (paper Table 11: 8 layers, 4 heads, 256 dim).

    Sequence = [obs token, time token, 5 noisy-action tokens], full attention
    (no causality here — denoising a chunk is not a temporal process).
    eps is read out at the action-token positions.
    """

    def __init__(self, act_dim: int = 2, horizon: int = 5, cond_dim: int = 384,
                 dim: int = 256, depth: int = 8, heads: int = 4):
        super().__init__()
        self.horizon = horizon
        self.act_in = nn.Linear(act_dim, dim)
        self.cond_in = nn.Linear(cond_dim, dim)
        self.time_in = nn.Sequential(nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, dim))
        self.pos_emb = nn.Parameter(torch.randn(1, horizon + 2, dim) * 0.02)
        self.blocks = nn.ModuleList(Block(dim, heads) for _ in range(depth))
        self.norm = nn.LayerNorm(dim)
        self.out = nn.Linear(dim, act_dim)
        S = horizon + 2
        self.register_buffer("mask", torch.ones(S, S, dtype=torch.bool), persistent=False)

    def forward(self, noisy: torch.Tensor, t: torch.Tensor, cond: torch.Tensor):
        """(B,5,2) noisy chunk, (B,) t, (B,384) obs readout -> (B,5,2) eps."""
        toks = torch.cat([
            self.cond_in(cond)[:, None],                       # (B,1,dim)
            self.time_in(sinusoidal_emb(t, self.pos_emb.shape[-1]))[:, None],
            self.act_in(noisy),                                # (B,5,dim)
        ], dim=1) + self.pos_emb
        for blk in self.blocks:
            toks = blk(toks, self.mask)
        return self.out(self.norm(toks[:, 2:]))                # action positions only


# ------------------------------------------------------------ full policy --
class PatchPolicy(nn.Module):
    """Trunk over dense patch tokens -> readout -> DDPM head over action chunks."""

    def __init__(self, obs_horizon: int = 2, act_horizon: int = 5,
                 act_dim: int = 2, n_diffusion_steps: int = 100,
                 views: int = 1):
        super().__init__()
        self.act_horizon, self.act_dim = act_horizon, act_dim
        self.views = views
        self.trunk = PatchTrunk(
            dim=384, depth=8, heads=6, T=obs_horizon, P=256, views=views
        )
        self.head = DiffusionHead(act_dim=act_dim, horizon=act_horizon, cond_dim=384)
        self.n_diffusion_steps = n_diffusion_steps
        self._sched = None  # built lazily on the right device

    def sched(self, device) -> DDPMSchedule:
        if self._sched is None or self._sched.betas.device != torch.device(device):
            self._sched = DDPMSchedule(self.n_diffusion_steps, device=device)
        return self._sched

    def loss(self, feats: torch.Tensor, actions: torch.Tensor) -> torch.Tensor:
        """feats (B,T,256,384), actions (B,5,2) in [-1,1] -> scalar MSE."""
        sched = self.sched(feats.device)
        cond = self.trunk(feats)[:, -1]  # last frame's readout has seen it all
        t = torch.randint(0, sched.n_steps, (len(actions),), device=feats.device)
        eps = torch.randn_like(actions)
        noisy = sched.add_noise(actions, eps, t)
        return nn.functional.mse_loss(self.head(noisy, t, cond), eps)

    @torch.no_grad()
    def act(self, feats: torch.Tensor) -> torch.Tensor:
        """feats (B,T,256,384) -> (B,5,2) action chunk in [-1,1]."""
        cond = self.trunk(feats)[:, -1]
        a = self.sched(feats.device).sample(self.head, cond, (self.act_horizon, self.act_dim))
        return a.clamp(-1, 1)


# ------------------------------------------------------------------- EMA ---
class EMA:
    """Exponential moving average of weights (Diffusion Policy's warmup rule:
    decay ramps 0 -> ~0.9999 so early garbage weights wash out quickly)."""

    def __init__(self, model: nn.Module, power: float = 0.75, max_decay: float = 0.9999):
        self.power, self.max_decay = power, max_decay
        self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()}

    @torch.no_grad()
    def update(self, model: nn.Module, step: int):
        decay = min(self.max_decay, 1 - (1 + step) ** -self.power)
        for k, v in model.state_dict().items():
            if v.dtype.is_floating_point:
                self.shadow[k].lerp_(v.detach(), 1 - decay)  # s = d*s + (1-d)*v
            else:
                self.shadow[k].copy_(v)


if __name__ == "__main__":
    torch.manual_seed(0)
    sched = DDPMSchedule(100)
    assert sched.alpha_bar[0] > 0.99 and sched.alpha_bar[-1] < 0.01, "schedule endpoints off"
    assert (sched.alpha_bar.diff() < 0).all(), "alpha_bar must decay monotonically"

    policy = PatchPolicy()
    n = sum(p.numel() for p in policy.parameters())
    feats, acts = torch.randn(4, 2, 256, 384), torch.rand(4, 5, 2) * 2 - 1
    loss = policy.loss(feats, acts)
    loss.backward()  # verify grads flow through trunk AND head
    grads = [p.grad for p in policy.parameters() if p.grad is not None]
    chunk = policy.act(feats[:2])
    assert chunk.shape == (2, 5, 2) and chunk.abs().max() <= 1
    print(f"params {n/1e6:.1f}M | loss {loss.item():.3f} | {len(grads)} tensors got grads")
    print("diffusion self-test passed ✓")