File size: 4,870 Bytes
51ce5b3 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | #!/usr/bin/env python
"""Probe the nvidia/physical-ai collection and emit a streaming specs YAML.
Walks the robotics repos, finds every LeRobot v2.x (sub)dataset that
HubEpisodeStream can consume, assigns embodiment ids BY ROBOT FAMILY (not per
dataset — 127 ids would defeat the embedding) and weights proportional to
sqrt(frames) (plain proportional lets the two biggest GR1 tasks dominate the
mixture; sqrt keeps diversity without starving the big sources).
v3.0 repos (BridgeData2_LeRobot_v3, LIBERO_LeRobot_v3) are NOT emitted — they
need the v3 shard path (build_shards.py); both were already in the C-scaled
pretraining mixture anyway.
Usage:
python make_stream_specs.py --out configs/physical_ai_stream.yaml
"""
from __future__ import annotations
import argparse
import json
import math
import os
import urllib.request
from concurrent.futures import ThreadPoolExecutor
FAMILIES = [ # ordered: first match wins; ids continue after C-scaled's 0..9
("gr1", 10), ("bimanual_panda_gripper", 11), ("bimanual_panda_hand", 12),
("single_panda_gripper", 13), ("sim_behavior_r1_pro", 14), ("unitree_g1", 15),
("g1", 15),
]
ROOTS = [
("nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim", None), # nested
("nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim", "LeRobot"), # nested under LeRobot/
("nvidia/PhysicalAI-GR00T-Tuned-Tasks", None), # nested
("nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1", None), # nested
("nvidia/GR00T-N1.7-AppleToPlate", ""), # root-level v2.1
]
def hf_json(url: str, token: str | None):
req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"} if token else {})
try:
return json.load(urllib.request.urlopen(req, timeout=60))
except Exception:
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="configs/physical_ai_stream.yaml")
ap.add_argument("--hz", type=float, default=10.0)
args = ap.parse_args()
token = None
tp = os.path.expanduser("~/.cache/huggingface/token")
if os.path.exists(tp):
token = open(tp).read().strip()
jobs = [] # (repo_id, prefix)
for repo, sub in ROOTS:
if sub == "": # root-level dataset
jobs.append((repo, ""))
continue
base = f"https://huggingface.co/api/datasets/{repo}/tree/main" + (f"/{sub}" if sub else "")
tree = hf_json(base, token) or []
for e in tree:
if e.get("type") == "directory":
jobs.append((repo, e["path"]))
def probe(job):
repo, pfx = job
p = f"{pfx}/" if pfx else ""
info = hf_json(f"https://huggingface.co/datasets/{repo}/resolve/main/{p}meta/info.json", token)
if not info or "total_frames" not in info:
return None
ver = str(info.get("codebase_version", ""))
if not ver.startswith("v2"):
return None # v3 needs the shard path, v1 untested
return dict(
repo_id=repo, prefix=pfx, frames=info["total_frames"], fps=info["fps"],
state_dim=info["features"]["observation.state"]["shape"][0],
action_dim=info["features"]["action"]["shape"][0],
ncam=sum(1 for v in info["features"].values() if v.get("dtype") == "video"),
)
with ThreadPoolExecutor(16) as ex:
found = [r for r in ex.map(probe, jobs) if r]
max_state = max(r["state_dim"] for r in found)
max_action = max(r["action_dim"] for r in found)
total = sum(r["frames"] for r in found)
lines = [
"# Autogenerated by make_stream_specs.py — streaming mixture over the",
"# LeRobot-v2.x part of the nvidia/physical-ai collection.",
f"# {len(found)} datasets, {total:,} source frames.",
f"# max_state_dim needed: {max_state}; max_action_dim needed: {max_action}.",
"datasets:",
]
for r in sorted(found, key=lambda x: (-x["frames"])):
name = r["prefix"] or r["repo_id"].split("/")[-1]
emb = next((i for k, i in FAMILIES if name.startswith(k) or f"/{k}" in name.lower() or k in name.lower()), 15)
w = math.sqrt(r["frames"])
lines += [
f" - repo_id: {r['repo_id']}",
f" prefix: \"{r['prefix']}\"",
f" weight: {w:.1f} # sqrt({r['frames']:,} fr); fps={r['fps']} state={r['state_dim']} act={r['action_dim']} cams={r['ncam']}",
f" embodiment_id: {emb}",
f" hz: {args.hz}",
]
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
with open(args.out, "w") as f:
f.write("\n".join(lines) + "\n")
print(f"{len(found)} датасетов, {total:,} кадров -> {args.out}")
print(f"max_state_dim={max_state} max_action_dim={max_action}")
if __name__ == "__main__":
main()
|