File size: 5,173 Bytes
07c6d07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Pre-encode per-finger tactile camera images through SITR and write the
features back into each episode HDF5 as ``/sitr/<f>`` (T, 768) float32.
The reasoning: SITR is a ~86M-param ViT-base. Running it once per frame is
cheap (~5–10 ms) but re-running it inside every ACT training step is wasteful.
This script does one forward pass per (episode × finger × frame), stores the
cls-token features in the HDF5, and lets ACT consume them via the existing
``--tactile-keys sitr_0 … sitr_4 --tactile-dim 3840`` machinery.
Source images: ``/images/raw_<f>`` (T, 240, 320) uint8 grayscale. Replicated
to 3 channels, resized to 224×224, normalized with the stats SITR was trained
under (``gsrl/dataloaders.py``: sample_mu / sample_std).
Output: ``/sitr/<f>`` (T, 768) float32, gzip-compressed.
"""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
import h5py
import numpy as np
import torch
import torch.nn.functional as F
GSRL_ROOT = "/home/allenbi/projects_25/gsrl"
SITR_CKPT = f"{GSRL_ROOT}/checkpoints/SITR_B18.pth"
# SITR training normalization (gsrl/dataloaders.py:14-15). These are applied
# *after* ToTensor() (i.e. on inputs already scaled to [0, 1]).
SITR_MEAN = (-1.2223, -1.8114, -1.7090)
SITR_STD = (11.7932, 12.7956, 13.6452)
N_FINGERS = 5
SITR_FEAT_DIM = 768
def load_sitr(device: torch.device) -> torch.nn.Module:
sys.path.insert(0, GSRL_ROOT)
from models.networks import SITR_base # type: ignore
model = SITR_base(num_calibration=0) # inference without calibration frames
ckpt = torch.load(SITR_CKPT, map_location="cpu")
state_dict = ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
missing, unexpected = model.load_state_dict(state_dict, strict=False)
print(f"[sitr] loaded {SITR_CKPT} missing={len(missing)} unexpected={len(unexpected)}")
return model.to(device).eval()
def preprocess_batch(images_uint8: np.ndarray, device: torch.device) -> torch.Tensor:
"""images_uint8: (B, H, W) uint8 grayscale → (B, 3, 224, 224) float32 normalized."""
x = torch.from_numpy(images_uint8).to(device).float() / 255.0 # (B, H, W)
x = x.unsqueeze(1).expand(-1, 3, -1, -1) # (B, 3, H, W)
x = F.interpolate(x, size=(224, 224), mode="bilinear", align_corners=False)
mean = torch.tensor(SITR_MEAN, device=device).view(1, 3, 1, 1)
std = torch.tensor(SITR_STD, device=device).view(1, 3, 1, 1)
return (x - mean) / std
@torch.no_grad()
def encode_episode(hf: h5py.File, model: torch.nn.Module, device: torch.device,
batch_size: int, overwrite: bool) -> int:
"""Encode all 5 finger streams and write /sitr/<f>. Returns number of features written."""
T = hf["/qpos"].shape[0]
n_done = 0
sitr_group = hf.require_group("sitr")
for f in range(N_FINGERS):
key = str(f)
if key in sitr_group and not overwrite:
print(f" finger {f}: /sitr/{f} exists, skipping")
continue
if key in sitr_group:
del sitr_group[key]
raw_ds = hf[f"/images/raw_{f}"]
out = np.zeros((T, SITR_FEAT_DIM), dtype=np.float32)
for start in range(0, T, batch_size):
end = min(start + batch_size, T)
batch_np = raw_ds[start:end] # (b, 240, 320) uint8
batch = preprocess_batch(batch_np, device) # (b, 3, 224, 224)
feats = model.forward_encoder(batch, c=None) # (b, N+1, 768)
out[start:end] = feats[:, 0, :].cpu().numpy() # cls token
sitr_group.create_dataset(key, data=out, compression="gzip", compression_opts=4)
n_done += 1
return n_done
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--episodes-dir", type=Path, default=Path("episodes/screwing"))
ap.add_argument("--episode-glob", type=str, default="episode_*.hdf5")
ap.add_argument("--batch-size", type=int, default=32)
ap.add_argument("--device", type=str,
default="cuda" if torch.cuda.is_available() else "cpu")
ap.add_argument("--overwrite", action="store_true",
help="Re-encode and overwrite existing /sitr/<f> datasets.")
args = ap.parse_args()
files = sorted(args.episodes_dir.glob(args.episode_glob))
if not files:
print(f"No HDF5 under {args.episodes_dir} matching {args.episode_glob}", file=sys.stderr)
sys.exit(1)
device = torch.device(args.device)
model = load_sitr(device)
t0 = time.time()
total_written = 0
for i, path in enumerate(files):
with h5py.File(path, "r+") as hf:
T = hf["/qpos"].shape[0]
t_ep = time.time()
n = encode_episode(hf, model, device, args.batch_size, args.overwrite)
dt = time.time() - t_ep
print(f"[{i+1}/{len(files)}] {path.name} T={T} wrote {n} streams ({dt:.1f}s)",
flush=True)
total_written += n
print(f"\ndone in {time.time() - t0:.1f}s; {total_written} /sitr/<f> datasets total.")
if __name__ == "__main__":
main()
|