#!/usr/bin/env python3 """ train_nexar_head.py ═══════════════════════════════════════════════════════════════════════════════ Train a binary collision-risk head on cached Qwen3-VL-4B CoT+BeliefToken per-frame visual features (Nexar-only). Input caches (from make_nexar_belief_cache.py): data/belief_cache_nexar_qwen3vl4b/train.pt data/belief_cache_nexar_qwen3vl4b/val.pt Cache layout: beliefs_frame [N, T, D] fp16 valid_frames [N, T] bool beliefs_text [N, D] fp16 labels [N] int64 (0 safe / 1 collision) meta dict (video_ids, hidden_dim, n_frames, ...) Head: A small Transformer encoder over the T frame embeddings + a mean-pool over valid frames, followed by a 2-layer MLP classifier. Keeps the temporal axis (LKAlert's main inductive bias) but stays tiny so 1200 clips don't overfit. Outputs: checkpoints/Nexar/qwen3vl4b_head/best.pt (head weights + meta) checkpoints/Nexar/qwen3vl4b_head/train_log.json Usage ───── python -m training.Policy.train_nexar_head \ --train_cache data/belief_cache_nexar_qwen3vl4b/train.pt \ --val_cache data/belief_cache_nexar_qwen3vl4b/val.pt \ --out_dir checkpoints/Nexar/qwen3vl4b_head \ --epochs 30 --batch_size 64 --lr 3e-4 """ from __future__ import annotations import argparse import json import logging import math from pathlib import Path from typing import Dict import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Policy.train_nexar_head") def binary_ap(y_true: np.ndarray, y_score: np.ndarray) -> float: from sklearn.metrics import average_precision_score if (y_true == 1).sum() == 0 or (y_true == 0).sum() == 0: return float("nan") return float(average_precision_score(y_true, y_score)) def binary_auc(y_true: np.ndarray, y_score: np.ndarray) -> float: from sklearn.metrics import roc_auc_score if (y_true == 1).sum() == 0 or (y_true == 0).sum() == 0: return float("nan") return float(roc_auc_score(y_true, y_score)) class NexarHead(nn.Module): """Temporal encoder over T frame beliefs + binary classifier.""" def __init__(self, hidden_dim: int, proj_dim: int = 512, n_layers: int = 2, n_heads: int = 8, dropout: float = 0.2): super().__init__() self.proj = nn.Linear(hidden_dim, proj_dim) enc_layer = nn.TransformerEncoderLayer( d_model=proj_dim, nhead=n_heads, dim_feedforward=proj_dim * 4, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) self.encoder = nn.TransformerEncoder(enc_layer, num_layers=n_layers) self.cls = nn.Sequential( nn.LayerNorm(proj_dim), nn.Linear(proj_dim, proj_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(proj_dim, 1), ) def forward(self, frames: torch.Tensor, valid: torch.Tensor) -> torch.Tensor: """ frames: [B, T, D] fp valid: [B, T] bool returns: [B] logit """ x = self.proj(frames) # [B, T, P] key_padding_mask = ~valid # True = pad x = self.encoder(x, src_key_padding_mask=key_padding_mask) denom = valid.sum(dim=1, keepdim=True).clamp(min=1).float() pooled = (x * valid.unsqueeze(-1).float()).sum(dim=1) / denom return self.cls(pooled).squeeze(-1) def load_cache(path: str | Path): d = torch.load(path, map_location="cpu", weights_only=False) return d def main(): ap = argparse.ArgumentParser("train_nexar_head") ap.add_argument("--train_cache", required=True) ap.add_argument("--val_cache", required=True) ap.add_argument("--out_dir", required=True) ap.add_argument("--proj_dim", type=int, default=512) ap.add_argument("--n_layers", type=int, default=2) ap.add_argument("--n_heads", type=int, default=8) ap.add_argument("--dropout", type=float, default=0.2) ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--weight_decay", type=float, default=1e-4) ap.add_argument("--epochs", type=int, default=30) ap.add_argument("--batch_size", type=int, default=64) ap.add_argument("--warmup_frac", type=float, default=0.1) ap.add_argument("--pos_weight", type=float, default=0.0, help=">0 overrides auto-balance; 0 = auto (neg/pos ratio)") ap.add_argument("--patience", type=int, default=8, help="Early stop after N epochs w/o val AP improvement") ap.add_argument("--seed", type=int, default=0) args = ap.parse_args() torch.manual_seed(args.seed) np.random.seed(args.seed) out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) logger.info(f"loading train cache: {args.train_cache}") tr = load_cache(args.train_cache) logger.info(f"loading val cache: {args.val_cache}") va = load_cache(args.val_cache) D = int(tr["meta"]["hidden_dim"]) T = int(tr["meta"]["n_frames"]) assert va["meta"]["hidden_dim"] == D and va["meta"]["n_frames"] == T, \ "train/val cache hidden_dim or n_frames mismatch" tr_x = tr["beliefs_frame"].float() # [N, T, D] tr_v = tr["valid_frames"].bool() tr_y = tr["labels"].long() va_x = va["beliefs_frame"].float() va_v = va["valid_frames"].bool() va_y = va["labels"].long() # drop any samples with label -1 (shouldn't happen on train/val) tr_keep = tr_y >= 0; va_keep = va_y >= 0 tr_x, tr_v, tr_y = tr_x[tr_keep], tr_v[tr_keep], tr_y[tr_keep] va_x, va_v, va_y = va_x[va_keep], va_v[va_keep], va_y[va_keep] n_pos = int((tr_y == 1).sum()); n_neg = int((tr_y == 0).sum()) logger.info(f"train: {len(tr_y)} (pos={n_pos} neg={n_neg})") logger.info(f"val: {len(va_y)} (pos={int((va_y==1).sum())} neg={int((va_y==0).sum())})") pos_weight_val = args.pos_weight if args.pos_weight > 0 else (n_neg / max(n_pos, 1)) pos_weight = torch.tensor([pos_weight_val], dtype=torch.float32, device="cuda") logger.info(f"pos_weight = {pos_weight_val:.3f}") tr_ds = TensorDataset(tr_x, tr_v, tr_y) va_ds = TensorDataset(va_x, va_v, va_y) tr_dl = DataLoader(tr_ds, batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True, drop_last=False) va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True, drop_last=False) model = NexarHead(hidden_dim=D, proj_dim=args.proj_dim, n_layers=args.n_layers, n_heads=args.n_heads, dropout=args.dropout).to("cuda") logger.info(f"params: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M") opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.999)) total_steps = args.epochs * len(tr_dl) warmup_steps = max(1, int(total_steps * args.warmup_frac)) def lr_at(step): if step < warmup_steps: return step / warmup_steps p = (step - warmup_steps) / max(1, total_steps - warmup_steps) return 0.5 * (1 + math.cos(math.pi * p)) sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_at) log = {"epochs": []} best_ap = -1.0; best_state = None; best_epoch = -1; no_imp = 0 for epoch in range(args.epochs): model.train() tr_loss_sum = 0.0; n = 0 for xb, vb, yb in tr_dl: xb = xb.to("cuda", non_blocking=True) vb = vb.to("cuda", non_blocking=True) yb = yb.to("cuda", non_blocking=True).float() logits = model(xb, vb) loss = F.binary_cross_entropy_with_logits(logits, yb, pos_weight=pos_weight) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step(); sched.step() tr_loss_sum += float(loss.item()) * xb.size(0); n += xb.size(0) tr_loss = tr_loss_sum / max(n, 1) model.eval() va_logits = []; va_labels = [] with torch.no_grad(): for xb, vb, yb in va_dl: xb = xb.to("cuda"); vb = vb.to("cuda") va_logits.append(model(xb, vb).cpu()) va_labels.append(yb) va_logits = torch.cat(va_logits).numpy() va_labels = torch.cat(va_labels).numpy() va_prob = 1 / (1 + np.exp(-va_logits)) ap = binary_ap(va_labels, va_prob) auc = binary_auc(va_labels, va_prob) lr_now = opt.param_groups[0]["lr"] logger.info(f"ep{epoch:02d} tr_loss={tr_loss:.4f} " f"val AP={ap:.4f} val AUC={auc:.4f} lr={lr_now:.2e}") log["epochs"].append({"epoch": epoch, "tr_loss": tr_loss, "val_ap": ap, "val_auc": auc, "lr": lr_now}) if ap > best_ap: best_ap = ap; best_epoch = epoch; no_imp = 0 best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} else: no_imp += 1 if no_imp >= args.patience: logger.info(f"early stop at epoch {epoch} (no val AP improvement for {args.patience} epochs)") break if best_state is None: raise SystemExit("no best state recorded; training failed") meta_out = { "hidden_dim": D, "n_frames": T, "proj_dim": args.proj_dim, "n_layers": args.n_layers, "n_heads": args.n_heads, "dropout": args.dropout, "best_epoch": best_epoch, "best_val_ap": best_ap, "train_cache": str(args.train_cache), "val_cache": str(args.val_cache), } torch.save({"state_dict": best_state, "meta": meta_out}, out_dir / "best.pt") with open(out_dir / "train_log.json", "w") as f: json.dump({"log": log, "best": meta_out}, f, indent=2) logger.info(f"best val AP = {best_ap:.4f} @ epoch {best_epoch}") logger.info(f"saved -> {out_dir/'best.pt'}") if __name__ == "__main__": main()