training/Policy/_balance_eval.py

"""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,
    )