File size: 6,466 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """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,
)
|