training/Policy/train_nexar_head.py

#!/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()