tinyvla2/scripts/prepare_wds.py

#!/usr/bin/env python
"""One-time prep for the wds packs: tar offset index + quantile normalization stats.

The packs ship mean/std only. That is unusable for several of them — go-stanford's
second action channel has std 0.008, so an outlier normalizes to ~250 sigma. Same
failure mode the canonical pipeline already fixed by switching to q01/q99.
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

from tinyvla.data.wds_shards import WdsShardSource

ROOT = Path.home() / "tinyvla_data" / "wds"


def packs():
    for repo in sorted(ROOT.iterdir()):
        if not repo.is_dir():
            continue
        if (repo / "manifest.json").exists():
            yield repo
        else:  # nested (unitree)
            for sub in sorted(repo.iterdir()):
                if (sub / "manifest.json").exists():
                    yield sub


def main():
    n_samples = int(sys.argv[1]) if len(sys.argv) > 1 else 20000
    for p in packs():
        t0 = time.time()
        try:
            src = WdsShardSource(p)
            n = len(src)
            q = src.compute_quantile_stats(n_samples=n_samples)
            import numpy as np
            a01 = np.array(q["action"]["q01"]); a99 = np.array(q["action"]["q99"])
            print(f"{p.name[:42]:44} n={n:>8} adim={len(a01):>3} "
                  f"span=[{np.round(a01,2)[:3]} .. {np.round(a99,2)[:3]}] {time.time()-t0:.0f}s",
                  flush=True)
        except Exception as e:
            print(f"{p.name[:42]:44} FAIL {type(e).__name__}: {str(e)[:90]}", flush=True)
    print("PREP DONE", flush=True)


if __name__ == "__main__":
    main()