test_vision.py

"""
test_vision.py — unit tests for the M2 image encoder (pure CPU, no sim, no render).

Run: python test_vision.py
"""

from __future__ import annotations

import torch
import torch.nn as nn

from vision import (
    replace_bn_with_gn,
    assert_no_batchnorm,
    SpatialSoftmax,
    ResNet18SpatialSoftmax,
    MultiCameraEncoder,
)


def test_replace_bn_with_gn():
    import torchvision
    m = torchvision.models.resnet18(weights=None)
    # sanity: raw resnet18 HAS batchnorm
    n_bn_before = sum(isinstance(x, nn.BatchNorm2d) for x in m.modules())
    assert n_bn_before > 0
    replace_bn_with_gn(m)
    assert_no_batchnorm(m)  # raises if any remain
    n_gn = sum(isinstance(x, nn.GroupNorm) for x in m.modules())
    assert n_gn == n_bn_before, (n_gn, n_bn_before)
    print(f"  ok: replaced {n_bn_before} BatchNorm2d -> GroupNorm, none remain")


def test_spatial_softmax_correctness():
    # single hot pixel at (row=4, col=0) in a 5x5 map, one keypoint channel.
    ss = SpatialSoftmax(in_channels=1, num_kp=1, temperature=0.01)
    # make proj an identity-ish passthrough so the input activation IS the logit
    with torch.no_grad():
        ss.proj.weight.copy_(torch.ones_like(ss.proj.weight))
        ss.proj.bias.zero_()
    feat = torch.full((1, 1, 5, 5), -10.0)
    feat[0, 0, 4, 0] = 10.0  # bottom-left corner: y=+1, x=-1
    out = ss(feat)  # (1, 2) = [kp_x, kp_y]
    kp_x, kp_y = out[0, 0].item(), out[0, 1].item()
    assert abs(kp_x - (-1.0)) < 0.05, kp_x
    assert abs(kp_y - (1.0)) < 0.05, kp_y
    assert ss.out_dim == 2
    print(f"  ok: spatial softmax localizes hot pixel -> (x={kp_x:.3f}, y={kp_y:.3f})")


def test_multicam_shape_contract():
    enc = MultiCameraEncoder(["agentview", "robot0_eye_in_hand"], proprio_dim=9, num_kp=32)
    B, To = 3, 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)
    assert out.shape == (B, To * (64 * 2 + 9)), out.shape
    assert enc.cond_per_step == 137
    print(f"  ok: MultiCameraEncoder(2 cams, To=2, 76^2) -> {tuple(out.shape)}")


def test_gradients_flow():
    enc = ResNet18SpatialSoftmax(num_kp=16)
    x = torch.randn(2, 3, 76, 76, requires_grad=True)
    loss = enc(x).sum()
    loss.backward()
    n_with_grad = sum(1 for p in enc.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
    n_total = sum(1 for _ in enc.parameters())
    assert n_with_grad == n_total, f"{n_with_grad}/{n_total} params got grads"
    print(f"  ok: all {n_total} encoder params received gradients")


def test_eval_determinism():
    # GroupNorm (no running stats) => two eval passes are bit-identical.
    enc = ResNet18SpatialSoftmax(num_kp=32).eval()
    x = torch.randn(2, 3, 76, 76)
    with torch.no_grad():
        a = enc(x)
        b = enc(x)
    assert torch.equal(a, b), (a - b).abs().max().item()
    print("  ok: eval-mode forward is deterministic (no BN running-stat drift)")


if __name__ == "__main__":
    torch.manual_seed(0)
    tests = [
        ("replace_bn_with_gn", test_replace_bn_with_gn),
        ("spatial_softmax_correctness", test_spatial_softmax_correctness),
        ("multicam_shape_contract", test_multicam_shape_contract),
        ("gradients_flow", test_gradients_flow),
        ("eval_determinism", test_eval_determinism),
    ]
    for name, fn in tests:
        print(f"[{name}]")
        fn()
    print("\nALL VISION TESTS PASSED")