training/Policy/evaluate_policy.py

#!/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()