test/test_tactile_vae.py

"""Smoke/contract tests for tactile_vae.model.tactile_vae."""

from pathlib import Path
import sys
import tempfile

import torch

_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

from tactile_vae.model.tactile_vae import TactileVAE, VAELoss, load_pretrained


def _small_model() -> TactileVAE:
    return TactileVAE(
        img_size=64,
        patch_size=16,
        in_chans=3,
        embed_dim=96,
        encoder_depth=2,
        encoder_heads=4,
        decoder_embed_dim=96,
        decoder_depth=2,
        decoder_heads=4,
        latent_dim=48,
    )


def test_forward_contract_shapes():
    torch.manual_seed(0)
    model = _small_model()
    x = torch.randn(2, 3, 64, 64)

    out = model(x)
    assert out["x_hat"].shape == x.shape
    assert out["mu"].shape == (2, 48)
    assert out["logvar"].shape == (2, 48)
    assert out["z"].shape == (2, 48)
    assert out["pred_patches"].shape == (2, (64 // 16) ** 2, 16 * 16 * 3)


def test_losses_finite_and_backprop():
    torch.manual_seed(0)
    model = _small_model()
    criterion = VAELoss(beta=1e-3, recon_type="l1", ssim_weight=0.1)
    x = torch.randn(2, 3, 64, 64)

    out = model(x)
    losses = criterion(out["x_hat"], x, out["mu"], out["logvar"])
    for name, val in losses.items():
        assert torch.isfinite(val).all(), f"non-finite loss component: {name}"

    losses["total"].backward()
    grad_ok = any(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters() if p.requires_grad)
    assert grad_ok, "expected finite gradients after backward"


def test_tiny_overfit_recon_drops():
    torch.manual_seed(0)
    model = _small_model()
    criterion = VAELoss(beta=1e-3, recon_type="l1", ssim_weight=0.0)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)

    x = torch.tanh(torch.randn(8, 3, 64, 64))

    with torch.no_grad():
        out0 = model(x)
        start = criterion(out0["x_hat"], x, out0["mu"], out0["logvar"])["recon_total"].item()

    for _ in range(30):
        out = model(x)
        losses = criterion(out["x_hat"], x, out["mu"], out["logvar"])
        opt.zero_grad(set_to_none=True)
        losses["total"].backward()
        opt.step()

    with torch.no_grad():
        out1 = model(x)
        end = criterion(out1["x_hat"], x, out1["mu"], out1["logvar"])["recon_total"].item()

    assert end < start, f"expected recon to decrease; start={start:.4f}, end={end:.4f}"


def test_checkpoint_roundtrip_consistent_eval():
    torch.manual_seed(0)
    model = _small_model().eval()
    x = torch.randn(2, 3, 64, 64)

    with torch.no_grad():
        out_ref = model(x, sample=False)["x_hat"]

    with tempfile.TemporaryDirectory() as tmpdir:
        ckpt_path = Path(tmpdir) / "tactile_vae.pt"
        torch.save(model.state_dict(), ckpt_path)
        loaded = load_pretrained(
            checkpoint=str(ckpt_path),
            model_kwargs={
                "img_size": 64,
                "patch_size": 16,
                "in_chans": 3,
                "embed_dim": 96,
                "encoder_depth": 2,
                "encoder_heads": 4,
                "decoder_embed_dim": 96,
                "decoder_depth": 2,
                "decoder_heads": 4,
                "latent_dim": 48,
            },
            freeze=True,
            strict=True,
        ).eval()

        with torch.no_grad():
            out_loaded = loaded(x, sample=False)["x_hat"]

    torch.testing.assert_close(out_ref, out_loaded, atol=1e-6, rtol=1e-5)


if __name__ == "__main__":
    test_forward_contract_shapes()
    test_losses_finite_and_backprop()
    test_tiny_overfit_recon_drops()
    test_checkpoint_roundtrip_consistent_eval()
    print("All tactile VAE tests passed.")