"""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))