training/Policy/diagnose_nexar_behavior.py

#!/usr/bin/env python3
"""
Diagnostic: what is the policy head actually DOING on Nexar?

Three failure modes we want to distinguish:
  (a) model is asleep  — predicts SILENT even as collision approaches
  (b) stuck in OBSERVE — never commits to ALERT even at TTA < 0.5s
  (c) late ALERT       — ALERT fires only when TTA is very small (bad driver UX)

Output:
  • Overall predicted-class distribution on Nexar val
  • Per-TTA-bucket predicted distribution for ego_positive samples
    (shows how prediction evolves as collision nears)
  • First-ALERT-time statistics: the TTA at which ALERT is first predicted
"""
from __future__ import annotations
import argparse
import json
from collections import Counter
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from training.Policy.temporal_trainer import (
    TemporalPolicyDataset, TemporalPolicyModel, temporal_collate_fn,
)
from training.Policy.trajectory_trainer import (
    TrajectoryPolicyDataset, TrajectoryPolicyModel, trajectory_collate_fn,
)

ACTION = {0: "SIL", 1: "OBS", 2: "ALR"}


def load_ckpt(ckpt_dir, hidden_dim, seq_len):
    meta = json.loads((ckpt_dir / "policy_meta.json").read_text())
    v = meta.get("version", "")
    if "trajectory" in v or "v7" in v:
        m = TrajectoryPolicyModel(hidden_dim=hidden_dim, seq_len=seq_len,
                                  use_gru=meta.get("use_gru", True),
                                  belief_noise_std=0.0)
        m.load_policy_checkpoint(str(ckpt_dir))
        return m, True, meta
    m = TemporalPolicyModel(hidden_dim=hidden_dim, seq_len=seq_len)
    m.load_policy_checkpoint(str(ckpt_dir))
    return m, False, meta


@torch.no_grad()
def run(model, loader, is_traj):
    model.eval()
    probs = []
    for b in tqdm(loader, desc="infer", ncols=80):
        if is_traj:
            lo, _ = model(b["belief_seqs"], b["tta_mean_seqs"], b["tta_var_seqs"])
        else:
            lo = model(b["belief_seqs"], b["tta_mean_seqs"], b["tta_var_seqs"])
        probs.append(F.softmax(lo, dim=-1).cpu().numpy())
    return np.concatenate(probs, 0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="checkpoints/Policy/temporal_long_mono/best")
    ap.add_argument("--label_dir", default="data/policy_labels")
    ap.add_argument("--cache_dir", default="data/belief_cache")
    ap.add_argument("--batch_size", type=int, default=512)
    ap.add_argument("--alert_bias", type=float, default=0.0,
                    help="decision-time bias added to P(ALERT) before argmax")
    args = ap.parse_args()

    ckpt_dir = Path(args.ckpt)
    meta = json.loads((ckpt_dir / "policy_meta.json").read_text())
    seq_len = meta.get("seq_len", 8)
    is_traj_meta = "trajectory" in meta.get("version", "") or "v7" in meta.get("version", "")

    ds_cls  = TrajectoryPolicyDataset if is_traj_meta else TemporalPolicyDataset
    collate = trajectory_collate_fn  if is_traj_meta else temporal_collate_fn
    val_ds = ds_cls(
        manifests=[Path(args.label_dir) / "val.json"],
        split="val",
        belief_cache_path=Path(args.cache_dir) / "val.pt",
        seq_len=seq_len,
    )
    loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
                        collate_fn=collate, num_workers=2, pin_memory=True)

    hidden_dim = val_ds._cache["beliefs"].shape[-1]
    model, is_traj, _ = load_ckpt(ckpt_dir, hidden_dim, seq_len)
    probs = run(model, loader, is_traj)  # [N, 3]

    # Apply the same alert_bias used at deployment for argmax
    adj = probs.copy()
    adj[:, 2] += args.alert_bias
    pred = adj.argmax(axis=1)

    labels   = np.array([s["action_label"] for s in val_ds.samples])
    ttas     = np.array([s["tta_raw"]     for s in val_ds.samples])
    cats     = np.array([s["category"]    for s in val_ds.samples])
    sources  = np.array([s.get("source", "?") for s in val_ds.samples])
    videos   = np.array([s["video_id"]    for s in val_ds.samples])

    nexar = sources == "nexar"
    print(f"\n══ {ckpt_dir}  (alert_bias={args.alert_bias}) ══")
    print(f"Nexar val: {int(nexar.sum())} samples "
          f"({int((labels[nexar]==2).sum())} true ALERT, "
          f"{int((labels[nexar]==0).sum())} SILENT, "
          f"{int((labels[nexar]==1).sum())} OBSERVE)\n")

    # ── (1) global prediction distribution, Nexar only ───────────────────────
    def pct(mask_sub):
        n = int(mask_sub.sum())
        if n == 0: return "—"
        c = Counter(pred[mask_sub].tolist())
        return " ".join(f"{ACTION[k]}={c.get(k,0)/n*100:5.1f}%" for k in (0,1,2))

    print("──────── Nexar predicted-class mix by true label ────────")
    print(f"  SILENT  true (n={int((labels[nexar]==0).sum())}):  {pct(nexar & (labels==0))}")
    print(f"  OBSERVE true (n={int((labels[nexar]==1).sum())}):  {pct(nexar & (labels==1))}")
    print(f"  ALERT   true (n={int((labels[nexar]==2).sum())}):  {pct(nexar & (labels==2))}")
    print()

    # ── (2) per-TTA bucket for ego_positive (the "collision coming" samples) ─
    ego = nexar & (cats == "ego_positive") & (labels == 2)  # true ALERT
    print(f"──────── Nexar ego_positive ALERT ({int(ego.sum())} samples): prediction vs TTA ────────")
    print(f"  TTA=time-to-collision at obs window (seconds)")
    bins = [(0.0,0.5),(0.5,1.0),(1.0,1.5),(1.5,2.0),(2.0,3.0),(3.0,5.0),(5.0,99.0)]
    print(f"  {'TTA bucket (s)':<14}  {'n':>5}  "
          f"{'P(SILENT)':>10}  {'P(OBSERVE)':>11}  {'P(ALERT)':>10}  "
          f"{'pred: SIL / OBS / ALR':<28}")
    for lo, hi in bins:
        m = ego & (ttas >= lo) & (ttas < hi)
        n = int(m.sum())
        if n == 0:
            continue
        ps = probs[m].mean(axis=0)
        c = Counter(pred[m].tolist())
        mix = f"{c.get(0,0)/n*100:4.0f} / {c.get(1,0)/n*100:4.0f} / {c.get(2,0)/n*100:4.0f}"
        print(f"  [{lo:4.1f},{hi:4.1f})      {n:>5}  "
              f"{ps[0]:>10.3f}  {ps[1]:>11.3f}  {ps[2]:>10.3f}  {mix:<28}")
    print()

    # ── (3) per-video "first-ALERT TTA" distribution ─────────────────────────
    # For each Nexar ego-positive video, find the LATEST tta (= earliest time)
    # at which the model predicted ALERT. Reports lead time.
    ego_vids = np.unique(videos[nexar & (cats == "ego_positive")])
    first_alert_leads = []
    never_alert = 0
    for vid in ego_vids:
        m = (videos == vid) & (labels == 2)
        if not m.any():
            continue
        p = pred[m]
        t = ttas[m]
        if (p == 2).any():
            lead = float(t[p == 2].max())
            first_alert_leads.append(lead)
        else:
            never_alert += 1
    leads = np.array(first_alert_leads) if first_alert_leads else np.array([0.0])
    print(f"──────── Per-ego-collision-video: lead time of FIRST ALERT ────────")
    print(f"  {len(ego_vids)} unique ego collision videos")
    print(f"  videos where ALERT never fired : {never_alert}  ({never_alert/max(len(ego_vids),1)*100:.1f}%)")
    print(f"  videos where ALERT fired       : {len(first_alert_leads)}")
    if first_alert_leads:
        print(f"  lead time (seconds before collision):")
        print(f"    mean   = {leads.mean():.2f}s")
        print(f"    median = {np.median(leads):.2f}s")
        print(f"    p25    = {np.percentile(leads,25):.2f}s")
        print(f"    p75    = {np.percentile(leads,75):.2f}s")
        print(f"    max    = {leads.max():.2f}s")
    print()

    # ── (4) Ever-OBSERVE-stuck? videos that saw OBSERVE but never ALERT ──────
    obs_but_never_alert = 0
    for vid in ego_vids:
        m = (videos == vid) & (labels == 2)
        if not m.any(): continue
        p = pred[m]
        if (p == 1).any() and not (p == 2).any():
            obs_but_never_alert += 1
    print(f"──────── OBSERVE-stuck diagnosis ────────")
    print(f"  ego collision videos that predicted OBSERVE but NEVER ALERT: "
          f"{obs_but_never_alert}  ({obs_but_never_alert/max(len(ego_vids),1)*100:.1f}%)")


if __name__ == "__main__":
    main()