Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """Package recorded agent episodes into the Space's replay-player format. | |
| Input: run directories produced by an arena runner (uitars-webtask-eval's | |
| run_arena.py or this repo's harness) containing frames/, transcript.txt, and | |
| result.json. Frames are assumed to be saved 8x per step (video padding); | |
| one frame per step is kept, downscaled to JPEG. | |
| python3 scripts/build_replays.py --agent "UI-TARS-1.5-7B (4-bit)" RUN_DIR... | |
| Output: server/static/replays/<task>/step_NNN.jpg + replay.json, and a | |
| top-level manifest.json the landing/replay/play pages read. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import re | |
| import sys | |
| from pathlib import Path | |
| from PIL import Image | |
| OUT_ROOT = Path(__file__).resolve().parent.parent / "server" / "static" / "replays" | |
| FRAME_DUP = 8 # frames saved per step by the runners | |
| JPEG_WIDTH = 880 | |
| JPEG_QUALITY = 68 | |
| _STEP = re.compile(r"^\[step (\d+)\] (\w+) point=(.*?) end=.*?content=(.*)$") | |
| DIAGNOSIS = { | |
| "checkout-form": "Clean run: precise field clicks, correct typing, " | |
| "three wizard steps, submit.", | |
| "email-triage": "State-tracking failure: archives worked, but the model never " | |
| "perceived rows disappearing — it re-clicked fixed coordinates " | |
| "while the list reflowed underneath, archiving 7 extra emails " | |
| "while blaming 'an unresponsive system'.", | |
| "settings-panel": "Environment finding: headless chromium renders native " | |
| "<select> dropdowns outside the page — the timezone options " | |
| "are invisible to any screenshot agent. Dark Mode itself was " | |
| "correctly enabled and saved.", | |
| "causal-gate": "Causal reasoning: correct. Grounding: a near-miss. The model " | |
| "opened the dropdown, selected “Friend referral,” watched the " | |
| "referral field appear, and typed the right code — it got the " | |
| "causal precondition exactly right. Its final click aimed at " | |
| "“Save changes” landed about 12px above the button, on the " | |
| "divider line just above it, so the click never registered " | |
| "and nothing was saved.", | |
| "flash-offer": "Same 'distrusts its own successful action' pattern as " | |
| "email-triage, reproduced in an unrelated task. The model " | |
| "correctly named the live code at each glance as it rotated, " | |
| "but its first click genuinely focused the input — a subtle " | |
| "focus-ring change it misread as 'the page didn't respond,' " | |
| "so it re-clicked the same point instead of typing, until the " | |
| "harness's repeat-action guard aborted the episode.", | |
| } | |
| def parse_transcript(path: Path) -> list[dict]: | |
| steps, current = [], None | |
| for line in path.read_text().splitlines(): | |
| m = _STEP.match(line) | |
| if m: | |
| current = {"step": int(m.group(1)), "action": m.group(2), | |
| "point": m.group(3), "content": m.group(4), "thought": ""} | |
| steps.append(current) | |
| elif current is not None and line.strip().startswith("thought:"): | |
| current["thought"] = line.split("thought:", 1)[1].strip() | |
| return steps | |
| def action_label(s: dict) -> str: | |
| label = s["action"] | |
| if s["point"] not in ("None", ""): | |
| label += f" @ {s['point']}" | |
| if s["content"] not in ("None", "''", '""'): | |
| label += f" {s['content']}" | |
| return label | |
| def build_one(run_dir: Path, agent: str) -> dict: | |
| result = json.loads((run_dir / "result.json").read_text()) | |
| task = result["task"] | |
| steps = parse_transcript(run_dir / "transcript.txt") | |
| frames = sorted((run_dir / "frames").glob("frame_*.png")) | |
| out_dir = OUT_ROOT / task | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| step_entries = [] | |
| for s in steps: | |
| idx = s["step"] * FRAME_DUP | |
| if idx >= len(frames): | |
| break | |
| img = Image.open(frames[idx]).convert("RGB") | |
| img = img.resize((JPEG_WIDTH, round(img.height * JPEG_WIDTH / img.width))) | |
| name = f"step_{s['step']:03d}.jpg" | |
| img.save(out_dir / name, quality=JPEG_QUALITY, optimize=True) | |
| step_entries.append({"img": name, "action": action_label(s), | |
| "thought": s["thought"]}) | |
| # trailing "final state" frame, if recorded beyond the last step | |
| final_idx = len(steps) * FRAME_DUP | |
| if final_idx < len(frames): | |
| img = Image.open(frames[final_idx]).convert("RGB") | |
| img = img.resize((JPEG_WIDTH, round(img.height * JPEG_WIDTH / img.width))) | |
| img.save(out_dir / "final.jpg", quality=JPEG_QUALITY, optimize=True) | |
| step_entries.append({"img": "final.jpg", "action": "final state", | |
| "thought": ""}) | |
| entry = { | |
| "task": task, | |
| "agent": agent, | |
| "seed": result["seed"], | |
| "instruction": result["instruction"], | |
| "success": result["success"], | |
| "subgoals": result["subgoals"], | |
| "diagnosis": DIAGNOSIS.get(task, ""), | |
| "steps": step_entries, | |
| } | |
| (out_dir / "replay.json").write_text(json.dumps(entry, indent=1)) | |
| return entry | |
| def main() -> int: | |
| ap = argparse.ArgumentParser(description=__doc__) | |
| ap.add_argument("run_dirs", nargs="+", type=Path) | |
| ap.add_argument("--agent", default="UI-TARS-1.5-7B (4-bit)") | |
| args = ap.parse_args() | |
| manifest = [build_one(d, args.agent) for d in args.run_dirs] | |
| OUT_ROOT.mkdir(parents=True, exist_ok=True) | |
| # manifest omits per-step data to keep the landing page light | |
| slim = [{k: v for k, v in e.items() if k != "steps"} | {"n_steps": len(e["steps"])} | |
| for e in manifest] | |
| (OUT_ROOT / "manifest.json").write_text(json.dumps(slim, indent=1)) | |
| total = sum(f.stat().st_size for f in OUT_ROOT.rglob("*") if f.is_file()) | |
| print(f"built {len(manifest)} replays -> {OUT_ROOT} ({total/1e6:.1f} MB)") | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |