Download scripts/yam_agent_loop.py from yqi19/mp_yam_code: direct link, hf CLI and curl.
- Browser
- Download file 13.3 kB
-
https://huggingface.co/yqi19/mp_yam_code/resolve/main/scripts/yam_agent_loop.py
- Command line
-
hf download hf://yqi19/mp_yam_code/scripts/yam_agent_loop.py
-
curl -L -o yam_agent_loop.py https://huggingface.co/yqi19/mp_yam_code/resolve/main/scripts/yam_agent_loop.py
13.3 kB
| """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) | |