File size: 1,643 Bytes
51a2b1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | #!/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()
|