#!/usr/bin/env python3 """ Standalone policy evaluation — detailed action breakdown + confusion matrix. Usage: # Evaluate a trained PolicyHead checkpoint: python -m training.Policy.evaluate_policy \ --sft_checkpoint checkpoints/SFT/sft_v2/best \ --policy_checkpoint checkpoints/Policy/policy_warmstart_v1/best \ --label_dir data/policy_labels # Evaluate the SFT baseline (untrained PolicyHead, random init): python -m training.Policy.evaluate_policy \ --sft_checkpoint checkpoints/SFT/sft_v2/best \ --label_dir data/policy_labels # (omit --policy_checkpoint to test random-init baseline) """ from __future__ import annotations import argparse import json import logging from collections import defaultdict from pathlib import Path from typing import Dict, List, Optional import numpy as np import torch from torch.utils.data import DataLoader from tqdm import tqdm import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from .policy_model import PolicyModel, ACTION_NAMES, N_ACTIONS from .policy_dataset import PolicyDataset, policy_collate_fn from .warm_start_trainer import compute_policy_score, _ratio logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Policy.evaluate") @torch.no_grad() def evaluate(model: PolicyModel, loader: DataLoader) -> dict: model.eval() cat_preds: Dict[str, List[int]] = defaultdict(list) cat_labels: Dict[str, List[int]] = defaultdict(list) cat_ttas: Dict[str, List[float]] = defaultdict(list) for batch in tqdm(loader, desc="Evaluating"): if "beliefs" in batch: logits = model.forward_cached(batch["beliefs"], batch["tta_means"], batch["tta_vars"]) else: logits = model(batch["images"], batch["metadata"]) preds = logits.argmax(dim=-1).cpu().tolist() for p, l, tta, cat in zip( preds, batch["action_labels"].tolist(), batch["tta_raws"].tolist(), batch["categories"], ): cat_preds[cat].append(p) cat_labels[cat].append(l) cat_ttas[cat].append(tta) # ── per-category action distribution ────────────────────────────────────── sep = "=" * 68 print(f"\n{sep}") print(f" Action distribution (predicted)") print(f" {'Category':<22} {'SILENT':>8} {'OBSERVE':>8} {'ALERT':>8} {'N':>7}") print("-" * 68) for cat in sorted(cat_preds): ps = cat_preds[cat] n = len(ps) s = _ratio(sum(1 for p in ps if p == 0), n) o = _ratio(sum(1 for p in ps if p == 1), n) a = _ratio(sum(1 for p in ps if p == 2), n) print(f" {cat:<22} {s:>8.3f} {o:>8.3f} {a:>8.3f} {n:>7}") print(sep) # ── ego metrics ─────────────────────────────────────────────────────────── ego_ps = cat_preds.get("ego_positive", []) ego_ls = cat_labels.get("ego_positive", []) ego_ts = cat_ttas.get("ego_positive", []) alert_ps = [p for p, l in zip(ego_ps, ego_ls) if l == 2] obs_ps = [p for p, l in zip(ego_ps, ego_ls) if l == 1] ego_alert_recall = _ratio(sum(1 for p in alert_ps if p == 2), len(alert_ps)) ego_observe_rate = _ratio(sum(1 for p in obs_ps if p == 1), len(obs_ps)) # ── non-ego metrics ─────────────────────────────────────────────────────── ne_ps = cat_preds.get("non_ego", []) non_ego_alert_rate = _ratio(sum(1 for p in ne_ps if p == 2), len(ne_ps)) non_ego_noalert_rate = 1.0 - non_ego_alert_rate non_ego_observe_rate = _ratio(sum(1 for p in ne_ps if p == 1), len(ne_ps)) # ── safe-neg metrics ────────────────────────────────────────────────────── sn_ps = cat_preds.get("safe_neg", []) safe_neg_silent_rate = _ratio(sum(1 for p in sn_ps if p == 0), len(sn_ps)) safe_neg_alert_leak = _ratio(sum(1 for p in sn_ps if p == 2), len(sn_ps)) # ── overall ─────────────────────────────────────────────────────────────── all_p = [p for ps in cat_preds.values() for p in ps] all_l = [l for ls in cat_labels.values() for l in ls] overall_acc = _ratio(sum(p == l for p, l in zip(all_p, all_l)), len(all_p)) score = compute_policy_score( ego_alert_recall = ego_alert_recall, safe_neg_silent_rate = safe_neg_silent_rate, safe_neg_alert_rate = safe_neg_alert_leak, ) # ── ego TTA-bucket ALERT rate ───────────────────────────────────────────── tta_buckets = [ ("[1.5, 2.0)", 1.5, 2.0), ("[2.0, 3.0)", 2.0, 3.0), ("[3.0, 4.0)", 3.0, 4.0), ("[4.0, 5.0)", 4.0, 5.0), ("[5.5, 8.0]", 5.5, 8.01), ("(8.0, +∞) ", 8.0, 1e9), ] print("\n Ego TTA-bucket ALERT rate (timing calibration):") print(f" {'TTA range':<15} {'ALERT':>8} {'OBSERVE':>8} {'SILENT':>8} {'N':>6}") print("-" * 52) for bname, lo, hi in tta_buckets: bucket_ps = [p for p, tta in zip(ego_ps, ego_ts) if lo <= tta < hi] if not bucket_ps: continue n = len(bucket_ps) a = _ratio(sum(1 for p in bucket_ps if p == 2), n) o = _ratio(sum(1 for p in bucket_ps if p == 1), n) s = _ratio(sum(1 for p in bucket_ps if p == 0), n) print(f" TTA {bname:<11} {a:>8.3f} {o:>8.3f} {s:>8.3f} {n:>6}") # ── summary ─────────────────────────────────────────────────────────────── print(f"\n{sep}") print(f" SUMMARY") print(f" ego_alert_recall : {ego_alert_recall:.4f} " f"(n_label_ALERT={len(alert_ps)})") print(f" ego_observe_rate : {ego_observe_rate:.4f} " f"(n_label_OBSERVE={len(obs_ps)})") print(f" non_ego_noalert_rate : {non_ego_noalert_rate:.4f} " f"(non_ego_alert={non_ego_alert_rate:.4f}, n={len(ne_ps)})") print(f" non_ego_observe_rate : {non_ego_observe_rate:.4f}") print(f" safe_neg_silent_rate : {safe_neg_silent_rate:.4f} " f"(alert_leak={safe_neg_alert_leak:.4f}, n={len(sn_ps)})") print(f" overall_acc : {overall_acc:.4f}") print(f" ★ policy_score : {score:.4f}") # ── confusion matrix ────────────────────────────────────────────────────── conf = np.zeros((N_ACTIONS, N_ACTIONS), dtype=int) for p, l in zip(all_p, all_l): conf[l][p] += 1 print(f"\n Confusion matrix [row=true_label, col=prediction]:") header = " " + " " * 12 + " ".join(f"pred_{n:6s}" for n in ACTION_NAMES.values()) print(header) for i, n in enumerate(ACTION_NAMES.values()): row = " ".join(f"{conf[i][j]:12d}" for j in range(N_ACTIONS)) print(f" label_{n:8s} {row}") print(sep + "\n") return { "policy_score": score, "ego_alert_recall": ego_alert_recall, "ego_observe_rate": ego_observe_rate, "non_ego_noalert_rate": non_ego_noalert_rate, "non_ego_alert_rate": non_ego_alert_rate, "non_ego_observe_rate": non_ego_observe_rate, "safe_neg_silent_rate": safe_neg_silent_rate, "safe_neg_alert_leak": safe_neg_alert_leak, "overall_acc": overall_acc, "confusion_matrix": conf.tolist(), "n_ego_alert_windows": len(alert_ps), "n_ego_obs_windows": len(obs_ps), "n_non_ego": len(ne_ps), "n_safe_neg": len(sn_ps), } def main(): parser = argparse.ArgumentParser("evaluate_policy") parser.add_argument("--sft_checkpoint", required=True) parser.add_argument("--policy_checkpoint", default=None, help="Dir with policy_head.pt. Omit to test random-init baseline.") parser.add_argument("--label_dir", default="data/policy_labels") parser.add_argument("--split", default="val", choices=["train", "val"]) parser.add_argument("--batch_size", type=int, default=256) parser.add_argument("--belief_cache_dir", default=None, help="Dir with {split}.pt belief cache. Much faster than image mode.") parser.add_argument("--output_json", default=None) args = parser.parse_args() model = PolicyModel( sft_checkpoint_dir = args.sft_checkpoint, use_bf16 = True, ) if args.policy_checkpoint is not None: model.load_policy_checkpoint(args.policy_checkpoint) logger.info(f"Evaluating trained PolicyHead from: {args.policy_checkpoint}") else: logger.info("No policy_checkpoint provided — evaluating random-init PolicyHead (baseline).") belief_cache_path = None if args.belief_cache_dir is not None: belief_cache_path = Path(args.belief_cache_dir) / f"{args.split}.pt" ds = PolicyDataset( manifests = [Path(args.label_dir) / f"{args.split}.json"], split = args.split, belief_cache_path = belief_cache_path, ) loader = DataLoader( ds, batch_size=args.batch_size, shuffle=False, num_workers=2, collate_fn=policy_collate_fn, ) metrics = evaluate(model, loader) if args.output_json: with open(args.output_json, "w") as f: json.dump(metrics, f, indent=2) logger.info(f"Metrics saved → {args.output_json}") if __name__ == "__main__": main()