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