Initial upload of tactile_vae (code, model, config, inference)

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+inference/reconstruction_grid.png filter=lfs diff=lfs merge=lfs -text
+inference/vae_baseline_3/reconstruction_grid.png filter=lfs diff=lfs merge=lfs -text
+test/color_jitter_compare.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,25 @@
+# Large training data (16 GB parquet) — do not push
+data/
+
+# Python build / cache
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+build/
+dist/
+
+# Test caches
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+
+# Local outputs / regenerable artifacts
+runs/
+test_output/
+
+# OS / editor
+.DS_Store
+.vscode/
+.idea/
+*.swp

README.md

@@ -0,0 +1,20 @@
+# tactile_vae
+
+ViT-based tactile variational autoencoder.
+
+## Architecture
+- Encoder: `PatchEmbed + Transformer blocks + fixed sin-cos positional embedding`
+- Latent: `mu/logvar` + reparameterization
+- Decoder: latent-conditioned transformer patch decoder + unpatchify
+
+## Usage
+
+```python
+import torch
+from tactile_vae.model import TactileVAE, VAELoss
+
+model = TactileVAE()
+out = model(torch.randn(2, 3, 128, 128))
+loss_fn = VAELoss(beta=1.0)
+losses = loss_fn(out["x_hat"], torch.randn(2, 3, 128, 128), out["mu"], out["logvar"])
+```

__init__.py

@@ -0,0 +1,3 @@
+"""tactile_vae package."""
+
+__version__ = "0.0.1"

config/__init__.py

@@ -0,0 +1,32 @@
+"""Configuration dataclasses for tactile_vae."""
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TactileVAEModelConfig:
+    img_size: int = 128
+    patch_size: int = 16
+    in_chans: int = 3
+    embed_dim: int = 256
+    encoder_depth: int = 4
+    encoder_heads: int = 8
+    decoder_embed_dim: int = 192
+    decoder_depth: int = 4
+    decoder_heads: int = 8
+    mlp_ratio: float = 4.0
+    latent_dim: int = 128
+
+
+@dataclass
+class TactileVAELossConfig:
+    beta: float = 1.0
+    recon_type: str = "l1"
+    ssim_weight: float = 0.0
+    ssim_window_size: int = 11
+
+
+@dataclass
+class TactileVAETrainConfig:
+    model: TactileVAEModelConfig = field(default_factory=TactileVAEModelConfig)
+    loss: TactileVAELossConfig = field(default_factory=TactileVAELossConfig)

config/train_vae.yaml

@@ -0,0 +1,75 @@
+# Tactile VAE training config.
+#
+# Paths can be absolute or relative to the config file. The training script
+# resolves relative paths against the repo root (`tactile_world_model/`).
+seed: 0
+device: auto                      # auto | cuda | cuda:0 | cpu
+
+# Each run lives at <runs_root>/<run_id>. If that directory already contains
+# a ckpt_last.pt the trainer auto-resumes from it (override with --no-resume
+# or set train.resume_from explicitly). Leave run_id null to auto-generate
+# a timestamped one like "run_2026-05-16_05-30-12".
+runs_root: runs
+run_id: vae_baseline
+
+model:
+  img_size: 128
+  patch_size: 16
+  in_chans: 3
+  embed_dim: 384
+  encoder_depth: 3
+  encoder_heads: 6
+  decoder_embed_dim: 384
+  decoder_depth: 8
+  decoder_heads: 12
+  mlp_ratio: 4.0
+  latent_dim: 256
+
+loss:
+  beta: 1.0e-3
+  recon_type: l1                  # l1 | mse
+  perceptual_type: lpips           # ssim | lpips
+  ssim_weight: 0.1                # used for SSIM; also fallback for LPIPS if lpips_weight unset
+  ssim_window_size: 11
+  lpips_weight: 1               # used when perceptual_type = lpips
+
+data:
+  root: tactile_vae/data
+  image_size: 128
+  cache_files: 1
+  splits_path: tactile_vae/dataset/splits_subset.json
+  meta_columns: []                # leave empty unless you actually consume them
+  return_meta: false
+  color_jitter:                   # null to disable
+    brightness: 0.2
+    contrast: 0.2
+    saturation: 0.2
+    hue: 0.05
+
+optim:
+  lr: 1.0e-4
+  weight_decay: 0.05
+  betas: [0.9, 0.95]
+  eps: 1.0e-8
+
+scheduler:
+  type: cosine                    # cosine | constant
+  warmup_steps: 1000
+  min_lr_ratio: 0.1
+
+train:
+  batch_size: 128
+  num_workers: 8
+  epochs: 50000
+  max_steps: null                 # int to cap total steps; null = run all epochs
+  log_every: 50
+  val_every_steps: 2000
+  num_val_batches: 100            # cap validation batches per check
+  sample_every_steps: 2000
+  num_sample_images: 16
+  ckpt_every_steps: 5000
+  keep_last_ckpts: 3              # rotate periodic checkpoints; "best" is kept forever
+  gradient_clip_norm: 1.0
+  amp: true                       # mixed-precision on CUDA; ignored on CPU
+  amp_dtype: bf16                 # enforced to bf16 by train_vae.py
+  resume_from: null               # path to checkpoint to resume from

dataset/__init__.py

@@ -0,0 +1,19 @@
+from tactile_vae.dataset.dataset import (
+    DEFAULT_DATA_ROOT,
+    DEFAULT_FILE_GLOB,
+    META_COLUMNS,
+    ColorJitterConfig,
+    ParquetFileShuffleSampler,
+    TactileParquetDataset,
+    default_image_transform,
+)
+
+__all__ = [
+    "DEFAULT_DATA_ROOT",
+    "DEFAULT_FILE_GLOB",
+    "META_COLUMNS",
+    "ColorJitterConfig",
+    "ParquetFileShuffleSampler",
+    "TactileParquetDataset",
+    "default_image_transform",
+]

dataset/dataset.py

@@ -0,0 +1,343 @@
+"""Map-style dataset over the fota_unlabeled tactile parquets.
+
+Each parquet row holds a JPEG-encoded `image` blob plus per-frame metadata
+(`obj_name`, pose, force, etc.). The dataset enumerates all parquet files
+under `root`, indexes them by file metadata only (no full read at init),
+and decodes JPEGs on demand.
+
+For DataLoader use:
+  - Each worker keeps an LRU of `cache_files` parquet image-columns in memory
+    (PyArrow `ChunkedArray` of binary refs). Random reads within a cached
+    file are zero-copy slices; misses pay one full column read (~2 GB).
+  - Combine with `ParquetFileShuffleSampler` for shuffle without thrashing:
+    shuffles file order each epoch and shuffles indices inside each file,
+    so a worker only ever pulls from one cached file at a time.
+"""
+from __future__ import annotations
+
+import io
+import json
+from collections import OrderedDict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Iterator, Mapping, Sequence
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, Sampler
+from torchvision.transforms import ColorJitter
+
+DEFAULT_DATA_ROOT = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/data")
+DEFAULT_FILE_GLOB = "train-*-of-*.parquet"
+DEFAULT_SPLITS_PATH = Path(__file__).with_name("splits.json")
+VALID_SPLITS = ("train", "val", "test")
+
+# Metadata columns we expose by default — float32 scalars usable as conditioning.
+META_COLUMNS: tuple[str, ...] = (
+    "x_mm", "y_mm", "z_mm",
+    "quat_x", "quat_y", "quat_z", "quat_w",
+    "f_x", "f_y", "f_z",
+    "grid_z_max", "grid_z_mean",
+)
+
+
+@dataclass
+class ColorJitterConfig:
+    """Training-time color-jitter knobs (forwarded to `torchvision.transforms.ColorJitter`).
+
+    Each value is the magnitude of perturbation around the identity. `0.0`
+    disables that channel; when all four are zero the config is a no-op.
+    Defaults match a mild augmentation reasonable for tactile RGB frames.
+
+    Treat as part of the training config — eval/inference paths should pass
+    `color_jitter=None` to the dataset so augmentation is off.
+    """
+    brightness: float = 0.2
+    contrast: float = 0.2
+    saturation: float = 0.2
+    hue: float = 0.1
+
+    def is_noop(self) -> bool:
+        return not any((self.brightness, self.contrast, self.saturation, self.hue))
+
+    def build(self) -> ColorJitter | None:
+        if self.is_noop():
+            return None
+        return ColorJitter(
+            brightness=self.brightness,
+            contrast=self.contrast,
+            saturation=self.saturation,
+            hue=self.hue,
+        )
+
+
+def _resolve_color_jitter(
+    cfg: ColorJitterConfig | Mapping[str, float] | ColorJitter | None,
+) -> ColorJitter | None:
+    if cfg is None:
+        return None
+    if isinstance(cfg, ColorJitter):
+        return cfg
+    if isinstance(cfg, Mapping):
+        cfg = ColorJitterConfig(**cfg)
+    if isinstance(cfg, ColorJitterConfig):
+        return cfg.build()
+    raise TypeError(f"unsupported color_jitter type: {type(cfg).__name__}")
+
+
+def default_image_transform(
+    image_size: int = 128,
+    color_jitter: ColorJitterConfig | Mapping[str, float] | ColorJitter | None = None,
+) -> Callable[[Image.Image], torch.Tensor]:
+    """Resize → (optional) color-jitter → CHW float32 in [0, 1].
+
+    `color_jitter` accepts a `ColorJitterConfig`, a dict of its fields, an
+    already-built `torchvision.transforms.ColorJitter`, or `None` to disable.
+    Apply jitter on PIL (before tensor conversion) — torchvision's RNG is
+    per-worker in PyTorch DataLoaders, so per-sample augmentation is correct
+    with `num_workers > 0`.
+    """
+    jitter = _resolve_color_jitter(color_jitter)
+
+    def _tx(img: Image.Image) -> torch.Tensor:
+        if img.size != (image_size, image_size):
+            img = img.resize((image_size, image_size), Image.BILINEAR)
+        if jitter is not None:
+            img = jitter(img)
+        arr = np.array(img, dtype=np.uint8, copy=True)  # H, W, 3, writable
+        t = torch.from_numpy(arr).permute(2, 0, 1).contiguous().float().div_(255.0) # [0, 1]
+        return t
+    return _tx
+
+
+class TactileParquetDataset(Dataset):
+    """Lazy parquet dataset that yields decoded RGB tensors (+ optional metadata).
+
+    Args:
+        root: directory containing parquet shards.
+        file_glob: glob pattern relative to `root`. Defaults to
+            `train-*-of-*.parquet`; `.raw.parquet` shards are excluded.
+        image_size: output square size if `transform` is the default.
+        transform: callable(PIL.Image) -> Tensor. If None, uses
+            `default_image_transform(image_size, color_jitter=color_jitter)`.
+            Pass an explicit transform to fully override preprocessing;
+            `color_jitter` is then ignored.
+        color_jitter: training-time RGB augmentation, applied inside the
+            default transform. Accepts a `ColorJitterConfig`, a dict of its
+            fields, a prebuilt `torchvision.transforms.ColorJitter`, or
+            `None` (default — augmentation off). For eval/inference, leave
+            as `None`.
+        split: one of `"train"`, `"val"`, `"test"`, or `None`. If set, the
+            dataset only exposes rows assigned to that split by `splits_path`.
+        splits_path: JSON manifest produced by `make_splits.py`. Defaults to
+            `tactile_vae/dataset/splits.json`. Only consulted when
+            `split is not None`.
+        return_meta: if True, `__getitem__` returns `(image, meta_dict)` where
+            `meta_dict` has float32 tensors for the columns in `meta_columns`.
+        meta_columns: which metadata columns to surface.
+        cache_files: how many parquet image-columns to keep in memory per worker.
+    """
+
+    def __init__(
+        self,
+        root: str | Path = DEFAULT_DATA_ROOT,
+        file_glob: str = DEFAULT_FILE_GLOB,
+        image_size: int = 128,
+        transform: Callable[[Image.Image], torch.Tensor] | None = None,
+        color_jitter: ColorJitterConfig | Mapping[str, float] | ColorJitter | None = None,
+        split: str | None = None,
+        splits_path: str | Path | None = None,
+        return_meta: bool = False,
+        meta_columns: Sequence[str] = META_COLUMNS,
+        cache_files: int = 1,
+    ) -> None:
+        root = Path(root)
+        files = sorted(root.glob(file_glob))
+        files = [p for p in files if ".raw." not in p.name]
+        if not files:
+            raise FileNotFoundError(f"no parquet shards under {root} matching {file_glob!r}")
+
+        self.files: list[Path] = files
+        self.image_size = image_size
+        self.transform = transform or default_image_transform(
+            image_size, color_jitter=color_jitter
+        )
+        self.split = split
+        self.return_meta = return_meta
+        self.meta_columns: list[str] = list(meta_columns)
+        self.cache_files = max(1, int(cache_files))
+
+        # True per-file row counts from parquet metadata (independent of split).
+        file_row_counts: list[int] = [pq.ParquetFile(p).metadata.num_rows for p in files]
+
+        # Build `_samples`: (N, 2) int64 array of (file_idx, row_idx), sorted
+        # by file_idx then row_idx. Grouping by file keeps contiguous global
+        # indices on the same parquet — that's what the file-cache relies on.
+        if split is None:
+            per_file_rows: list[np.ndarray] = [
+                np.arange(n, dtype=np.int64) for n in file_row_counts
+            ]
+        else:
+            if split not in VALID_SPLITS:
+                raise ValueError(f"split must be one of {VALID_SPLITS}; got {split!r}")
+            spath = Path(splits_path) if splits_path is not None else DEFAULT_SPLITS_PATH
+            if not spath.exists():
+                raise FileNotFoundError(
+                    f"split={split!r} requested but {spath} not found. "
+                    "Generate it with `python tactile_vae/dataset/make_splits.py`."
+                )
+            with spath.open() as f:
+                manifest = json.load(f)
+            split_map: dict[str, list[int]] = manifest["splits"][split]
+            per_file_rows = []
+            for p, total_rows in zip(files, file_row_counts):
+                rows = split_map.get(p.name, [])
+                arr = np.asarray(rows, dtype=np.int64)
+                if arr.size and (arr.min() < 0 or arr.max() >= total_rows):
+                    raise ValueError(
+                        f"split manifest for {p.name} has row index out of range "
+                        f"[0, {total_rows}): min={int(arr.min())} max={int(arr.max())}"
+                    )
+                # Sort within file so cache reads progress monotonically.
+                per_file_rows.append(np.sort(arr))
+
+        chunks = [
+            np.stack([np.full(rows.shape[0], fi, dtype=np.int64), rows], axis=1)
+            for fi, rows in enumerate(per_file_rows)
+            if rows.shape[0] > 0
+        ]
+        self._samples: np.ndarray = (
+            np.concatenate(chunks, axis=0)
+            if chunks
+            else np.zeros((0, 2), dtype=np.int64)
+        )
+        per_file_counts = [rows.shape[0] for rows in per_file_rows]
+        self._per_file_offsets: np.ndarray = np.concatenate(
+            [[0], np.cumsum(per_file_counts, dtype=np.int64)]
+        )
+
+        # Per-worker caches; populated lazily.
+        self._image_cache: "OrderedDict[int, object]" = OrderedDict()
+        self._meta_cache: "OrderedDict[int, dict[str, np.ndarray]]" = OrderedDict()
+
+    def __len__(self) -> int:
+        return int(self._samples.shape[0])
+
+    @property
+    def num_files(self) -> int:
+        return len(self.files)
+
+    def file_lengths(self) -> list[int]:
+        """Number of samples drawn from each file under the current split."""
+        return np.diff(self._per_file_offsets).astype(int).tolist()
+
+    def sample_id(self, idx: int) -> str:
+        """Stable ID for sample `idx`: `<file-stem>:<row>` — matches `make_splits.py`."""
+        fi, ri = self._locate(idx)
+        return f"{self.files[fi].stem}:{ri}"
+
+    def __getstate__(self) -> dict:
+        # Strip caches before pickling to workers — pyarrow objects don't
+        # pickle cleanly and the cache should be per-worker anyway.
+        state = self.__dict__.copy()
+        state["_image_cache"] = OrderedDict()
+        state["_meta_cache"] = OrderedDict()
+        return state
+
+    def _locate(self, idx: int) -> tuple[int, int]:
+        n = self.__len__()
+        if idx < 0:
+            idx += n
+        if idx < 0 or idx >= n:
+            raise IndexError(idx)
+        fi, ri = self._samples[idx]
+        return int(fi), int(ri)
+
+    def _get_image_column(self, fi: int):
+        col = self._image_cache.get(fi)
+        if col is not None:
+            self._image_cache.move_to_end(fi)
+            return col
+        col = pq.read_table(self.files[fi], columns=["image"]).column("image")
+        self._image_cache[fi] = col
+        while len(self._image_cache) > self.cache_files:
+            self._image_cache.popitem(last=False)
+        return col
+
+    def _get_meta_columns(self, fi: int) -> dict[str, np.ndarray]:
+        if not self.meta_columns:
+            return {}
+        cached = self._meta_cache.get(fi)
+        if cached is not None:
+            self._meta_cache.move_to_end(fi)
+            return cached
+        tbl = pq.read_table(self.files[fi], columns=self.meta_columns)
+        cached = {name: tbl.column(name).to_numpy(zero_copy_only=False)
+                  for name in self.meta_columns}
+        self._meta_cache[fi] = cached
+        while len(self._meta_cache) > self.cache_files:
+            self._meta_cache.popitem(last=False)
+        return cached
+
+    def __getitem__(self, idx: int):
+        fi, ri = self._locate(idx)
+        col = self._get_image_column(fi)
+        raw = col[ri].as_py()
+        with Image.open(io.BytesIO(raw)) as im:
+            im = im.convert("RGB")
+            tensor = self.transform(im)
+
+        if not self.return_meta:
+            return tensor
+
+        meta_np = self._get_meta_columns(fi)
+        meta = {
+            name: torch.as_tensor(float(arr[ri]), dtype=torch.float32)
+            for name, arr in meta_np.items()
+        }
+        return tensor, meta
+
+
+class ParquetFileShuffleSampler(Sampler[int]):
+    """Shuffle that keeps each worker reading from one parquet at a time.
+
+    Shuffles file order per epoch, then yields indices inside each file in a
+    shuffled order. With `cache_files=1` on the dataset this means each
+    worker only ever pays one file-load cost per file.
+    """
+
+    def __init__(
+        self,
+        dataset: TactileParquetDataset,
+        seed: int = 0,
+        shuffle_files: bool = True,
+        shuffle_within_file: bool = True,
+    ) -> None:
+        self.dataset = dataset
+        self.seed = seed
+        self.shuffle_files = shuffle_files
+        self.shuffle_within_file = shuffle_within_file
+        self.epoch = 0
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __iter__(self) -> Iterator[int]:
+        rng = np.random.default_rng(self.seed + self.epoch)
+        n_files = self.dataset.num_files
+        file_order = rng.permutation(n_files) if self.shuffle_files else np.arange(n_files)
+        offsets = self.dataset._per_file_offsets
+        for fi in file_order:
+            start = int(offsets[fi])
+            end = int(offsets[fi + 1])
+            n = end - start
+            if n == 0:
+                continue  # this file contributes no samples to the current split
+            local = rng.permutation(n) if self.shuffle_within_file else np.arange(n)
+            for li in local:
+                yield int(start + li)

dataset/make_splits.py

@@ -0,0 +1,165 @@
+"""Generate a stable train/val/test split manifest for the fota_unlabeled parquets.
+
+Each sample is identified by `(parquet_filename, row_index)`. We hash that
+identity to assign a stable bucket independent of file order, so adding /
+reordering shards never reshuffles existing samples (only new files get
+fresh assignments).
+
+Run:
+    python tactile_vae/dataset/make_splits.py
+
+Default writes to `tactile_vae/dataset/splits.json` and looks like:
+
+    {
+      "seed": 42,
+      "ratios": {"train": 0.9, "val": 0.05, "test": 0.05},
+      "counts": {...}, "total": ...,
+      "files": ["train-00000-of-00008.parquet", ...],
+      "splits": {
+        "train": {"train-00000-of-00008.parquet": [0, 3, 4, ...], ...},
+        "val":   {...},
+        "test":  {...}
+      }
+    }
+
+Per-file row lists are sorted ascending so they're easy to diff across runs
+and cheap to bisect / slice from a parquet reader.
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import struct
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tactile_vae.dataset.dataset import DEFAULT_DATA_ROOT, DEFAULT_FILE_GLOB
+
+DEFAULT_OUTPUT = Path(__file__).with_name("splits.json")
+DEFAULT_RATIOS = (0.9, 0.05, 0.05)
+DEFAULT_SEED = 42
+SPLIT_NAMES = ("train", "val", "test")
+
+
+def _bucket_for(filename: str, row_idx: int, seed: int) -> float:
+    """Map a (filename, row) pair to a deterministic float in [0, 1).
+
+    Uses BLAKE2b for a stable cross-run hash that doesn't depend on Python's
+    PYTHONHASHSEED. Adding new files leaves existing samples in their bucket.
+    """
+    h = hashlib.blake2b(digest_size=8)
+    h.update(struct.pack("<I", seed))
+    h.update(filename.encode("utf-8"))
+    h.update(b"\x00")
+    h.update(struct.pack("<q", row_idx))
+    return int.from_bytes(h.digest(), "little") / 2**64
+
+
+def _assign(filename: str, n_rows: int, seed: int, cutoffs: tuple[float, float]) -> np.ndarray:
+    """Return a uint8 array of length `n_rows` with bucket {0,1,2} per row."""
+    # Vectorized blake2b is impractical, but for n ≤ ~100k per file the
+    # python loop is still fast (<1s/file) and keeps the bucketing portable.
+    out = np.empty(n_rows, dtype=np.uint8)
+    train_cut, val_cut = cutoffs
+    for i in range(n_rows):
+        u = _bucket_for(filename, i, seed)
+        if u < train_cut:
+            out[i] = 0  # train
+        elif u < val_cut:
+            out[i] = 1  # val
+        else:
+            out[i] = 2  # test
+    return out
+
+
+def make_splits(
+    root: Path = DEFAULT_DATA_ROOT,
+    file_glob: str = DEFAULT_FILE_GLOB,
+    ratios: tuple[float, float, float] = DEFAULT_RATIOS,
+    seed: int = DEFAULT_SEED,
+) -> dict:
+    if not (len(ratios) == 3 and abs(sum(ratios) - 1.0) < 1e-6):
+        raise ValueError(f"ratios must sum to 1.0; got {ratios} (sum={sum(ratios)})")
+    train_r, val_r, _ = ratios
+    cutoffs = (train_r, train_r + val_r)
+
+    files = sorted(p for p in root.glob(file_glob) if ".raw." not in p.name)
+    if not files:
+        raise FileNotFoundError(f"no parquet shards under {root} matching {file_glob!r}")
+
+    per_split: dict[str, dict[str, list[int]]] = {n: {} for n in SPLIT_NAMES}
+    counts = {n: 0 for n in SPLIT_NAMES}
+    total = 0
+    t0 = time.time()
+    for p in files:
+        n = pq.ParquetFile(p).metadata.num_rows
+        total += n
+        assignments = _assign(p.name, n, seed=seed, cutoffs=cutoffs)
+        for split_idx, split_name in enumerate(SPLIT_NAMES):
+            rows = np.flatnonzero(assignments == split_idx).tolist()
+            per_split[split_name][p.name] = rows
+            counts[split_name] += len(rows)
+        print(f"  {p.name}: rows={n:,} "
+              f"train={int((assignments==0).sum()):,} "
+              f"val={int((assignments==1).sum()):,} "
+              f"test={int((assignments==2).sum()):,}",
+              flush=True)
+    dt = time.time() - t0
+
+    print(f"\nAssignment done in {dt:.1f}s")
+    for n in SPLIT_NAMES:
+        print(f"  {n:5s}: {counts[n]:>8,} ({100*counts[n]/total:.2f}%)")
+    print(f"  total: {total:>8,}")
+
+    return {
+        "seed": seed,
+        "ratios": {"train": ratios[0], "val": ratios[1], "test": ratios[2]},
+        "counts": counts,
+        "total": total,
+        "source_root": str(root),
+        "source_glob": file_glob,
+        "files": [p.name for p in files],
+        "split_names": list(SPLIT_NAMES),
+        "splits": per_split,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+    p.add_argument("--root", type=Path, default=DEFAULT_DATA_ROOT)
+    p.add_argument("--glob", default=DEFAULT_FILE_GLOB)
+    p.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
+    p.add_argument("--seed", type=int, default=DEFAULT_SEED)
+    p.add_argument("--train", type=float, default=DEFAULT_RATIOS[0])
+    p.add_argument("--val", type=float, default=DEFAULT_RATIOS[1])
+    p.add_argument("--test", type=float, default=DEFAULT_RATIOS[2])
+    return p.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    manifest = make_splits(
+        root=args.root,
+        file_glob=args.glob,
+        ratios=(args.train, args.val, args.test),
+        seed=args.seed,
+    )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w") as f:
+        json.dump(manifest, f, separators=(",", ":"))
+    size_mb = args.output.stat().st_size / 1e6
+    print(f"\nwrote {args.output} ({size_mb:.2f} MB)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

inference/per_sample_metrics.json

@@ -0,0 +1,352 @@
+[
+  {
+    "subset_rank": 0,
+    "dataset_index": 28002,
+    "sample_id": "train-00000-of-00008:31065",
+    "mae": 0.01892464980483055,
+    "mse": 0.0015020258724689484
+  },
+  {
+    "subset_rank": 1,
+    "dataset_index": 10540,
+    "sample_id": "train-00000-of-00008:11713",
+    "mae": 0.018629208207130432,
+    "mse": 0.001927545526996255
+  },
+  {
+    "subset_rank": 2,
+    "dataset_index": 315096,
+    "sample_id": "train-00006-of-00008:40249",
+    "mae": 0.0188576839864254,
+    "mse": 0.001595273963175714
+  },
+  {
+    "subset_rank": 3,
+    "dataset_index": 316578,
+    "sample_id": "train-00006-of-00008:41869",
+    "mae": 0.01664964109659195,
+    "mse": 0.001574263907968998
+  },
+  {
+    "subset_rank": 4,
+    "dataset_index": 271509,
+    "sample_id": "train-00005-of-00008:37878",
+    "mae": 0.01655184105038643,
+    "mse": 0.001182803069241345
+  },
+  {
+    "subset_rank": 5,
+    "dataset_index": 319123,
+    "sample_id": "train-00006-of-00008:44695",
+    "mae": 0.0181032232940197,
+    "mse": 0.0014666281640529633
+  },
+  {
+    "subset_rank": 6,
+    "dataset_index": 33232,
+    "sample_id": "train-00000-of-00008:36906",
+    "mae": 0.019307494163513184,
+    "mse": 0.0023352960124611855
+  },
+  {
+    "subset_rank": 7,
+    "dataset_index": 190231,
+    "sample_id": "train-00004-of-00008:498",
+    "mae": 0.015174083411693573,
+    "mse": 0.000804901123046875
+  },
+  {
+    "subset_rank": 8,
+    "dataset_index": 29925,
+    "sample_id": "train-00000-of-00008:33188",
+    "mae": 0.019524700939655304,
+    "mse": 0.002352712210267782
+  },
+  {
+    "subset_rank": 9,
+    "dataset_index": 100407,
+    "sample_id": "train-00001-of-00008:38654",
+    "mae": 0.02003355696797371,
+    "mse": 0.0019040403421968222
+  },
+  {
+    "subset_rank": 10,
+    "dataset_index": 187441,
+    "sample_id": "train-00003-of-00008:42709",
+    "mae": 0.019151723012328148,
+    "mse": 0.0011700440663844347
+  },
+  {
+    "subset_rank": 11,
+    "dataset_index": 302683,
+    "sample_id": "train-00006-of-00008:26476",
+    "mae": 0.018105922266840935,
+    "mse": 0.0015790475299581885
+  },
+  {
+    "subset_rank": 12,
+    "dataset_index": 202331,
+    "sample_id": "train-00004-of-00008:13870",
+    "mae": 0.019837066531181335,
+    "mse": 0.0021240056958049536
+  },
+  {
+    "subset_rank": 13,
+    "dataset_index": 103226,
+    "sample_id": "train-00001-of-00008:41777",
+    "mae": 0.018462352454662323,
+    "mse": 0.0015763709088787436
+  },
+  {
+    "subset_rank": 14,
+    "dataset_index": 284692,
+    "sample_id": "train-00006-of-00008:6576",
+    "mae": 0.016698075458407402,
+    "mse": 0.0011528099421411753
+  },
+  {
+    "subset_rank": 15,
+    "dataset_index": 146700,
+    "sample_id": "train-00002-of-00008:43384",
+    "mae": 0.02269577607512474,
+    "mse": 0.001561635173857212
+  },
+  {
+    "subset_rank": 16,
+    "dataset_index": 15249,
+    "sample_id": "train-00000-of-00008:16924",
+    "mae": 0.01747109368443489,
+    "mse": 0.0012944282498210669
+  },
+  {
+    "subset_rank": 17,
+    "dataset_index": 8226,
+    "sample_id": "train-00000-of-00008:9149",
+    "mae": 0.01862267032265663,
+    "mse": 0.0015876510879024863
+  },
+  {
+    "subset_rank": 18,
+    "dataset_index": 241700,
+    "sample_id": "train-00005-of-00008:4679",
+    "mae": 0.017128579318523407,
+    "mse": 0.0009350610198453069
+  },
+  {
+    "subset_rank": 19,
+    "dataset_index": 249617,
+    "sample_id": "train-00005-of-00008:13464",
+    "mae": 0.018679119646549225,
+    "mse": 0.0016614518826827407
+  },
+  {
+    "subset_rank": 20,
+    "dataset_index": 6151,
+    "sample_id": "train-00000-of-00008:6832",
+    "mae": 0.019154969602823257,
+    "mse": 0.0019272298086434603
+  },
+  {
+    "subset_rank": 21,
+    "dataset_index": 201535,
+    "sample_id": "train-00004-of-00008:13012",
+    "mae": 0.02364823780953884,
+    "mse": 0.003448478877544403
+  },
+  {
+    "subset_rank": 22,
+    "dataset_index": 303654,
+    "sample_id": "train-00006-of-00008:27559",
+    "mae": 0.01817791536450386,
+    "mse": 0.001617558067664504
+  },
+  {
+    "subset_rank": 23,
+    "dataset_index": 237059,
+    "sample_id": "train-00004-of-00008:52418",
+    "mae": 0.017815029248595238,
+    "mse": 0.001500395592302084
+  },
+  {
+    "subset_rank": 24,
+    "dataset_index": 240896,
+    "sample_id": "train-00005-of-00008:3774",
+    "mae": 0.017452210187911987,
+    "mse": 0.001413716934621334
+  },
+  {
+    "subset_rank": 25,
+    "dataset_index": 95772,
+    "sample_id": "train-00001-of-00008:33454",
+    "mae": 0.01503452192991972,
+    "mse": 0.0008028325974009931
+  },
+  {
+    "subset_rank": 26,
+    "dataset_index": 225780,
+    "sample_id": "train-00004-of-00008:39911",
+    "mae": 0.020092152059078217,
+    "mse": 0.002289179712533951
+  },
+  {
+    "subset_rank": 27,
+    "dataset_index": 1992,
+    "sample_id": "train-00000-of-00008:2219",
+    "mae": 0.017521802335977554,
+    "mse": 0.0013774571707472205
+  },
+  {
+    "subset_rank": 28,
+    "dataset_index": 65230,
+    "sample_id": "train-00000-of-00008:72439",
+    "mae": 0.01493510790169239,
+    "mse": 0.000781821203418076
+  },
+  {
+    "subset_rank": 29,
+    "dataset_index": 339712,
+    "sample_id": "train-00007-of-00008:21662",
+    "mae": 0.020267298445105553,
+    "mse": 0.0018565801437944174
+  },
+  {
+    "subset_rank": 30,
+    "dataset_index": 46260,
+    "sample_id": "train-00000-of-00008:51312",
+    "mae": 0.01745472103357315,
+    "mse": 0.0011822863016277552
+  },
+  {
+    "subset_rank": 31,
+    "dataset_index": 12500,
+    "sample_id": "train-00000-of-00008:13870",
+    "mae": 0.017031168565154076,
+    "mse": 0.0013072905130684376
+  },
+  {
+    "subset_rank": 32,
+    "dataset_index": 1019,
+    "sample_id": "train-00000-of-00008:1144",
+    "mae": 0.01728898286819458,
+    "mse": 0.0012464504688978195
+  },
+  {
+    "subset_rank": 33,
+    "dataset_index": 348025,
+    "sample_id": "train-00007-of-00008:30906",
+    "mae": 0.014839177951216698,
+    "mse": 0.0006855515530332923
+  },
+  {
+    "subset_rank": 34,
+    "dataset_index": 114566,
+    "sample_id": "train-00002-of-00008:7678",
+    "mae": 0.018796399235725403,
+    "mse": 0.0015862470027059317
+  },
+  {
+    "subset_rank": 35,
+    "dataset_index": 65379,
+    "sample_id": "train-00000-of-00008:72604",
+    "mae": 0.013970417901873589,
+    "mse": 0.0007457585888914764
+  },
+  {
+    "subset_rank": 36,
+    "dataset_index": 361296,
+    "sample_id": "train-00007-of-00008:45623",
+    "mae": 0.02225802093744278,
+    "mse": 0.0027801282703876495
+  },
+  {
+    "subset_rank": 37,
+    "dataset_index": 229059,
+    "sample_id": "train-00004-of-00008:43552",
+    "mae": 0.019468698650598526,
+    "mse": 0.0020611113868653774
+  },
+  {
+    "subset_rank": 38,
+    "dataset_index": 150090,
+    "sample_id": "train-00003-of-00008:1264",
+    "mae": 0.02170371823012829,
+    "mse": 0.00205692695453763
+  },
+  {
+    "subset_rank": 39,
+    "dataset_index": 321279,
+    "sample_id": "train-00007-of-00008:1230",
+    "mae": 0.017519645392894745,
+    "mse": 0.0015829935437068343
+  },
+  {
+    "subset_rank": 40,
+    "dataset_index": 206314,
+    "sample_id": "train-00004-of-00008:18294",
+    "mae": 0.017531011253595352,
+    "mse": 0.0018101041205227375
+  },
+  {
+    "subset_rank": 41,
+    "dataset_index": 157328,
+    "sample_id": "train-00003-of-00008:9302",
+    "mae": 0.01673971116542816,
+    "mse": 0.000940843892749399
+  },
+  {
+    "subset_rank": 42,
+    "dataset_index": 271578,
+    "sample_id": "train-00005-of-00008:37956",
+    "mae": 0.017461789771914482,
+    "mse": 0.0012990653049200773
+  },
+  {
+    "subset_rank": 43,
+    "dataset_index": 235323,
+    "sample_id": "train-00004-of-00008:50492",
+    "mae": 0.01795295439660549,
+    "mse": 0.0015181120252236724
+  },
+  {
+    "subset_rank": 44,
+    "dataset_index": 3083,
+    "sample_id": "train-00000-of-00008:3430",
+    "mae": 0.015775706619024277,
+    "mse": 0.0009211775613948703
+  },
+  {
+    "subset_rank": 45,
+    "dataset_index": 111555,
+    "sample_id": "train-00002-of-00008:4339",
+    "mae": 0.01904463954269886,
+    "mse": 0.0021079955622553825
+  },
+  {
+    "subset_rank": 46,
+    "dataset_index": 195644,
+    "sample_id": "train-00004-of-00008:6481",
+    "mae": 0.0185789056122303,
+    "mse": 0.0015466390177607536
+  },
+  {
+    "subset_rank": 47,
+    "dataset_index": 249696,
+    "sample_id": "train-00005-of-00008:13548",
+    "mae": 0.018499009311199188,
+    "mse": 0.0016744902823120356
+  },
+  {
+    "subset_rank": 48,
+    "dataset_index": 208395,
+    "sample_id": "train-00004-of-00008:20588",
+    "mae": 0.01778639853000641,
+    "mse": 0.0018380036344751716
+  },
+  {
+    "subset_rank": 49,
+    "dataset_index": 179055,
+    "sample_id": "train-00003-of-00008:33409",
+    "mae": 0.01800130307674408,
+    "mse": 0.0015790577745065093
+  }
+]

inference/summary.json

@@ -0,0 +1,10 @@
+{
+  "config": "/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_4/config.snapshot.yaml",
+  "checkpoint": "/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_4/checkpoints/last.ckpt",
+  "split": "train",
+  "seed": 0,
+  "selected_num_samples": 50,
+  "mean_mae": 0.018208201732486485,
+  "mean_mse": 0.001575469592353329,
+  "grid_path": "/group2/ct/weihanx/tactile_world_model/tactile_vae/inference/reconstruction_grid.png"
+}

inference/vae_baseline_3/per_sample_metrics.json

@@ -0,0 +1,352 @@
+[
+  {
+    "subset_rank": 0,
+    "dataset_index": 28002,
+    "sample_id": "train-00000-of-00008:31065",
+    "mae": 0.0166653823107481,
+    "mse": 0.0014663146575912833
+  },
+  {
+    "subset_rank": 1,
+    "dataset_index": 10540,
+    "sample_id": "train-00000-of-00008:11713",
+    "mae": 0.017510399222373962,
+    "mse": 0.0017775435699149966
+  },
+  {
+    "subset_rank": 2,
+    "dataset_index": 315096,
+    "sample_id": "train-00006-of-00008:40249",
+    "mae": 0.017776841297745705,
+    "mse": 0.00159529410302639
+  },
+  {
+    "subset_rank": 3,
+    "dataset_index": 316578,
+    "sample_id": "train-00006-of-00008:41869",
+    "mae": 0.01655813306570053,
+    "mse": 0.0015516902785748243
+  },
+  {
+    "subset_rank": 4,
+    "dataset_index": 271509,
+    "sample_id": "train-00005-of-00008:37878",
+    "mae": 0.01444108597934246,
+    "mse": 0.0008555233362130821
+  },
+  {
+    "subset_rank": 5,
+    "dataset_index": 319123,
+    "sample_id": "train-00006-of-00008:44695",
+    "mae": 0.018781110644340515,
+    "mse": 0.0018363429699093103
+  },
+  {
+    "subset_rank": 6,
+    "dataset_index": 33232,
+    "sample_id": "train-00000-of-00008:36906",
+    "mae": 0.01676507666707039,
+    "mse": 0.001682119444012642
+  },
+  {
+    "subset_rank": 7,
+    "dataset_index": 190231,
+    "sample_id": "train-00004-of-00008:498",
+    "mae": 0.014347122982144356,
+    "mse": 0.0007567598950117826
+  },
+  {
+    "subset_rank": 8,
+    "dataset_index": 29925,
+    "sample_id": "train-00000-of-00008:33188",
+    "mae": 0.016472000628709793,
+    "mse": 0.0017040781676769257
+  },
+  {
+    "subset_rank": 9,
+    "dataset_index": 100407,
+    "sample_id": "train-00001-of-00008:38654",
+    "mae": 0.017876047641038895,
+    "mse": 0.0018246680265292525
+  },
+  {
+    "subset_rank": 10,
+    "dataset_index": 187441,
+    "sample_id": "train-00003-of-00008:42709",
+    "mae": 0.016988320276141167,
+    "mse": 0.0009907050989568233
+  },
+  {
+    "subset_rank": 11,
+    "dataset_index": 302683,
+    "sample_id": "train-00006-of-00008:26476",
+    "mae": 0.018481288105249405,
+    "mse": 0.0022351769730448723
+  },
+  {
+    "subset_rank": 12,
+    "dataset_index": 202331,
+    "sample_id": "train-00004-of-00008:13870",
+    "mae": 0.017929738387465477,
+    "mse": 0.001798365032300353
+  },
+  {
+    "subset_rank": 13,
+    "dataset_index": 103226,
+    "sample_id": "train-00001-of-00008:41777",
+    "mae": 0.017578527331352234,
+    "mse": 0.0016123488312587142
+  },
+  {
+    "subset_rank": 14,
+    "dataset_index": 284692,
+    "sample_id": "train-00006-of-00008:6576",
+    "mae": 0.015600966289639473,
+    "mse": 0.00086554343579337
+  },
+  {
+    "subset_rank": 15,
+    "dataset_index": 146700,
+    "sample_id": "train-00002-of-00008:43384",
+    "mae": 0.021709434688091278,
+    "mse": 0.0013935761526226997
+  },
+  {
+    "subset_rank": 16,
+    "dataset_index": 15249,
+    "sample_id": "train-00000-of-00008:16924",
+    "mae": 0.014922507107257843,
+    "mse": 0.0007923931698314846
+  },
+  {
+    "subset_rank": 17,
+    "dataset_index": 8226,
+    "sample_id": "train-00000-of-00008:9149",
+    "mae": 0.016841208562254906,
+    "mse": 0.0015739825321361423
+  },
+  {
+    "subset_rank": 18,
+    "dataset_index": 241700,
+    "sample_id": "train-00005-of-00008:4679",
+    "mae": 0.01727406494319439,
+    "mse": 0.0009527048096060753
+  },
+  {
+    "subset_rank": 19,
+    "dataset_index": 249617,
+    "sample_id": "train-00005-of-00008:13464",
+    "mae": 0.0178667064756155,
+    "mse": 0.0016339406138285995
+  },
+  {
+    "subset_rank": 20,
+    "dataset_index": 6151,
+    "sample_id": "train-00000-of-00008:6832",
+    "mae": 0.018926797434687614,
+    "mse": 0.0019104413222521544
+  },
+  {
+    "subset_rank": 21,
+    "dataset_index": 201535,
+    "sample_id": "train-00004-of-00008:13012",
+    "mae": 0.020377418026328087,
+    "mse": 0.0026058475486934185
+  },
+  {
+    "subset_rank": 22,
+    "dataset_index": 303654,
+    "sample_id": "train-00006-of-00008:27559",
+    "mae": 0.0186694897711277,
+    "mse": 0.002249116078019142
+  },
+  {
+    "subset_rank": 23,
+    "dataset_index": 237059,
+    "sample_id": "train-00004-of-00008:52418",
+    "mae": 0.01686570607125759,
+    "mse": 0.0015954857226461172
+  },
+  {
+    "subset_rank": 24,
+    "dataset_index": 240896,
+    "sample_id": "train-00005-of-00008:3774",
+    "mae": 0.01646217331290245,
+    "mse": 0.0014256120193749666
+  },
+  {
+    "subset_rank": 25,
+    "dataset_index": 95772,
+    "sample_id": "train-00001-of-00008:33454",
+    "mae": 0.014364289119839668,
+    "mse": 0.0007188305607996881
+  },
+  {
+    "subset_rank": 26,
+    "dataset_index": 225780,
+    "sample_id": "train-00004-of-00008:39911",
+    "mae": 0.018801426514983177,
+    "mse": 0.002039157785475254
+  },
+  {
+    "subset_rank": 27,
+    "dataset_index": 1992,
+    "sample_id": "train-00000-of-00008:2219",
+    "mae": 0.018728084862232208,
+    "mse": 0.0016581187956035137
+  },
+  {
+    "subset_rank": 28,
+    "dataset_index": 65230,
+    "sample_id": "train-00000-of-00008:72439",
+    "mae": 0.013995742425322533,
+    "mse": 0.0007463646470569074
+  },
+  {
+    "subset_rank": 29,
+    "dataset_index": 339712,
+    "sample_id": "train-00007-of-00008:21662",
+    "mae": 0.020348068326711655,
+    "mse": 0.0021553956903517246
+  },
+  {
+    "subset_rank": 30,
+    "dataset_index": 46260,
+    "sample_id": "train-00000-of-00008:51312",
+    "mae": 0.017040027305483818,
+    "mse": 0.001189670292660594
+  },
+  {
+    "subset_rank": 31,
+    "dataset_index": 12500,
+    "sample_id": "train-00000-of-00008:13870",
+    "mae": 0.017332643270492554,
+    "mse": 0.0015696667833253741
+  },
+  {
+    "subset_rank": 32,
+    "dataset_index": 1019,
+    "sample_id": "train-00000-of-00008:1144",
+    "mae": 0.01800120621919632,
+    "mse": 0.0015177734894677997
+  },
+  {
+    "subset_rank": 33,
+    "dataset_index": 348025,
+    "sample_id": "train-00007-of-00008:30906",
+    "mae": 0.01363951526582241,
+    "mse": 0.0006664479151368141
+  },
+  {
+    "subset_rank": 34,
+    "dataset_index": 114566,
+    "sample_id": "train-00002-of-00008:7678",
+    "mae": 0.01782606542110443,
+    "mse": 0.0016285644378513098
+  },
+  {
+    "subset_rank": 35,
+    "dataset_index": 65379,
+    "sample_id": "train-00000-of-00008:72604",
+    "mae": 0.012702388688921928,
+    "mse": 0.0006791208870708942
+  },
+  {
+    "subset_rank": 36,
+    "dataset_index": 361296,
+    "sample_id": "train-00007-of-00008:45623",
+    "mae": 0.01820572093129158,
+    "mse": 0.0015046261250972748
+  },
+  {
+    "subset_rank": 37,
+    "dataset_index": 229059,
+    "sample_id": "train-00004-of-00008:43552",
+    "mae": 0.018051331862807274,
+    "mse": 0.0017623631283640862
+  },
+  {
+    "subset_rank": 38,
+    "dataset_index": 150090,
+    "sample_id": "train-00003-of-00008:1264",
+    "mae": 0.0181649811565876,
+    "mse": 0.0011993813095614314
+  },
+  {
+    "subset_rank": 39,
+    "dataset_index": 321279,
+    "sample_id": "train-00007-of-00008:1230",
+    "mae": 0.016737129539251328,
+    "mse": 0.0016499034827575088
+  },
+  {
+    "subset_rank": 40,
+    "dataset_index": 206314,
+    "sample_id": "train-00004-of-00008:18294",
+    "mae": 0.015606882981956005,
+    "mse": 0.0015939919976517558
+  },
+  {
+    "subset_rank": 41,
+    "dataset_index": 157328,
+    "sample_id": "train-00003-of-00008:9302",
+    "mae": 0.017554117366671562,
+    "mse": 0.0011766867246478796
+  },
+  {
+    "subset_rank": 42,
+    "dataset_index": 271578,
+    "sample_id": "train-00005-of-00008:37956",
+    "mae": 0.01577179506421089,
+    "mse": 0.0010084250243380666
+  },
+  {
+    "subset_rank": 43,
+    "dataset_index": 235323,
+    "sample_id": "train-00004-of-00008:50492",
+    "mae": 0.016777435317635536,
+    "mse": 0.0015847641043365002
+  },
+  {
+    "subset_rank": 44,
+    "dataset_index": 3083,
+    "sample_id": "train-00000-of-00008:3430",
+    "mae": 0.015516860410571098,
+    "mse": 0.0011327709071338177
+  },
+  {
+    "subset_rank": 45,
+    "dataset_index": 111555,
+    "sample_id": "train-00002-of-00008:4339",
+    "mae": 0.0198848657310009,
+    "mse": 0.002476822817698121
+  },
+  {
+    "subset_rank": 46,
+    "dataset_index": 195644,
+    "sample_id": "train-00004-of-00008:6481",
+    "mae": 0.016649916768074036,
+    "mse": 0.001522408565506339
+  },
+  {
+    "subset_rank": 47,
+    "dataset_index": 249696,
+    "sample_id": "train-00005-of-00008:13548",
+    "mae": 0.017414681613445282,
+    "mse": 0.0016105175018310547
+  },
+  {
+    "subset_rank": 48,
+    "dataset_index": 208395,
+    "sample_id": "train-00004-of-00008:20588",
+    "mae": 0.016060233116149902,
+    "mse": 0.0016828887164592743
+  },
+  {
+    "subset_rank": 49,
+    "dataset_index": 179055,
+    "sample_id": "train-00003-of-00008:33409",
+    "mae": 0.017118770629167557,
+    "mse": 0.0016486085951328278
+  }
+]

inference/vae_baseline_3/summary.json

@@ -0,0 +1,10 @@
+{
+  "config": "/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_3/config.snapshot.yaml",
+  "checkpoint": "/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_3/checkpoints/last.ckpt",
+  "split": "train",
+  "seed": 0,
+  "selected_num_samples": 50,
+  "mean_mae": 0.017119634542614223,
+  "mean_mse": 0.0014961768814828248,
+  "grid_path": "/group2/ct/weihanx/tactile_world_model/tactile_vae/inference/vae_baseline_3/reconstruction_grid.png"
+}

model/__init__.py

@@ -0,0 +1,15 @@
+from tactile_vae.model.tactile_vae import (
+    BetaVAELoss,
+    SSIMLoss,
+    TactileVAE,
+    VAELoss,
+    load_pretrained,
+)
+
+__all__ = [
+    "TactileVAE",
+    "VAELoss",
+    "BetaVAELoss",
+    "SSIMLoss",
+    "load_pretrained",
+]

model/cosmos_tokenizer.py

@@ -0,0 +1,136 @@
+"""Cosmos CV4x8x8 tokenizer / detokenizer for single-frame RGB images.
+
+Wraps the Cosmos JIT encoder and decoder with a clean encode/decode interface.
+Input images are expected as (B, 3, H, W) float32 in [-1, 1].
+Latents have shape (B, 16, 1, H/8, W/8) – the temporal dim is always 1 for
+single frames, matching the CV (Causal Video) 4×8×8 compression scheme.
+
+Usage:
+    tokenizer = CosmosTokenizer(ckpt_dir, device)
+    z     = tokenizer(x, mode="encode")   # or tokenizer.encode(x)
+    x_hat = tokenizer(z, mode="decode")   # or tokenizer.decode(z)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+_DEFAULT_CKPT_DIR = Path(
+    "/group2/ct/weihanx/tactile_world_model/tactile_wm/pretrained_models/"
+    "Cosmos-0.1-Tokenizer-CV4x8x8"
+)
+
+
+class CosmosTokenizer(nn.Module):
+    def __init__(
+        self,
+        ckpt_dir: str | Path = _DEFAULT_CKPT_DIR,
+        device: str | torch.device = "cuda",
+        dtype: str = "bfloat16",
+    ):
+        super().__init__()
+        self.device = torch.device(device)
+        self.dtype = getattr(torch, dtype)
+        ckpt_dir = Path(ckpt_dir)
+        for name in ("encoder.jit", "decoder.jit"):
+            if not (ckpt_dir / name).exists():
+                raise FileNotFoundError(f"Cosmos checkpoint not found: {ckpt_dir / name}")
+        self.encoder = torch.jit.load(str(ckpt_dir / "encoder.jit"), map_location=self.device).eval()
+        self.decoder = torch.jit.load(str(ckpt_dir / "decoder.jit"), map_location=self.device).eval()
+
+    def to(self, device):
+        super().to(device)
+        self.device = torch.device(device)
+        self.encoder = self.encoder.to(device)
+        self.decoder = self.decoder.to(device)
+        return self
+
+    @staticmethod
+    def _extract(obj) -> torch.Tensor:
+        if isinstance(obj, torch.Tensor):
+            return obj
+        if isinstance(obj, (tuple, list)):
+            for item in obj:
+                t = CosmosTokenizer._extract(item)
+                if isinstance(t, torch.Tensor):
+                    return t
+        raise TypeError(f"no tensor in model output: {type(obj)!r}")
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Tokenize (B, 3, H, W) float32 [-1,1] → latent (B, 16, 1, H/8, W/8)."""
+        video = x.to(device=self.device, dtype=self.dtype).unsqueeze(2)  # (B,3,1,H,W)
+        return self._extract(self.encoder(video))
+
+    @torch.no_grad()
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """Detokenize latent (B, 16, 1, H/8, W/8) → (B, 3, H, W) float32 [-1,1]."""
+        z = z.to(device=self.device, dtype=self.dtype)
+        recon = self._extract(self.decoder(z))  # (B, 3, 1, H, W)
+        return recon[:, :, 0].float()
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, mode: str) -> torch.Tensor:
+        if mode == "encode":
+            return self.encode(x)
+        if mode == "decode":
+            return self.decode(x)
+        raise ValueError(f"mode must be 'encode' or 'decode', got {mode!r}")
+
+
+if __name__ == "__main__":
+    _EPISODE_PATH = Path("/group2/ct/weihanx/tactile_world_model/mode1_v1/0323_episode_000.pt")
+    _OUT_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/test_output/cosmos")
+    _SAMPLE_INDICES = [0, 100, 500, 1000, 2000]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"device: {device}")
+
+    # ── Load tokenizer ────────────────────────────────────────────────────────
+    tokenizer = CosmosTokenizer(ckpt_dir=_DEFAULT_CKPT_DIR, device=device)
+    print(f"encoder loaded: {_DEFAULT_CKPT_DIR / 'encoder.jit'}")
+    print(f"decoder loaded: {_DEFAULT_CKPT_DIR / 'decoder.jit'}")
+
+    # ── Load episode frames ───────────────────────────────────────────────────
+    ep = torch.load(str(_EPISODE_PATH), map_location="cpu", weights_only=False)
+    views = ep["view"]  # (T, 3, H, W) uint8
+    frames_u8 = views[_SAMPLE_INDICES]  # (N, 3, H, W) uint8
+    frames = frames_u8.float() / 127.5 - 1.0  # (N, 3, H, W) float32 in [-1, 1]
+    print(f"\nepisode frames: {tuple(frames.shape)}")
+
+    # ── Tokenize ──────────────────────────────────────────────────────────────
+    z = tokenizer.encode(frames)
+    print(f"latent shape:   {tuple(z.shape)}")
+
+    # ── Detokenize ────────────────────────────────────────────────────────────
+    x_hat = tokenizer.decode(z)
+    print(f"recon shape:    {tuple(x_hat.shape)}")
+
+    # ── Save panels and print PSNR ────────────────────────────────────────────
+    _OUT_DIR.mkdir(parents=True, exist_ok=True)
+    psnrs = []
+    for i, idx in enumerate(_SAMPLE_INDICES):
+        orig_np = ((frames[i].permute(1, 2, 0).clamp(-1, 1) + 1) * 127.5).byte().numpy()
+        recon_np = ((x_hat[i].permute(1, 2, 0).clamp(-1, 1) + 1) * 127.5).byte().cpu().numpy()
+        diff_np = np.abs(orig_np.astype(int) - recon_np.astype(int)).astype(np.uint8)
+        mse = float(((orig_np.astype(float) - recon_np.astype(float)) ** 2).mean())
+        psnr = 10 * np.log10(255.0 ** 2 / mse) if mse > 0 else float("inf")
+        psnrs.append(psnr)
+        print(f"  frame {idx:5d}  PSNR={psnr:.2f} dB")
+
+        h, w = orig_np.shape[:2]
+        panel = Image.new("RGB", (3 * w + 16, h), (20, 20, 20))
+        panel.paste(Image.fromarray(orig_np), (0, 0))
+        panel.paste(Image.fromarray(recon_np), (w + 8, 0))
+        panel.paste(Image.fromarray(diff_np), (2 * w + 16, 0))
+        panel.save(_OUT_DIR / f"cosmos_frame_{idx:05d}_panel.png")
+        Image.fromarray(orig_np).save(_OUT_DIR / f"cosmos_frame_{idx:05d}_input.png")
+        Image.fromarray(recon_np).save(_OUT_DIR / f"cosmos_frame_{idx:05d}_recon.png")
+
+    print(f"\nmean PSNR: {np.mean(psnrs):.2f} dB  (over {len(psnrs)} frames)")
+    print(f"panels saved to {_OUT_DIR}")

model/tactile_vae.py

@@ -0,0 +1,503 @@
+"""ViT-based tactile VAE (no MAE masking).
+
+Architecture:
+- Encoder: PatchEmbed + Transformer blocks + fixed sin-cos positional embedding.
+- Latent: mu/logvar heads + reparameterization.
+- Decoder: latent-conditioned transformer decoder that predicts image patches.
+- Reconstruction: unpatchify patch predictions into image space.
+
+Training objective: regular VAE loss (reconstruction + beta * KL), with optional SSIM.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional
+from PIL import Image
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.vision_transformer import Block, PatchEmbed
+
+
+DEFAULT_TACTILE_VAE_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_wm/pretrained_models")
+DEFAULT_CHECKPOINT_NAME = "ckpt_best.pt"
+
+
+def _get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / (10000**omega)
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)
+
+
+def _get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
+    assert embed_dim % 2 == 0
+    emb_h = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray:
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = _get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def unpatchify(pred_patches: torch.Tensor, patch_size: int, in_chans: int) -> torch.Tensor:
+    """Convert patch predictions (B, L, p*p*C) to images (B, C, H, W)."""
+    h = w = int(pred_patches.shape[1] ** 0.5)
+    assert h * w == pred_patches.shape[1], "number of patches must be a square"
+    x = pred_patches.reshape(pred_patches.shape[0], h, w, patch_size, patch_size, in_chans)
+    x = torch.einsum("nhwpqc->nchpwq", x)
+    return x.reshape(pred_patches.shape[0], in_chans, h * patch_size, h * patch_size)
+
+
+class ViTEncoder(nn.Module):
+    """PatchEmbed + transformer blocks + fixed sin-cos positional embeddings."""
+
+    def __init__(
+        self,
+        img_size: int,
+        patch_size: int,
+        in_chans: int,
+        embed_dim: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        norm_layer: type[nn.Module] = nn.LayerNorm,
+    ):
+        super().__init__()
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)
+        self.blocks = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer) for _ in range(depth)]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        self._initialize_weights()
+
+    def _initialize_weights(self) -> None:
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1], int(self.patch_embed.num_patches**0.5), cls_token=True
+        )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        torch.nn.init.normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patch_embed(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        for blk in self.blocks:
+            x = blk(x)
+        return self.norm(x)
+
+
+class ViTDecoder(nn.Module):
+    """Latent-conditioned transformer decoder that predicts image patches."""
+
+    def __init__(
+        self,
+        img_size: int,
+        patch_size: int,
+        in_chans: int,
+        latent_dim: int,
+        embed_dim: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        norm_layer: type[nn.Module] = nn.LayerNorm,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.num_patches = (img_size // patch_size) ** 2
+
+        self.z_token = nn.Linear(latent_dim, embed_dim)
+        self.patch_tokens = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=False)
+
+        self.blocks = nn.ModuleList(
+            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer) for _ in range(depth)]
+        )
+        self.norm = norm_layer(embed_dim)
+        self.pred = nn.Linear(embed_dim, patch_size * patch_size * in_chans)
+
+        self._initialize_weights()
+
+    def _initialize_weights(self) -> None:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=False)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        torch.nn.init.normal_(self.patch_tokens, std=0.02)
+        self.apply(ViTEncoder._init_weights)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        b = z.shape[0]
+        ztok = self.z_token(z).unsqueeze(1)
+        ptok = self.patch_tokens.expand(b, -1, -1)
+        x = ztok + ptok
+        x = x + self.pos_embed
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return self.pred(x)
+
+
+class TactileVAE(nn.Module):
+    """Regular ViT-based VAE for tactile image reconstruction."""
+
+    def __init__(
+        self,
+        img_size: int = 128,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 256,
+        encoder_depth: int = 4,
+        encoder_heads: int = 8,
+        decoder_embed_dim: int = 192,
+        decoder_depth: int = 4,
+        decoder_heads: int = 8,
+        mlp_ratio: float = 4.0,
+        latent_dim: int = 128,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+
+        self.encoder = ViTEncoder(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            depth=encoder_depth,
+            num_heads=encoder_heads,
+            mlp_ratio=mlp_ratio,
+        )
+
+        self.mu_head = nn.Linear(embed_dim, latent_dim)
+        self.logvar_head = nn.Linear(embed_dim, latent_dim)
+
+        self.decoder = ViTDecoder(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            latent_dim=latent_dim,
+            embed_dim=decoder_embed_dim,
+            depth=decoder_depth,
+            num_heads=decoder_heads,
+            mlp_ratio=mlp_ratio,
+        )
+
+    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
+        enc_tokens = self.encoder(x)
+        cls = enc_tokens[:, 0]
+        mu = self.mu_head(cls)
+        logvar = self.logvar_head(cls)
+        return mu, logvar, {"enc_tokens": enc_tokens}
+
+    @staticmethod
+    def reparameterize(mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def decode(self, z: torch.Tensor, enc_ctx: Optional[dict[str, torch.Tensor]] = None) -> torch.Tensor:
+        del enc_ctx  # decoder is latent-conditioned only for regular VAE.
+        pred_patches = self.decoder(z)
+        return unpatchify(pred_patches, patch_size=self.patch_size, in_chans=self.in_chans)
+
+    def reconstruct(self, x: torch.Tensor, use_mean: bool = True) -> torch.Tensor:
+        mu, logvar, enc_ctx = self.encode(x)
+        z = mu if use_mean else self.reparameterize(mu, logvar)
+        return self.decode(z, enc_ctx=enc_ctx)
+
+    def forward(self, x: torch.Tensor, sample: bool = True) -> dict[str, torch.Tensor]:
+        mu, logvar, enc_ctx = self.encode(x)
+        z = self.reparameterize(mu, logvar) if sample else mu
+        pred_patches = self.decoder(z)
+        x_hat = unpatchify(pred_patches, patch_size=self.patch_size, in_chans=self.in_chans)
+        return {
+            "x_hat": x_hat,
+            "mu": mu,
+            "logvar": logvar,
+            "z": z,
+            "pred_patches": pred_patches,
+            "enc_ctx": enc_ctx,
+        }
+
+
+class SSIMLoss(nn.Module):
+    """Simple differentiable SSIM loss (1 - SSIM mean)."""
+
+    def __init__(self, window_size: int = 11, channels: int = 3):
+        super().__init__()
+        self.window_size = window_size
+        self.channels = channels
+        self.padding = window_size // 2
+        self.register_buffer(
+            "kernel",
+            torch.ones((channels, 1, window_size, window_size), dtype=torch.float32)
+            / (window_size * window_size),
+            persistent=False,
+        )
+
+    def _filter(self, x: torch.Tensor) -> torch.Tensor:
+        if x.shape[1] == self.channels:
+            kernel = self.kernel.to(device=x.device, dtype=x.dtype)
+        else:
+            kernel = torch.ones(
+                (x.shape[1], 1, self.window_size, self.window_size),
+                device=x.device,
+                dtype=x.dtype,
+            ) / (self.window_size * self.window_size)
+        return F.conv2d(x, kernel, padding=self.padding, groups=x.shape[1])
+
+    def forward(self, x_hat: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        c1 = 0.01**2
+        c2 = 0.03**2
+        mu_x = self._filter(x)
+        mu_y = self._filter(x_hat)
+        sigma_x = self._filter(x * x) - mu_x * mu_x
+        sigma_y = self._filter(x_hat * x_hat) - mu_y * mu_y
+        sigma_xy = self._filter(x * x_hat) - mu_x * mu_y
+        ssim_map = ((2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)) / (
+            (mu_x * mu_x + mu_y * mu_y + c1) * (sigma_x + sigma_y + c2) + 1e-8
+        )
+        return 1.0 - ssim_map.mean()
+
+
+class VAELoss(nn.Module):
+    """Regular VAE loss: reconstruction + beta * KL."""
+
+    def __init__(
+        self,
+        beta: float = 1.0,
+        recon_type: str = "l1",
+        ssim_weight: float = 0.0,
+        ssim_window_size: int = 11,
+    ):
+        super().__init__()
+        if recon_type not in {"l1", "mse"}:
+            raise ValueError(f"recon_type must be 'l1' or 'mse', got: {recon_type}")
+        self.beta = beta
+        self.recon_type = recon_type
+        self.ssim_weight = ssim_weight
+        self.ssim_loss = SSIMLoss(window_size=ssim_window_size)
+
+    def forward(
+        self,
+        x_hat: torch.Tensor,
+        x: torch.Tensor,
+        mu: torch.Tensor,
+        logvar: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        recon = F.l1_loss(x_hat, x) if self.recon_type == "l1" else F.mse_loss(x_hat, x)
+        ssim_term = self.ssim_loss(x_hat, x) if self.ssim_weight > 0 else x_hat.new_zeros(())
+        recon_total = recon + self.ssim_weight * ssim_term
+        kl = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
+        total = recon_total + self.beta * kl
+        return {
+            "total": total,
+            "recon": recon,
+            "ssim": ssim_term,
+            "recon_total": recon_total,
+            "kl": kl,
+        }
+
+
+class BetaVAELoss(VAELoss):
+    """Backward-compatible alias for VAELoss."""
+
+
+class TactileVAEWrapper(nn.Module):
+    """Inference wrapper around TactileVAE with a mode-based interface.
+
+    Usage:
+        wrapper = TactileVAEWrapper(ckpt_path, device)
+        z     = wrapper(x, mode="encode")   # (B, latent_dim)
+        x_hat = wrapper(z, mode="decode")   # (B, C, H, W)
+    Or equivalently:
+        z     = wrapper.encode(x)
+        x_hat = wrapper.decode(z)
+    """
+
+    def __init__(
+        self,
+        ckpt_path: str | Path,
+        device: str | torch.device = "cpu",
+        model_kwargs: Optional[dict[str, Any]] = None,
+    ):
+        super().__init__()
+        self.device = torch.device(device)
+        self.vae = self._load(ckpt_path, model_kwargs or {})
+
+    def _load(self, ckpt_path: str | Path, model_kwargs: dict) -> TactileVAE:
+        ckpt_path = Path(ckpt_path)
+        if not ckpt_path.exists():
+            raise FileNotFoundError(f"TactileVAE checkpoint not found: {ckpt_path}")
+        state = torch.load(str(ckpt_path), map_location=self.device, weights_only=False)
+        state_dict = _unwrap_state(state)
+        vae = TactileVAE(**model_kwargs)
+        vae.load_state_dict(state_dict, strict=True)
+        vae.eval().to(self.device)
+        vae.requires_grad_(False)
+        return vae
+
+    def to(self, device):
+        super().to(device)
+        self.device = torch.device(device)
+        self.vae = self.vae.to(device)
+        return self
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (B, C, H, W) float in [-1, 1]. Returns z: (B, latent_dim) using mu."""
+        mu, _logvar, _ctx = self.vae.encode(x.to(self.device))
+        return mu
+
+    @torch.no_grad()
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """z: (B, latent_dim). Returns x_hat: (B, C, H, W) float in [-1, 1]."""
+        return self.vae.decode(z.to(self.device))
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, mode: str) -> torch.Tensor:
+        if mode == "encode":
+            return self.encode(x)
+        if mode == "decode":
+            return self.decode(x)
+        raise ValueError(f"mode must be 'encode' or 'decode', got {mode!r}")
+
+
+def _resolve_checkpoint(checkpoint: Optional[str | Path], vae_dir: str | Path) -> Path:
+    if checkpoint is None:
+        return Path(vae_dir) / DEFAULT_CHECKPOINT_NAME
+    p = Path(checkpoint)
+    return p if p.is_absolute() else Path(vae_dir) / p
+
+
+def _unwrap_state(state: Any) -> dict[str, torch.Tensor]:
+    if isinstance(state, dict):
+        if "state_dict" in state and isinstance(state["state_dict"], dict):
+            return state["state_dict"]
+        if "model" in state and isinstance(state["model"], dict):
+            return state["model"]
+        return state
+    raise TypeError(f"Unsupported checkpoint payload type: {type(state)!r}")
+
+
+def load_pretrained(
+    checkpoint: Optional[str | Path] = None,
+    vae_dir: str | Path = DEFAULT_TACTILE_VAE_DIR,
+    map_location: str | torch.device = "cpu",
+    freeze: bool = True,
+    strict: bool = True,
+    model_kwargs: Optional[dict[str, Any]] = None,
+) -> TactileVAE:
+    ckpt_path = _resolve_checkpoint(checkpoint, vae_dir)
+    if not ckpt_path.exists():
+        raise FileNotFoundError(f"Tactile VAE checkpoint not found at: {ckpt_path}")
+
+    state = torch.load(str(ckpt_path), map_location=map_location)
+    state_dict = _unwrap_state(state)
+
+    model = TactileVAE(**(model_kwargs or {}))
+    model.load_state_dict(state_dict, strict=strict)
+
+    if freeze:
+        model.eval()
+        for p in model.parameters():
+            p.requires_grad_(False)
+    return model
+
+
+if __name__ == "__main__":
+    # Architecture roundtrip test (no real checkpoint required).
+    # To test with a real checkpoint, set CKPT_PATH below.
+    _CKPT_PATH = Path("/group2/ct/weihanx/tactile_world_model/tactile_wm/pretrained_models/ckpt_best.pt")
+    _OUT_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/test_output")
+    _EPISODE_PATH = Path("/group2/ct/weihanx/tactile_world_model/mode1_v1/0323_episode_000.pt")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"device: {device}")
+
+    # ── 1. Architecture test with random weights ─────────────────────────────
+    print("\n[1] Architecture roundtrip with random weights")
+    vae = TactileVAE().eval().to(device)
+    x_rand = torch.randn(2, 3, 128, 128, device=device)
+    with torch.no_grad():
+        out = vae(x_rand)
+    print(f"  input:  {tuple(x_rand.shape)}")
+    print(f"  z:      {tuple(out['z'].shape)}")
+    print(f"  x_hat:  {tuple(out['x_hat'].shape)}")
+
+    # Separate encode → decode
+    with torch.no_grad():
+        mu, logvar, _ = vae.encode(x_rand)
+        x_hat2 = vae.decode(mu)
+    print(f"  encode → z (mu): {tuple(mu.shape)}")
+    print(f"  decode → x_hat:  {tuple(x_hat2.shape)}")
+
+    # ── 2. TactileVAEWrapper with real checkpoint (if available) ──────────────
+    print(f"\n[2] TactileVAEWrapper from checkpoint: {_CKPT_PATH}")
+    if not _CKPT_PATH.exists():
+        print("  checkpoint not found — skipping pretrained test")
+    else:
+        ep = torch.load(str(_EPISODE_PATH), map_location="cpu", weights_only=False)
+        views = ep["view"]  # (T, 3, H, W) uint8
+        sample_indices = [0, 100, 500, 1000, 2000]
+        frames_u8 = views[sample_indices]  # (N, 3, H, W)
+        frames = frames_u8.float() / 127.5 - 1.0  # [-1, 1]
+
+        wrapper = TactileVAEWrapper(str(_CKPT_PATH), device=device)
+        z = wrapper.encode(frames)
+        x_hat = wrapper.decode(z)
+        print(f"  frames: {tuple(frames.shape)}  z: {tuple(z.shape)}  x_hat: {tuple(x_hat.shape)}")
+
+        _OUT_DIR.mkdir(parents=True, exist_ok=True)
+        for i, idx in enumerate(sample_indices):
+            orig_np = ((frames[i].permute(1, 2, 0).clamp(-1, 1) + 1) * 127.5).byte().numpy()
+            recon_np = ((x_hat[i].permute(1, 2, 0).clamp(-1, 1) + 1) * 127.5).byte().cpu().numpy()
+            diff_np = (np.abs(orig_np.astype(int) - recon_np.astype(int))).astype(np.uint8)
+            h, w = orig_np.shape[:2]
+            panel = Image.new("RGB", (3 * w + 16, h), (20, 20, 20))
+            panel.paste(Image.fromarray(orig_np), (0, 0))
+            panel.paste(Image.fromarray(recon_np), (w + 8, 0))
+            panel.paste(Image.fromarray(diff_np), (2 * w + 16, 0))
+            panel.save(_OUT_DIR / f"vae_frame_{idx:05d}_panel.png")
+            mse = float(((orig_np.astype(float) - recon_np.astype(float)) ** 2).mean())
+            psnr = 10 * np.log10(255.0 ** 2 / mse) if mse > 0 else float("inf")
+            print(f"  frame {idx:5d}  PSNR={psnr:.2f} dB")
+        print(f"  saved panels to {_OUT_DIR}")
+
+    print("\nAll tests passed.")

pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tactile_vae"
+version = "0.0.1"
+description = "ViT-based tactile variational autoencoder"
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.1",
+    "torchvision>=0.16",
+    "timm>=1.0",
+    "numpy",
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["tactile_vae*"]
+exclude = ["data*"]

script/inference.py

@@ -0,0 +1,234 @@
+"""General Lightning-based inference script for TactileVAE.
+
+Features:
+- Load any Lightning `.ckpt` checkpoint.
+- Load any config YAML.
+- Randomly select `N` samples from any split (`train` / `val` / `test`).
+- Run reconstruction inference and save metrics + visualization.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import yaml
+from PIL import Image
+from torch.utils.data import DataLoader, Subset
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tactile_vae.dataset import TactileParquetDataset
+from tactile_vae.model import TactileVAE
+
+DEFAULT_CONFIG = Path("/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_3/config.snapshot.yaml")
+DEFAULT_CKPT = Path("/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_3/checkpoints/last.ckpt")
+DEFAULT_OUT_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/inference/vae_baseline_3")
+
+
+def _resolve_path(p: str | Path) -> Path:
+    path = Path(p)
+    return path if path.is_absolute() else (_REPO_ROOT / path).resolve()
+
+
+def load_config(path: Path) -> dict:
+    with path.open() as f:
+        cfg = yaml.safe_load(f)
+    if not isinstance(cfg, dict):
+        raise ValueError(f"invalid config: {path}")
+    cfg["data"]["root"] = str(_resolve_path(cfg["data"]["root"]))
+    if cfg["data"].get("splits_path"):
+        cfg["data"]["splits_path"] = str(_resolve_path(cfg["data"]["splits_path"]))
+    return cfg
+
+
+def pick_device(spec: str) -> torch.device:
+    if spec == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(spec)
+
+
+class InferenceModule(pl.LightningModule):
+    """Minimal LightningModule used for strict Lightning checkpoint loading."""
+
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        self.model = TactileVAE(**config["model"])
+
+    def forward(self, x, **kw):
+        return self.model(x, **kw)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", type=Path, default=DEFAULT_CONFIG, help="config yaml")
+    p.add_argument("--ckpt", type=Path, default=DEFAULT_CKPT, help="Lightning checkpoint .ckpt")
+    p.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR, help="output directory")
+    p.add_argument("--split", type=str, default="test", choices=["train", "val", "test"])
+    p.add_argument("--num-samples", type=int, default=50, help="number of random samples from the split")
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--num-workers", type=int, default=0)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--device", type=str, default="auto", help="auto / cuda / cpu / cuda:0 ...")
+    p.add_argument("--max-grid", type=int, default=16, help="max samples shown in saved reconstruction grid")
+    return p.parse_args()
+
+
+def build_dataset(cfg: dict, split: str) -> TactileParquetDataset:
+    dcfg = cfg["data"]
+    return TactileParquetDataset(
+        root=dcfg["root"],
+        split=split,
+        splits_path=dcfg.get("splits_path"),
+        image_size=dcfg["image_size"],
+        cache_files=dcfg.get("cache_files", 1),
+        color_jitter=None,
+    )
+
+
+def select_subset(ds: TactileParquetDataset, n: int, seed: int) -> tuple[Subset, list[int]]:
+    n = min(max(1, int(n)), len(ds))
+    rng = np.random.default_rng(seed)
+    idx = rng.choice(len(ds), size=n, replace=False).tolist()
+    return Subset(ds, idx), idx
+
+
+@torch.no_grad()
+def run_inference(
+    module: InferenceModule,
+    ds: TactileParquetDataset,
+    subset_idx: list[int],
+    loader: DataLoader,
+    device: torch.device,
+) -> tuple[list[dict[str, Any]], float, float, list[tuple[torch.Tensor, torch.Tensor]]]:
+    module.eval().to(device)
+    per_sample: list[dict[str, Any]] = []
+    vis_pairs: list[tuple[torch.Tensor, torch.Tensor]] = []
+    mae_total = 0.0
+    mse_total = 0.0
+    n_total = 0
+
+    cursor = 0
+    for x in loader:
+        x = x.to(device, non_blocking=True)
+        out = module.model(x, sample=False)
+        x_hat = out["x_hat"]
+
+        abs_err = (x - x_hat).abs().mean(dim=(1, 2, 3))
+        sq_err = ((x - x_hat) ** 2).mean(dim=(1, 2, 3))
+        bs = x.shape[0]
+
+        for i in range(bs):
+            gidx = subset_idx[cursor + i]
+            sample_id = ds.sample_id(gidx)
+            mae_i = float(abs_err[i].item())
+            mse_i = float(sq_err[i].item())
+            per_sample.append(
+                {
+                    "subset_rank": cursor + i,
+                    "dataset_index": int(gidx),
+                    "sample_id": sample_id,
+                    "mae": mae_i,
+                    "mse": mse_i,
+                }
+            )
+            vis_pairs.append((x[i].detach().cpu(), x_hat[i].detach().cpu()))
+            mae_total += mae_i
+            mse_total += mse_i
+            n_total += 1
+        cursor += bs
+
+    mae_mean = mae_total / max(1, n_total)
+    mse_mean = mse_total / max(1, n_total)
+    return per_sample, mae_mean, mse_mean, vis_pairs
+
+
+def save_grid(pairs: list[tuple[torch.Tensor, torch.Tensor]], out_path: Path, n_show: int, image_size: int) -> None:
+    n = min(n_show, len(pairs))
+    if n <= 0:
+        return
+    h = w = int(image_size)
+    canvas = np.zeros((2 * h, n * w, 3), dtype=np.uint8)
+    for i in range(n):
+        src, rec = pairs[i]
+        src_np = (src.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+        rec_np = (rec.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+        canvas[:h, i * w : (i + 1) * w] = src_np
+        canvas[h:, i * w : (i + 1) * w] = rec_np
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    Image.fromarray(canvas).save(out_path)
+
+
+def main() -> None:
+    args = parse_args()
+    cfg = load_config(args.config)
+    device = pick_device(args.device)
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    print(f"config: {args.config}")
+    print(f"ckpt: {args.ckpt}")
+    print(f"split: {args.split}")
+    print(f"num_samples: {args.num_samples}")
+    print(f"device: {device}")
+    print(f"out_dir: {args.out_dir}")
+
+    ds = build_dataset(cfg, split=args.split)
+    subset, subset_idx = select_subset(ds, args.num_samples, args.seed)
+    print(f"split_size={len(ds)}  selected={len(subset_idx)}")
+    print(f"preview_sample_ids={[ds.sample_id(i) for i in subset_idx[:5]]}")
+
+    loader = DataLoader(
+        subset,
+        batch_size=min(max(1, args.batch_size), len(subset)),
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+        drop_last=False,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    module = InferenceModule.load_from_checkpoint(
+        str(args.ckpt),
+        config=cfg,
+        strict=True,
+        map_location="cpu",
+    )
+
+    per_sample, mae_mean, mse_mean, vis_pairs = run_inference(
+        module=module, ds=ds, subset_idx=subset_idx, loader=loader, device=device
+    )
+
+    grid_path = args.out_dir / "reconstruction_grid.png"
+    save_grid(vis_pairs, out_path=grid_path, n_show=args.max_grid, image_size=cfg["data"]["image_size"])
+
+    summary = {
+        "config": str(args.config),
+        "checkpoint": str(args.ckpt),
+        "split": args.split,
+        "seed": args.seed,
+        "selected_num_samples": len(subset_idx),
+        "mean_mae": mae_mean,
+        "mean_mse": mse_mean,
+        "grid_path": str(grid_path),
+    }
+    with (args.out_dir / "summary.json").open("w") as f:
+        json.dump(summary, f, indent=2)
+    with (args.out_dir / "per_sample_metrics.json").open("w") as f:
+        json.dump(per_sample, f, indent=2)
+
+    print(f"mean_mae={mae_mean:.6f}  mean_mse={mse_mean:.6f}")
+    print(f"saved: {args.out_dir / 'summary.json'}")
+    print(f"saved: {args.out_dir / 'per_sample_metrics.json'}")
+    print(f"saved: {grid_path}")
+
+
+if __name__ == "__main__":
+    main()

script/train_vae.py

@@ -0,0 +1,807 @@
+"""Train TactileVAE on the fota_unlabeled parquet dataset.
+
+Run:
+    python tactile_vae/script/train_vae.py --config tactile_vae/config/train_vae.yaml
+
+Each run lives in `<runs_root>/<run_id>/`. Re-launching with the same
+`run_id` auto-resumes from `ckpt_last.pt` in that directory (override with
+`--no-resume`, or `--resume-from <path>` to resume from a specific file).
+
+Writes into `<runs_root>/<run_id>/`:
+  - metrics.csv          per-step training + per-eval validation metrics
+  - samples/step_*.png   input vs. reconstruction grid (val-set images)
+  - ckpt_last.pt         most recent checkpoint
+  - ckpt_step_*.pt       periodic checkpoints (rotated; keep_last_ckpts)
+  - ckpt_best.pt         lowest monitored validation metric (default: val/total)
+  - run.log              stdout mirror
+  - config.snapshot.yaml the resolved config (first launch — preserved on resume)
+
+Checkpoints are saved as:
+    {"state_dict": ..., "optimizer": ..., "scaler": ..., "scheduler": ...,
+     "step": int, "epoch": int, "config": dict, "best_val_recon": float}
+which `tactile_vae.model.load_pretrained` can re-open via its `state_dict` key.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import datetime as dt
+import math
+import os
+import random
+import signal
+import sys
+import time
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import yaml
+from PIL import Image
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
+try:
+    import wandb  # optional — only imported if WANDB_PROJECT is set in env.
+except ImportError:  # pragma: no cover - wandb is optional
+    wandb = None  # type: ignore[assignment]
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tactile_vae.dataset import (
+    ColorJitterConfig,
+    ParquetFileShuffleSampler,
+    TactileParquetDataset,
+)
+from tactile_vae.model import TactileVAE, VAELoss
+
+
+# ---------------------------------------------------------------------------
+# Utility: config loading, path resolution, logging
+# ---------------------------------------------------------------------------
+
+def _resolve_path(p: str | None) -> Path | None:
+    if p is None:
+        return None
+    path = Path(p)
+    return path if path.is_absolute() else (_REPO_ROOT / path).resolve()
+
+
+def _autogenerate_run_id() -> str:
+    return "run_" + dt.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+
+def load_config(path: Path) -> dict:
+    with path.open() as f:
+        cfg = yaml.safe_load(f)
+
+    # `run_id` is the run's identity; auto-generate if missing.
+    if not cfg.get("run_id"):
+        cfg["run_id"] = _autogenerate_run_id()
+
+    # Derived `output_dir` = <runs_root>/<run_id>. An explicit `output_dir`
+    # in the YAML (legacy) is honored verbatim.
+    if cfg.get("output_dir"):
+        cfg["output_dir"] = str(_resolve_path(cfg["output_dir"]))
+    else:
+        runs_root = _resolve_path(cfg.get("runs_root", "runs"))
+        cfg["output_dir"] = str(runs_root / cfg["run_id"])
+
+    cfg["data"]["root"] = str(_resolve_path(cfg["data"]["root"]))
+    if cfg["data"].get("splits_path"):
+        cfg["data"]["splits_path"] = str(_resolve_path(cfg["data"]["splits_path"]))
+    if cfg["train"].get("resume_from"):
+        cfg["train"]["resume_from"] = str(_resolve_path(cfg["train"]["resume_from"]))
+    return cfg
+
+
+def _maybe_autoresume(cfg: dict, *, allow_autoresume: bool) -> None:
+    """If the run dir already has ckpt_last.pt and the user didn't pin
+    `resume_from`, auto-resume from it. Mutates `cfg["train"]` in place."""
+    if cfg["train"].get("resume_from") or not allow_autoresume:
+        return
+    last = Path(cfg["output_dir"]) / "ckpt_last.pt"
+    if last.exists():
+        cfg["train"]["resume_from"] = str(last)
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def pick_device(spec: str) -> torch.device:
+    if spec == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(spec)
+
+
+def init_wandb(config: dict, output_dir: Path) -> Any:
+    """Initialize a wandb run if `WANDB_PROJECT` is set in the environment.
+
+    Run id / name default to the training-script `run_id` so that re-launching
+    with the same run_id continues the same wandb run (`resume="allow"`).
+    Returns the wandb run handle, or None when wandb is unavailable / disabled.
+    """
+    if wandb is None or not os.environ.get("WANDB_PROJECT"):
+        return None
+    run = wandb.init(
+        project=os.environ["WANDB_PROJECT"],
+        entity=os.environ.get("WANDB_ENTITY"),
+        id=os.environ.get("WANDB_RUN_ID") or config["run_id"],
+        name=os.environ.get("WANDB_NAME") or config["run_id"],
+        resume="allow",
+        config=config,
+        mode=os.environ.get("WANDB_MODE", "online"),
+        dir=str(output_dir),
+    )
+    return run
+
+
+class TeeLogger:
+    """stdout that also appends to a file."""
+
+    def __init__(self, path: Path):
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self._fh = path.open("a", buffering=1)
+        self._stdout = sys.stdout
+
+    def write(self, msg: str) -> None:
+        self._stdout.write(msg)
+        self._fh.write(msg)
+
+    def flush(self) -> None:
+        self._stdout.flush()
+        self._fh.flush()
+
+
+# ---------------------------------------------------------------------------
+# Data + model + optim builders
+# ---------------------------------------------------------------------------
+
+def build_datasets(data_cfg: dict) -> tuple[TactileParquetDataset, TactileParquetDataset]:
+    common = dict(
+        root=data_cfg["root"],
+        image_size=data_cfg["image_size"],
+        cache_files=data_cfg.get("cache_files", 1),
+        splits_path=data_cfg.get("splits_path"),
+        return_meta=data_cfg.get("return_meta", False),
+    )
+    if data_cfg.get("meta_columns"):
+        common["meta_columns"] = data_cfg["meta_columns"]
+
+    jitter_cfg = data_cfg.get("color_jitter")
+    color_jitter = ColorJitterConfig(**jitter_cfg) if jitter_cfg else None
+
+    train_ds = TactileParquetDataset(split="train", color_jitter=color_jitter, **common)
+    val_ds = TactileParquetDataset(split="val", color_jitter=None, **common)
+    return train_ds, val_ds
+
+
+def build_model(model_cfg: dict) -> TactileVAE:
+    return TactileVAE(**model_cfg)
+
+
+class ConfigurablePerceptualVAELoss(nn.Module):
+    """VAE loss with configurable perceptual term: SSIM or LPIPS."""
+
+    def __init__(self, loss_cfg: dict):
+        super().__init__()
+        self.perceptual_type = str(loss_cfg.get("perceptual_type", "ssim")).lower()
+        if self.perceptual_type not in {"ssim", "lpips"}:
+            raise ValueError(
+                f"loss.perceptual_type must be one of [ssim, lpips], got: {self.perceptual_type!r}"
+            )
+        self.aux_key = self.perceptual_type
+        self.ssim_impl: VAELoss | None = None
+        self.lpips_impl: nn.Module | None = None
+
+        if self.perceptual_type == "ssim":
+            self.ssim_impl = VAELoss(**loss_cfg)
+        else:
+            self.beta = float(loss_cfg.get("beta", 1e-3))
+            self.recon_type = str(loss_cfg.get("recon_type", "l1")).lower()
+            self.lpips_weight = float(loss_cfg.get("lpips_weight", loss_cfg.get("ssim_weight", 0.1)))
+            try:
+                import lpips  # type: ignore
+            except ImportError as exc:  # pragma: no cover - depends on runtime env
+                raise ImportError(
+                    "LPIPS loss requested but `lpips` is not installed. "
+                    "Install with: pip install lpips"
+                ) from exc
+            self.lpips_impl = lpips.LPIPS(net="alex")
+            self.lpips_impl.eval()
+            for p in self.lpips_impl.parameters():
+                p.requires_grad = False
+
+    def forward(self, x_hat: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, logvar: torch.Tensor) -> dict[str, torch.Tensor]:
+        if self.perceptual_type == "ssim":
+            assert self.ssim_impl is not None
+            return self.ssim_impl(x_hat, x, mu, logvar)
+
+        if self.recon_type == "l1":
+            recon = F.l1_loss(x_hat, x)
+        elif self.recon_type == "mse":
+            recon = F.mse_loss(x_hat, x)
+        else:
+            raise ValueError(f"loss.recon_type must be one of [l1, mse], got: {self.recon_type!r}")
+
+        # LPIPS expects inputs in [-1, 1].
+        with torch.amp.autocast(device_type=x_hat.device.type, enabled=False):
+            x_hat_lp = (2.0 * x_hat.float()) - 1.0
+            x_lp = (2.0 * x.float()) - 1.0
+            assert self.lpips_impl is not None
+            lpips_val = self.lpips_impl(x_hat_lp, x_lp).mean()
+        recon_total = recon + self.lpips_weight * lpips_val
+        kl = (-0.5 * (1 + logvar - mu.pow(2) - logvar.exp())).mean()
+        total = recon_total + self.beta * kl
+        return {
+            "total": total,
+            "recon": recon,
+            "recon_total": recon_total,
+            "lpips": lpips_val,
+            "kl": kl,
+        }
+
+
+def build_loss(loss_cfg: dict) -> nn.Module:
+    return ConfigurablePerceptualVAELoss(loss_cfg)
+
+
+def build_optimizer(params, optim_cfg: dict) -> torch.optim.Optimizer:
+    return torch.optim.AdamW(
+        params,
+        lr=optim_cfg["lr"],
+        weight_decay=optim_cfg.get("weight_decay", 0.0),
+        betas=tuple(optim_cfg.get("betas", (0.9, 0.95))),
+        eps=optim_cfg.get("eps", 1e-8),
+    )
+
+
+def lr_at_step(step: int, base_lr: float, total_steps: int, sched_cfg: dict) -> float:
+    warmup = int(sched_cfg.get("warmup_steps", 0))
+    sched = sched_cfg.get("type", "constant")
+    if step < warmup:
+        return base_lr * (step + 1) / max(1, warmup)
+    if sched == "constant":
+        return base_lr
+    if sched == "cosine":
+        min_ratio = float(sched_cfg.get("min_lr_ratio", 0.1))
+        # Cosine from base_lr → base_lr * min_ratio over the remaining steps.
+        progress = (step - warmup) / max(1, total_steps - warmup)
+        progress = min(max(progress, 0.0), 1.0)
+        cos = 0.5 * (1.0 + math.cos(math.pi * progress))
+        return base_lr * (min_ratio + (1 - min_ratio) * cos)
+    raise ValueError(f"unknown scheduler type: {sched}")
+
+
+# ---------------------------------------------------------------------------
+# Training utilities: validation, sampling, checkpoints
+# ---------------------------------------------------------------------------
+
+@dataclass
+class MetricAccum:
+    sum: float = 0.0
+    n: int = 0
+
+    def add(self, v: float, count: int = 1) -> None:
+        self.sum += v * count
+        self.n += count
+
+    def mean(self) -> float:
+        return self.sum / self.n if self.n else float("nan")
+
+
+@torch.no_grad()
+def run_validation(
+    model: TactileVAE,
+    criterion: nn.Module,
+    loader: DataLoader,
+    device: torch.device,
+    max_batches: int,
+) -> dict[str, float]:
+    model.eval()
+    accs: dict[str, MetricAccum] = {}
+    for i, batch in enumerate(loader):
+        if i >= max_batches:
+            break
+        x = batch.to(device, non_blocking=True)
+        out = model(x, sample=False)
+        losses = criterion(out["x_hat"], x, out["mu"], out["logvar"])
+        bs = x.shape[0]
+        for k, v in losses.items():
+            if k not in accs:
+                accs[k] = MetricAccum()
+            accs[k].add(v.item(), bs)
+    model.train()
+    return {f"val/{k}": a.mean() for k, a in accs.items()}
+
+
+def _to_uint8_hwc(t: torch.Tensor) -> np.ndarray:
+    arr = t.detach().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
+    return (arr * 255).astype(np.uint8)
+
+
+@torch.no_grad()
+def save_sample_grid(
+    model: TactileVAE,
+    val_ds: TactileParquetDataset,
+    device: torch.device,
+    out_path: Path,
+    n: int,
+    rng_state: np.random.Generator,
+) -> tuple[list[np.ndarray], list[np.ndarray]]:
+    """Sample `n` images from val, run reconstruction, save a top=target/bottom=recon grid.
+
+    Returns (targets, reconstructions) as lists of HWC uint8 arrays for wandb logging.
+    """
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    indices = rng_state.choice(len(val_ds), size=n, replace=False).tolist()
+    imgs = torch.stack([val_ds[i] for i in indices]).to(device, non_blocking=True)
+    model.eval()
+    recon = model(imgs, sample=False)["x_hat"]
+    model.train()
+
+    targets = [_to_uint8_hwc(imgs[i]) for i in range(n)]
+    recons  = [_to_uint8_hwc(recon[i]) for i in range(n)]
+
+    # Local PNG: top row = target, bottom row = reconstruction.
+    h = w = val_ds.image_size
+    canvas = np.zeros((2 * h, n * w, 3), dtype=np.uint8)
+    for i in range(n):
+        canvas[:h, i * w : (i + 1) * w] = targets[i]
+        canvas[h:, i * w : (i + 1) * w] = recons[i]
+    Image.fromarray(canvas).save(out_path)
+
+    return targets, recons
+
+
+def save_checkpoint(
+    path: Path,
+    *,
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    scaler: torch.amp.GradScaler | None,
+    step: int,
+    epoch: int,
+    config: dict,
+    best_val_metric: float,
+    best_metric_name: str,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload: dict[str, Any] = {
+        "state_dict": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "step": step,
+        "epoch": epoch,
+        "config": config,
+        "best_val_metric": best_val_metric,
+        "best_metric_name": best_metric_name,
+        # Backward compatibility for older checkpoints/resume logic.
+        "best_val_recon": best_val_metric,
+    }
+    if scaler is not None:
+        payload["scaler"] = scaler.state_dict()
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    torch.save(payload, tmp)
+    os.replace(tmp, path)
+
+
+def rotate_periodic_ckpts(out_dir: Path, keep: int) -> None:
+    ckpts = sorted(out_dir.glob("ckpt_step_*.pt"))
+    while len(ckpts) > keep:
+        ckpts.pop(0).unlink(missing_ok=True)
+
+
+# ---------------------------------------------------------------------------
+# Main training loop
+# ---------------------------------------------------------------------------
+
+def train(config: dict) -> None:
+    set_seed(config["seed"])
+    device = pick_device(config["device"])
+    out_dir = Path(config["output_dir"])
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    sys.stdout = TeeLogger(out_dir / "run.log")  # type: ignore[assignment]
+    print(f"== Tactile VAE training ==")
+    print(f"run_id: {config['run_id']}    device: {device}")
+    print(f"output_dir: {out_dir}")
+    if config["train"].get("resume_from"):
+        print(f"resume_from: {config['train']['resume_from']}")
+
+    wandb_run = init_wandb(config, out_dir)
+    if wandb_run is not None:
+        print(f"wandb: project={os.environ.get('WANDB_PROJECT')} "
+              f"run_id={wandb_run.id} url={wandb_run.url}")
+    else:
+        print("wandb: disabled (set WANDB_PROJECT to enable)")
+
+    # Snapshot the resolved config on first launch; preserve the original on
+    # resume so the config used to start the run isn't silently overwritten.
+    snap = out_dir / "config.snapshot.yaml"
+    if not snap.exists():
+        with snap.open("w") as f:
+            yaml.safe_dump(config, f, sort_keys=False)
+
+    train_ds, val_ds = build_datasets(config["data"])
+    print(f"datasets: train={len(train_ds):,}  val={len(val_ds):,}")
+
+    tcfg = config["train"]
+    train_sampler = ParquetFileShuffleSampler(train_ds, seed=config["seed"])
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=tcfg["batch_size"],
+        sampler=train_sampler,
+        num_workers=tcfg["num_workers"],
+        pin_memory=device.type == "cuda",
+        drop_last=True,
+        persistent_workers=tcfg["num_workers"] > 0,
+        prefetch_factor=2 if tcfg["num_workers"] > 0 else None,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=tcfg["batch_size"],
+        shuffle=False,
+        num_workers=max(2, tcfg["num_workers"] // 2),
+        pin_memory=device.type == "cuda",
+        drop_last=False,
+    )
+    steps_per_epoch = len(train_loader)
+    total_steps = (
+        tcfg["max_steps"]
+        if tcfg.get("max_steps")
+        else steps_per_epoch * tcfg["epochs"]
+    )
+    print(f"steps/epoch={steps_per_epoch:,}  total_steps={total_steps:,}")
+
+    model = build_model(config["model"]).to(device)
+    criterion = build_loss(config["loss"]).to(device)
+    optimizer = build_optimizer(model.parameters(), config["optim"])
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"model: {model.__class__.__name__} params={n_params:,}")
+
+    use_amp = bool(tcfg.get("amp", False)) and device.type == "cuda"
+    amp_dtype_cfg = str(tcfg.get("amp_dtype", "bf16")).lower()
+    if amp_dtype_cfg not in {"bf16", "bfloat16"}:
+        print(f"[info] overriding train.amp_dtype={amp_dtype_cfg!r} to 'bf16' (enforced)")
+    amp_dtype = torch.bfloat16
+    if not use_amp:
+        amp_dtype = torch.float32
+
+    # Enforced bf16 path: no GradScaler.
+    scaler = None
+
+    # Graceful shutdown on preemption/cancel: write ckpt_last then exit.
+    stop_requested = False
+
+    def _request_stop(signum: int, _frame) -> None:
+        nonlocal stop_requested
+        stop_requested = True
+        try:
+            sig_name = signal.Signals(signum).name
+        except ValueError:
+            sig_name = str(signum)
+        print(f"[signal] received {sig_name}; stopping after current step and saving ckpt_last.pt")
+
+    prev_sigterm = signal.getsignal(signal.SIGTERM)
+    prev_sigint = signal.getsignal(signal.SIGINT)
+    signal.signal(signal.SIGTERM, _request_stop)
+    signal.signal(signal.SIGINT, _request_stop)
+
+    # ----- resume ---------------------------------------------------------
+    step = 0
+    epoch_start = 0
+    best_metric_name = str(tcfg.get("best_metric", "val/total"))
+    best_val_metric = float("inf")
+    if tcfg.get("resume_from"):
+        ckpt = torch.load(tcfg["resume_from"], map_location=device)
+        model.load_state_dict(ckpt["state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer"])
+        if scaler is not None and "scaler" in ckpt:
+            scaler.load_state_dict(ckpt["scaler"])
+        step = int(ckpt.get("step", 0))
+        epoch_start = int(ckpt.get("epoch", 0))
+        best_val_metric = float(
+            ckpt.get("best_val_metric", ckpt.get("best_val_recon", float("inf")))
+        )
+        best_metric_name = str(ckpt.get("best_metric_name", best_metric_name))
+        print(f"resumed from {tcfg['resume_from']} @ step={step} epoch={epoch_start}")
+
+    # ----- bookkeeping ----------------------------------------------------
+    metrics_csv = out_dir / "metrics.csv"
+    new_csv = not metrics_csv.exists()
+    csv_fh = metrics_csv.open("a", newline="", buffering=1)
+    csv_writer = csv.writer(csv_fh)
+    if new_csv:
+        aux_metric_name = str(getattr(criterion, "aux_key", "ssim"))
+        csv_writer.writerow(
+            ["step", "epoch", "lr", "split",
+             "loss_total", "recon", "recon_total", aux_metric_name, "kl", "throughput"]
+        )
+
+    aux_metric_name = str(getattr(criterion, "aux_key", "ssim"))
+    metric_keys = ("total", "recon", "recon_total", aux_metric_name, "kl")
+    running = {k: deque(maxlen=tcfg["log_every"]) for k in metric_keys}
+    grad_norm_running: deque[float] = deque(maxlen=tcfg["log_every"])
+    sample_rng = np.random.default_rng(config["seed"] + 1)
+    t_window = time.time()
+    samples_in_window = 0
+
+    base_lr = config["optim"]["lr"]
+    model.train()
+
+    # ----- training loop --------------------------------------------------
+    done = False
+    try:
+        for epoch in range(epoch_start, tcfg["epochs"]):
+            train_sampler.set_epoch(epoch)
+            epoch_accs = {k: MetricAccum() for k in metric_keys}
+            epoch_samples = 0
+            for batch in train_loader:
+                if step >= total_steps or stop_requested:
+                    done = True
+                    break
+
+                lr = lr_at_step(step, base_lr, total_steps, config["scheduler"])
+                for g in optimizer.param_groups:
+                    g["lr"] = lr
+
+                x = batch.to(device, non_blocking=True)
+                optimizer.zero_grad(set_to_none=True)
+
+                with torch.amp.autocast(device.type, dtype=amp_dtype, enabled=use_amp):
+                    out = model(x)
+                    losses = criterion(out["x_hat"], x, out["mu"], out["logvar"])
+                    loss = losses["total"]
+
+                if not torch.isfinite(loss).item():
+                    print(
+                        f"[warn] non-finite loss at step={step + 1}, epoch={epoch}; "
+                        "skipping optimizer step"
+                    )
+                    optimizer.zero_grad(set_to_none=True)
+                    continue
+
+                grad_norm_value: float | None = None
+                if scaler is not None:
+                    scaler.scale(loss).backward()
+                    scaler.unscale_(optimizer)
+                    if tcfg.get("gradient_clip_norm"):
+                        gn = torch.nn.utils.clip_grad_norm_(
+                            model.parameters(),
+                            tcfg["gradient_clip_norm"],
+                        )
+                        grad_norm_value = float(gn.item())
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    loss.backward()
+                    if tcfg.get("gradient_clip_norm"):
+                        gn = torch.nn.utils.clip_grad_norm_(
+                            model.parameters(),
+                            tcfg["gradient_clip_norm"],
+                        )
+                        grad_norm_value = float(gn.item())
+                    optimizer.step()
+
+                bs = x.shape[0]
+                for k, dq in running.items():
+                    dq.append(losses[k].item())
+                for k, acc in epoch_accs.items():
+                    acc.add(losses[k].item(), bs)
+                if grad_norm_value is not None and math.isfinite(grad_norm_value):
+                    grad_norm_running.append(grad_norm_value)
+                samples_in_window += bs
+                epoch_samples += bs
+                step += 1
+
+                # ----- step-level logging -----------------------------------
+                if step % tcfg["log_every"] == 0 or step == 1:
+                    now = time.time()
+                    throughput = samples_in_window / max(1e-6, now - t_window)
+                    means = {k: sum(dq) / len(dq) for k, dq in running.items()}
+                    print(
+                        f"step {step:>7} | ep {epoch:>3} | lr {lr:.2e} | "
+                        f"total {means['total']:.4f} recon {means['recon']:.4f} "
+                        f"{aux_metric_name} {means[aux_metric_name]:.4f} kl {means['kl']:.4f} "
+                        f"gclip {((sum(grad_norm_running) / len(grad_norm_running)) if grad_norm_running else float('nan')):.3f} | "
+                        f"{throughput:.0f} img/s"
+                    )
+                    csv_writer.writerow([
+                        step, epoch, f"{lr:.6g}", "train",
+                        f"{means['total']:.6g}", f"{means['recon']:.6g}",
+                        f"{means['recon_total']:.6g}", f"{means[aux_metric_name]:.6g}",
+                        f"{means['kl']:.6g}", f"{throughput:.1f}",
+                    ])
+                    if wandb_run is not None:
+                        wandb_run.log({
+                            "train/total": means["total"],
+                            "train/recon": means["recon"],
+                            "train/recon_total": means["recon_total"],
+                            f"train/{aux_metric_name}": means[aux_metric_name],
+                            "train/kl": means["kl"],
+                            "train/throughput_img_per_s": throughput,
+                            "train/lr": lr,
+                            "epoch": epoch,
+                        }, step=step)
+                    t_window = now
+                    samples_in_window = 0
+
+                # ----- validation -------------------------------------------
+                if step % tcfg["val_every_steps"] == 0:
+                    vmetrics = run_validation(
+                        model, criterion, val_loader, device,
+                        max_batches=tcfg["num_val_batches"],
+                    )
+                    print(
+                        f"  [val @ step {step}]  "
+                        + "  ".join(f"{k.split('/')[-1]}={v:.4f}" for k, v in vmetrics.items())
+                    )
+                    csv_writer.writerow([
+                        step, epoch, f"{lr:.6g}", "val",
+                        f"{vmetrics['val/total']:.6g}", f"{vmetrics['val/recon']:.6g}",
+                        f"{vmetrics['val/recon_total']:.6g}", f"{vmetrics[f'val/{aux_metric_name}']:.6g}",
+                        f"{vmetrics['val/kl']:.6g}", "",
+                    ])
+                    if wandb_run is not None:
+                        wandb_run.log(vmetrics, step=step)
+                    if best_metric_name not in vmetrics:
+                        raise KeyError(
+                            f"train.best_metric={best_metric_name!r} not found in validation metrics "
+                            f"{sorted(vmetrics.keys())}"
+                        )
+                    if vmetrics[best_metric_name] < best_val_metric:
+                        best_val_metric = vmetrics[best_metric_name]
+                        save_checkpoint(
+                            out_dir / "ckpt_best.pt",
+                            model=model, optimizer=optimizer, scaler=scaler,
+                            step=step, epoch=epoch, config=config,
+                            best_val_metric=best_val_metric,
+                            best_metric_name=best_metric_name,
+                        )
+                        print(
+                            f"  -> new best {best_metric_name}={best_val_metric:.4f}, "
+                            "saved ckpt_best.pt"
+                        )
+
+                # ----- sample grid ------------------------------------------
+                if step % tcfg["sample_every_steps"] == 0:
+                    sample_path = out_dir / "samples" / f"step_{step:07d}.png"
+                    targets, recons = save_sample_grid(
+                        model, val_ds, device,
+                        out_path=sample_path,
+                        n=tcfg["num_sample_images"],
+                        rng_state=sample_rng,
+                    )
+                    if wandb_run is not None:
+                        wandb_run.log({
+                            "samples/target": [
+                                wandb.Image(tgt, caption=f"sample {i}") for i, tgt in enumerate(targets)
+                            ],
+                            "samples/reconstruction": [
+                                wandb.Image(rec, caption=f"sample {i}") for i, rec in enumerate(recons)
+                            ],
+                        }, step=step)
+
+                # ----- periodic checkpoint ----------------------------------
+                if step % tcfg["ckpt_every_steps"] == 0:
+                    save_checkpoint(
+                        out_dir / f"ckpt_step_{step:07d}.pt",
+                        model=model, optimizer=optimizer, scaler=scaler,
+                        step=step, epoch=epoch, config=config,
+                        best_val_metric=best_val_metric,
+                        best_metric_name=best_metric_name,
+                    )
+                    save_checkpoint(
+                        out_dir / "ckpt_last.pt",
+                        model=model, optimizer=optimizer, scaler=scaler,
+                        step=step, epoch=epoch, config=config,
+                        best_val_metric=best_val_metric,
+                        best_metric_name=best_metric_name,
+                    )
+                    rotate_periodic_ckpts(out_dir, tcfg["keep_last_ckpts"])
+                    print(f"  saved ckpt_step_{step:07d}.pt")
+
+            # ----- epoch-end logging ----------------------------------------
+            if epoch_samples > 0:
+                epoch_means = {k: acc.mean() for k, acc in epoch_accs.items()}
+                print(
+                    f"[epoch {epoch} end @ step {step}]  "
+                    + "  ".join(f"{k}={v:.4f}" for k, v in epoch_means.items())
+                )
+                csv_writer.writerow([
+                    step, epoch, f"{lr:.6g}", "train_epoch",
+                    f"{epoch_means['total']:.6g}", f"{epoch_means['recon']:.6g}",
+                    f"{epoch_means['recon_total']:.6g}", f"{epoch_means[aux_metric_name]:.6g}",
+                    f"{epoch_means['kl']:.6g}", "",
+                ])
+                if wandb_run is not None:
+                    wandb_run.log(
+                        {f"epoch_train/{k}": v for k, v in epoch_means.items()} | {"epoch": epoch},
+                        step=step,
+                    )
+
+            if done:
+                break
+    finally:
+        signal.signal(signal.SIGTERM, prev_sigterm)
+        signal.signal(signal.SIGINT, prev_sigint)
+
+    save_checkpoint(
+        out_dir / "ckpt_last.pt",
+        model=model, optimizer=optimizer, scaler=scaler,
+        step=step, epoch=epoch, config=config,
+        best_val_metric=best_val_metric,
+        best_metric_name=best_metric_name,
+    )
+    csv_fh.close()
+    if wandb_run is not None:
+        wandb_run.summary["best_val_metric"] = best_val_metric
+        wandb_run.summary["best_metric_name"] = best_metric_name
+        wandb_run.summary["final_step"] = step
+        wandb_run.finish()
+    print(f"done. step={step} best_{best_metric_name}={best_val_metric:.4f}")
+
+
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--config",
+        type=Path,
+        default=Path(__file__).resolve().parents[1] / "config" / "train_vae.yaml",
+    )
+    p.add_argument("--run-id", type=str, default=None,
+                   help="override run_id from the config (output_dir becomes runs_root/<run-id>)")
+    p.add_argument("--output-dir", type=str, default=None,
+                   help="override output_dir directly (bypasses runs_root/run_id derivation)")
+    p.add_argument("--resume-from", type=str, default=None,
+                   help="resume from a specific checkpoint path (overrides auto-resume)")
+    p.add_argument("--no-resume", action="store_true",
+                   help="start fresh even if <output_dir>/ckpt_last.pt exists")
+    return p.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+
+    # Read the YAML, then apply CLI overrides BEFORE deriving output_dir so
+    # `--run-id` re-points the run directory.
+    with args.config.open() as f:
+        raw_cfg = yaml.safe_load(f)
+    if args.run_id:
+        raw_cfg["run_id"] = args.run_id
+    if args.output_dir:
+        raw_cfg["output_dir"] = args.output_dir   # bypasses runs_root/run_id
+
+    # Funnel through the normal loader to resolve paths + autogen run_id.
+    tmp = args.config.parent / f".__cli_override_{os.getpid()}.yaml"
+    try:
+        with tmp.open("w") as f:
+            yaml.safe_dump(raw_cfg, f, sort_keys=False)
+        cfg = load_config(tmp)
+    finally:
+        tmp.unlink(missing_ok=True)
+
+    if args.resume_from:
+        cfg["train"]["resume_from"] = str(_resolve_path(args.resume_from))
+    _maybe_autoresume(cfg, allow_autoresume=not args.no_resume)
+
+    train(cfg)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

script/train_vae.sh

@@ -0,0 +1,184 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=tactile_vae
+#SBATCH --partition=ct
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:h100:1
+#SBATCH --requeue
+#SBATCH --output=/group2/ct/weihanx/tactile_world_model/slurm-logs/tactile_vae.%j.log
+#SBATCH --error=/group2/ct/weihanx/tactile_world_model/slurm-logs/tactile_vae.%j.log
+
+# Train tactile_vae.model.TactileVAE on the fota_unlabeled parquet dataset.
+#
+# Each run lives at <RUNS_DIR>/<RUN_ID>/. Re-launching with the same RUN_ID
+# auto-resumes from ckpt_last.pt; wandb keeps the same run id, so metrics
+# append to the same dashboard.
+#
+# Usage (sbatch):  sbatch tactile_vae/script/train_vae.sh <run_id> [config.yaml]
+# Usage (local):   ./tactile_vae/script/train_vae.sh        <run_id> [config.yaml]
+#
+# Diagnostics: set DEBUG=1 to enable `set -x` command tracing.
+#   sbatch --export=ALL,DEBUG=1 tactile_vae/script/train_vae.sh <run_id>
+
+# Force unbuffered stdout/stderr so the slurm log shows progress live, not in
+# one giant flush at the end. (Without this, NFS-backed log files can look
+# completely empty for minutes while bash + python buffer output.)
+exec 1> >(stdbuf -oL -eL cat) 2>&1
+
+set -euo pipefail
+[[ "${DEBUG:-0}" == "1" ]] && set -x
+
+# ============================================================
+#  Inputs (positional)
+# ============================================================
+if [[ $# -lt 1 ]]; then
+  echo "Usage: $0 <run_id> [config.yaml]" >&2
+  echo "  run_id   : required. Both the output subdir name and the wandb run id." >&2
+  echo "  config   : optional. Defaults to tactile_vae/config/train_vae.yaml." >&2
+  exit 2
+fi
+RUN_ID="$1"
+CONFIG="${2:-tactile_vae/config/train_vae.yaml}"
+
+# ============================================================
+#  Paths
+# ============================================================
+WORKDIR="/group2/ct/weihanx/tactile_world_model"
+# Keep this aligned with train_vae.yaml default `runs_root: runs`.
+RUNS_DIR="$WORKDIR/tactile_world_model/runs"
+RUN_DIR="$RUNS_DIR/$RUN_ID"
+DATA_DIR="$WORKDIR/tactile_vae/data"
+SPLITS_PATH="$WORKDIR/tactile_vae/dataset/splits.json"
+
+# Conda env with all required deps installed. Override via env var if you
+# prefer a different env (e.g. CONDA_ENV=samaudio311 sbatch ...).
+#   torch torchvision timm numpy pyarrow PIL pyyaml wandb
+# `twm` is the project's standard env (matches tactile_jepa training).
+CONDA_ENV="${CONDA_ENV:-twm}"
+
+mkdir -p "$WORKDIR/slurm-logs"
+mkdir -p "$RUNS_DIR"
+umask 027
+
+# ============================================================
+#  Print startup info IMMEDIATELY — before any heavy operation
+#  (conda activate / python imports) so the slurm log is never
+#  silent for more than a fraction of a second.
+# ============================================================
+echo "=== Tactile VAE training ==="
+echo "Host:       $(hostname)"
+echo "Job ID:     ${SLURM_JOB_ID:-N/A}"
+echo "Start time: $(date)"
+echo "Run ID:     $RUN_ID"
+echo "Workdir:    $WORKDIR"
+echo "Config:     $CONFIG"
+echo "Run dir:    $RUN_DIR"
+echo "Conda env:  $CONDA_ENV"
+echo
+
+# ============================================================
+#  Environment knobs
+# ============================================================
+export OMP_NUM_THREADS=8
+export MKL_NUM_THREADS=8
+export TOKENIZERS_PARALLELISM="false"
+export PYTHONFAULTHANDLER=1
+export PYTHONUNBUFFERED=1   # ensures `print()` in Python flushes per line
+
+# ============================================================
+#  Weights & Biases (mirrors jepa_training.sh — same account)
+# ============================================================
+export WANDB_API_KEY="76cdc4261bf436617e661171fd41d80403e69e9b"
+export WANDB_ENTITY="weihanx-university-of-michigan"
+export WANDB_USERNAME="weihanx@umich.edu"
+export WANDB_PROJECT="tactile_vae"
+export WANDB_MODE="online"
+export WANDB_RUN_ID="$RUN_ID"
+export WANDB_NAME="$RUN_ID"
+export WANDB_SERVICE_WAIT=300
+export WANDB_INIT_TIMEOUT=300
+export WANDB_START_METHOD="thread"
+export WANDB_CONSOLE="wrap"
+# Keep wandb metadata/cache off network storage to speed init/resume.
+export WANDB_DIR="${WANDB_DIR:-/tmp/$USER/wandb/$RUN_ID}"
+export WANDB_CACHE_DIR="${WANDB_CACHE_DIR:-/tmp/$USER/wandb-cache}"
+export WANDB_DATA_DIR="${WANDB_DATA_DIR:-/tmp/$USER/wandb-data}"
+mkdir -p "$WANDB_DIR" "$WANDB_CACHE_DIR" "$WANDB_DATA_DIR"
+
+# Debug knob: disable wandb entirely to isolate startup stalls.
+# Default is enabled; set DISABLE_WANDB=1 to disable.
+DISABLE_WANDB="${DISABLE_WANDB:-0}"
+if [[ "$DISABLE_WANDB" == "1" ]]; then
+  unset WANDB_PROJECT WANDB_ENTITY WANDB_API_KEY WANDB_USERNAME
+  export WANDB_MODE="disabled"
+fi
+
+echo "--- Wandb ---"
+echo "  project=${WANDB_PROJECT:-<disabled>}  entity=${WANDB_ENTITY:-<disabled>}"
+echo "  run_id=$WANDB_RUN_ID  name=$WANDB_NAME  mode=$WANDB_MODE"
+echo "  dir=$WANDB_DIR"
+echo "  cache_dir=$WANDB_CACHE_DIR"
+echo "  data_dir=$WANDB_DATA_DIR"
+if [[ -n "${WANDB_API_KEY:-}" ]]; then
+  echo "  api_key=${WANDB_API_KEY:0:10}...${WANDB_API_KEY: -4}"
+else
+  echo "  api_key=<disabled>"
+fi
+echo
+
+# ============================================================
+#  Sanity checks (cheap; before conda activate)
+# ============================================================
+if [[ ! -f "$WORKDIR/$CONFIG" ]] && [[ ! -f "$CONFIG" ]]; then
+  echo "ERROR: config not found: $CONFIG (or $WORKDIR/$CONFIG)" >&2
+  exit 2
+fi
+if [[ ! -d "$DATA_DIR" ]]; then
+  echo "ERROR: data dir does not exist: $DATA_DIR" >&2
+  exit 2
+fi
+if [[ ! -f "$SPLITS_PATH" ]]; then
+  echo "ERROR: splits manifest not found: $SPLITS_PATH" >&2
+  echo "  Generate it with: python tactile_vae/dataset/make_splits.py" >&2
+  exit 2
+fi
+
+if [[ -f "$RUN_DIR/ckpt_last.pt" ]]; then
+  echo "Resume:     auto-resume from $RUN_DIR/ckpt_last.pt"
+else
+  echo "Resume:     fresh run (no $RUN_DIR/ckpt_last.pt)"
+fi
+echo
+
+# ============================================================
+#  Resolve Python interpreter
+# ============================================================
+# Fast path: call env python directly to avoid expensive `conda activate`
+# startup on busy shared filesystems. Fallback to full activation if needed.
+PYTHON_BIN="${PYTHON_BIN:-$HOME/miniconda3/envs/$CONDA_ENV/bin/python}"
+if [[ -x "$PYTHON_BIN" ]]; then
+  echo "[$(date +%H:%M:%S)] using env python directly: $PYTHON_BIN"
+else
+  echo "[$(date +%H:%M:%S)] env python not found; falling back to conda activate..."
+  source ~/miniconda3/etc/profile.d/conda.sh
+  echo "[$(date +%H:%M:%S)] activating $CONDA_ENV..."
+  conda activate "$CONDA_ENV"
+  PYTHON_BIN="$(which python)"
+  echo "[$(date +%H:%M:%S)] env activated. python = $PYTHON_BIN"
+fi
+echo "[$(date +%H:%M:%S)] GPU(s): ${CUDA_VISIBLE_DEVICES:-$(nvidia-smi -L 2>/dev/null | head -1 || echo none)}"
+echo
+
+# ============================================================
+#  Launch training (`-u` is also forced by PYTHONUNBUFFERED above)
+# ============================================================
+cd "$WORKDIR"
+echo "[$(date +%H:%M:%S)] launching trainer..."
+"$PYTHON_BIN" -u tactile_vae/script/train_vae.py \
+    --config "$CONFIG" \
+    --run-id "$RUN_ID"
+
+echo
+echo "[$(date +%H:%M:%S)] Finished."
+echo "Run dir contents:"
+ls -lh "$RUN_DIR" 2>/dev/null || echo "  (empty)"

script/train_vae_pl.py

@@ -0,0 +1,624 @@
+"""Train TactileVAE with PyTorch Lightning.
+
+Run:
+    python tactile_vae/script/train_vae_pl.py --config tactile_vae/config/train_vae.yaml
+
+Same YAML config format as train_vae.py.
+
+Checkpoints written to <output_dir>/:
+  ckpt_best.pt / ckpt_last.pt / ckpt_step_*.pt  — original format (TactileVAEWrapper compat)
+  checkpoints/last.ckpt                          — Lightning format (full resume with trainer state)
+"""
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import math
+import os
+import random
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import yaml
+from PIL import Image
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.loggers import CSVLogger
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import DataLoader
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tactile_vae.dataset import ColorJitterConfig, ParquetFileShuffleSampler, TactileParquetDataset
+from tactile_vae.model import TactileVAE, VAELoss
+
+
+# ---------------------------------------------------------------------------
+# Utilities (same as train_vae.py)
+# ---------------------------------------------------------------------------
+
+def _resolve_path(p: str | None) -> Path | None:
+    if p is None:
+        return None
+    path = Path(p)
+    return path if path.is_absolute() else (_REPO_ROOT / path).resolve()
+
+
+def _autogenerate_run_id() -> str:
+    return "run_" + dt.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+
+
+def load_config(path: Path) -> dict:
+    with path.open() as f:
+        cfg = yaml.safe_load(f)
+    if not cfg.get("run_id"):
+        cfg["run_id"] = _autogenerate_run_id()
+    if cfg.get("output_dir"):
+        cfg["output_dir"] = str(_resolve_path(cfg["output_dir"]))
+    else:
+        runs_root = _resolve_path(cfg.get("runs_root", "runs"))
+        cfg["output_dir"] = str(runs_root / cfg["run_id"])
+    cfg["data"]["root"] = str(_resolve_path(cfg["data"]["root"]))
+    if cfg["data"].get("splits_path"):
+        cfg["data"]["splits_path"] = str(_resolve_path(cfg["data"]["splits_path"]))
+    if cfg["train"].get("resume_from"):
+        cfg["train"]["resume_from"] = str(_resolve_path(cfg["train"]["resume_from"]))
+    return cfg
+
+
+def _maybe_autoresume(cfg: dict, *, allow_autoresume: bool) -> None:
+    if cfg["train"].get("resume_from") or not allow_autoresume:
+        return
+    # Prefer Lightning checkpoint for full state restore (step count, optimizer, etc.)
+    last_ckpt = Path(cfg["output_dir"]) / "checkpoints" / "last.ckpt"
+    if last_ckpt.exists():
+        cfg["train"]["resume_from"] = str(last_ckpt)
+        return
+    last_pt = Path(cfg["output_dir"]) / "ckpt_last.pt"
+    if last_pt.exists():
+        cfg["train"]["resume_from"] = str(last_pt)
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def build_datasets(data_cfg: dict) -> tuple[TactileParquetDataset, TactileParquetDataset]:
+    common = dict(
+        root=data_cfg["root"],
+        image_size=data_cfg["image_size"],
+        cache_files=data_cfg.get("cache_files", 1),
+        splits_path=data_cfg.get("splits_path"),
+        return_meta=data_cfg.get("return_meta", False),
+    )
+    if data_cfg.get("meta_columns"):
+        common["meta_columns"] = data_cfg["meta_columns"]
+    jitter_cfg = data_cfg.get("color_jitter")
+    color_jitter = ColorJitterConfig(**jitter_cfg) if jitter_cfg else None
+    train_ds = TactileParquetDataset(split="train", color_jitter=color_jitter, **common)
+    val_ds = TactileParquetDataset(split="val", color_jitter=None, **common)
+    return train_ds, val_ds
+
+
+def lr_at_step(step: int, base_lr: float, total_steps: int, sched_cfg: dict) -> float:
+    warmup = int(sched_cfg.get("warmup_steps", 0))
+    sched = sched_cfg.get("type", "constant")
+    if step < warmup:
+        return base_lr * (step + 1) / max(1, warmup)
+    if sched == "constant":
+        return base_lr
+    if sched == "cosine":
+        min_ratio = float(sched_cfg.get("min_lr_ratio", 0.1))
+        progress = (step - warmup) / max(1, total_steps - warmup)
+        progress = min(max(progress, 0.0), 1.0)
+        return base_lr * (min_ratio + (1 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress)))
+    raise ValueError(f"unknown scheduler type: {sched!r}")
+
+
+# ---------------------------------------------------------------------------
+# LightningModule
+# ---------------------------------------------------------------------------
+
+
+class ConfigurablePerceptualVAELoss(nn.Module):
+    """VAE loss with configurable perceptual term: SSIM or LPIPS."""
+
+    def __init__(self, loss_cfg: dict):
+        super().__init__()
+        self.perceptual_type = str(loss_cfg.get("perceptual_type", "ssim")).lower()
+        if self.perceptual_type not in {"ssim", "lpips"}:
+            raise ValueError(
+                f"loss.perceptual_type must be one of [ssim, lpips], got: {self.perceptual_type!r}"
+            )
+        self.aux_key = self.perceptual_type
+        self.ssim_impl: VAELoss | None = None
+        self.lpips_impl: nn.Module | None = None
+
+        if self.perceptual_type == "ssim":
+            self.ssim_impl = VAELoss(**loss_cfg)
+        else:
+            self.beta = float(loss_cfg.get("beta", 1e-3))
+            self.recon_type = str(loss_cfg.get("recon_type", "l1")).lower()
+            self.lpips_weight = float(loss_cfg.get("lpips_weight", loss_cfg.get("ssim_weight", 0.1)))
+            try:
+                import lpips  # type: ignore
+            except ImportError as exc:  # pragma: no cover - depends on runtime env
+                raise ImportError(
+                    "LPIPS loss requested but `lpips` is not installed. "
+                    "Install with: pip install lpips"
+                ) from exc
+            self.lpips_impl = lpips.LPIPS(net="alex")
+            self.lpips_impl.eval()
+            for p in self.lpips_impl.parameters():
+                p.requires_grad = False
+
+    def forward(self, x_hat: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, logvar: torch.Tensor) -> dict[str, torch.Tensor]:
+        if self.perceptual_type == "ssim":
+            assert self.ssim_impl is not None
+            return self.ssim_impl(x_hat, x, mu, logvar)
+
+        if self.recon_type == "l1":
+            recon = F.l1_loss(x_hat, x)
+        elif self.recon_type == "mse":
+            recon = F.mse_loss(x_hat, x)
+        else:
+            raise ValueError(f"loss.recon_type must be one of [l1, mse], got: {self.recon_type!r}")
+
+        # LPIPS expects inputs in [-1, 1], and is more stable in fp32.
+        with torch.amp.autocast(device_type=x_hat.device.type, enabled=False):
+            x_hat_lp = (2.0 * x_hat.float()) - 1.0
+            x_lp = (2.0 * x.float()) - 1.0
+            assert self.lpips_impl is not None
+            lpips_val = self.lpips_impl(x_hat_lp, x_lp).mean()
+        recon_total = recon + self.lpips_weight * lpips_val
+        kl = (-0.5 * (1 + logvar - mu.pow(2) - logvar.exp())).mean()
+        total = recon_total + self.beta * kl
+        return {
+            "total": total,
+            "recon": recon,
+            "recon_total": recon_total,
+            "lpips": lpips_val,
+            "kl": kl,
+        }
+
+
+class TactileVAEModule(pl.LightningModule):
+    def __init__(self, config: dict, *, step_offset: int = 0, total_steps: int = 0):
+        super().__init__()
+        self.config = config
+        self.step_offset = int(step_offset)
+        self.total_steps = int(total_steps)
+        self.model = TactileVAE(**config["model"])
+        self.criterion = ConfigurablePerceptualVAELoss(config["loss"])
+
+    def forward(self, x, **kw):
+        return self.model(x, **kw)
+
+    def training_step(self, batch, batch_idx):
+        x = batch
+        out = self.model(x)
+        losses = self.criterion(out["x_hat"], x, out["mu"], out["logvar"])
+        if not torch.isfinite(losses["total"]).item():
+            print(
+                f"[warn] non-finite loss at step={self.trainer.global_step + self.step_offset + 1}, "
+                f"epoch={self.trainer.current_epoch}; skipping optimizer step"
+            )
+            return None
+        self.log("train/total", losses["total"], prog_bar=True, on_step=True, on_epoch=False, batch_size=x.shape[0])
+        self.log_dict(
+            {f"train/{k}": v for k, v in losses.items() if k != "total"},
+            on_step=True, on_epoch=False, batch_size=x.shape[0],
+        )
+        return losses["total"]
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        x = batch
+        out = self.model(x, sample=False)
+        losses = self.criterion(out["x_hat"], x, out["mu"], out["logvar"])
+        self.log_dict(
+            {f"val/{k}": v for k, v in losses.items()},
+            on_step=False, on_epoch=True, batch_size=x.shape[0],
+        )
+
+    def configure_optimizers(self):
+        optim_cfg = self.config["optim"]
+        optimizer = torch.optim.AdamW(
+            self.model.parameters(),
+            lr=optim_cfg["lr"],
+            weight_decay=optim_cfg.get("weight_decay", 0.0),
+            betas=tuple(optim_cfg.get("betas", (0.9, 0.95))),
+            eps=optim_cfg.get("eps", 1e-8),
+        )
+        base_lr = float(optim_cfg["lr"])
+        sched_cfg = self.config["scheduler"]
+        scheduler = LambdaLR(
+            optimizer,
+            lr_lambda=lambda step: lr_at_step(
+                step + self.step_offset, base_lr, self.total_steps, sched_cfg
+            ) / base_lr,
+        )
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
+        }
+
+
+# ---------------------------------------------------------------------------
+# LightningDataModule
+# ---------------------------------------------------------------------------
+
+class TactileVAEDataModule(pl.LightningDataModule):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        self.train_ds: TactileParquetDataset | None = None
+        self.val_ds: TactileParquetDataset | None = None
+        self.train_sampler: ParquetFileShuffleSampler | None = None
+
+    def setup(self, stage: str | None = None):
+        if self.train_ds is not None:
+            return
+        self.train_ds, self.val_ds = build_datasets(self.config["data"])
+        self.train_sampler = ParquetFileShuffleSampler(self.train_ds, seed=self.config["seed"])
+
+    def train_dataloader(self):
+        tcfg = self.config["train"]
+        return DataLoader(
+            self.train_ds,
+            batch_size=tcfg["batch_size"],
+            sampler=self.train_sampler,
+            num_workers=tcfg["num_workers"],
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=tcfg["num_workers"] > 0,
+            prefetch_factor=2 if tcfg["num_workers"] > 0 else None,
+        )
+
+    def val_dataloader(self):
+        tcfg = self.config["train"]
+        return DataLoader(
+            self.val_ds,
+            batch_size=tcfg["batch_size"],
+            shuffle=False,
+            num_workers=max(2, tcfg["num_workers"] // 2),
+            pin_memory=True,
+            drop_last=False,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Callbacks
+# ---------------------------------------------------------------------------
+
+class SetEpochCallback(pl.Callback):
+    """Keeps ParquetFileShuffleSampler epoch-aware for proper per-epoch shuffling."""
+
+    def __init__(self, *, epoch_offset: int = 0):
+        self.epoch_offset = int(epoch_offset)
+
+    def on_train_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        dm = trainer.datamodule
+        if hasattr(dm, "train_sampler") and hasattr(dm.train_sampler, "set_epoch"):
+            dm.train_sampler.set_epoch(trainer.current_epoch + self.epoch_offset)
+
+
+class SampleGridCallback(pl.Callback):
+    """Saves a top=original / bottom=reconstruction image grid every N steps."""
+
+    def __init__(self, config: dict, *, step_offset: int = 0):
+        self.sample_every = config["train"]["sample_every_steps"]
+        self.n = config["train"]["num_sample_images"]
+        self.out_dir = Path(config["output_dir"]) / "samples"
+        self.rng = np.random.default_rng(config["seed"] + 1)
+        self.step_offset = int(step_offset)
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        effective_step = trainer.global_step + self.step_offset
+        if effective_step > 0 and effective_step % self.sample_every == 0:
+            self._save_grid(trainer, pl_module, effective_step)
+
+    @torch.no_grad()
+    def _save_grid(self, trainer, pl_module, step):
+        val_ds = trainer.datamodule.val_ds
+        device = pl_module.device
+        self.out_dir.mkdir(parents=True, exist_ok=True)
+        idx = self.rng.choice(len(val_ds), size=self.n, replace=False).tolist()
+        imgs = torch.stack([val_ds[i] for i in idx]).to(device)
+        pl_module.eval()
+        recon = pl_module.model(imgs, sample=False)["x_hat"]
+        pl_module.train()
+        h = w = val_ds.image_size
+        canvas = np.zeros((2 * h, self.n * w, 3), dtype=np.uint8)
+        for i in range(self.n):
+            orig = (imgs[i].cpu().clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+            rec = (recon[i].cpu().clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+            canvas[:h, i * w:(i + 1) * w] = orig
+            canvas[h:, i * w:(i + 1) * w] = rec
+        Image.fromarray(canvas).save(self.out_dir / f"step_{step:07d}.png")
+
+
+class CompatCheckpointCallback(pl.Callback):
+    """Saves ckpt_last.pt / ckpt_step_*.pt / ckpt_best.pt in the original format
+    so that TactileVAEWrapper.load_pretrained keeps working unchanged."""
+
+    def __init__(
+        self,
+        config: dict,
+        *,
+        step_offset: int = 0,
+        epoch_offset: int = 0,
+        initial_best_val_metric: float = float("inf"),
+    ):
+        self.config = config
+        self.out_dir = Path(config["output_dir"])
+        self.ckpt_every = config["train"]["ckpt_every_steps"]
+        self.keep_last = config["train"]["keep_last_ckpts"]
+        self.best_metric = config["train"].get("best_metric", "val/total")
+        self.best_val_metric = float(initial_best_val_metric)
+        self.step_offset = int(step_offset)
+        self.epoch_offset = int(epoch_offset)
+
+    def _build_payload(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> dict:
+        # LightningModule.state_dict() prefixes all keys with "model." — strip it.
+        sd = {k[len("model."):]: v for k, v in pl_module.state_dict().items() if k.startswith("model.")}
+        return {
+            "state_dict": sd,
+            "optimizer": trainer.optimizers[0].state_dict(),
+            "step": trainer.global_step + self.step_offset,
+            "epoch": trainer.current_epoch + self.epoch_offset,
+            "config": self.config,
+            "best_val_metric": self.best_val_metric,
+            "best_metric_name": self.best_metric,
+            "best_val_recon": self.best_val_metric,  # backward compat key
+        }
+
+    def _save(self, path: Path, trainer, pl_module) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = path.with_suffix(path.suffix + ".tmp")
+        torch.save(self._build_payload(trainer, pl_module), tmp)
+        os.replace(tmp, path)
+
+    def _rotate(self) -> None:
+        ckpts = sorted(self.out_dir.glob("ckpt_step_*.pt"))
+        while len(ckpts) > self.keep_last:
+            ckpts.pop(0).unlink(missing_ok=True)
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        effective_step = trainer.global_step + self.step_offset
+        if effective_step > 0 and effective_step % self.ckpt_every == 0:
+            self._save(self.out_dir / f"ckpt_step_{effective_step:07d}.pt", trainer, pl_module)
+            self._save(self.out_dir / "ckpt_last.pt", trainer, pl_module)
+            self._rotate()
+            print(f"  saved ckpt_step_{effective_step:07d}.pt")
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        val = float(trainer.callback_metrics.get(self.best_metric, float("inf")))
+        if val < self.best_val_metric:
+            self.best_val_metric = val
+            self._save(self.out_dir / "ckpt_best.pt", trainer, pl_module)
+            print(f"  -> new best {self.best_metric}={val:.4f}, saved ckpt_best.pt")
+
+    def on_train_end(self, trainer, pl_module):
+        self._save(self.out_dir / "ckpt_last.pt", trainer, pl_module)
+
+
+class CompatResumeStateCallback(pl.Callback):
+    """Loads optimizer state from compat .pt resume checkpoints."""
+
+    def __init__(self, optim_state: dict[str, Any] | None):
+        self.optim_state = optim_state
+
+    def on_fit_start(self, trainer, pl_module):
+        if self.optim_state is None:
+            return
+        if not trainer.optimizers:
+            return
+        trainer.optimizers[0].load_state_dict(self.optim_state)
+        print("loaded optimizer state from compat checkpoint")
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--config", type=Path,
+                   default=Path(__file__).resolve().parents[1] / "config" / "train_vae.yaml")
+    p.add_argument("--run-id", type=str, default=None,
+                   help="override run_id (output_dir = runs_root/<run-id>)")
+    p.add_argument("--output-dir", type=str, default=None,
+                   help="override output_dir directly")
+    p.add_argument("--resume-from", type=str, default=None,
+                   help="path to .ckpt (Lightning) or .pt (compat) checkpoint")
+    p.add_argument("--no-resume", action="store_true",
+                   help="start fresh even if ckpt_last.pt / last.ckpt exists")
+    return p.parse_args()
+
+
+def _init_loggers(cfg: dict, out_dir: Path) -> list[Any]:
+    loggers: list[Any] = [CSVLogger(str(out_dir), name="", version="")]
+    if os.environ.get("WANDB_PROJECT"):
+        try:
+            from pytorch_lightning.loggers import WandbLogger
+            loggers.append(WandbLogger(
+                project=os.environ["WANDB_PROJECT"],
+                entity=os.environ.get("WANDB_ENTITY"),
+                id=os.environ.get("WANDB_RUN_ID") or cfg["run_id"],
+                name=os.environ.get("WANDB_NAME") or cfg["run_id"],
+                save_dir=str(out_dir),
+                config=cfg,
+            ))
+        except ImportError:
+            print("wandb not available — logging disabled")
+    return loggers
+
+
+def _build_trainer(
+    cfg: dict,
+    *,
+    callbacks: list[pl.Callback],
+    loggers: list[Any],
+    precision: str,
+    resume_from: str | None,
+    resume_step_offset: int,
+    total_steps: int,
+) -> pl.Trainer:
+    tcfg = cfg["train"]
+    trainer_kwargs: dict[str, Any] = {
+        "accelerator": "auto",
+        "devices": 1,
+        "precision": precision,
+        "callbacks": callbacks,
+        "logger": loggers,
+        "limit_val_batches": tcfg["num_val_batches"],
+        "val_check_interval": tcfg["val_every_steps"],
+        "check_val_every_n_epoch": None,  # step-based only; disable epoch-end validation
+        "log_every_n_steps": tcfg["log_every"],
+        "gradient_clip_val": tcfg.get("gradient_clip_norm") or None,
+        "num_sanity_val_steps": 0,
+        "default_root_dir": str(cfg["output_dir"]),
+    }
+    if resume_from and Path(resume_from).suffix != ".ckpt":
+        remaining_steps = max(0, total_steps - resume_step_offset)
+        trainer_kwargs["max_steps"] = remaining_steps
+        print(f"compat resume remaining_steps={remaining_steps}")
+    elif tcfg.get("max_steps"):
+        trainer_kwargs["max_steps"] = tcfg["max_steps"]
+    else:
+        trainer_kwargs["max_epochs"] = tcfg["epochs"]
+    return pl.Trainer(**trainer_kwargs)
+
+
+def main(cfg: dict) -> None:
+    set_seed(cfg["seed"])
+    out_dir = Path(cfg["output_dir"])
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Resume bookkeeping for compat .pt checkpoints.
+    resume_step_offset = 0
+    resume_epoch_offset = 0
+    resume_optimizer_state: dict[str, Any] | None = None
+    resume_best_val_metric = float("inf")
+    resume_from = cfg["train"].get("resume_from")
+    if resume_from and Path(resume_from).suffix != ".ckpt":
+        compat = torch.load(str(resume_from), map_location="cpu", weights_only=False)
+        resume_step_offset = int(compat.get("step", 0))
+        resume_epoch_offset = int(compat.get("epoch", 0))
+        resume_optimizer_state = compat.get("optimizer")
+        resume_best_val_metric = float(
+            compat.get("best_val_metric", compat.get("best_val_recon", float("inf")))
+        )
+
+    snap = out_dir / "config.snapshot.yaml"
+    if not snap.exists():
+        with snap.open("w") as f:
+            yaml.safe_dump(cfg, f, sort_keys=False)
+
+    # Build data module early so SampleGridCallback can reference val_ds via trainer.datamodule
+    datamodule = TactileVAEDataModule(cfg)
+    datamodule.setup()
+    print(f"datasets: train={len(datamodule.train_ds):,}  val={len(datamodule.val_ds):,}")
+
+    tcfg = cfg["train"]
+    steps_per_epoch = len(datamodule.train_dataloader())
+    total_steps = tcfg["max_steps"] if tcfg.get("max_steps") else steps_per_epoch * tcfg["epochs"]
+    print(f"steps/epoch={steps_per_epoch:,}  total_steps={total_steps:,}")
+    module = TactileVAEModule(cfg, step_offset=resume_step_offset, total_steps=total_steps)
+    n_params = sum(p.numel() for p in module.model.parameters())
+    print(f"model: {module.model.__class__.__name__}  params={n_params:,}")
+
+    # Precision
+    use_amp = bool(tcfg.get("amp", False))
+    if use_amp:
+        amp_dtype = str(tcfg.get("amp_dtype", "bf16")).lower()
+        if amp_dtype not in {"bf16", "bfloat16"}:
+            print(f"[info] overriding train.amp_dtype={amp_dtype!r} to 'bf16' (enforced)")
+        precision = "bf16-mixed"
+    else:
+        precision = "32"
+
+    loggers = _init_loggers(cfg, out_dir)
+
+    callbacks = [
+        SetEpochCallback(epoch_offset=resume_epoch_offset),
+        SampleGridCallback(cfg, step_offset=resume_step_offset),
+        CompatCheckpointCallback(
+            cfg,
+            step_offset=resume_step_offset,
+            epoch_offset=resume_epoch_offset,
+            initial_best_val_metric=resume_best_val_metric,
+        ),
+        CompatResumeStateCallback(resume_optimizer_state),
+        ModelCheckpoint(
+            dirpath=str(out_dir / "checkpoints"),
+            filename="last",
+            save_last=True,
+            save_top_k=0,
+            every_n_train_steps=tcfg["ckpt_every_steps"],
+        ),
+    ]
+
+    trainer = _build_trainer(
+        cfg,
+        callbacks=callbacks,
+        loggers=loggers,
+        precision=precision,
+        resume_from=resume_from,
+        resume_step_offset=resume_step_offset,
+        total_steps=total_steps,
+    )
+
+    # Resume: .ckpt = full Lightning resume; .pt = model+optimizer+offsets compat resume.
+    ckpt_path: str | None = None
+    if resume_from:
+        rf = Path(resume_from)
+        if rf.suffix == ".ckpt":
+            ckpt_path = str(rf)
+            print(f"resuming (Lightning): {rf}")
+        else:
+            ckpt = torch.load(str(rf), map_location="cpu", weights_only=False)
+            module.model.load_state_dict(ckpt["state_dict"])
+            print(
+                f"resuming (compat): {rf}  "
+                f"step={resume_step_offset} epoch={resume_epoch_offset}"
+            )
+
+    trainer.fit(module, datamodule=datamodule, ckpt_path=ckpt_path)
+    print(f"done. global_step={trainer.global_step}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    with args.config.open() as f:
+        raw_cfg = yaml.safe_load(f)
+    if args.run_id:
+        raw_cfg["run_id"] = args.run_id
+    if args.output_dir:
+        raw_cfg["output_dir"] = args.output_dir
+
+    tmp = args.config.parent / f".__cli_override_{os.getpid()}.yaml"
+    try:
+        with tmp.open("w") as f:
+            yaml.safe_dump(raw_cfg, f, sort_keys=False)
+        cfg = load_config(tmp)
+    finally:
+        tmp.unlink(missing_ok=True)
+
+    if args.resume_from:
+        cfg["train"]["resume_from"] = str(_resolve_path(args.resume_from))
+    _maybe_autoresume(cfg, allow_autoresume=not args.no_resume)
+
+    main(cfg)

script/train_vae_pl.sh

@@ -0,0 +1,173 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=tactile_vae_pl
+#SBATCH --partition=sharedp
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=2
+#SBATCH --gres=gpu:h100:4
+#SBATCH --requeue
+#SBATCH --exclude=mfmc10
+#SBATCH --output=/group2/ct/weihanx/tactile_world_model/slurm-logs/tactile_vae_pl.%j.log
+#SBATCH --error=/group2/ct/weihanx/tactile_world_model/slurm-logs/tactile_vae_pl.%j.log
+
+# Train TactileVAE with PyTorch Lightning (train_vae_pl.py).
+#
+# Each run lives at <RUNS_DIR>/<RUN_ID>/. Re-launching with the same RUN_ID
+# auto-resumes from checkpoints/last.ckpt (Lightning) or ckpt_last.pt (compat).
+# wandb keeps the same run id so metrics append to the same dashboard.
+#
+# Usage (sbatch):  sbatch tactile_vae/script/train_vae_pl.sh <run_id> [config.yaml]
+# Usage (local):   ./tactile_vae/script/train_vae_pl.sh        <run_id> [config.yaml]
+#
+# Diagnostics: set DEBUG=1 to enable `set -x` command tracing.
+#   sbatch --export=ALL,DEBUG=1 tactile_vae/script/train_vae_pl.sh <run_id>
+
+exec 1> >(stdbuf -oL -eL cat) 2>&1
+
+set -euo pipefail
+[[ "${DEBUG:-0}" == "1" ]] && set -x
+
+# ============================================================
+#  Inputs (positional)
+# ============================================================
+if [[ $# -lt 1 ]]; then
+  echo "Usage: $0 <run_id> [config.yaml]" >&2
+  echo "  run_id   : required. Both the output subdir name and the wandb run id." >&2
+  echo "  config   : optional. Defaults to tactile_vae/config/train_vae.yaml." >&2
+  exit 2
+fi
+RUN_ID="$1"
+CONFIG="${2:-tactile_vae/config/train_vae.yaml}"
+
+# ============================================================
+#  Paths
+# ============================================================
+WORKDIR="/group2/ct/weihanx/tactile_world_model"
+RUNS_DIR="$WORKDIR/runs"
+RUN_DIR="$RUNS_DIR/$RUN_ID"
+DATA_DIR="$WORKDIR/tactile_vae/data"
+SPLITS_PATH="$WORKDIR/tactile_vae/dataset/splits_subset.json"
+
+CONDA_ENV="${CONDA_ENV:-twm}"
+
+mkdir -p "$WORKDIR/slurm-logs"
+mkdir -p "$RUNS_DIR"
+umask 027
+
+# ============================================================
+#  Print startup info
+# ============================================================
+echo "=== Tactile VAE (PyTorch Lightning) ==="
+echo "Host:       $(hostname)"
+echo "Job ID:     ${SLURM_JOB_ID:-N/A}"
+echo "Start time: $(date)"
+echo "Run ID:     $RUN_ID"
+echo "Workdir:    $WORKDIR"
+echo "Config:     $CONFIG"
+echo "Run dir:    $RUN_DIR"
+echo "Conda env:  $CONDA_ENV"
+echo
+
+# ============================================================
+#  Environment knobs
+# ============================================================
+export OMP_NUM_THREADS=8
+export MKL_NUM_THREADS=8
+export TOKENIZERS_PARALLELISM="false"
+export PYTHONFAULTHANDLER=1
+export PYTHONUNBUFFERED=1
+
+# ============================================================
+#  Weights & Biases
+# ============================================================
+export WANDB_API_KEY="76cdc4261bf436617e661171fd41d80403e69e9b"
+export WANDB_ENTITY="weihanx-university-of-michigan"
+export WANDB_USERNAME="weihanx@umich.edu"
+export WANDB_PROJECT="tactile_vae"
+export WANDB_MODE="online"
+export WANDB_RUN_ID="$RUN_ID"
+export WANDB_NAME="$RUN_ID"
+export WANDB_SERVICE_WAIT=300
+export WANDB_INIT_TIMEOUT=300
+export WANDB_START_METHOD="thread"
+export WANDB_CONSOLE="wrap"
+export WANDB_DIR="${WANDB_DIR:-/tmp/$USER/wandb/$RUN_ID}"
+export WANDB_CACHE_DIR="${WANDB_CACHE_DIR:-/tmp/$USER/wandb-cache}"
+export WANDB_DATA_DIR="${WANDB_DATA_DIR:-/tmp/$USER/wandb-data}"
+mkdir -p "$WANDB_DIR" "$WANDB_CACHE_DIR" "$WANDB_DATA_DIR"
+
+DISABLE_WANDB="${DISABLE_WANDB:-0}"
+if [[ "$DISABLE_WANDB" == "1" ]]; then
+  unset WANDB_PROJECT WANDB_ENTITY WANDB_API_KEY WANDB_USERNAME
+  export WANDB_MODE="disabled"
+fi
+
+echo "--- Wandb ---"
+echo "  project=${WANDB_PROJECT:-<disabled>}  entity=${WANDB_ENTITY:-<disabled>}"
+echo "  run_id=$WANDB_RUN_ID  name=$WANDB_NAME  mode=$WANDB_MODE"
+echo "  dir=$WANDB_DIR"
+echo "  cache_dir=$WANDB_CACHE_DIR"
+echo "  data_dir=$WANDB_DATA_DIR"
+if [[ -n "${WANDB_API_KEY:-}" ]]; then
+  echo "  api_key=${WANDB_API_KEY:0:10}...${WANDB_API_KEY: -4}"
+else
+  echo "  api_key=<disabled>"
+fi
+echo
+
+# ============================================================
+#  Sanity checks
+# ============================================================
+if [[ ! -f "$WORKDIR/$CONFIG" ]] && [[ ! -f "$CONFIG" ]]; then
+  echo "ERROR: config not found: $CONFIG (or $WORKDIR/$CONFIG)" >&2
+  exit 2
+fi
+if [[ ! -d "$DATA_DIR" ]]; then
+  echo "ERROR: data dir does not exist: $DATA_DIR" >&2
+  exit 2
+fi
+if [[ ! -f "$SPLITS_PATH" ]]; then
+  echo "ERROR: splits manifest not found: $SPLITS_PATH" >&2
+  echo "  Generate it with: python tactile_vae/dataset/make_splits.py" >&2
+  exit 2
+fi
+
+# Report resume state (Lightning ckpt takes priority over compat ckpt)
+if [[ -f "$RUN_DIR/checkpoints/last.ckpt" ]]; then
+  echo "Resume:     auto-resume from $RUN_DIR/checkpoints/last.ckpt (Lightning)"
+elif [[ -f "$RUN_DIR/ckpt_last.pt" ]]; then
+  echo "Resume:     auto-resume from $RUN_DIR/ckpt_last.pt (compat)"
+else
+  echo "Resume:     fresh run (no checkpoint found)"
+fi
+echo
+
+# ============================================================
+#  Resolve Python interpreter
+# ============================================================
+PYTHON_BIN="${PYTHON_BIN:-$HOME/miniconda3/envs/$CONDA_ENV/bin/python}"
+if [[ -x "$PYTHON_BIN" ]]; then
+  echo "[$(date +%H:%M:%S)] using env python directly: $PYTHON_BIN"
+else
+  echo "[$(date +%H:%M:%S)] env python not found; falling back to conda activate..."
+  source ~/miniconda3/etc/profile.d/conda.sh
+  echo "[$(date +%H:%M:%S)] activating $CONDA_ENV..."
+  conda activate "$CONDA_ENV"
+  PYTHON_BIN="$(which python)"
+  echo "[$(date +%H:%M:%S)] env activated. python = $PYTHON_BIN"
+fi
+echo "[$(date +%H:%M:%S)] GPU(s): ${CUDA_VISIBLE_DEVICES:-$(nvidia-smi -L 2>/dev/null | head -1 || echo none)}"
+echo
+
+# ============================================================
+#  Launch training
+# ============================================================
+cd "$WORKDIR"
+echo "[$(date +%H:%M:%S)] launching trainer (PyTorch Lightning)..."
+"$PYTHON_BIN" -u tactile_vae/script/train_vae_pl.py \
+    --config "$CONFIG" \
+    --run-id "$RUN_ID"
+
+echo
+echo "[$(date +%H:%M:%S)] Finished."
+echo "Run dir contents:"
+ls -lh "$RUN_DIR" 2>/dev/null || echo "  (empty)"

test/__init__.py

@@ -0,0 +1 @@
+"""Tests for tactile_vae."""

test/test_raw_parquet.py

@@ -0,0 +1,130 @@
+"""Sanity checks for the decoded raw-RGB parquet at
+`tactile_vae/data/train-00000-of-00008.raw.parquet`.
+
+We don't reload the JPEG source to compare pixel-for-pixel — we only check
+that the file is well-formed: row count matches the source shard, every
+image blob is `H*W*3` bytes, and the on-disk size is consistent with the
+sum of expected per-row sizes.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pyarrow.compute as pc
+import pyarrow.parquet as pq
+import pytest
+
+DATA_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/data")
+JPEG_PATH = DATA_DIR / "train-00000-of-00008.parquet"
+RAW_PATH = DATA_DIR / "train-00000-of-00008.raw.parquet"
+
+EXPECTED_ROWS = 72859           # from the JPEG source shard
+EXPECTED_FORMAT = "raw_rgb_uint8"
+EXPECTED_CHANNELS = 3
+EXPECTED_HW = (480, 640)        # all rows in this shard share these dims
+
+# Loose tolerance for total file size vs. payload size (parquet overhead).
+SIZE_TOLERANCE_BYTES = 64 * 1024 * 1024     # 64 MB header/footer/metadata slack
+
+
+@pytest.fixture(scope="module")
+def pf() -> pq.ParquetFile:
+    if not RAW_PATH.exists():
+        pytest.skip(f"{RAW_PATH} not present")
+    return pq.ParquetFile(RAW_PATH)
+
+
+def test_row_count_matches_source(pf: pq.ParquetFile) -> None:
+    assert pf.metadata.num_rows == EXPECTED_ROWS
+    if JPEG_PATH.exists():
+        src_rows = pq.ParquetFile(JPEG_PATH).metadata.num_rows
+        assert pf.metadata.num_rows == src_rows, (
+            f"raw rows {pf.metadata.num_rows} != source rows {src_rows}"
+        )
+
+
+def test_schema_has_expected_columns(pf: pq.ParquetFile) -> None:
+    names = set(pf.schema_arrow.names)
+    for required in ("image", "image_format", "height", "width"):
+        assert required in names, f"missing column: {required}"
+
+
+def test_image_format_string_is_raw(pf: pq.ParquetFile) -> None:
+    fmts = pf.read(columns=["image_format"]).column("image_format").to_pylist()
+    unique = set(fmts)
+    assert unique == {EXPECTED_FORMAT}, f"unexpected image_format values: {unique}"
+
+
+def test_blob_lengths_match_height_width(pf: pq.ParquetFile) -> None:
+    """Every `image` blob must be exactly H*W*C bytes."""
+    tbl = pf.read(columns=["image", "height", "width"])
+    images = tbl.column("image")
+    heights = tbl.column("height").to_numpy()
+    widths = tbl.column("width").to_numpy()
+
+    # Length-per-row across chunks. Avoid combine_chunks(): a 67 GB binary
+    # column overflows the 32-bit offset limit of pa.binary().
+    lengths = pc.binary_length(images).to_numpy(zero_copy_only=False).astype(np.int64)
+
+    expected = heights.astype(np.int64) * widths.astype(np.int64) * EXPECTED_CHANNELS
+    mismatch = np.flatnonzero(lengths != expected)
+    assert mismatch.size == 0, (
+        f"{mismatch.size} rows have wrong blob size; "
+        f"first bad row {int(mismatch[0])}: got {int(lengths[mismatch[0]])} "
+        f"expected {int(expected[mismatch[0]])}"
+    )
+
+
+def test_first_row_decodes_to_valid_image(pf: pq.ParquetFile) -> None:
+    """Spot-check: row 0 reshapes to (H, W, 3) uint8, plausible value range."""
+    batch = next(pf.iter_batches(batch_size=1, columns=["image", "height", "width"]))
+    raw = batch["image"][0].as_py()
+    h = int(batch["height"][0].as_py())
+    w = int(batch["width"][0].as_py())
+
+    assert (h, w) == EXPECTED_HW, f"unexpected (h,w)=({h},{w})"
+    assert len(raw) == h * w * EXPECTED_CHANNELS
+
+    arr = np.frombuffer(raw, dtype=np.uint8).reshape(h, w, EXPECTED_CHANNELS)
+    assert arr.dtype == np.uint8
+    assert arr.min() >= 0 and arr.max() <= 255
+    # A real tactile image has > 1 unique value — guard against accidental zero-fill.
+    assert int(arr.std()) > 0
+
+
+def test_total_loaded_payload_bytes(pf: pq.ParquetFile) -> None:
+    """Sum of (h*w*3) over all rows should approximate the on-disk file size,
+    within a small slack for parquet headers/footers/metadata."""
+    tbl = pf.read(columns=["height", "width"])
+    heights = tbl.column("height").to_numpy().astype(np.int64)
+    widths = tbl.column("width").to_numpy().astype(np.int64)
+    payload_bytes = int((heights * widths * EXPECTED_CHANNELS).sum())
+
+    on_disk = RAW_PATH.stat().st_size
+    overhead = on_disk - payload_bytes
+    assert payload_bytes == EXPECTED_ROWS * EXPECTED_HW[0] * EXPECTED_HW[1] * EXPECTED_CHANNELS
+    assert 0 < overhead < SIZE_TOLERANCE_BYTES, (
+        f"unexpected file overhead: payload={payload_bytes:,}B "
+        f"on_disk={on_disk:,}B overhead={overhead:,}B "
+        f"(tolerance {SIZE_TOLERANCE_BYTES:,}B)"
+    )
+
+    print(
+        f"\n  raw parquet OK: {EXPECTED_ROWS:,} rows × "
+        f"{EXPECTED_HW[0]}×{EXPECTED_HW[1]}×{EXPECTED_CHANNELS}"
+        f" = {payload_bytes/1e9:.2f} GB payload, "
+        f"{on_disk/1e9:.2f} GB on disk "
+        f"({overhead/1e6:.1f} MB overhead)"
+    )
+
+
+if __name__ == "__main__":
+    pf_ = pq.ParquetFile(RAW_PATH)
+    test_row_count_matches_source(pf_)
+    test_schema_has_expected_columns(pf_)
+    test_image_format_string_is_raw(pf_)
+    test_blob_lengths_match_height_width(pf_)
+    test_first_row_decodes_to_valid_image(pf_)
+    test_total_loaded_payload_bytes(pf_)
+    print("All raw-parquet tests passed.")

test/test_tactile_vae.py

@@ -0,0 +1,126 @@
+"""Smoke/contract tests for tactile_vae.model.tactile_vae."""
+
+from pathlib import Path
+import sys
+import tempfile
+
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tactile_vae.model.tactile_vae import TactileVAE, VAELoss, load_pretrained
+
+
+def _small_model() -> TactileVAE:
+    return TactileVAE(
+        img_size=64,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=96,
+        encoder_depth=2,
+        encoder_heads=4,
+        decoder_embed_dim=96,
+        decoder_depth=2,
+        decoder_heads=4,
+        latent_dim=48,
+    )
+
+
+def test_forward_contract_shapes():
+    torch.manual_seed(0)
+    model = _small_model()
+    x = torch.randn(2, 3, 64, 64)
+
+    out = model(x)
+    assert out["x_hat"].shape == x.shape
+    assert out["mu"].shape == (2, 48)
+    assert out["logvar"].shape == (2, 48)
+    assert out["z"].shape == (2, 48)
+    assert out["pred_patches"].shape == (2, (64 // 16) ** 2, 16 * 16 * 3)
+
+
+def test_losses_finite_and_backprop():
+    torch.manual_seed(0)
+    model = _small_model()
+    criterion = VAELoss(beta=1e-3, recon_type="l1", ssim_weight=0.1)
+    x = torch.randn(2, 3, 64, 64)
+
+    out = model(x)
+    losses = criterion(out["x_hat"], x, out["mu"], out["logvar"])
+    for name, val in losses.items():
+        assert torch.isfinite(val).all(), f"non-finite loss component: {name}"
+
+    losses["total"].backward()
+    grad_ok = any(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters() if p.requires_grad)
+    assert grad_ok, "expected finite gradients after backward"
+
+
+def test_tiny_overfit_recon_drops():
+    torch.manual_seed(0)
+    model = _small_model()
+    criterion = VAELoss(beta=1e-3, recon_type="l1", ssim_weight=0.0)
+    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    x = torch.tanh(torch.randn(8, 3, 64, 64))
+
+    with torch.no_grad():
+        out0 = model(x)
+        start = criterion(out0["x_hat"], x, out0["mu"], out0["logvar"])["recon_total"].item()
+
+    for _ in range(30):
+        out = model(x)
+        losses = criterion(out["x_hat"], x, out["mu"], out["logvar"])
+        opt.zero_grad(set_to_none=True)
+        losses["total"].backward()
+        opt.step()
+
+    with torch.no_grad():
+        out1 = model(x)
+        end = criterion(out1["x_hat"], x, out1["mu"], out1["logvar"])["recon_total"].item()
+
+    assert end < start, f"expected recon to decrease; start={start:.4f}, end={end:.4f}"
+
+
+def test_checkpoint_roundtrip_consistent_eval():
+    torch.manual_seed(0)
+    model = _small_model().eval()
+    x = torch.randn(2, 3, 64, 64)
+
+    with torch.no_grad():
+        out_ref = model(x, sample=False)["x_hat"]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        ckpt_path = Path(tmpdir) / "tactile_vae.pt"
+        torch.save(model.state_dict(), ckpt_path)
+        loaded = load_pretrained(
+            checkpoint=str(ckpt_path),
+            model_kwargs={
+                "img_size": 64,
+                "patch_size": 16,
+                "in_chans": 3,
+                "embed_dim": 96,
+                "encoder_depth": 2,
+                "encoder_heads": 4,
+                "decoder_embed_dim": 96,
+                "decoder_depth": 2,
+                "decoder_heads": 4,
+                "latent_dim": 48,
+            },
+            freeze=True,
+            strict=True,
+        ).eval()
+
+        with torch.no_grad():
+            out_loaded = loaded(x, sample=False)["x_hat"]
+
+    torch.testing.assert_close(out_ref, out_loaded, atol=1e-6, rtol=1e-5)
+
+
+if __name__ == "__main__":
+    test_forward_contract_shapes()
+    test_losses_finite_and_backprop()
+    test_tiny_overfit_recon_drops()
+    test_checkpoint_roundtrip_consistent_eval()
+    print("All tactile VAE tests passed.")

test/visualize_color_jitter.py

@@ -0,0 +1,85 @@
+"""Visual sanity check: 10 samples with vs. without color jitter.
+
+Writes `color_jitter_compare.png` next to this script. Layout:
+  row 0 — base image (no augmentation)
+  row 1 — color-jittered image (training-time augmentation)
+  row 2 — |jittered − base| amplified for visibility
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tactile_vae.dataset import ColorJitterConfig, TactileParquetDataset
+
+N_SAMPLES = 10
+SEED = 0
+JITTER = ColorJitterConfig(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.05)
+OUT_PNG = Path(__file__).with_name("color_jitter_compare.png")
+
+
+def _to_hwc(t: torch.Tensor) -> np.ndarray:
+    return t.detach().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
+
+
+def main() -> None:
+    ds_plain = TactileParquetDataset(image_size=128, color_jitter=None)
+    ds_jit = TactileParquetDataset(image_size=128, color_jitter=JITTER)
+
+    # Sample 10 distinct indices across the dataset for variety.
+    n = len(ds_plain)
+    rng = np.random.default_rng(SEED)
+    indices = sorted(rng.choice(n, size=N_SAMPLES, replace=False).tolist())
+    print(f"Comparing {N_SAMPLES} samples (indices: {indices[:3]}... of {n:,})")
+    print(f"Jitter config: {JITTER}")
+
+    torch.manual_seed(SEED)  # fix jitter RNG so the PNG is reproducible
+
+    fig, axes = plt.subplots(3, N_SAMPLES, figsize=(2 * N_SAMPLES, 6.6))
+
+    for col, idx in enumerate(indices):
+        base = ds_plain[idx]                 # (3,128,128) float in [0,1]
+        jittered = ds_jit[idx]
+        base_hwc = _to_hwc(base)
+        jit_hwc = _to_hwc(jittered)
+
+        # Amplify residual so visible per-pixel changes are visible at a glance.
+        diff = np.abs(jit_hwc - base_hwc)
+        scale = max(diff.max(), 1e-6)
+        diff_vis = np.clip(diff / scale, 0, 1)
+
+        axes[0, col].imshow(base_hwc)
+        axes[0, col].set_title(f"idx={idx}", fontsize=8)
+        axes[1, col].imshow(jit_hwc)
+        axes[1, col].set_title(f"Δmean={float(jit_hwc.mean()-base_hwc.mean()):+.3f}",
+                               fontsize=8)
+        axes[2, col].imshow(diff_vis)
+        axes[2, col].set_title(f"|Δ|max={float(diff.max()):.3f}", fontsize=8)
+        for r in range(3):
+            axes[r, col].set_xticks([])
+            axes[r, col].set_yticks([])
+
+    axes[0, 0].set_ylabel("no jitter", fontsize=10)
+    axes[1, 0].set_ylabel("with jitter", fontsize=10)
+    axes[2, 0].set_ylabel("|Δ| (norm)", fontsize=10)
+
+    fig.suptitle(
+        f"Color jitter comparison (b={JITTER.brightness}, c={JITTER.contrast}, "
+        f"s={JITTER.saturation}, h={JITTER.hue})",
+        fontsize=11,
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.96))
+    fig.savefig(OUT_PNG, dpi=120)
+    print(f"wrote {OUT_PNG} ({OUT_PNG.stat().st_size/1024:.1f} KB)")
+
+
+if __name__ == "__main__":
+    main()