#!/usr/bin/env python3 """Compare Robometer vs TrackFlow at discriminating SUCCESS from FAILURE videos. Per-task (run one task per process so the 3 tasks parallelize across GPUs). - TrackFlow score = DISTANCE-WEIGHTED progress (mode="progress_gated": progress * exp(-matched_dist/k_pixel)) — NOT the distance-blind progress_dtw. - Robometer score = max SUCCESS value over the video (its success flag), plus max progress reported for reference. Metric: ROC-AUC (Mann-Whitney) of success-vs-failure separation per scorer. Run (pwm env). One task per GPU: TF_GPU=0 CUDA_VISIBLE_DEVICES=0 python compare_robo_trackflow.py --task Open-Drawer TF_GPU=1 CUDA_VISIBLE_DEVICES=1 python compare_robo_trackflow.py --task Push-Box TF_GPU=2 CUDA_VISIBLE_DEVICES=2 python compare_robo_trackflow.py --task Stack-Can Robometer server shared on localhost:8801. """ import sys, os, json, glob, argparse, traceback import numpy as np RL = "/home/hongyi/PWM/policy_rl/rl" for p in [RL, f"{RL}/old/online"]: if p not in sys.path: sys.path.insert(0, p) from rewards import RobometerScorer, TrackFlowScorer M = "/home/hongyi/PWM/reward_model/3pointr_models" ODH = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Open-Drawer_halffail" PBH = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Push-Box_halffail" SC = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can" N_EACH = 10 TF_GPU = os.environ.get("TF_GPU", "0") K_PIXEL = 50.0 def clean(paths): return [p for p in paths if not any(t in os.path.basename(p) for t in ("_mask", "_pred", "target", "achieved", "_actions"))] def pick(globpat, n): return clean(sorted(glob.glob(globpat)))[:n] def stackcan_ckpt(): for c in ("stack_can_canplate_150pts_384_flow_backbone_final.pth", "stack_can_canplate_150pts_384_flow_backbone.pth"): if os.path.exists(f"{M}/{c}"): return f"{M}/{c}" raise FileNotFoundError("no stack_can_canplate ckpt") def TASKS(): return { "Open-Drawer": { "success": pick(f"{ODH}/success_*.mp4", N_EACH), "failure": pick(f"{ODH}/failure_*.mp4", N_EACH), "robo_task": "A humanoid robot opens a drawer with its dexterous hands.", "tf": dict(pointr_ckpt=f"{M}/open_drawer_2d_wb_150pts_flow_backbone.pth", prompt="wooden item", prompt_fallbacks=["wooden block", "the wooden block", "wooden box", "wooden drawer"], subprocess_gpu=TF_GPU), "note": "success+failure both ACT rollouts (clean, same domain)", }, "Push-Box": { "success": pick("/data0/hongyi/reward_model_data/track_cache/demos_push_box/tableHF_ep0*/input_short.mp4", N_EACH), "failure": pick(f"{PBH}/failure_*.mp4", N_EACH), "robo_task": "A humanoid robot pushes a box to a marker.", "tf": dict(pointr_ckpt=f"{M}/push_box_2d_150pts_384_flow_backbone.pth", prompt="cheez-it box", prompt_fallbacks=["the box", "red box", "cereal box"], subprocess_gpu=TF_GPU), "note": "CAVEAT: success=teleop demos, failure=ACT rollouts (domain confound)", }, "Stack-Can": { "success": pick(f"{SC}/tableHF_ep[0-9][0-9][0-9].mp4", N_EACH), "failure": pick(f"{SC}/table2_failure_ep*.mp4", N_EACH), "robo_task": "A humanoid robot stacks a can onto the plate.", "tf": dict(pointr_ckpt=stackcan_ckpt(), prompt="red can", prompt_fallbacks=["can", "orange can"], mark_prompt="white round plate", mark_prompt_fallbacks=["white circle", "white disc", "white object on the table"], subprocess_gpu=TF_GPU), "note": "CAVEAT: success=table1 teleop, failure=table2 ACT rollouts (table+domain confound)", }, } def auc(s, f): s, f = np.asarray(s), np.asarray(f) if len(s) == 0 or len(f) == 0: return float("nan") wins = sum((sv > fv) + 0.5 * (sv == fv) for sv in s for fv in f) return float(wins) / (len(s) * len(f)) def main(): ap = argparse.ArgumentParser() ap.add_argument("--task", required=True) a = ap.parse_args() cfg = TASKS()[a.task] out_dir = "/home/hongyi/PWM/reward_model/comparison" os.makedirs(out_dir, exist_ok=True) out = f"{out_dir}/robo_vs_trackflow_{a.task}.json" robo = RobometerScorer(host="localhost", port=8801) if hasattr(robo, "wait_until_ready"): robo.wait_until_ready(timeout=600) tf = TrackFlowScorer(**cfg["tf"], verbose=False) print(f"\n===== {a.task} (TF_GPU={TF_GPU}) ===== {cfg['note']}", flush=True) print(f" {len(cfg['success'])} success, {len(cfg['failure'])} failure", flush=True) rows = [] for label, vids in (("success", cfg["success"]), ("failure", cfg["failure"])): for v in vids: rec = {"video": os.path.basename(v), "label": label} try: rp, rs = robo.score(v, task=cfg["robo_task"]) rec["robo_prog_max"] = float(np.asarray(rp).max()) if len(rp) else 0.0 rec["robo_succ_max"] = float(np.asarray(rs).max()) if len(rs) else 0.0 except Exception as e: rec["robo_succ_max"] = None; rec["robo_err"] = str(e)[:120] try: tp = tf.score(v, mode="progress_gated", k_pixel=K_PIXEL) # DISTANCE-WEIGHTED rec["track_max"] = float(np.asarray(tp).max()) if len(tp) else 0.0 except Exception as e: rec["track_max"] = None; rec["track_err"] = str(e)[:120] rows.append(rec) print(f" [{label[:4]}] {rec['video'][:40]:40s} " f"robo_succ={rec.get('robo_succ_max')} track={rec.get('track_max')}", flush=True) def vals(key, lab): return [r[key] for r in rows if r["label"] == lab and r.get(key) is not None] agg = {} for scorer, key in (("robometer_success", "robo_succ_max"), ("robometer_progress", "robo_prog_max"), ("trackflow_gated", "track_max")): s, f = vals(key, "success"), vals(key, "failure") agg[scorer] = {"mean_success": round(float(np.mean(s)), 3) if s else None, "mean_failure": round(float(np.mean(f)), 3) if f else None, "auc": round(auc(s, f), 3), "n_succ": len(s), "n_fail": len(f)} result = {"task": a.task, "note": cfg["note"], "k_pixel": K_PIXEL, "tf_mode": "progress_gated", "agg": agg, "rows": rows} with open(out, "w") as fp: json.dump(result, fp, indent=2) print(f" --> robometer(success) AUC={agg['robometer_success']['auc']} " f"trackflow(gated) AUC={agg['trackflow_gated']['auc']} -> {out}", flush=True) if __name__ == "__main__": try: main() except Exception: traceback.print_exc()