scripts/yam_agent_loop.py

"""Agent-in-the-loop task pipeline: RUN -> EVALUATE (VLM) -> UPDATE POLICY -> PROPOSE new tasks.

The demo scripts already report a numeric EPISODE_RESULT, but that number only checks the final
object pose -- it cannot see that a cup landed upside-down, that a basket is sunk into the table,
or that the "lifted" object was actually dragged. So each run is also rendered to a contact sheet
and handed to a vision agent, and the agent's verdict is what drives the next iteration.

Three stages, each usable on its own:

  evaluate  render a contact sheet per video + emit the reviewer prompts, and record verdicts
  update    turn verdicts into concrete parameter edits for the next run of that task
  propose   suggest new tasks from the asset library that are unlike the ones already covered

    python scripts/yam_agent_loop.py evaluate --videos outputs/tasks
    python scripts/yam_agent_loop.py update   --verdicts outputs/verify/verdicts.json
    python scripts/yam_agent_loop.py propose  --n 10

The vision call itself is deliberately NOT hardcoded to one provider: `evaluate` writes a
`review_requests.json` that a harness (Claude Code subagents, or any VLM endpoint) consumes, and
`update` reads back a `verdicts.json`. That keeps credentials out of this repo -- nothing here
uploads your renders anywhere.
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
VERIFY = REPO/"outputs"/"verify"

# ---------------------------------------------------------------- knowledge learned so far
# Each entry: the visual symptom a reviewer can see -> the parameter that actually fixes it.
# These were all derived from real failures in this suite, not invented.
FIX_RULES = [
    {"symptom": "object upside down or lying on its side at the start",
     "cause": "RoboTwin GLBs are authored Y-up, so an identity spawn quaternion lands them rolled",
     "fix": "pass a +90 deg roll: --block_rpy 90,0,0 / --pot_rpy 90,0,0 / --container_rpy 90,0,0"},
    {"symptom": "container is enormous, fills the table, or object never reaches its rim",
     "cause": "the GLB->USD converter reads model_data scale but does not bake it in",
     "fix": "pass the asset's model_data scale explicitly via --container_scale"},
    {"symptom": "arm never moves; tracking error stays huge from the first frame",
     "cause": "a tall container was spawned on top of the arm's home pose and blocks it",
     "fix": "move the container away from the home eef (world ~(0.08,-0.19)) or shrink it"},
    {"symptom": "jaws close fully (finger separation -> 0) and the object stays put",
     "cause": "clamp force too high for a thin-walled vessel, or grasp height on a taper",
     "fix": "lower YAM_GRIP_EFFORT to ~40-55, or grasp at the rim with --grasp_top"},
    {"symptom": "object slips out mid-carry",
     "cause": "clamp force too low for a heavy/large object",
     "fix": "raise YAM_GRIP_EFFORT to ~75-85"},
    {"symptom": "jaws stall but the object does not move (contact with the wrong part)",
     "cause": "jaw axis parallel to a protruding handle instead of across it",
     "fix": "rotate the wrist: --jaw x (or --grip_tilt for an ear handle)"},
    {"symptom": "two-arm lift tilts to one side",
     "cause": "the mesh origin is not centred between the two handles",
     "fix": "give each arm its own offset: --half_y_l / --half_y_r"},
    {"symptom": "second object lands on the first and rolls out of the container",
     "cause": "both objects released at the same drop point",
     "fix": "spread the drop points along x (already automatic in yam_multi_pick)"},
    {"symptom": "descend error > ~4 cm, arm visibly short of the object",
     "cause": "target outside the reliable reach envelope",
     "fix": "move the object within ~0.35 m of the arm root at (-0.2, +-0.2)"},
    {"symptom": "object visibly dragged rather than lifted",
     "cause": "grasp height below the object centre, or grip on a narrowing taper",
     "fix": "re-seat the object on the table first, then grasp at its measured mid-height"},
]

REVIEW_PROMPT = """You are visually verifying a robot manipulation video.

Read the contact sheet {sheet} (7 frames sampled across the episode, left-to-right, top-to-bottom,
in time order). Judge from the IMAGES ONLY -- do not read logs.

Claimed task: {claim}

Report:
- VERDICT: CONFIRMED / PARTIAL / REFUTED
- what the final frames actually show (where did each object end up?)
- visual oddities: object upside-down or on its side, sunk into or floating above a surface,
  landed outside its container, gripper closed beside the object, one arm idle in a two-arm task,
  container tipped, arm contorted, object never left the table.
Be skeptical and concrete."""


def sheet_for(video: Path) -> Path:
    return VERIFY/f"{video.stem}.png"


def make_sheet(video: Path, n_tiles: int = 7) -> Path | None:
    """Render a contact sheet so a vision model can judge the episode from one image."""
    try:
        import imageio.v2 as iio
        from PIL import Image
    except ImportError:
        print("[loop] need imageio + pillow", file=sys.stderr); return None
    try:
        r = iio.get_reader(str(video)); n = r.count_frames()
    except Exception as e:
        print(f"[loop] cannot read {video}: {e}"); return None
    idx = [int(n*t) for t in [i/(n_tiles-1)*0.95+0.02 for i in range(n_tiles)]]
    tiles = [Image.fromarray(r.get_data(min(i, n-1))).resize((360, 272)) for i in idx]
    cols = 4; rows = (len(tiles)+cols-1)//cols
    sheet = Image.new("RGB", (360*cols, 272*rows))
    for k, im in enumerate(tiles):
        sheet.paste(im, (360*(k % cols), 272*(k//cols)))
    VERIFY.mkdir(parents=True, exist_ok=True)
    out = sheet_for(video)
    sheet.save(out)
    return out


def cmd_evaluate(args):
    """Render sheets and emit one review request per video for the vision agents."""
    videos = sorted(Path(args.videos).glob("*.mp4"))
    claims = {}
    cfile = REPO/"tasks"/"claims.json"
    if cfile.exists():
        claims = json.loads(cfile.read_text())
    requests = []
    for v in videos:
        s = make_sheet(v)
        if s is None:
            continue
        claim = claims.get(v.stem, f"the task named {v.stem}")
        requests.append({"video": str(v.relative_to(REPO)), "sheet": str(s.relative_to(REPO)),
                         "claim": claim,
                         "prompt": REVIEW_PROMPT.format(sheet=s.relative_to(REPO), claim=claim)})
        print(f"[loop] sheet {s.relative_to(REPO)}")
    out = VERIFY/"review_requests.json"
    out.write_text(json.dumps(requests, indent=2))
    print(f"\n[loop] {len(requests)} review request(s) -> {out.relative_to(REPO)}")
    print("[loop] hand each `prompt` to a vision agent, then write its verdict into "
          "outputs/verify/verdicts.json as {task: {verdict, notes}}")
    return 0


def suggest_fixes(notes: str) -> list[dict]:
    """Map a reviewer's free-text observation onto the parameter that actually fixes it."""
    low = notes.lower()
    hits = []
    for rule in FIX_RULES:
        key = rule["symptom"].split()[0:3]
        probes = {
            "upside": ["upside", "inverted", "on its side", "sideways", "tipped over"],
            "container": ["enormous", "huge", "too big", "fills the table", "oversized"],
            "arm": ["never moves", "does not move", "blocked", "stuck at home"],
            "jaws": ["closes fully", "closed beside", "closes beside", "misses the object", "no contact"],
            "object": ["slips", "dropped mid", "falls out", "drags", "dragged", "never left the table"],
            "two-arm": ["tilt", "tilted", "one side", "one arm idle"],
            "second": ["rolls out", "lands on the first", "bounced out", "outside the container"],
            "descend": ["short of", "out of reach", "cannot reach", "too far"],
        }.get(key[0], [])
        if any(p in low for p in probes):
            hits.append(rule)
    return hits


def cmd_update(args):
    """Turn reviewer verdicts into the concrete next action for each task."""
    vfile = Path(args.verdicts)
    if not vfile.exists():
        raise SystemExit(f"no verdicts at {vfile}; run `evaluate` and collect agent verdicts first")
    verdicts = json.loads(vfile.read_text())
    plan = {}
    for task, v in verdicts.items():
        verdict = str(v.get("verdict", "")).upper()
        notes = str(v.get("notes", ""))
        if verdict == "CONFIRMED":
            plan[task] = {"action": "keep", "why": "reviewer confirmed the task visually"}
            continue
        fixes = suggest_fixes(notes)
        plan[task] = {
            "action": "rerun_with_changes" if fixes else "needs_human_look",
            "verdict": verdict,
            "reviewer_notes": notes,
            "changes": [{"symptom": f["symptom"], "cause": f["cause"], "fix": f["fix"]} for f in fixes],
        }
    out = VERIFY/"update_plan.json"
    out.write_text(json.dumps(plan, indent=2))
    for task, p in plan.items():
        print(f"\n== {task}: {p['action']}")
        for c in p.get("changes", []):
            print(f"   - {c['symptom']}\n     -> {c['fix']}")
    print(f"\n[loop] plan -> {out.relative_to(REPO)}")
    return 0


# ---------------------------------------------------------------- task proposal
SKILL_AXES = ["pick-place", "insertion", "stacking", "semantic-selection", "tool-use",
              "bimanual-lift", "bimanual-handover", "pouring", "pushing", "sweeping",
              "hanging", "weighing", "constrained-pull"]


def covered_tasks() -> dict[str, list[str]]:
    out = {}
    tdir = REPO/"tasks"
    if not tdir.exists():
        return out
    try:
        import yaml
    except ImportError:
        return out
    for f in tdir.glob("*.yaml"):
        spec = yaml.safe_load(f.read_text())
        out[spec.get("name", f.stem)] = spec.get("tags", [])
    return out


def cmd_propose(args):
    """Propose new tasks that are UNLIKE what is already covered, from assets actually present."""
    usd_root = Path(os.environ.get("ROBOTWIN_USD", "/home/yu/internship_yu/robotwin_usd"))
    have = sorted({p.parent.name for p in usd_root.glob("*/*.usd")}) if usd_root.exists() else []
    covered = covered_tasks()
    used_tags = {t for tags in covered.values() for t in tags}
    gaps = [s for s in SKILL_AXES if s not in used_tags]

    catalogue = {
        "pouring": ("091_kettle", "002_bowl", "pour beads from the kettle into a bowl by rotating the wrist"),
        "sweeping": ("083_brush", "082_smallshovel", "sweep loose beads into a dustpan with the brush"),
        "weighing": ("035_apple", "072_electronicscale", "place the fruit onto the electronic scale"),
        "hanging": ("039_mug", "040_rack", "hang the mug on the rack"),
        "tool-use": ("117_whiteboard-eraser", None, "wipe a marked region with the eraser"),
        "constrained-pull": (None, None, "pull a wooden bar straight out of a slotted cube"),
        "bimanual-handover": ("110_basket", None, "left arm passes the basket to the right arm"),
        "pushing": (None, None, "push a block to a target region without grasping it"),
        "insertion": (None, None, "drop a ring over a standing post"),
        "stacking": ("108_block", None, "build a 3-block tower on a marked square"),
        "semantic-selection": ("006_hamburg", "110_basket", "put only the food items in the basket"),
        "bimanual-lift": ("060_kitchenpot", None, "both arms lift a wide object by its two handles"),
        "pick-place": ("021_cup", "062_plasticbox", "move the cup into a plastic box"),
    }

    print(f"[loop] {len(covered)} task(s) registered; {len(have)} RoboTwin categories converted")
    print(f"[loop] skill axes not yet covered: {', '.join(gaps) if gaps else '(none)'}\n")
    proposals = []
    for skill in (gaps or SKILL_AXES):
        obj, tgt, desc = catalogue.get(skill, (None, None, f"a {skill} task"))
        missing = [a for a in (obj, tgt) if a and a not in have]
        proposals.append({
            "skill": skill, "description": desc,
            "objects": [a for a in (obj, tgt) if a],
            "assets_ready": not missing,
            "convert_first": missing,
        })
        if len(proposals) >= args.n:
            break
    for p in proposals:
        flag = "" if p["assets_ready"] else f"  (convert first: {', '.join(p['convert_first'])})"
        print(f"  [{p['skill']}] {p['description']}{flag}")
    out = VERIFY/"proposals.json"
    VERIFY.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(proposals, indent=2))
    print(f"\n[loop] proposals -> {out.relative_to(REPO)}")
    return 0


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("evaluate"); p.add_argument("--videos", default=str(REPO/"outputs"/"tasks"))
    p = sub.add_parser("update"); p.add_argument("--verdicts", default=str(VERIFY/"verdicts.json"))
    p = sub.add_parser("propose"); p.add_argument("--n", type=int, default=10)
    args = ap.parse_args()
    return {"evaluate": cmd_evaluate, "update": cmd_update, "propose": cmd_propose}[args.cmd](args)


if __name__ == "__main__":
    sys.exit(main() or 0)