Qistee/HTT integration: feature_policy.py + README (validated bit-identical)

qistee_integration/README.md

@@ -0,0 +1,47 @@
+# T3 / SITR inference in the Qistee/HTT robot pipeline
+
+`2026-08-12` · validated end-to-end on the main workstation: FeaturePolicy
+action chunks are BIT-IDENTICAL (max diff 0.0) to the tactile_act reference
+inference on the same inputs; TorchScript features match stored training
+features at cos ≥ 0.99997.
+
+## Why sitr/t3 could not run in the existing pipeline
+
+`htt_policy.py` is hard-wired to the in-graph HTT checkpoint:
+1. `checkpoint_config()` hardcodes `tactile_dim=960` — T3/SITR are 3840;
+2. `HTTPolicy` always builds `TactileImageBranch` and loads `strict=True` —
+   T3/SITR checkpoints have no `tactile_branch.*` weights;
+3. its preprocessing (bg-subtract + 2-frame pairs) is the HTT convention —
+   T3/SITR features come from raw single frames with their own normalizations.
+
+## Fix (3 files, no third-party code needed — only torch)
+
+1. Copy `feature_policy.py` (this folder) into the Qistee/HTT repo root.
+2. Download the TorchScript encoders from
+   hf.co/AllenBi21/tactile-act-screw-baselines → `exported/`:
+   `t3_feature_encoder.ts.pt`, `sitr_feature_encoder.ts.pt` (~345 MB each;
+   encoder-specific normalization baked into the graphs).
+3. Checkpoints: `screw_t3_50k/act_best.pt`, `screw_sitr5_50k/act_best.pt`
+   (same HF repo).
+
+Robot loop (mirror inference.py's sharpa mode, minus background capture —
+T3/SITR need NO background, NO frame pairs):
+
+```python
+from feature_policy import load_feature_policy, TSFeatureEncoder
+policy = load_feature_policy("checkpoint/t3_act_best.pt", device)
+encoder = TSFeatureEncoder("exported/t3_feature_encoder.ts.pt", device)
+...
+feats = encoder(hand.raw_frame())              # (5,240,320) -> (1,3840)
+chunk = policy.predict_action_chunk(qpos_t, feats)[0]
+```
+
+Keep 30 Hz and the 30-step chunk queue exactly as in the HTT path. Finger
+order 0..4 must match the camera order used in training (same as HTT's).
+
+## Sanity check before robot use
+
+Load an episode HDF5 from hf datasets AllenBi21/tactile-act-screw-episodes,
+feed frame t's five raw images + qpos, and confirm the first action of the
+chunk changes smoothly across consecutive t (and features vs stored /t3//sitr
+groups have cosine > 0.999 — same test as tools/live_features.py --smoke).

qistee_integration/feature_policy.py

@@ -0,0 +1,102 @@
+"""ACT policies conditioned on EXTERNAL tactile features (T3 / SITR).
+
+Counterpart of ``htt_policy.py`` for the ``screw_t3_50k`` / ``screw_sitr5_50k``
+checkpoints (hf.co/AllenBi21/tactile-act-screw-baselines). Differences vs the
+HTT policy:
+
+  * tactile_dim = 3840 (5 fingers x 768), read from the checkpoint config —
+    NOT hardcoded;
+  * there is NO in-graph tactile branch: features come from a TorchScript
+    encoder (``t3_feature_encoder.ts.pt`` / ``sitr_feature_encoder.ts.pt``,
+    same HF repo, ``exported/``) — no third-party code needed, only torch;
+  * NO background subtraction and NO frame pairs: both encoders consume raw
+    single frames (their authors' released conventions; normalization
+    constants are baked inside the TorchScript graphs).
+
+Robot loop (mirrors inference.py's sharpa mode, minus background capture):
+
+    policy = load_feature_policy("checkpoint/act_best.pt", device)
+    encoder = TSFeatureEncoder("exported/t3_feature_encoder.ts.pt", device)
+    ...
+    feats = encoder(hand.raw_frame())                # (5,240,320) uint8 -> (3840,)
+    chunk = policy.predict_action_chunk(qpos_t, feats)
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from act_policy.config import ACTConfig
+from act_policy.model import ACT
+from act_policy.policy import _Normalizer
+
+FINGER_COUNT = 5
+FEATURE_DIM = 768
+
+
+class FeaturePolicy(nn.Module):
+    """Inference-only ACT whose state-dict names match the feature checkpoints
+    (same layout as HTTPolicy minus ``tactile_branch``)."""
+
+    def __init__(self, config: ACTConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.model = ACT(config)
+        self.state_norm = _Normalizer(config.state_dim)
+        self.tactile_norm = _Normalizer(config.tactile_dim)
+        self.action_norm = _Normalizer(config.action_dim)
+        self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+        self.register_buffer("image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    @torch.inference_mode()
+    def predict_action_chunk(self, qpos: Tensor, tactile_features: Tensor) -> Tensor:
+        """qpos (B,22); tactile_features (B,3840) from the TorchScript encoder."""
+        self.eval()
+        batch = {
+            "state": self.state_norm.normalize(qpos),
+            "tactile": self.tactile_norm.normalize(tactile_features),
+        }
+        actions, _ = self.model(batch)
+        return self.action_norm.unnormalize(actions)
+
+
+def load_feature_policy(checkpoint: str, device: torch.device) -> FeaturePolicy:
+    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    cfg_dict = dict(payload["config"])
+    fields = {f for f in ACTConfig.__dataclass_fields__}  # type: ignore[attr-defined]
+    cfg_dict = {k: v for k, v in cfg_dict.items() if k in fields}
+    cfg_dict["image_keys"] = []
+    config = ACTConfig(**cfg_dict)
+    if config.tactile_dim != FINGER_COUNT * FEATURE_DIM:
+        raise ValueError(
+            f"checkpoint tactile_dim={config.tactile_dim}, expected "
+            f"{FINGER_COUNT * FEATURE_DIM} — is this a T3/SITR checkpoint? "
+            "(HTT in-graph checkpoints belong to htt_policy.py)")
+    policy = FeaturePolicy(config)
+    policy.load_state_dict(payload["model"], strict=True)
+    return policy.to(device).eval()
+
+
+class TSFeatureEncoder:
+    """TorchScript T3/SITR feature encoder. Input: (5,H,W) uint8/float raw
+    grayscale fingertip frames (fingers 0..4). Output: (1,3840) float32.
+    Preprocessing here = /255 -> 3ch -> bilinear 224; the encoder-specific
+    normalization is baked into the TorchScript graph."""
+
+    def __init__(self, path: str, device: torch.device) -> None:
+        self.module = torch.jit.load(path, map_location=device).eval()
+        self.device = device
+
+    @torch.inference_mode()
+    def __call__(self, frames) -> Tensor:
+        x = torch.as_tensor(frames, device=self.device).float()
+        if x.ndim != 3 or x.shape[0] != FINGER_COUNT:
+            raise ValueError(f"expected (5,H,W) raw frames, got {tuple(x.shape)}")
+        if x.max() > 1.0:
+            x = x / 255.0
+        x = x.unsqueeze(1).expand(-1, 3, -1, -1)
+        x = F.interpolate(x, size=(224, 224), mode="bilinear", align_corners=False)
+        feats = self.module(x)                     # (5, 768), finger order kept
+        return feats.reshape(1, FINGER_COUNT * FEATURE_DIM)