| """Scale up: run many randomized episodes per task and keep only the ones that succeeded. |
| |
| Each episode is a fresh seed, so the scene randomization differs every time; the scripted solver |
| plays it and the task's own conditions decide whether it counts. Failed episodes are discarded |
| (or kept under --keep-failures for debugging), which is what makes this a demonstration |
| generator rather than a batch runner. |
| |
| python scripts/generate.py --episodes 20 # every task |
| python scripts/generate.py --tasks grape_box,stack_blocks --episodes 50 |
| python scripts/generate.py --tasks passing --episodes 100 # only tasks known to work |
| python scripts/generate.py --episodes 20 --resume # skip episodes already done |
| |
| Writes: |
| <out>/<task>/ep<seed>.mp4 video, successes only unless --keep-failures |
| <out>/<task>/ep<seed>.json per-episode result + the exact randomized placements |
| <out>/manifest.json running totals and success rate per task |
| |
| One episode is one simulator launch, so this is slow and long-running by nature; it prints a |
| line per episode and is safe to interrupt and --resume. |
| """ |
| import argparse |
| import json |
| import os |
| import re |
| import subprocess |
| import sys |
| import time |
|
|
| REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--tasks", default="all", |
| help="'all', 'passing', or a comma-separated list of task names") |
| ap.add_argument("--episodes", type=int, default=10, help="episodes per task") |
| ap.add_argument("--start-seed", type=int, default=100, |
| help="seeds run from here upward; keep clear of the 0-9 used for development") |
| ap.add_argument("--out", default="data") |
| ap.add_argument("--keep-failures", action="store_true") |
| ap.add_argument("--resume", action="store_true", help="skip episodes whose json already exists") |
| ap.add_argument("--timeout", type=int, default=1500) |
| ap.add_argument("--python", default=sys.executable) |
| a = ap.parse_args() |
|
|
|
|
| def registered(): |
| out = subprocess.run([a.python, "scripts/yam_task.py", "--list"], cwd=REPO, |
| capture_output=True, text=True).stdout |
| return [l.split()[0] for l in out.splitlines()[1:] if l.strip() and not l.startswith(("name", "-")) |
| and len(l.split()) > 2] |
|
|
|
|
| def known_passing(): |
| """Tasks whose most recent development run succeeded -- the ones worth spending seeds on.""" |
| good = [] |
| for t in registered(): |
| log = os.path.join(REPO, "outputs", "tasks", f"v3_{t}.log") |
| if os.path.exists(log) and "EPISODE_RESULT: SUCCESS" in open(log, errors="ignore").read(): |
| good.append(t) |
| return good |
|
|
|
|
| tasks = (registered() if a.tasks == "all" else |
| known_passing() if a.tasks == "passing" else |
| [t.strip() for t in a.tasks.split(",") if t.strip()]) |
| if not tasks: |
| raise SystemExit("no tasks selected") |
|
|
| out_root = os.path.join(REPO, a.out) |
| os.makedirs(out_root, exist_ok=True) |
| manifest_path = os.path.join(out_root, "manifest.json") |
| manifest = json.load(open(manifest_path)) if os.path.exists(manifest_path) else {} |
|
|
| env = dict(os.environ) |
| env.setdefault("OMNI_KIT_ACCEPT_EULA", "YES") |
| env.setdefault("PYTHONUNBUFFERED", "1") |
|
|
| print(f"[gen] {len(tasks)} task(s) x {a.episodes} episodes, seeds " |
| f"{a.start_seed}..{a.start_seed+a.episodes-1} -> {out_root}", flush=True) |
| t_start = time.time() |
|
|
| for task in tasks: |
| tdir = os.path.join(out_root, task) |
| os.makedirs(tdir, exist_ok=True) |
| rec = manifest.setdefault(task, {"attempted": 0, "succeeded": 0, "episodes": []}) |
|
|
| for i in range(a.episodes): |
| seed = a.start_seed+i |
| meta_p = os.path.join(tdir, f"ep{seed}.json") |
| if a.resume and os.path.exists(meta_p): |
| continue |
| vid = os.path.join(tdir, f"ep{seed}.mp4") |
| t0 = time.time() |
| try: |
| p = subprocess.run( |
| [a.python, "scripts/yam_task.py", "--task", task, "--seed", str(seed), |
| "--video", vid, |
| "--kit_args=--/rtx/verifyDriverVersion/enabled=false"], |
| cwd=REPO, env=env, capture_output=True, text=True, timeout=a.timeout) |
| log = p.stdout+p.stderr |
| except subprocess.TimeoutExpired: |
| log = "TIMEOUT" |
| ok = "EPISODE_RESULT: SUCCESS" in log |
| |
| m = re.search(r"placements=(\{.*?\})\s*$", log, re.M) |
| meta = {"task": task, "seed": seed, "success": ok, |
| "seconds": round(time.time()-t0, 1), |
| "placements": m.group(1) if m else None, |
| "conditions": re.findall(r"^\[env\] (PASS|FAIL) (.+)$", log, re.M)} |
| json.dump(meta, open(meta_p, "w"), indent=1) |
| if not ok and not a.keep_failures and os.path.exists(vid): |
| os.remove(vid) |
| rec["attempted"] += 1 |
| rec["succeeded"] += int(ok) |
| rec["episodes"].append({"seed": seed, "success": ok}) |
| json.dump(manifest, open(manifest_path, "w"), indent=1) |
| print(f"[gen] {task:<18} seed {seed:<4} {'ok ' if ok else 'FAIL'} " |
| f"{meta['seconds']:>5.1f}s ({rec['succeeded']}/{rec['attempted']})", flush=True) |
|
|
| print(f"\n[gen] done in {(time.time()-t_start)/60:.1f} min") |
| tot_a = sum(r["attempted"] for r in manifest.values()) |
| tot_s = sum(r["succeeded"] for r in manifest.values()) |
| for t, r in sorted(manifest.items()): |
| if r["attempted"]: |
| print(f" {t:<20} {r['succeeded']:>4}/{r['attempted']:<4} " |
| f"{100*r['succeeded']/r['attempted']:5.1f}%") |
| print(f" {'TOTAL':<20} {tot_s:>4}/{tot_a:<4} " |
| f"{(100*tot_s/tot_a if tot_a else 0):5.1f}% manifest -> {manifest_path}") |
|
|
