Blame for tools/live_features.py

"""Live per-frame tactile features for the T3 / SITR ACT policies.

Produces the exact 3840-d vector (5 fingers x 768, finger order 0..4) that
`screw_t3_50k` / `screw_sitr5_50k` were trained on — same preprocessing and
constants as the offline encoders (tools/encode_t3.py, tools/encode_sitr.py).

Robot-loop usage:
    from tools.live_features import LiveTactileEncoder
    enc = LiveTactileEncoder("t3", device="cuda")     # or "sitr"
    feat = enc(frames5)      # frames5: (5, H, W) uint8 grayscale, fingers 0..4
    action = infer.step({"state": qpos22, "tactile": feat})

Third-party dependencies (NOT bundled):
  t3   — repo + t3_medium weights, loaded through tactile_fusion's
         load_baselines._load_t3_raw (set TACTILE_FUSION_ROOT if not default)
  sitr — gsrl repo + checkpoints/SITR_B18.pth (set GSRL_ROOT if not default)

Self-test (compares against features stored in an episode HDF5):
    python -m tools.live_features --kind t3   --episode episodes/screwing/episode_000_screwing_0.hdf5
    python -m tools.live_features --kind sitr --episode episodes/screwing/episode_000_screwing_0.hdf5
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

TACTILE_FUSION_ROOT = os.environ.get(
    "TACTILE_FUSION_ROOT", "/home/allenbi/projects_25/tactile_fusion")
GSRL_ROOT = os.environ.get("GSRL_ROOT", "/home/allenbi/projects_25/gsrl")

N_FINGERS = 5
FEAT_DIM = 768

# T3: per-dataset stats computed ONCE on the screwing corpus (encode_t3.py).
# Fixed at training time — never recompute at deployment.
T3_MEAN, T3_STD = 0.49611, 0.18685
# SITR: training-time stats from gsrl/dataloaders.py, applied on [0,1] input.
SITR_MEAN = (-1.2223, -1.8114, -1.7090)
SITR_STD = (11.7932, 12.7956, 13.6452)


def _strip_shadowing_paths():
    """tactile_act's `data/` package shadows tactile_fusion's — drop it before
    importing tactile_fusion modules (same fix as encode_t3.py)."""
    sys.path[:] = [p for p in sys.path
                   if "tactile_act" not in Path(p or ".").resolve().as_posix()]
    for name in list(sys.modules):
        if name == "data" or name.startswith("data."):
            del sys.modules[name]


class LiveTactileEncoder:
    def __init__(self, kind: str, device: str = "cuda"):
        assert kind in ("t3", "sitr"), kind
        self.kind = kind
        self.device = torch.device(device)
        if kind == "t3":
            self._init_t3()
        else:
            self._init_sitr()

    # ------------------------------------------------------------- T3
    def _init_t3(self):
        _strip_shadowing_paths()
        if TACTILE_FUSION_ROOT not in sys.path:
            sys.path.insert(0, TACTILE_FUSION_ROOT)
        cwd = os.getcwd()
        os.chdir(TACTILE_FUSION_ROOT)   # relative third_party paths
        try:
            from load_baselines import _load_t3_raw  # type: ignore
            model, domain = _load_t3_raw(modality="9dtact")
        finally:
            os.chdir(cwd)
        model = model.to(self.device).eval()
        self._encoder = model.encoders[domain]
        self._trunk = model.trunk
        self._model = model  # keep alive
        m = torch.full((1, 3, 1, 1), T3_MEAN, device=self.device)
        s = torch.full((1, 3, 1, 1), T3_STD, device=self.device)
        self._norm = (m, s)

    # ------------------------------------------------------------- SITR
    def _init_sitr(self):
        if GSRL_ROOT not in sys.path:
            sys.path.insert(0, GSRL_ROOT)
        from models.networks import SITR_base  # type: ignore
        model = SITR_base(num_calibration=0)
        ckpt = torch.load(os.path.join(GSRL_ROOT, "checkpoints/SITR_B18.pth"),
                          map_location="cpu")
        sd = ckpt.get("state_dict", ckpt)
        model.load_state_dict(sd, strict=False)
        self._model = model.to(self.device).eval()
        m = torch.tensor(SITR_MEAN, device=self.device).view(1, 3, 1, 1)
        s = torch.tensor(SITR_STD, device=self.device).view(1, 3, 1, 1)
        self._norm = (m, s)

    # ------------------------------------------------------------- common
    def _preprocess(self, frames: np.ndarray) -> torch.Tensor:
        """(B, H, W) uint8 grayscale -> (B, 3, 224, 224) normalized float."""
        x = torch.from_numpy(np.ascontiguousarray(frames)).to(self.device)
        x = x.float() / 255.0
        x = x.unsqueeze(1).expand(-1, 3, -1, -1)
        x = F.interpolate(x, size=(224, 224), mode="bilinear",
                          align_corners=False)
        m, s = self._norm
        return (x - m) / s

    @torch.no_grad()
    def encode_frames(self, frames: np.ndarray) -> np.ndarray:
        """(B, H, W) uint8 -> (B, 768) float32."""
        x = self._preprocess(frames)
        if self.kind == "t3":
            tokens = self._encoder(x)
            tokens = self._trunk(tokens)
            tokens = F.layer_norm(tokens, (tokens.shape[-1],))
            out = tokens[:, 0, :]
        else:
            out = self._model.forward_encoder(x, c=None)[:, 0, :]
        return out.float().cpu().numpy()

    def __call__(self, frames5: np.ndarray) -> np.ndarray:
        """(5, H, W) uint8, fingers 0..4 -> (3840,) float32 (order preserved)."""
        assert frames5.shape[0] == N_FINGERS, frames5.shape
        return self.encode_frames(frames5).reshape(-1)


# ------------------------------------------------------------------ self-test
def _selftest(kind: str, episode: str, device: str, n_frames: int = 4) -> int:
    import h5py
    enc = LiveTactileEncoder(kind, device=device)
    worst = 1.0
    with h5py.File(episode, "r") as hf:
        T = hf["/qpos"].shape[0]
        idx = np.linspace(0, T - 1, n_frames).astype(int)
        for t in idx:
            frames5 = np.stack([hf[f"/images/raw_{f}"][t]
                                for f in range(N_FINGERS)])
            live = enc(frames5).reshape(N_FINGERS, FEAT_DIM)
            stored = np.stack([hf[f"/{kind}/{f}"][t] for f in range(N_FINGERS)])
            cos = np.sum(live * stored, 1) / (
                np.linalg.norm(live, axis=1) * np.linalg.norm(stored, axis=1))
            worst = min(worst, float(cos.min()))
            print(f"t={t:4d}  cos(live, stored) per finger: "
                  + " ".join(f"{c:.4f}" for c in cos))
    ok = worst > 0.999
    print(f"[selftest {kind}] worst cosine {worst:.5f} -> "
          + ("PASS" if ok else "FAIL (feature mismatch — do not deploy)"))
    return 0 if ok else 1


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--kind", choices=["t3", "sitr"], required=True)
    ap.add_argument("--episode", required=True)
    ap.add_argument("--device",
                    default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()
    sys.exit(_selftest(args.kind, args.episode, args.device))