scripts/yam_task_runner.py

"""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"

# skill name -> the script that implements it. A skill is "how to do it"; a task is "what,
# with which objects, and what counts as done".
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  # noqa: E402
    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)