"""Phase D shared utility — Universal balance gate. Used by: - Head-RL trainers (train_head_{dpo,kto,ppo}.py) for in-loop validation - tools/balance_gate_eval.py for PASS/FAIL aggregation Universal balance gate (NeurIPS / VLAlert-X paper requirement): r_OBSERVE >= 0.20 AND r_ALERT >= 0.70 AND r_SILENT >= 0.85 AND AP(ALERT) >= 0.85 AND AUROC(HAZARD) >= 0.60 AND FP rate on safe_neg <= 0.15 """ from __future__ import annotations import sys from pathlib import Path from typing import Optional import numpy as np import torch import torch.nn.functional as F from sklearn.metrics import (average_precision_score, roc_auc_score, confusion_matrix) ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(ROOT)) GATE = { "r_OBSERVE_min": 0.20, "r_ALERT_min": 0.70, "r_SILENT_min": 0.85, "AP_alert_min": 0.85, "AUROC_hazard_min": 0.60, "FP_safe_neg_max": 0.15, } @torch.no_grad() def predict_val_probs(policy, danger_head, val_cache, device, batch_size=256): """Forward all val samples through DangerHead + PolicyHead. Returns softmax probs [N, 3] as a numpy array. """ policy.eval(); danger_head.eval() N = len(val_cache["tick_action"]) out = np.zeros((N, 3), dtype=np.float32) for i in range(0, N, batch_size): bc = val_cache["belief_content"][i:i+batch_size].to(device, dtype=torch.float32) v = val_cache["valid_frames"][i:i+batch_size].to(device) pp = val_cache["policy_position"][i:i+batch_size].to(device, dtype=torch.float32) prev = torch.full((pp.shape[0],), 3, dtype=torch.long, device=device) dh_out = danger_head(bc, valid_frames=v) logits = policy(pp, dh_out["perception_summary"], dh_out["per_frame"], prev, valid_frames=v) out[i:i+pp.shape[0]] = F.softmax(logits.float(), dim=-1).cpu().numpy() return out def decode_argmax(probs): return probs.argmax(axis=-1) def decode_threshold(probs, tau_obs: float = 0.20, tau_alert: float = 0.40): """OBSERVE-first decoder: predict OBS if P(OBS) > tau_obs AND P(ALR) < tau_alert, else argmax. Used at eval time after calibrating tau.""" pred = probs.argmax(axis=-1) obs_gate = (probs[:, 1] > tau_obs) & (probs[:, 2] < tau_alert) pred[obs_gate] = 1 return pred def compute_gate_metrics(probs, tick_action, category=None, source=None, tta_raw=None, decode_mode="argmax", tau_obs=0.20, tau_alert=0.40): """Compute all metrics needed for the universal balance gate. probs: [N, 3] float — softmax over (SILENT, OBSERVE, ALERT) tick_action: [N] int — ground-truth tick action category: [N] str (optional) — 'safe_neg', 'ego_positive', 'non_ego' Returns a dict of {metric_name: value, "PASS_gate": bool, ...} """ y_3 = np.asarray(tick_action) N = len(y_3) if decode_mode == "threshold": pred = decode_threshold(probs, tau_obs=tau_obs, tau_alert=tau_alert) else: pred = decode_argmax(probs) cm = confusion_matrix(y_3, pred, labels=[0, 1, 2]) rec = cm.diagonal() / cm.sum(axis=1).clip(min=1) r_sil, r_obs, r_alr = float(rec[0]), float(rec[1]), float(rec[2]) P_alert = probs[:, 2] P_hazard = 1.0 - probs[:, 0] # ALERT-binary: positive iff tick_action == ALERT (deployment metric) y_alert = (y_3 == 2).astype(int) # HAZARD-binary: positive iff tick_action != SILENT (capability metric) y_hazard = (y_3 != 0).astype(int) ap_alert = float(average_precision_score(y_alert, P_alert)) au_alert = float(roc_auc_score(y_alert, P_alert)) ap_hazard = float(average_precision_score(y_hazard, P_hazard)) au_hazard = float(roc_auc_score(y_hazard, P_hazard)) # FP rate on safe_neg fp_safe_neg = float("nan") if category is not None: cat_arr = np.asarray(category) sn_mask = (cat_arr == "safe_neg") if sn_mask.sum() > 0: fp_safe_neg = float((pred[sn_mask] == 2).mean()) # Balanced accuracy (mean of recalls) val_bal = (r_sil + r_obs + r_alr) / 3.0 # Composite metric (used for best-ckpt selection during training) composite = 0.4 * val_bal + 0.3 * ap_alert + 0.3 * au_hazard # PASS gate passes = ( r_obs >= GATE["r_OBSERVE_min"] and r_alr >= GATE["r_ALERT_min"] and r_sil >= GATE["r_SILENT_min"] and ap_alert >= GATE["AP_alert_min"] and au_hazard >= GATE["AUROC_hazard_min"] and (np.isnan(fp_safe_neg) or fp_safe_neg <= GATE["FP_safe_neg_max"]) ) return { "N": N, "r_SILENT": r_sil, "r_OBSERVE": r_obs, "r_ALERT": r_alr, "val_balanced_acc": val_bal, "AP_alert": ap_alert, "AUROC_alert": au_alert, "AP_hazard": ap_hazard, "AUROC_hazard": au_hazard, "FP_safe_neg": fp_safe_neg, "composite": composite, "argmax_dist": np.bincount(pred, minlength=3).tolist(), "tick_action_dist": np.bincount(y_3, minlength=3).tolist(), "PASS_gate": bool(passes), "decode_mode": decode_mode, "tau_obs": tau_obs, "tau_alert": tau_alert, } def format_gate_row(m: dict, tag: str = "") -> str: """One-line summary string for logging.""" pass_str = "PASS" if m["PASS_gate"] else "FAIL" fp = m["FP_safe_neg"] fp_s = f"{fp:.3f}" if not np.isnan(fp) else "N/A" return (f"[{pass_str}] {tag} r_SIL={m['r_SILENT']:.3f} r_OBS={m['r_OBSERVE']:.3f} " f"r_ALR={m['r_ALERT']:.3f} AP_alr={m['AP_alert']:.4f} " f"AUR_haz={m['AUROC_hazard']:.4f} FP_safe={fp_s} " f"composite={m['composite']:.4f}") def evaluate_policy_on_val(policy, danger_head, val_cache, device, batch_size=256, decode_mode="argmax", tau_obs=0.20, tau_alert=0.40): """Convenience: forward + gate metrics in one call.""" probs = predict_val_probs(policy, danger_head, val_cache, device, batch_size) return compute_gate_metrics( probs, tick_action=val_cache["tick_action"].numpy(), category=val_cache.get("category", None), source=val_cache.get("source", None), tta_raw=(val_cache.get("tick_tta_raw", None).numpy() if val_cache.get("tick_tta_raw", None) is not None else None), decode_mode=decode_mode, tau_obs=tau_obs, tau_alert=tau_alert, )