| """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_MEAN, T3_STD = 0.49611, 0.18685 |
| |
| 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() |
|
|
| |
| 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) |
| try: |
| from load_baselines import _load_t3_raw |
| 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 |
| 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) |
|
|
| |
| def _init_sitr(self): |
| if GSRL_ROOT not in sys.path: |
| sys.path.insert(0, GSRL_ROOT) |
| from models.networks import SITR_base |
| 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) |
|
|
| |
| 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) |
|
|
|
|
| |
| 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)) |
|
|
