Upload tinyvla2/src/tinyvla/data/wds_shards.py with huggingface_hub

tinyvla2/src/tinyvla/data/wds_shards.py

@@ -0,0 +1,282 @@
+"""Map-style source over prebuilt WebDataset tar shards (AlexWortega/wds-*).
+
+Layout produced by the packer, uniform across all six repos:
+
+    <root>/manifest.json          norm_stats, embodiment_id, chunk, dims, shard list
+    <root>/tasks.json             {task_index: instruction string}
+    <root>/shard-XXXXX.tar        <name>_<ep>_<frame>.cam0.jpg
+                                  <name>_<ep>_<frame>.cam1.jpg
+                                  <name>_<ep>_<frame>.meta.npz
+
+`meta.npz` carries `state`, `action_chunk` (chunk, action_dim), `action_is_pad`,
+`action_dim`, `state_dim`, `task_index`, `embodiment_id`.
+
+Why map-style and not IterableDataset: the training loop mixes these with
+`CanonicalSource` inside `WeightedMixtureDataset`, which is map-style and seeds
+sampling per index. tar members have fixed offsets, so once the offsets are
+indexed a sample is a plain seek+read — no need to rewrite the loader as a
+stream, and matched-pairs sampling across variants keeps working.
+
+The index (one pass over the tar headers, no payload read) is cached next to the
+shards, because building it for the 635-shard unitree pack takes minutes.
+
+Emits exactly the schema CanonicalSource emits, so it drops into the existing
+mixture unchanged. Actions here are NATIVE (unitree joint positions 7-28d, nav
+waypoints/step-deltas 3d), not the 7-dim canonical EE space — the per-sample
+`action_dim_mask` is what lets one head serve both.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import pickle
+import tarfile
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+INDEX_VERSION = 1
+
+
+def _pad_strict(x: torch.Tensor, dim: int, what: str) -> torch.Tensor:
+    """Zero-pad the last dim, REFUSING to truncate.
+
+    Silently cutting a 28-dim G1 action down to 8 is exactly the class of
+    convention bug this project keeps paying for, so it raises instead.
+    """
+    d = x.shape[-1]
+    if d > dim:
+        raise ValueError(
+            f"{what} has {d} dims but max is {dim} — raise max_{what}_dim in the "
+            f"policy config instead of truncating"
+        )
+    return F.pad(x, (0, dim - d)) if d < dim else x
+
+
+class WdsShardSource(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        root: str | Path,
+        embodiment_id: int | None = None,
+        chunk: int = 50,
+        image_size: int = 256,
+        max_state_dim: int = 32,
+        max_action_dim: int = 32,
+        morphology: torch.Tensor | None = None,
+        robot_prompt: str | None = None,
+        normalize: bool = True,
+        emit_latent_image: bool = False,
+    ):
+        self.root = Path(root).expanduser()
+        self.chunk = chunk
+        self.image_size = image_size
+        self.max_state_dim = max_state_dim
+        self.max_action_dim = max_action_dim
+        self.morphology = morphology
+        self.robot_prompt = robot_prompt
+        self.normalize = normalize
+        self.emit_latent_image = emit_latent_image
+
+        self.manifest = json.loads((self.root / "manifest.json").read_text())
+        self.name = self.manifest.get("name", self.root.name)
+        self.action_space = self.manifest.get("action_space", "native")
+        # config may override the packed id (three nav repos ship the same id 16
+        # for three different robots)
+        self.embodiment_id = int(
+            embodiment_id if embodiment_id is not None else self.manifest["embodiment_id"]
+        )
+
+        tasks_path = self.root / "tasks.json"
+        self._tasks = (
+            {int(k): v for k, v in json.loads(tasks_path.read_text()).items()}
+            if tasks_path.exists()
+            else {}
+        )
+
+        ns = self.manifest["norm_stats"]
+        self._s_mu, self._s_sd = self._stat(ns.get("observation.state"))
+        self._a_mu, self._a_sd = self._stat(ns.get("action"))
+
+        self._index = self._load_or_build_index()
+
+        # The packs ship mean/std only, which is unusable here: go-stanford's
+        # second action channel has std 0.008, so a single outlier normalizes to
+        # ~250 sigma (observed range [-121, 243]). Same reason canonical actions
+        # use quantiles. Prefer cached q01/q99 -> [-1, 1]; fall back to mean/std
+        # with a hard clip so a missing stats file degrades instead of exploding.
+        qs = self.root / "_qstats.json"
+        self.q_stats = json.loads(qs.read_text()) if qs.exists() else None
+        if self.q_stats is not None:
+            self._a_span, self._a_mid = self._span_mid(self.q_stats["action"])
+            self._s_span, self._s_mid = self._span_mid(self.q_stats["state"])
+        self._fh: dict[str, object] = {}  # per-worker file handles, opened lazily
+
+    # ------------------------------------------------------------------ setup
+
+    @staticmethod
+    def _span_mid(q):
+        q01 = torch.tensor(q["q01"], dtype=torch.float32)
+        q99 = torch.tensor(q["q99"], dtype=torch.float32)
+        floor = 0.01 * torch.median(torch.abs(torch.cat([q01, q99])) + 1e-6)
+        return torch.clamp(q99 - q01, min=float(floor) if float(floor) > 0 else 1e-3), 0.5 * (q01 + q99)
+
+    def compute_quantile_stats(self, n_samples: int = 40000, seed: int = 0) -> dict:
+        """q01/q99 per dim over sampled meta.npz, written to _qstats.json."""
+        rng = np.random.default_rng(seed)
+        n = len(self)
+        pick = rng.choice(n, size=min(n_samples, n), replace=False)
+        acts, states = [], []
+        for i in pick:
+            si, _, _, _, _, mo, ms = (int(v) for v in self._index["entries"][i])
+            meta = np.load(io.BytesIO(self._read(si, mo, ms)), allow_pickle=False)
+            acts.append(meta["action_chunk"].astype(np.float32))
+            states.append(meta["state"].astype(np.float32))
+        a = np.concatenate(acts, 0)
+        st = np.stack(states, 0)
+        out = {
+            "action": {"q01": np.quantile(a, 0.01, 0).tolist(), "q99": np.quantile(a, 0.99, 0).tolist()},
+            "state": {"q01": np.quantile(st, 0.01, 0).tolist(), "q99": np.quantile(st, 0.99, 0).tolist()},
+            "n": int(len(pick)),
+        }
+        (self.root / "_qstats.json").write_text(json.dumps(out))
+        return out
+
+    @staticmethod
+    def _stat(s):
+        if not s:
+            return None, None
+        mu = torch.tensor(s["mean"], dtype=torch.float32)
+        sd = torch.tensor(s["std"], dtype=torch.float32)
+        # scand/tartandrive ship an all-zero state with zero std; clamping keeps
+        # the division finite and the (already zero) channel stays zero
+        return mu, sd.clamp(min=1e-6)
+
+    def _load_or_build_index(self):
+        cache = self.root / f"_offsets_v{INDEX_VERSION}.pkl"
+        if cache.exists():
+            return pickle.loads(cache.read_bytes())
+
+        shards = sorted(self.root.glob("shard-*.tar"))
+        if not shards:
+            raise FileNotFoundError(f"no shard-*.tar under {self.root}")
+
+        entries = []  # (shard_idx, cam0_off, cam0_sz, cam1_off, cam1_sz, meta_off, meta_sz)
+        for si, sp in enumerate(shards):
+            by_key: dict[str, dict] = {}
+            with tarfile.open(sp) as tf:
+                for m in tf:  # streaming header scan, payload never read
+                    if not m.isfile():
+                        continue
+                    key, _, ext = m.name.partition(".")
+                    by_key.setdefault(key, {})[ext] = (m.offset_data, m.size)
+            for key in sorted(by_key):
+                parts = by_key[key]
+                if "cam0.jpg" not in parts or "meta.npz" not in parts:
+                    continue  # incomplete sample, drop it rather than guess
+                c0 = parts["cam0.jpg"]
+                c1 = parts.get("cam1.jpg", (-1, 0))
+                mo = parts["meta.npz"]
+                entries.append((si, c0[0], c0[1], c1[0], c1[1], mo[0], mo[1]))
+
+        idx = {
+            "shards": [str(p) for p in shards],
+            "entries": np.asarray(entries, dtype=np.int64),
+        }
+        cache.write_bytes(pickle.dumps(idx))
+        return idx
+
+    # ------------------------------------------------------------------ access
+
+    def __len__(self):
+        return len(self._index["entries"])
+
+    def _read(self, shard_idx: int, off: int, size: int) -> bytes:
+        path = self._index["shards"][shard_idx]
+        fh = self._fh.get(path)
+        if fh is None:
+            fh = open(path, "rb")  # noqa: SIM115 — kept open per worker on purpose
+            self._fh[path] = fh
+        fh.seek(off)
+        return fh.read(size)
+
+    def _img(self, blob: bytes) -> torch.Tensor:
+        from PIL import Image
+
+        img = Image.open(io.BytesIO(blob)).convert("RGB")
+        x = torch.from_numpy(np.asarray(img, dtype=np.uint8).copy()).permute(2, 0, 1).float() / 255.0
+        if x.shape[-2:] != (self.image_size, self.image_size):
+            x = F.interpolate(x[None], size=(self.image_size, self.image_size),
+                              mode="bilinear", align_corners=False)[0]
+        return x
+
+    def __getitem__(self, i: int) -> dict:
+        si, c0o, c0s, c1o, c1s, mo, ms = (int(v) for v in self._index["entries"][i])
+
+        cam0 = self._img(self._read(si, c0o, c0s))
+        # single-camera packs (all five nav sets) still ship a cam1.jpg member,
+        # but it is a ZERO-BYTE placeholder — presence is not enough, check size
+        if c1o >= 0 and c1s > 0:
+            cam1 = self._img(self._read(si, c1o, c1s))
+            cam_mask = torch.tensor([True, True])
+        else:
+            cam1 = torch.zeros_like(cam0)
+            cam_mask = torch.tensor([True, False])
+
+        meta = np.load(io.BytesIO(self._read(si, mo, ms)), allow_pickle=False)
+        state = torch.from_numpy(meta["state"].astype(np.float32))
+        action = torch.from_numpy(meta["action_chunk"].astype(np.float32))
+        is_pad = torch.from_numpy(meta["action_is_pad"].astype(bool))
+
+        if self.normalize:
+            if self.q_stats is not None:
+                state = (2.0 * (state - self._s_mid) / self._s_span).clamp(-3, 3)
+                action = (2.0 * (action - self._a_mid) / self._a_span).clamp(-3, 3)
+            else:
+                if self._s_mu is not None and len(state):
+                    n = min(len(state), len(self._s_mu))
+                    state = torch.cat([(state[:n] - self._s_mu[:n]) / self._s_sd[:n], state[n:]])
+                    state = state.clamp(-10, 10)
+                if self._a_mu is not None:
+                    n = min(action.shape[-1], len(self._a_mu))
+                    action = torch.cat(
+                        [(action[..., :n] - self._a_mu[:n]) / self._a_sd[:n], action[..., n:]], dim=-1
+                    ).clamp(-10, 10)
+
+        a_dim = action.shape[-1]
+        mask = torch.zeros(self.max_action_dim, dtype=torch.bool)
+        mask[: min(a_dim, self.max_action_dim)] = True
+
+        if action.shape[0] < self.chunk:
+            pad = self.chunk - action.shape[0]
+            action = torch.cat([action, action[-1:].expand(pad, -1)], 0)
+            is_pad = torch.cat([is_pad, torch.ones(pad, dtype=torch.bool)])
+        action, is_pad = action[: self.chunk], is_pad[: self.chunk]
+
+        task = self._tasks.get(int(meta["task_index"]), "")
+        if self.robot_prompt:
+            task = f"{self.robot_prompt} {task}".strip()
+
+        out = {
+            "observation.images.cam0": cam0,
+            "observation.images.cam1": cam1,
+            "camera_mask": cam_mask,
+            "observation.state": _pad_strict(state, self.max_state_dim, "state"),
+            "action": _pad_strict(action, self.max_action_dim, "action"),
+            "action_dim_mask": mask,
+            "action_is_pad": is_pad,
+            "embodiment_id": torch.tensor(self.embodiment_id, dtype=torch.long),
+            "task": task,
+        }
+        if self.morphology is not None:
+            out["morphology"] = self.morphology
+        # Batches mix wds packs with CanonicalSource, so every optional key the
+        # other source emits must be emitted here too or default_collate raises
+        # on whichever sample lacks it. CanonicalSource always emits
+        # spatial_cell (-1 when unlabelled) for exactly this reason.
+        out["spatial_cell"] = torch.tensor(-1, dtype=torch.long)
+        if self.emit_latent_image:
+            out["latent_image"] = cam0  # no episode history here: no staleness aug
+        return out