screw_sitr5_ingraph_50k: in-graph fine-tuned SITR ACT (snapshot step 7200) + TS encoder export + code (sitr branch, bf16 flag, feature_policy in-graph loading)

act_policy/policy.py

@@ -16,6 +16,7 @@ from torch import Tensor, nn
 
 from .config import ACTConfig
 from .model import ACT, FlowMatchingHead
+from .sitr_branch import SITRBranch
 from .tactile_fusion_branch import TactileFusionBranch
 
 
@@ -74,10 +75,15 @@ class ACTPolicy(nn.Module):
             self.tactile_norm = _Normalizer(config.tactile_dim)
         self.action_norm = _Normalizer(config.action_dim)
 
-        # In-graph tactile-fusion encoder (joint-trained or frozen).
-        self.tactile_branch: TactileFusionBranch | None = None
+        # In-graph tactile encoder (joint-trained or frozen). Modality "sitr"
+        # gets the raw-single-frame SITR branch; anything else the
+        # tactile-fusion (HTT) branch.
+        self.tactile_branch: nn.Module | None = None
         if config.tactile_image_keys:
-            self.tactile_branch = TactileFusionBranch(config)
+            if config.tactile_encoder_modality == "sitr":
+                self.tactile_branch = SITRBranch(config)
+            else:
+                self.tactile_branch = TactileFusionBranch(config)
 
         # ImageNet stats for the torchvision-pretrained ResNet backbone.
         self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))

act_policy/sitr_branch.py

@@ -0,0 +1,88 @@
+"""In-graph SITR encoder branch for ACT.
+
+Counterpart of ``TactileFusionBranch`` for ``tactile_encoder_modality="sitr"``.
+SITR's released convention: raw single frames, no background subtraction, no
+frame pairs. The per-episode bg tensor is still accepted (and ignored) so the
+dataset/policy interface stays identical to the HTT in-graph path.
+
+Preprocessing mirrors tools/encode_sitr.py: /255 -> replicate to 3 channels ->
+bilinear 224x224 -> normalize with SITR's training stats -> forward_encoder,
+cls token. With ``tactile_image_group="tactile_imgs"`` the frames arrive
+pre-resized to 224x224 and the resize is a no-op.
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ACTConfig
+
+GSRL_ROOT = os.environ.get("GSRL_ROOT", "/home/allenbi/projects_25/gsrl")
+DEFAULT_SITR_CKPT = os.path.join(GSRL_ROOT, "checkpoints/SITR_B18.pth")
+
+# SITR training normalization (gsrl/dataloaders.py), applied on [0,1] input.
+SITR_MEAN = (-1.2223, -1.8114, -1.7090)
+SITR_STD = (11.7932, 12.7956, 13.6452)
+
+
+class SITRBranch(nn.Module):
+    """SITR_base (ViT-B/16) run per finger frame, cls-token readout.
+
+    Inputs at forward:
+        frames: (B, n_fingers, T, H, W) uint8 grayscale — only frames[:, :, -1]
+                (the current frame) is encoded; pass ``tactile_image_pair=1``.
+        bg:     ignored (SITR consumes raw frames).
+
+    Output: (B, n_fingers * 768) float32 tactile vector.
+    """
+
+    def __init__(self, config: ACTConfig):
+        super().__init__()
+        if not config.tactile_image_keys:
+            raise ValueError("SITRBranch requires non-empty tactile_image_keys.")
+        if int(config.tactile_encoder_feature_dim) != 768:
+            raise ValueError(
+                "SITR_base emits 768-d cls features; set --tactile-encoder-feature-dim 768."
+            )
+        if GSRL_ROOT not in sys.path:
+            sys.path.insert(0, GSRL_ROOT)
+        from models.networks import SITR_base  # noqa: E402
+
+        model = SITR_base(num_calibration=0)  # calibration-free, as in encode_sitr.py
+        ckpt_path = config.tactile_encoder_ckpt or DEFAULT_SITR_CKPT
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        state = ckpt.get("state_dict", ckpt) if isinstance(ckpt, dict) else ckpt
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        print(f"[sitr-branch] loaded {ckpt_path}  missing={len(missing)}  unexpected={len(unexpected)}")
+        # forward_encoder never touches the decoder/contrastive heads; keep them
+        # out of the optimizer.
+        for p in model.decoder_pred.parameters():
+            p.requires_grad = False
+        for p in model.contrastive_head.parameters():
+            p.requires_grad = False
+        self.model = model
+
+        self.register_buffer("image_mean", torch.tensor(SITR_MEAN).view(1, 3, 1, 1))
+        self.register_buffer("image_std", torch.tensor(SITR_STD).view(1, 3, 1, 1))
+        self.target_hw = tuple(config.tactile_encoder_target_hw)
+        self.feature_dim = int(config.tactile_encoder_feature_dim)
+        self.n_fingers = len(config.tactile_image_keys)
+
+        if config.tactile_encoder_freeze:
+            for p in self.model.parameters():
+                p.requires_grad = False
+
+    def forward(self, frames: torch.Tensor, bg: torch.Tensor | None = None) -> torch.Tensor:
+        B, n_fingers, T, H, W = frames.shape
+        x = frames[:, :, -1]                                        # (B, nF, H, W)
+        x = (x.float() if x.dtype == torch.uint8 else x) / 255.0
+        x = x.reshape(B * n_fingers, 1, H, W).expand(-1, 3, -1, -1)
+        if (H, W) != self.target_hw:
+            x = F.interpolate(x, size=self.target_hw, mode="bilinear", align_corners=False)
+        x = (x - self.image_mean) / self.image_std
+        f = self.model.forward_encoder(x, c=None)[:, 0, :]          # cls (B*nF, 768)
+        return f.reshape(B, n_fingers * self.feature_dim)

qistee_integration/feature_policy.py

@@ -75,7 +75,14 @@ def load_feature_policy(checkpoint: str, device: torch.device) -> FeaturePolicy:
             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)
+    # Feature checkpoints (screw_t3_50k / screw_sitr5_50k) have no tactile_branch.
+    # In-graph checkpoints (screw_sitr5_ingraph_50k) carry the fine-tuned encoder
+    # under tactile_branch.* — at deploy those weights live in the run's exported
+    # TorchScript encoder instead, so drop them here. Their tactile_norm was never
+    # initialized (identity), hence the same predict path is exact for both.
+    state = {k: v for k, v in payload["model"].items()
+             if not k.startswith("tactile_branch.")}
+    policy.load_state_dict(state, strict=True)
     return policy.to(device).eval()
 
 

screw_sitr5_ingraph_50k/DEPLOY.md

@@ -0,0 +1,62 @@
+# screw_sitr5_ingraph_50k — SITR fine-tuned IN-GRAPH ACT (screw task)
+
+`2026-08-12` · **snapshot of a still-running training** — `act_best.pt` here is
+the best-val checkpoint at global step 7,200 (train-val L1 0.2652, already below
+every other SITR/HTT condition's all-time best). It will be replaced by the
+final best when the run stops (patience 5 × 200 steps on val).
+
+## What it is
+
+Same 50K protocol as the other `screw_*_50k` runs (bs=32, lr 2e-5 warmup 5K →
+cosine 2e-6, chunk 30, CVAE, val = last 2 episodes), but the SITR encoder
+(`SITR_base`, ViT-B/16, calibration-free) runs **inside the graph** on raw
+single frames and is **fine-tuned jointly** (0.1× LR, i.e. 2e-6). Training
+forward/backward ran under bf16 autocast; the checkpoint is ordinary fp32.
+Naive/full setting: all 5 fingers, full 768-d cls per finger → 3840-d tactile.
+
+Unlike the frozen-feature `screw_sitr5_50k`, the encoder weights here differ
+from `SITR_B18.pth` — you MUST use this folder's TorchScript export, not
+`exported/sitr_feature_encoder.ts.pt`.
+
+## Files
+
+| File | Meaning |
+|---|---|
+| `act_best.pt` | full checkpoint (model incl. fine-tuned encoder + config + optimizer + scheduler) |
+| `sitr_ingraph_feature_encoder.ts.pt` | TorchScript of the FINE-TUNED encoder; contract identical to `exported/*.ts.pt`: input (B,3,224,224) float in [0,1], SITR normalization baked in, output (B,768) cls |
+| `args.json` | training invocation |
+
+## Robot loop (identical to the T3/SITR feature policies)
+
+Raw frames only — NO background capture, NO frame pairs, finger order 0..4
+(camera order, same as training):
+
+```python
+from feature_policy import load_feature_policy, TSFeatureEncoder  # repo ≥ this commit
+policy  = load_feature_policy("screw_sitr5_ingraph_50k/act_best.pt", device)
+encoder = TSFeatureEncoder("screw_sitr5_ingraph_50k/sitr_ingraph_feature_encoder.ts.pt", device)
+...
+feats = encoder(hand.raw_frame())              # (5,240,320) uint8 -> (1,3840)
+chunk = policy.predict_action_chunk(qpos_t, feats)[0]
+```
+
+Requires the updated `qistee_integration/feature_policy.py` (this commit): it
+drops the `tactile_branch.*` weights when loading in-graph checkpoints (they
+live in the TS encoder at deploy). This is exact, not approximate: in-graph
+checkpoints never initialize `tactile_norm` (identity), and the split path was
+validated bit-identical (feature maxdiff 0.0, action-chunk maxdiff 0.0 vs the
+reference in-graph `ACTPolicy` on training frames).
+
+Alternative (tactile_act repo on the machine): `ACTInference.from_checkpoint`
+loads `act_best.pt` directly — needs `act_policy/sitr_branch.py` (this commit)
+plus the gsrl repo for the `SITR_base` class definition and `SITR_B18.pth`
+(constructor init only; weights are immediately overwritten from the ACT
+checkpoint). The TorchScript path avoids the third-party dependency entirely.
+
+## Caveat carried over from the feature-SITR analysis
+
+Fine-tuning does NOT add a reference-frame mechanism: the encoder still
+consumes raw frames, so fingertip-camera appearance drift between the training
+data and the rig shifts features with nothing to anchor them (HTT's bg-capture
+anchors it; SITR-in-graph has no bg). Same deployment interpretation guardrails
+as `screw_sitr5_50k` / `screw_t3_50k` (see FEATURE_PREPROCESSING.md).

screw_sitr5_ingraph_50k/act_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe4b6e0fda53a53c8a2e1060bdc548edffbb8106304cd2a7da5aa5898d3dc937
+size 1537168989

screw_sitr5_ingraph_50k/args.json

@@ -0,0 +1,74 @@
+{
+  "data": "./episodes/screwing",
+  "smoke": false,
+  "episode_glob": "episode_*.hdf5",
+  "state_dim": 22,
+  "action_dim": 22,
+  "image_keys": [],
+  "image_shape": [
+    3,
+    240,
+    320
+  ],
+  "tactile_dim": 0,
+  "tactile_keys": [],
+  "tactile_image_keys": [
+    "raw_0",
+    "raw_1",
+    "raw_2",
+    "raw_3",
+    "raw_4"
+  ],
+  "tactile_image_pair": 1,
+  "tactile_image_group": "tactile_imgs",
+  "tactile_bg_group": "tactile_bg_224",
+  "tactile_encoder_modality": "sitr",
+  "tactile_encoder_ckpt": "/home/allenbi/projects_25/gsrl/checkpoints/SITR_B18.pth",
+  "tactile_encoder_pretrain_config": "/home/allenbi/projects_25/tactile_fusion/config/model/pretrain.yaml",
+  "tactile_encoder_feature_dim": 768,
+  "tactile_encoder_target_hw": [
+    224,
+    224
+  ],
+  "tactile_encoder_lr_scale": 0.1,
+  "freeze_tactile_encoder": false,
+  "state_history": 1,
+  "tactile_history": 1,
+  "delta_actions": false,
+  "state_indices": null,
+  "action_indices": null,
+  "chunk_size": 30,
+  "n_action_steps": 30,
+  "no_vae": false,
+  "use_flow_matching": false,
+  "fm_hidden_dim": 512,
+  "fm_n_mlp_layers": 3,
+  "fm_time_embed_dim": 128,
+  "fm_inference_steps": 10,
+  "lr": 2e-05,
+  "lr_backbone": 1e-05,
+  "weight_decay": 0.0001,
+  "lr_warmup_steps": 5000,
+  "lr_warmup_init": 1e-06,
+  "lr_cosine": true,
+  "lr_min": 2e-06,
+  "batch_size": 32,
+  "num_workers": 4,
+  "bf16": true,
+  "steps": 50000,
+  "log_every": 200,
+  "save_every": 10000,
+  "out": "./checkpoints",
+  "device": "cuda",
+  "resume": null,
+  "run_name": "screw_sitr5_ingraph_50k",
+  "tb_dir": "./runs",
+  "no_tb": false,
+  "no_stats": false,
+  "stats_batch_size": 32,
+  "stats_num_workers": 4,
+  "val_episodes": 2,
+  "val_every": 200,
+  "val_batches": 0,
+  "patience": 5
+}

screw_sitr5_ingraph_50k/sitr_ingraph_feature_encoder.ts.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9b05621e199620a49c63d6ceadd3d64cd7f3b19360af0f50f1f0a034b342f9e
+size 346358387

train.py

@@ -65,7 +65,9 @@ def parse_args() -> argparse.Namespace:
                    help="HDF5 group for per-episode bg images. Use 'tactile_bg_224' alongside "
                         "the resized tactile_imgs.")
     p.add_argument("--tactile-encoder-modality", type=str, default="9dtact",
-                   choices=["9dtact", "gsmini"], help="Tactile-fusion modality.")
+                   choices=["9dtact", "gsmini", "sitr"],
+                   help="In-graph tactile encoder: tactile-fusion modality, or 'sitr' "
+                        "for the raw-single-frame SITR branch (SITR_base, 768-d cls).")
     p.add_argument("--tactile-encoder-ckpt", type=str,
                    default="/home/allenbi/projects_25/tactile_fusion/checkpoints/pretrain_joint_alignscratch/"
                            "pretrain_multimodal_20260521_005341/checkpoint_step_34000.pth")
@@ -123,6 +125,9 @@ def parse_args() -> argparse.Namespace:
 
     p.add_argument("--batch-size", type=int, default=8)
     p.add_argument("--num-workers", type=int, default=4)
+    p.add_argument("--bf16", action="store_true",
+                   help="Run the training forward/backward under torch.autocast(bfloat16). "
+                        "Weights, optimizer, checkpoints and validation stay fp32.")
     p.add_argument("--steps", type=int, default=100_000)
     p.add_argument("--log-every", type=int, default=50)
     p.add_argument("--save-every", type=int, default=5000)
@@ -480,7 +485,9 @@ def main() -> None:
             if step >= args.steps or early_stop:
                 break
             batch = move_to(batch, args.device)
-            loss, info = policy(batch)
+            with torch.autocast(device_type="cuda" if "cuda" in args.device else "cpu",
+                                dtype=torch.bfloat16, enabled=args.bf16):
+                loss, info = policy(batch)
             optimizer.zero_grad(set_to_none=True)
             loss.backward()
             optimizer.step()