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