vision.py

"""
vision.py — image observation encoder for Diffusion Policy (M2).

Follows the Diffusion Policy CNN-vision recipe (Chi et al. 2023):
  - ResNet-18 backbone, trained FROM SCRATCH (no ImageNet pretraining).
  - Every BatchNorm replaced by GroupNorm. BatchNorm + EMA + temporally
    correlated batches silently degrades DP; GroupNorm avoids running-stat drift.
  - Global average pool / fc replaced by a SpatialSoftmax head that returns
    expected keypoint coordinates (learned attention over the feature map).
  - One separate encoder per camera (no weight sharing), features concatenated
    with proprioception to form the U-Net's global conditioning vector.

The diffusion core (diffusion_policy.py) is unchanged: MultiCameraEncoder just
produces a (B, To * cond_per_step) global_cond in place of LowDimEncoder.
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


# --------------------------------------------------------------------------- #
# BatchNorm -> GroupNorm surgery
# --------------------------------------------------------------------------- #
def replace_bn_with_gn(module: nn.Module, features_per_group: int = 16) -> nn.Module:
    """Recursively replace every nn.BatchNorm2d with a GroupNorm of matching channels.

    num_groups = max(1, C // features_per_group). Mutates `module` in place and
    also returns it for convenience.
    """
    for name, child in module.named_children():
        if isinstance(child, nn.BatchNorm2d):
            C = child.num_features
            setattr(
                module, name,
                nn.GroupNorm(num_groups=max(1, C // features_per_group), num_channels=C),
            )
        else:
            replace_bn_with_gn(child, features_per_group)
    return module


def assert_no_batchnorm(module: nn.Module) -> None:
    """Raise AssertionError listing any remaining BatchNorm submodule paths."""
    offenders = [
        name for name, m in module.named_modules()
        if isinstance(m, nn.modules.batchnorm._BatchNorm)
    ]
    assert not offenders, f"BatchNorm found (breaks EMA/DP): {offenders}"


# --------------------------------------------------------------------------- #
# Spatial softmax head
# --------------------------------------------------------------------------- #
class SpatialSoftmax(nn.Module):
    """Spatial-softmax keypoint head (Finn et al. 2016; DP vision encoder).

    Projects the feature map to `num_kp` channels, softmaxes each over the
    spatial grid, and returns the expected (x, y) coordinate of each keypoint in
    normalized [-1, 1] image coordinates.

    forward: (B, C, H, W) -> (B, 2*num_kp)
    """

    def __init__(self, in_channels: int, num_kp: int = 32, temperature: float = 1.0):
        super().__init__()
        self.num_kp = num_kp
        self.proj = nn.Conv2d(in_channels, num_kp, kernel_size=1)
        self.register_buffer("temperature", torch.tensor(float(temperature)))
        self._grid_hw: tuple[int, int] | None = None  # cache guard for pos buffers

    def _ensure_grid(self, H: int, W: int, device, dtype) -> None:
        if self._grid_hw == (H, W):
            return
        ys, xs = torch.meshgrid(
            torch.linspace(-1.0, 1.0, H, device=device, dtype=dtype),
            torch.linspace(-1.0, 1.0, W, device=device, dtype=dtype),
            indexing="ij",
        )
        # (H*W,) flattened grids, registered as buffers so .to(device) follows the module
        self.register_buffer("pos_x", xs.reshape(-1), persistent=False)
        self.register_buffer("pos_y", ys.reshape(-1), persistent=False)
        self._grid_hw = (H, W)

    def forward(self, feat: torch.Tensor) -> torch.Tensor:
        B, C, H, W = feat.shape
        self._ensure_grid(H, W, feat.device, feat.dtype)
        logits = self.proj(feat).reshape(B, self.num_kp, H * W)  # (B, K, H*W)
        attn = F.softmax(logits / self.temperature, dim=-1)      # (B, K, H*W)
        kp_x = (attn * self.pos_x).sum(dim=-1)                   # (B, K)
        kp_y = (attn * self.pos_y).sum(dim=-1)                   # (B, K)
        return torch.cat([kp_x, kp_y], dim=-1)                   # (B, 2K)

    @property
    def out_dim(self) -> int:
        return 2 * self.num_kp


# --------------------------------------------------------------------------- #
# Per-camera ResNet-18 + spatial softmax
# --------------------------------------------------------------------------- #
class ResNet18SpatialSoftmax(nn.Module):
    """ResNet-18 (from scratch, GroupNorm) truncated to conv features + SpatialSoftmax.

    forward: (B, 3, h, w) -> (B, 2*num_kp).  For 76x76 input the conv trunk emits
    (B, 512, 3, 3); coarse spatial resolution is expected and fine for keypoints.
    """

    def __init__(self, num_kp: int = 32):
        super().__init__()
        import torchvision

        m = torchvision.models.resnet18(weights=None)  # no pretraining (DP convention)
        # keep conv1 .. layer4, drop avgpool + fc
        self.backbone = nn.Sequential(*list(m.children())[:-2])
        replace_bn_with_gn(self.backbone)
        self.head = SpatialSoftmax(512, num_kp)
        assert_no_batchnorm(self)

    def forward(self, img: torch.Tensor) -> torch.Tensor:
        return self.head(self.backbone(img))

    @property
    def out_dim(self) -> int:
        return self.head.out_dim


# --------------------------------------------------------------------------- #
# Multi-camera encoder (produces the U-Net global conditioning)
# --------------------------------------------------------------------------- #
class MultiCameraEncoder(nn.Module):
    """Encode {camera: (B, To, 3, h, w)} + proprio (B, To, P) -> (B, To*cond_per_step).

    Separate ResNet per camera (DP convention). Per timestep the conditioning is
    [cam0_kp | cam1_kp | ... | proprio], flattened over To to match the low-dim
    encoder's (B, To*obs_dim) global_cond contract.
    """

    def __init__(self, cameras: list[str], proprio_dim: int = 9, num_kp: int = 32):
        super().__init__()
        self.cameras = list(cameras)
        self.proprio_dim = proprio_dim
        self.encoders = nn.ModuleDict(
            {cam: ResNet18SpatialSoftmax(num_kp) for cam in self.cameras}
        )
        self.cond_per_step = 2 * num_kp * len(self.cameras) + proprio_dim
        assert_no_batchnorm(self)

    def forward(self, images: dict[str, torch.Tensor], proprio: torch.Tensor) -> torch.Tensor:
        # images[cam]: (B, To, 3, h, w) ; proprio: (B, To, P)
        B, To = proprio.shape[0], proprio.shape[1]
        feats = []
        for cam in self.cameras:
            x = images[cam]
            b, t = x.shape[0], x.shape[1]
            x = x.reshape(b * t, *x.shape[2:])          # (B*To, 3, h, w)
            f = self.encoders[cam](x)                   # (B*To, 2K)
            feats.append(f.reshape(b, t, -1))           # (B, To, 2K)
        feats.append(proprio)                           # (B, To, P)
        cond = torch.cat(feats, dim=-1)                 # (B, To, cond_per_step)
        return cond.reshape(B, To * self.cond_per_step)

    @property
    def global_cond_dim_per_step(self) -> int:
        return self.cond_per_step


if __name__ == "__main__":
    import torchvision
    enc = MultiCameraEncoder(["agentview", "robot0_eye_in_hand"], proprio_dim=9, num_kp=32)
    assert_no_batchnorm(enc)
    B, To = 4, 2
    imgs = {c: torch.randn(B, To, 3, 76, 76) for c in enc.cameras}
    prop = torch.randn(B, To, 9)
    out = enc(imgs, prop)
    expected = To * (64 * 2 + 9)
    assert out.shape == (B, expected), (out.shape, expected)
    print(f"MultiCameraEncoder OK: -> {tuple(out.shape)}  (cond_per_step={enc.cond_per_step})")
    print("ALL CHECKS PASSED")