| """Registry + runner for the YAML-declared YAM tasks. |
| |
| A task is a YAML file under ``tasks/`` (see ``tasks/fruit_basket.yaml``). This module turns it |
| into a command line for the skill script that implements it, so adding a task means writing a |
| YAML file -- not another 400-line demo script. Modelled on RoboLab's task registration: a task |
| declares its scene, its instruction variants, and its success conditions; the conditions come |
| from ``bimanual.yam_tasks.conditions``. |
| |
| python scripts/yam_task_runner.py list # registry table |
| python scripts/yam_task_runner.py show fruit_basket # resolved command + conditions |
| python scripts/yam_task_runner.py run fruit_basket # run it, print PASS/FAIL + video |
| python scripts/yam_task_runner.py run --all --tag robotwin-asset |
| |
| No Isaac import here: the runner only builds and launches commands, so ``list``/``show`` are |
| instant and usable without a GPU. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import yaml |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| TASK_DIR = REPO/"tasks" |
| PY = os.environ.get("YAM_PYTHON", sys.executable) |
| KIT_ARGS = "--kit_args=--/rtx/verifyDriverVersion/enabled=false" |
|
|
| |
| |
| SKILLS = { |
| "pick_place": "scripts/yam_grasp_prm.py", |
| "insert": "scripts/yam_grasp_prm.py", |
| "multi_pick": "scripts/yam_multi_pick.py", |
| "stack": "scripts/yam_stack_target.py", |
| "dual_arm_lift": "scripts/yam_pot_lift.py", |
| "tool_use": "scripts/yam_tool_tasks.py", |
| "sort": "scripts/yam_cubes.py", |
| } |
|
|
|
|
| def load_tasks() -> dict[str, dict]: |
| tasks = {} |
| for f in sorted(TASK_DIR.glob("*.yaml")): |
| spec = yaml.safe_load(f.read_text()) |
| spec["_file"] = str(f.relative_to(REPO)) |
| tasks[spec["name"]] = spec |
| return tasks |
|
|
|
|
| def _xy(v): |
| return f"{float(v[0])},{float(v[1])}" |
|
|
|
|
| def build_command(spec: dict) -> list[str]: |
| """Resolve a task spec into an executable command line.""" |
| skill = spec["skill"] |
| if skill not in SKILLS: |
| raise SystemExit(f"unknown skill {skill!r}; known: {sorted(SKILLS)}") |
| cmd = [PY, str(REPO/SKILLS[skill]), "--headless"] |
| scene = spec.get("scene", {}) |
| params = spec.get("params", {}) |
| objects = scene.get("objects", []) |
| container = scene.get("container") |
| video = spec.get("video", f"outputs/tasks/{spec['name']}.mp4") |
|
|
| if skill == "multi_pick": |
| pick = params.get("pick", [o["name"] for o in objects if o.get("role") == "target"]) |
| leave = params.get("leave", [o["name"] for o in objects if o.get("role") == "distractor"]) |
| by_name = {o["name"]: o for o in objects} |
| cmd += ["--pick", ",".join(pick), "--leave", ",".join(leave), |
| "--pick_xy", ";".join(_xy(by_name[n]["xy"]) for n in pick)] |
| if leave: |
| cmd += ["--leave_xy", ";".join(_xy(by_name[n]["xy"]) for n in leave)] |
| if container: |
| cmd += ["--basket_xy", _xy(container["xy"]), |
| "--container", container["usd"], |
| "--container_scale", str(container.get("scale", 1.0))] |
| if container.get("on_top"): |
| cmd += ["--on_top"] |
| if params.get("label"): |
| cmd += ["--label", params["label"]] |
|
|
| elif skill in ("pick_place", "insert"): |
| obj = objects[0] |
| cmd += ["--obj", obj["name"], "--obj_xy", _xy(obj["xy"])] |
| if container: |
| cmd += ["--box_xy", _xy(container["xy"])] |
| if container.get("usd"): |
| cmd += ["--container", container["usd"], |
| "--container_scale", str(container.get("scale", 1.0))] |
| else: |
| cmd += ["--basket"] |
| if skill == "insert": |
| cmd += ["--insert", "--basket"] |
| if params.get("hole"): |
| cmd += [f"--hole={params['hole']}"] |
| if params.get("grasp_top"): |
| cmd += [f"--grasp_top={params['grasp_top']}"] |
| if params.get("jaw"): |
| cmd += ["--jaw", params["jaw"]] |
|
|
| elif skill == "stack": |
| cmd += ["--blocks", ",".join(params["blocks"]), |
| "--block_xy", ";".join(_xy(o["xy"]) for o in objects if o.get("role") != "target_region"), |
| "--target_xy", _xy(scene["target_region"]["xy"]), |
| "--target_size", str(scene["target_region"].get("size", 0.11))] |
|
|
| elif skill == "dual_arm_lift": |
| cmd += ["--obj", objects[0]["name"], "--pot_xy", _xy(objects[0]["xy"])] |
| if params.get("pot_rpy"): |
| cmd += ["--pot_rpy", params["pot_rpy"]] |
| if params.get("lift"): |
| cmd += ["--lift", str(params["lift"])] |
|
|
| elif skill == "tool_use": |
| cmd += ["--task", params["mode"]] |
| if objects: |
| cmd += ["--obj", objects[0]["name"], "--obj_xy", _xy(objects[0]["xy"])] |
| if params.get("tool"): |
| cmd += ["--tool", params["tool"]] |
| if scene.get("target_region"): |
| cmd += ["--target_xy", _xy(scene["target_region"]["xy"]), |
| "--target_size", str(scene["target_region"].get("size", 0.10))] |
|
|
| elif skill == "sort": |
| cmd += ["--task", params.get("mode", "sort")] |
|
|
| if "episode" in params: |
| cmd += ["--episode", str(params["episode"])] |
| cmd += ["--video", video, KIT_ARGS] |
| return cmd |
|
|
|
|
| def env_for(spec: dict) -> dict: |
| """Extra env vars a task needs (RoboTwin asset dir, runtime-registered objects).""" |
| env = dict(os.environ) |
| env.setdefault("OMNI_KIT_ACCEPT_EULA", "YES") |
| env.setdefault("OMNI_KIT_ALLOW_ROOT", "1") |
| env.setdefault("PYTHONUNBUFFERED", "1") |
| env.setdefault("ROBOTWIN_USD", os.environ.get("ROBOTWIN_USD", "/home/yu/internship_yu/robotwin_usd")) |
| extra = spec.get("scene", {}).get("register_objects") |
| if extra: |
| env["YAM_RW_OBJECTS"] = ",".join( |
| f"{o['name']}={o['usd']}:{o.get('scale', 0.05)}:{o.get('mass', 0.10)}" for o in extra) |
| return env |
|
|
|
|
| def describe_conditions(spec: dict) -> list[str]: |
| out = [] |
| for c in spec.get("success", []): |
| (kind, kwargs), = c.items() |
| out.append(f"{kind}({', '.join(f'{k}={v}' for k, v in kwargs.items())})") |
| return out |
|
|
|
|
| def cmd_list(tasks, args): |
| tag = args.tag |
| rows = [(n, s.get("skill", "?"), ",".join(s.get("tags", [])), s.get("video", "")) |
| for n, s in sorted(tasks.items()) if not tag or tag in s.get("tags", [])] |
| w = max([len(r[0]) for r in rows]+[4]) |
| print(f"{'name'.ljust(w)} {'skill'.ljust(14)} tags") |
| print("-"*(w+70)) |
| for n, sk, tg, _ in rows: |
| print(f"{n.ljust(w)} {sk.ljust(14)} {tg}") |
| print(f"\n{len(rows)} task(s) in {TASK_DIR.relative_to(REPO)}/") |
|
|
|
|
| def cmd_show(tasks, args): |
| spec = tasks[args.name] |
| print(f"# {spec['name']} ({spec['_file']})") |
| print(f"instruction: {spec.get('instruction', {}).get('default','')}") |
| print(f"skill: {spec['skill']} -> {SKILLS.get(spec['skill'])}") |
| print(f"tags: {', '.join(spec.get('tags', []))}") |
| print("success conditions:") |
| for c in describe_conditions(spec): |
| print(f" - {c}") |
| print("\ncommand:") |
| print(" " + " ".join(build_command(spec))) |
| extra = env_for(spec).get("YAM_RW_OBJECTS") |
| if extra: |
| print(f"\nYAM_RW_OBJECTS={extra}") |
|
|
|
|
| def verify_run(spec: dict, video: Path, verifier: str | None) -> dict: |
| """Agent-in-the-loop check: the numeric EPISODE_RESULT only inspects the final object pose, |
| so it happily passes an episode where the cup landed upside-down or the basket is sunk into |
| the table. Render a contact sheet and have a vision agent judge what actually happened. |
| |
| `verifier` is a shell command template receiving {sheet} and {claim}; it must print JSON |
| {"verdict": ..., "notes": ...}. With no verifier the request is written out for a harness |
| (e.g. a Claude Code subagent) to pick up -- nothing is uploaded anywhere by default. |
| """ |
| sys.path.insert(0, str(REPO/"scripts")) |
| from yam_agent_loop import make_sheet, REVIEW_PROMPT |
| sheet = make_sheet(video) |
| if sheet is None: |
| return {"verdict": "NO_VIDEO", "notes": "no frames rendered"} |
| claim = spec.get("instruction", {}).get("specific") or spec.get("instruction", {}).get("default", spec["name"]) |
| prompt = REVIEW_PROMPT.format(sheet=sheet.relative_to(REPO), claim=claim) |
| if not verifier: |
| req = REPO/"outputs"/"verify"/f"{spec['name']}_request.json" |
| req.write_text(json.dumps({"sheet": str(sheet.relative_to(REPO)), "claim": claim, |
| "prompt": prompt}, indent=2)) |
| return {"verdict": "PENDING_AGENT", "sheet": str(sheet.relative_to(REPO)), |
| "request": str(req.relative_to(REPO))} |
| out = subprocess.run(verifier.format(sheet=str(sheet), claim=json.dumps(claim)), |
| shell=True, capture_output=True, text=True, cwd=REPO) |
| try: |
| return json.loads(out.stdout.strip().splitlines()[-1]) |
| except Exception: |
| return {"verdict": "UNPARSED", "notes": (out.stdout or out.stderr)[-400:]} |
|
|
|
|
| def cmd_run(tasks, args): |
| names = sorted(tasks) if args.all else [args.name] |
| if args.tag: |
| names = [n for n in names if args.tag in tasks[n].get("tags", [])] |
| results = {} |
| for n in names: |
| spec = tasks[n] |
| cmd = build_command(spec) |
| log = REPO/"outputs"/"tasks"/f"{n}.log" |
| log.parent.mkdir(parents=True, exist_ok=True) |
| (REPO/"outputs"/"verify").mkdir(parents=True, exist_ok=True) |
| print(f"\n===== RUN {n} =====\n {' '.join(cmd)}", flush=True) |
| with open(log, "w") as fh: |
| rc = subprocess.call(cmd, cwd=REPO, env=env_for(spec), stdout=fh, stderr=subprocess.STDOUT) |
| text = log.read_text() |
| verdict = "UNKNOWN" |
| for line in text.splitlines(): |
| if "EPISODE_RESULT" in line: |
| verdict = "SUCCESS" if "SUCCESS" in line else "FAIL" |
| entry = {"numeric": verdict, "exit": rc, "video": spec.get("video")} |
| if not args.no_verify: |
| v = REPO/spec.get("video", f"outputs/tasks/{n}.mp4") |
| entry["visual"] = verify_run(spec, v, args.verifier) if v.exists() else {"verdict": "NO_VIDEO"} |
| results[n] = entry |
| print(f" exit={rc} numeric={verdict} visual={entry.get('visual', {}).get('verdict', 'skipped')}" |
| f" video={spec.get('video')}", flush=True) |
| (REPO/"outputs"/"verify"/"run_results.json").write_text(json.dumps(results, indent=2)) |
| print("\n===== SUMMARY (numeric result + agent visual verdict) =====") |
| for n, e in results.items(): |
| vis = e.get("visual", {}).get("verdict", "skipped") |
| print(f" {e['numeric'].ljust(8)} | {str(vis).ljust(14)} {n} -> {e['video']}") |
| print("\nagreeing SUCCESS+CONFIRMED is the only state that counts as done; " |
| "feed the rest to `yam_agent_loop.py update` for the concrete parameter fix.") |
| return 0 if all(e["numeric"] == "SUCCESS" for e in results.values()) else 1 |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| p = sub.add_parser("list"); p.add_argument("--tag", default="") |
| p = sub.add_parser("show"); p.add_argument("name") |
| p = sub.add_parser("run"); p.add_argument("name", nargs="?"); p.add_argument("--all", action="store_true") |
| p.add_argument("--tag", default="") |
| p.add_argument("--no-verify", action="store_true", help="skip the agent visual check") |
| p.add_argument("--verifier", default=os.environ.get("YAM_VERIFIER"), |
| help="shell command template with {sheet} and {claim}; must print JSON " |
| "{verdict, notes}. Omit to just emit the review request for an agent.") |
| args = ap.parse_args() |
| tasks = load_tasks() |
| if args.cmd == "list": |
| return cmd_list(tasks, args) |
| if args.cmd == "show": |
| return cmd_show(tasks, args) |
| if args.cmd == "run": |
| if not args.all and not args.name: |
| raise SystemExit("give a task name or --all") |
| return cmd_run(tasks, args) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main() or 0) |
|
|
