training/Policy/train_danger_head.py

"""VLAlert-X v2 Phase 3 — train Danger Head on dual-stream cache.

Per-frame BCE on continuous danger label + clip-level BCE on max-frame target.
5 seeds × 50 epochs with cosine LR + early stop on best val AUC.

Usage:
    python -m training.Policy.train_danger_head \
        --train_cache data/belief_cache_v2/sft_x_v2__train.pt \
        --val_cache   data/belief_cache_v2/sft_x_v2__val.pt \
        --out_dir     checkpoints/danger_v2 \
        --epochs 50 --seed 0
"""
from __future__ import annotations

import argparse
import json
import logging
import math
import random
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score, average_precision_score
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from lkalert.models.danger_head import DangerHead, danger_loss

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("train_danger_v2")


def set_seed(s: int) -> None:
    random.seed(s)
    np.random.seed(s)
    torch.manual_seed(s)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(s)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def slice_belief_layers(belief: torch.Tensor,
                          subset: Optional[List[int]],
                          n_total_layers: int = 4) -> torch.Tensor:
    """Slice the concat'd belief tensor [..., n_total_layers * D_per] along the
    last dim, keeping only the layers indexed by `subset` (in {0..n_total_layers-1}).

    Layer order in the cache (per make_cache_x_v2.py:298-302) matches the
    `belief_layers` tuple stored in cache metadata — typically (20, 24, 28, 32),
    so index 0 = lowest layer, index 3 = highest layer.

    If `subset is None` or equals the full set, returns input unchanged.
    """
    if subset is None or sorted(subset) == list(range(n_total_layers)):
        return belief
    D_total = belief.shape[-1]
    if D_total % n_total_layers != 0:
        raise ValueError(
            f"belief_content last-dim {D_total} not divisible by "
            f"n_total_layers={n_total_layers}; cache may have a different layer count.")
    D_per = D_total // n_total_layers
    bad = [i for i in subset if not (0 <= i < n_total_layers)]
    if bad:
        raise ValueError(f"subset indices out of range [0,{n_total_layers}): {bad}")
    parts = [belief[..., i * D_per:(i + 1) * D_per] for i in sorted(subset)]
    return torch.cat(parts, dim=-1)


class BeliefCacheDataset(Dataset):
    """Loads belief_content (and optionally policy_position) + danger labels."""

    def __init__(self, cache_path: Path,
                 belief_layer_subset: Optional[List[int]] = None):
        d = torch.load(cache_path, weights_only=False, map_location="cpu")
        belief_raw = d["belief_content"].float()                 # [N, 8, D_total]
        # Per-layer slice if subset given (cheap; no cache re-extraction).
        # Cache metadata records the layer indices these dims came from.
        cache_layers = d.get("belief_layers", [20, 24, 28, 32])
        n_total = len(cache_layers)
        self.belief_content = slice_belief_layers(
            belief_raw, belief_layer_subset, n_total_layers=n_total)
        self.belief_layer_subset = belief_layer_subset
        self.cache_layers = list(cache_layers)
        self.valid_frames   = d["valid_frames"]                 # [N, 8] bool
        self.danger_pf      = d["danger_pf"].float()            # [N, 8] in [0,1]
        self.tick_action    = d["tick_action"].long()           # [N]
        self.actions_pf     = d["actions_pf"].long()            # [N, 8]
        self.n = self.belief_content.shape[0]
        if belief_layer_subset is not None:
            kept = [cache_layers[i] for i in sorted(belief_layer_subset)]
            logger.info(f"  loaded {cache_path}  N={self.n}  "
                        f"belief={tuple(self.belief_content.shape)}  "
                        f"(subset idx={sorted(belief_layer_subset)} → layers={kept})")
        else:
            logger.info(f"  loaded {cache_path}  N={self.n}  "
                        f"belief={tuple(self.belief_content.shape)}")

    def __len__(self):
        return self.n

    def __getitem__(self, i):
        return {
            "belief_content":  self.belief_content[i],
            "valid_frames":    self.valid_frames[i],
            "danger_pf":       self.danger_pf[i],
            "tick_action":     self.tick_action[i],
            "actions_pf":      self.actions_pf[i],
        }


def collate(batch):
    return {k: torch.stack([b[k] for b in batch]) for k in batch[0]}


@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    all_pf_pred, all_pf_target, all_pf_mask = [], [], []
    all_clip_pred, all_clip_target = [], []
    for b in loader:
        bc   = b["belief_content"].to(device)
        v    = b["valid_frames"].to(device)
        d    = b["danger_pf"].to(device)
        out  = model(bc, valid_frames=v)
        all_pf_pred.append(out["per_frame"].cpu().numpy())
        all_pf_target.append(d.cpu().numpy())
        all_pf_mask.append(v.cpu().numpy())
        all_clip_pred.append(out["clip"].cpu().numpy())
        all_clip_target.append(d.max(dim=1).values.cpu().numpy())
    pf_p = np.concatenate(all_pf_pred).flatten()
    pf_t = np.concatenate(all_pf_target).flatten()
    pf_m = np.concatenate(all_pf_mask).flatten()
    pf_p, pf_t = pf_p[pf_m], pf_t[pf_m]
    # binarize target at 0.5 for AUC/AP (continuous label, threshold mid)
    pf_t_bin = (pf_t >= 0.5).astype(np.int32)
    clip_p = np.concatenate(all_clip_pred)
    clip_t = np.concatenate(all_clip_target)
    clip_t_bin = (clip_t >= 0.5).astype(np.int32)
    metrics = {
        "per_frame_mse": float(((pf_p - pf_t) ** 2).mean()),
        "per_frame_mae": float(np.abs(pf_p - pf_t).mean()),
        "clip_mse":      float(((clip_p - clip_t) ** 2).mean()),
    }
    # AUC / AP defined only if both classes present; wrap in try/except to
    # tolerate sklearn edge-cases (rare IndexError seen on N~233k flat arrays).
    if 0 < pf_t_bin.sum() < len(pf_t_bin):
        try:
            metrics["per_frame_auc"] = float(roc_auc_score(pf_t_bin, pf_p))
        except (IndexError, ValueError) as e:
            logger.warning(f"per_frame_auc failed: {e}; falling back to AP")
        try:
            metrics["per_frame_ap"] = float(average_precision_score(pf_t_bin, pf_p))
        except (IndexError, ValueError) as e:
            logger.warning(f"per_frame_ap failed: {e}")
    if 0 < clip_t_bin.sum() < len(clip_t_bin):
        try:
            metrics["clip_auc"] = float(roc_auc_score(clip_t_bin, clip_p))
        except (IndexError, ValueError) as e:
            logger.warning(f"clip_auc failed: {e}")
        try:
            metrics["clip_ap"] = float(average_precision_score(clip_t_bin, clip_p))
        except (IndexError, ValueError) as e:
            logger.warning(f"clip_ap failed: {e}")
    # If per-frame AUC and AP both failed (sklearn edge case on N~233k),
    # fall back to clip_auc → clip_ap → -MSE so training can still save ckpts.
    if "per_frame_auc" not in metrics:
        if "per_frame_ap" in metrics:
            metrics["per_frame_auc"] = metrics["per_frame_ap"]
        elif "clip_auc" in metrics:
            metrics["per_frame_auc"] = metrics["clip_auc"]
        elif "clip_ap" in metrics:
            metrics["per_frame_auc"] = metrics["clip_ap"]
        else:
            # last resort: invert MSE (smaller = better → larger surrogate)
            metrics["per_frame_auc"] = 1.0 - min(1.0, metrics.get("per_frame_mse", 1.0))
    return metrics


def train(args):
    set_seed(args.seed)
    args.out_dir = Path(args.out_dir)
    args.out_dir.mkdir(parents=True, exist_ok=True)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Parse layer subset CSV ("0,1,2,3") into a list[int], or None for full.
    subset: Optional[List[int]] = None
    if args.belief_layer_subset and args.belief_layer_subset.lower() != "all":
        subset = sorted({int(x) for x in args.belief_layer_subset.split(",")
                          if x.strip()})

    train_ds = BeliefCacheDataset(args.train_cache, belief_layer_subset=subset)
    val_ds   = BeliefCacheDataset(args.val_cache,   belief_layer_subset=subset)
    in_dim = int(train_ds.belief_content.shape[-1])
    logger.info(f"  in_dim (BELIEF_CONTENT) = {in_dim}")

    train_loader = DataLoader(train_ds, batch_size=args.batch_size,
                                shuffle=True, num_workers=2,
                                collate_fn=collate, pin_memory=True)
    val_loader   = DataLoader(val_ds, batch_size=args.batch_size,
                                shuffle=False, num_workers=2,
                                collate_fn=collate, pin_memory=True)

    model = DangerHead(in_dim=in_dim, hidden=args.hidden,
                         k_queries=args.k_queries,
                         dropout=args.dropout).to(device)
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f"  model params: {n_params/1e6:.2f} M")

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr,
                              weight_decay=args.weight_decay)
    n_steps = math.ceil(len(train_loader) * args.epochs)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=n_steps)

    best_metric = -1.0
    best_epoch  = -1
    epochs_no_improve = 0
    log: List[Dict] = []

    for ep in range(args.epochs):
        model.train()
        running = 0.0
        running_frame = 0.0
        running_clip = 0.0
        n_batch = 0
        pbar = tqdm(train_loader, ncols=80, desc=f"ep{ep}")
        for b in pbar:
            bc   = b["belief_content"].to(device, non_blocking=True)
            v    = b["valid_frames"].to(device, non_blocking=True)
            d    = b["danger_pf"].to(device, non_blocking=True)
            out  = model(bc, valid_frames=v)
            losses = danger_loss(out, d, valid_frames=v,
                                   w_clip=args.w_clip)
            losses["loss"].backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
            running       += losses["loss"].item()
            running_frame += losses["frame_loss"].item()
            running_clip  += losses["clip_loss"].item()
            n_batch += 1
            pbar.set_postfix(loss=running / max(1, n_batch),
                              lr=sched.get_last_lr()[0])

        val_metrics = evaluate(model, val_loader, device)
        record = {
            "epoch":             ep,
            "train_loss":        running / max(1, n_batch),
            "train_frame_loss":  running_frame / max(1, n_batch),
            "train_clip_loss":   running_clip / max(1, n_batch),
            "val":               val_metrics,
        }
        log.append(record)
        # Selection metric: "higher is better" by default; for MSE-family metrics
        # we negate so the same > comparison works.
        sel = args.selection_metric
        raw = val_metrics.get(sel, None)
        if raw is None:
            score = val_metrics.get("per_frame_auc", 0.0)
        elif sel.endswith("_mse") or sel.endswith("_mae"):
            score = -float(raw)  # lower is better
        else:
            score = float(raw)
        logger.info(f"[ep {ep}] " + json.dumps({
            "train_loss": f"{record['train_loss']:.4f}",
            **{k: f"{v:.4f}" if isinstance(v, float) else v
                 for k, v in val_metrics.items()}}))

        if score > best_metric:
            best_metric = score
            best_epoch  = ep
            epochs_no_improve = 0
            torch.save({"model": model.state_dict(),
                         "args": vars(args),
                         "epoch": ep,
                         "val_metrics": val_metrics,
                         "in_dim": in_dim,
                         "belief_layer_subset": subset,
                         "cache_layers": train_ds.cache_layers},
                         args.out_dir / "best.pt")
        else:
            epochs_no_improve += 1
            if epochs_no_improve >= args.patience:
                logger.info(f"[stop] no improvement for {args.patience} epochs")
                break

    (args.out_dir / "training_log.json").write_text(json.dumps(log, indent=2))
    logger.info(f"[done] best val_per_frame_auc = {best_metric:.4f} @ epoch {best_epoch}")
    logger.info(f"       ckpt: {args.out_dir / 'best.pt'}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train_cache", type=Path,
                    default=ROOT / "data/belief_cache_v2/sft_x_v2__train.pt")
    ap.add_argument("--val_cache",   type=Path,
                    default=ROOT / "data/belief_cache_v2/sft_x_v2__val.pt")
    ap.add_argument("--out_dir",     type=Path, required=True)
    ap.add_argument("--epochs",      type=int, default=50)
    ap.add_argument("--batch_size",  type=int, default=128)
    ap.add_argument("--lr",          type=float, default=3e-4)
    ap.add_argument("--weight_decay", type=float, default=1e-4)
    ap.add_argument("--hidden",      type=int, default=512)
    ap.add_argument("--k_queries",   type=int, default=4)
    ap.add_argument("--dropout",     type=float, default=0.2)
    ap.add_argument("--w_clip",      type=float, default=0.5)
    ap.add_argument("--patience",    type=int, default=10)
    ap.add_argument("--seed",        type=int, default=0)
    ap.add_argument("--belief_layer_subset", type=str, default="",
                    help="CSV of layer indices into cache_layers (e.g. '0,1,2,3' "
                          "or '3' or '2,3'); empty/'all' = use full concat. "
                          "For default cache (20,24,28,32): idx 0=L20, 3=L32.")
    ap.add_argument("--selection_metric", type=str, default="per_frame_auc",
                    help="Val metric used for best-ckpt selection. Options: "
                          "per_frame_auc, per_frame_ap, clip_auc, clip_ap, "
                          "per_frame_mse, per_frame_mae, clip_mse (mse/mae are "
                          "negated internally for higher-is-better).")
    args = ap.parse_args()
    train(args)


if __name__ == "__main__":
    main()