wds mixture: unitree+nav converters (gnm/vamos), wds source branch in train_fast, AMP autoselect (V100 fp16), embedding expand 16->32, push/fetch/validate tooling, DATA.md
b152c62 verified | #!/usr/bin/env python | |
| """One-time shard builder: LeRobot v3 video dataset -> sequential JPEG shards. | |
| Why this exists | |
| --------------- | |
| Training reads ONE random frame per sample. On a LeRobot v3 repo that means an | |
| h264 seek into a ~195 MB mp4 holding ~10 concatenated episodes, per camera, per | |
| sample. That is the bottleneck the project notes already call out ("video decode | |
| is the training bottleneck, use 12+ dataloader workers"), and no amount of GPU | |
| fixes it. | |
| This pass decodes every video ONCE, sequentially (the cheap direction), at a | |
| target rate, resizes to the training resolution, and writes JPEG bytes into | |
| parquet shards. Training then does sequential reads + JPEG decode, which is | |
| ~100x cheaper per frame and trivially parallel. | |
| Layout produced in --out: | |
| shards/shard-{i:05d}.parquet episode_index, frame_index, cam0, cam1 (JPEG bytes) | |
| frames.parquet episode_index, frame_index, state[D], action[D], task_index | |
| tasks.json task_index -> instruction string | |
| stats.json per-source mean/std for state and action | |
| manifest.json fps, image_size, dims, camera keys, counts | |
| Usage (parallel across boxes: each takes a disjoint episode range): | |
| python build_shards.py --repo-id zekaiwang/trex_dataset --out /data/trex_shards \ | |
| --hz 10 --image-size 256 --part 0 --num-parts 4 | |
| Verified against zekaiwang/trex_dataset metadata: episodes are contiguous and | |
| ordered inside every mp4, frame counts match the timestamps exactly, and the | |
| declared fps (30) is the real one. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import io | |
| import json | |
| import os | |
| from pathlib import Path | |
| import numpy as np | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| # --------------------------------------------------------------------- metadata | |
| def load_metadata(repo_id: str, revision: str | None): | |
| from huggingface_hub import hf_hub_download | |
| def grab(path): | |
| return hf_hub_download(repo_id, path, repo_type="dataset", revision=revision) | |
| info = json.loads(Path(grab("meta/info.json")).read_text()) | |
| from huggingface_hub import HfApi | |
| api = HfApi() | |
| files = api.list_repo_files(repo_id, repo_type="dataset", revision=revision) | |
| ep_files = sorted(f for f in files if f.startswith("meta/episodes/") and f.endswith(".parquet")) | |
| tables = [pq.read_table(grab(f)) for f in ep_files] | |
| episodes = pa.concat_tables(tables) if len(tables) > 1 else tables[0] | |
| return info, episodes, files | |
| def pick_cameras(info: dict, cam0: str | None, cam1: str | None): | |
| """Mirror CanonicalSource's cam0/cam1 choice unless the caller overrides it. | |
| CanonicalSource sorts the image keys, takes PRIMARY_RE (or the first key) as | |
| cam0 and the first WRIST_RE match as cam1. On trex that resolves to | |
| head_left / left_wrist — right_wrist and all 20 tactile streams are dropped. | |
| Made explicit here rather than left to a regex. | |
| """ | |
| import re | |
| PRIMARY_RE = re.compile(r"top|front|base|main|scene|image_0$|images\.image$|image2$") | |
| WRIST_RE = re.compile(r"wrist|gripper|hand|image_1$") | |
| keys = sorted(k for k, v in info["features"].items() if v.get("dtype") == "video") | |
| if cam0 is None: | |
| cam0 = next((k for k in keys if PRIMARY_RE.search(k)), keys[0]) | |
| if cam1 is None: | |
| cam1 = next((k for k in keys if k != cam0 and WRIST_RE.search(k)), None) | |
| assert cam0 in keys, f"{cam0} not among {keys}" | |
| assert cam1 is None or cam1 in keys, f"{cam1} not among {keys}" | |
| return cam0, cam1 | |
| # ------------------------------------------------------------------ video decode | |
| class VideoReader: | |
| """Opens one mp4 at a time, seeks to an episode, decodes its frames forward. | |
| Episodes inside a file are contiguous and time-ordered (verified), so walking | |
| episodes in order keeps this effectively sequential: one seek per episode, | |
| costing at most one extra GOP. | |
| """ | |
| def __init__(self, repo_id, revision, key, cache_dir, evict: bool): | |
| self.repo_id, self.revision, self.key = repo_id, revision, key | |
| self.cache_dir, self.evict = cache_dir, evict | |
| self._path = None | |
| self._ident = None | |
| self._container = None | |
| self._stream = None | |
| def _open(self, chunk_index: int, file_index: int): | |
| ident = (chunk_index, file_index) | |
| if ident == self._ident: | |
| return | |
| self._close(drop=True) | |
| from huggingface_hub import hf_hub_download | |
| rel = f"videos/{self.key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4" | |
| self._path = hf_hub_download( | |
| self.repo_id, rel, repo_type="dataset", revision=self.revision, cache_dir=self.cache_dir | |
| ) | |
| import av | |
| self._container = av.open(self._path) | |
| self._stream = self._container.streams.video[0] | |
| self._stream.thread_type = "AUTO" # frame+slice threading; the whole point | |
| self._ident = ident | |
| def _close(self, drop: bool): | |
| if self._container is not None: | |
| self._container.close() | |
| self._container = None | |
| if drop and self.evict and self._path and os.path.exists(self._path): | |
| real = os.path.realpath(self._path) # hf cache stores a symlink to the blob | |
| for p in {self._path, real}: | |
| try: | |
| os.remove(p) | |
| except OSError: | |
| pass | |
| self._path = self._ident = self._stream = None | |
| def close(self): | |
| self._close(drop=True) | |
| def read_episode(self, chunk_index, file_index, from_ts, n_frames, fps, stride): | |
| """Yield (frame_index_in_episode, PIL.Image) for every stride-th frame.""" | |
| self._open(chunk_index, file_index) | |
| tb = self._stream.time_base | |
| target = int(from_ts / tb) | |
| self._container.seek(target, stream=self._stream, backward=True, any_frame=False) | |
| wanted = set(range(0, n_frames, stride)) | |
| eps = 0.5 / fps # half a frame of slack against float timestamps | |
| got = 0 | |
| for frame in self._container.decode(self._stream): | |
| if frame.pts is None: | |
| continue | |
| t = float(frame.pts * tb) | |
| if t < from_ts - eps: | |
| continue | |
| idx = int(round((t - from_ts) * fps)) | |
| if idx >= n_frames: | |
| break | |
| if idx in wanted: | |
| wanted.discard(idx) | |
| got += 1 | |
| yield idx, frame.to_image() | |
| if not wanted: | |
| break | |
| if wanted: | |
| raise RuntimeError( | |
| f"{self.key} chunk{chunk_index}/file{file_index} @{from_ts:.3f}: " | |
| f"decoded {got}, still missing {len(wanted)} frames" | |
| ) | |
| def encode_jpeg(img, size: int, quality: int) -> bytes: | |
| """Resize to (size, size) and JPEG-encode. | |
| Plain resize, no crop: this reproduces CanonicalSource._img, which squashes | |
| the aspect ratio with F.interpolate. Matching it matters more than being | |
| geometrically nicer, otherwise the shards are off-distribution for the | |
| pretrained checkpoint. | |
| """ | |
| from PIL import Image | |
| if img.size != (size, size): | |
| img = img.resize((size, size), Image.BILINEAR) | |
| buf = io.BytesIO() | |
| img.save(buf, format="JPEG", quality=quality, subsampling=0) | |
| return buf.getvalue() | |
| # ------------------------------------------------------------------------ state | |
| def load_frame_table(repo_id, revision, files, episodes, keep_mask, stride, cache_dir, | |
| skip_bad=False): | |
| bad_rows: set[int] = set() | |
| """Pull state/action/task for the kept frames out of the data parquets.""" | |
| from huggingface_hub import hf_hub_download | |
| ep = {c: np.asarray(episodes[c]) for c in | |
| ["episode_index", "length", "dataset_from_index", "data/chunk_index", "data/file_index"]} | |
| want = {} # (chunk, file) -> list of episode rows | |
| for i in np.nonzero(keep_mask)[0]: | |
| want.setdefault((int(ep["data/chunk_index"][i]), int(ep["data/file_index"][i])), []).append(i) | |
| out_ep, out_fr, out_state, out_action, out_task = [], [], [], [], [] | |
| tasks: dict[str, int] = {} | |
| for (ci, fi), rows in sorted(want.items()): | |
| rel = f"data/chunk-{ci:03d}/file-{fi:03d}.parquet" | |
| if rel not in files: | |
| raise FileNotFoundError(rel) | |
| path = hf_hub_download(repo_id, rel, repo_type="dataset", revision=revision, cache_dir=cache_dir) | |
| t = pq.read_table(path, columns=["episode_index", "frame_index", "observation.state", "action", "task_index"]) | |
| d = {c: t[c].to_numpy(zero_copy_only=False) for c in t.column_names} | |
| for r in rows: | |
| e = int(ep["episode_index"][r]) | |
| sel = np.nonzero(d["episode_index"] == e)[0] | |
| if len(sel) != int(ep["length"][r]): | |
| if skip_bad: | |
| # unitree v3 repos ship a few episodes whose rows are | |
| # missing from the indexed file — drop them consistently | |
| # from BOTH the frame table and the video stage | |
| print(f" SKIP episode {e}: parquet has {len(sel)} rows, meta says {ep['length'][r]}") | |
| bad_rows.add(int(r)) | |
| continue | |
| raise RuntimeError(f"episode {e}: parquet has {len(sel)} rows, meta says {ep['length'][r]}") | |
| order = np.argsort(d["frame_index"][sel]) | |
| sel = sel[order][::stride] | |
| out_ep.append(np.full(len(sel), e, dtype=np.int32)) | |
| out_fr.append(np.arange(len(sel), dtype=np.int32)) | |
| out_state.append(np.stack([np.asarray(x, dtype=np.float32) for x in d["observation.state"][sel]])) | |
| out_action.append(np.stack([np.asarray(x, dtype=np.float32) for x in d["action"][sel]])) | |
| task = episodes["tasks"][r].as_py() | |
| task = task[0] if isinstance(task, list) and task else "" | |
| out_task.append(np.full(len(sel), tasks.setdefault(task, len(tasks)), dtype=np.int32)) | |
| del t, d | |
| return ( | |
| np.concatenate(out_ep), np.concatenate(out_fr), | |
| np.concatenate(out_state), np.concatenate(out_action), | |
| np.concatenate(out_task), tasks, bad_rows, | |
| ) | |
| # ------------------------------------------------------------------------- main | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--repo-id", default="zekaiwang/trex_dataset") | |
| ap.add_argument("--revision", default=None) | |
| ap.add_argument("--out", type=Path, required=True) | |
| ap.add_argument("--hz", type=float, default=10.0, help="target rate; 0 = keep native fps") | |
| ap.add_argument("--image-size", type=int, default=256) | |
| ap.add_argument("--jpeg-quality", type=int, default=92) | |
| ap.add_argument("--rows-per-shard", type=int, default=20000) | |
| ap.add_argument("--cam0", default=None, help="override the primary camera key") | |
| ap.add_argument("--cam1", default=None, help="override the wrist camera key") | |
| ap.add_argument("--part", type=int, default=0) | |
| ap.add_argument("--num-parts", type=int, default=1) | |
| ap.add_argument("--max-episodes", type=int, default=None, help="smoke test on the first N") | |
| ap.add_argument("--cache-dir", default=None) | |
| ap.add_argument("--keep-videos", action="store_true", help="do NOT delete mp4s after use") | |
| ap.add_argument("--skip-bad-episodes", action="store_true", | |
| help="drop episodes whose parquet rows disagree with meta instead of failing") | |
| args = ap.parse_args() | |
| info, episodes, files = load_metadata(args.repo_id, args.revision) | |
| fps = float(info["fps"]) | |
| stride = 1 if args.hz <= 0 else max(1, int(round(fps / args.hz))) | |
| eff_hz = fps / stride | |
| cam0, cam1 = pick_cameras(info, args.cam0, args.cam1) | |
| n_ep = episodes.num_rows | |
| lengths = np.asarray(episodes["length"]) | |
| # sanity: the LIBERO lesson — never trust declared fps | |
| for key in (cam0, cam1): | |
| if key is None: | |
| continue | |
| f = np.asarray(episodes[f"videos/{key}/from_timestamp"]) | |
| t = np.asarray(episodes[f"videos/{key}/to_timestamp"]) | |
| measured = lengths / (t - f) | |
| if not np.allclose(measured, fps, atol=1e-3): | |
| raise SystemExit( | |
| f"declared fps={fps} but {key} measures {measured.min():.3f}..{measured.max():.3f} — " | |
| "fix the rate before building shards" | |
| ) | |
| order = np.arange(n_ep) if args.max_episodes is None else np.arange(min(n_ep, args.max_episodes)) | |
| mine = order[args.part :: args.num_parts] | |
| keep = np.zeros(n_ep, dtype=bool) | |
| keep[mine] = True | |
| print(f"repo {args.repo_id}: {n_ep} episodes, {lengths.sum():,} frames @ {fps} Hz") | |
| print(f"cam0={cam0}\ncam1={cam1}") | |
| print(f"stride {stride} -> {eff_hz:g} Hz, part {args.part}/{args.num_parts}: {len(mine)} episodes") | |
| args.out.mkdir(parents=True, exist_ok=True) | |
| (args.out / "shards").mkdir(exist_ok=True) | |
| print("collecting state/action ...") | |
| ep_i, fr_i, state, action, task_i, tasks, bad_rows = load_frame_table( | |
| args.repo_id, args.revision, files, episodes, keep, stride, args.cache_dir, | |
| skip_bad=args.skip_bad_episodes, | |
| ) | |
| if bad_rows: | |
| mine = np.asarray([r for r in mine if int(r) not in bad_rows]) | |
| print(f" dropped {len(bad_rows)} bad episodes, {len(mine)} remain") | |
| print(f" {len(ep_i):,} frames kept, state {state.shape[1]}d, action {action.shape[1]}d, {len(tasks)} tasks") | |
| readers = { | |
| "cam0": VideoReader(args.repo_id, args.revision, cam0, args.cache_dir, not args.keep_videos), | |
| "cam1": None if cam1 is None else VideoReader(args.repo_id, args.revision, cam1, args.cache_dir, not args.keep_videos), | |
| } | |
| schema = pa.schema([ | |
| ("episode_index", pa.int32()), ("frame_index", pa.int32()), | |
| ("cam0", pa.binary()), ("cam1", pa.binary()), | |
| ]) | |
| buf = {k: [] for k in ["episode_index", "frame_index", "cam0", "cam1"]} | |
| shard_id, written, in_shard = args.part * 100000, 0, 0 | |
| writer = None | |
| def flush(final=False): | |
| nonlocal writer, shard_id, buf, written, in_shard | |
| if buf["episode_index"]: | |
| if writer is None: | |
| path = args.out / "shards" / f"shard-{shard_id:05d}.parquet" | |
| writer = pq.ParquetWriter(path, schema, compression="zstd", compression_level=1) | |
| writer.write_table(pa.table(buf, schema=schema)) | |
| written += len(buf["episode_index"]) | |
| in_shard += len(buf["episode_index"]) | |
| buf = {k: [] for k in buf} | |
| if writer is not None and (final or in_shard >= args.rows_per_shard): | |
| writer.close() | |
| writer = None | |
| shard_id += 1 | |
| in_shard = 0 | |
| import time | |
| t0 = time.time() | |
| for n, r in enumerate(mine): | |
| e = int(episodes["episode_index"][r].as_py()) | |
| L = int(lengths[r]) | |
| frames = {} | |
| for slot, key in (("cam0", cam0), ("cam1", cam1)): | |
| if key is None: | |
| continue | |
| rd = readers[slot] | |
| ci = int(episodes[f"videos/{key}/chunk_index"][r].as_py()) | |
| fi = int(episodes[f"videos/{key}/file_index"][r].as_py()) | |
| ts = float(episodes[f"videos/{key}/from_timestamp"][r].as_py()) | |
| frames[slot] = { | |
| i // stride: encode_jpeg(img, args.image_size, args.jpeg_quality) | |
| for i, img in rd.read_episode(ci, fi, ts, L, fps, stride) | |
| } | |
| n_kept = len(range(0, L, stride)) | |
| blank = b"" | |
| for j in range(n_kept): | |
| buf["episode_index"].append(e) | |
| buf["frame_index"].append(j) | |
| buf["cam0"].append(frames["cam0"][j]) | |
| buf["cam1"].append(frames.get("cam1", {}).get(j, blank)) | |
| if len(buf["episode_index"]) >= 2048: | |
| flush() | |
| if (n + 1) % 50 == 0: | |
| el = time.time() - t0 | |
| print(f" {n+1}/{len(mine)} episodes, {written:,} frames, " | |
| f"{written/el:.0f} fr/s, eta {(len(mine)-n-1)*el/(n+1)/60:.1f} min", flush=True) | |
| flush(final=True) | |
| for rd in readers.values(): | |
| if rd is not None: | |
| rd.close() | |
| # frames.parquet holds only the kept rows, in shard order | |
| pq.write_table( | |
| pa.table({ | |
| "episode_index": pa.array(ep_i, pa.int32()), | |
| "frame_index": pa.array(fr_i, pa.int32()), | |
| "state": pa.array(list(state), pa.list_(pa.float32(), state.shape[1])), | |
| "action": pa.array(list(action), pa.list_(pa.float32(), action.shape[1])), | |
| "task_index": pa.array(task_i, pa.int32()), | |
| }), | |
| args.out / f"frames-part{args.part:02d}.parquet", | |
| compression="zstd", | |
| ) | |
| (args.out / f"tasks-part{args.part:02d}.json").write_text( | |
| json.dumps({v: k for k, v in tasks.items()}, ensure_ascii=False) | |
| ) | |
| (args.out / f"stats-part{args.part:02d}.json").write_text(json.dumps({ | |
| "observation.state": {"mean": state.mean(0).tolist(), "std": state.std(0).tolist()}, | |
| "action": {"mean": action.mean(0).tolist(), "std": action.std(0).tolist()}, | |
| "count": int(len(ep_i)), | |
| })) | |
| (args.out / f"manifest-part{args.part:02d}.json").write_text(json.dumps({ | |
| "repo_id": args.repo_id, "revision": args.revision, | |
| "source_fps": fps, "stride": stride, "fps": eff_hz, | |
| "image_size": args.image_size, "jpeg_quality": args.jpeg_quality, | |
| "cam0": cam0, "cam1": cam1, | |
| "state_dim": int(state.shape[1]), "action_dim": int(action.shape[1]), | |
| "episodes": int(len(mine)), "frames": int(written), | |
| }, indent=2)) | |
| print(f"done: {written:,} frames in {(time.time()-t0)/60:.1f} min -> {args.out}") | |
| if __name__ == "__main__": | |
| main() | |