training/danger/train_danger_v3_hazard.py

"""Phase G.0c — Re-train DangerHead with 8-way hazard auxiliary head.

Joint loss = existing alert-binary BCE (per-frame + clip) +
             0.3 · CE(hazard_logits, hazard_target)

Hazard targets come from `data/policy_labels/hazard_categories_*.json`
built by `tools/build_hazard_labels.py`.

Output: checkpoints/danger_v3_hazard/best.pt
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

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

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

N_HAZARDS = 8


class HazardDataset(Dataset):
    def __init__(self, cache_path: Path, hazard_path: Path):
        self.cache = torch.load(cache_path, weights_only=False, map_location="cpu")
        hz = json.loads(hazard_path.read_text())
        # Index hazard labels by video_id (parallel to cache['ids'])
        ids_to_h = dict(zip(hz["ids"], hz["labels"]))
        self.hazard = torch.tensor(
            [ids_to_h.get(vid, 7) for vid in self.cache["ids"]],
            dtype=torch.long)
        logger.info(f"  loaded {cache_path.name}: N={len(self.cache['ids'])} "
                     f"hazard dist={torch.bincount(self.hazard, minlength=N_HAZARDS).tolist()}")

    def __len__(self):
        return len(self.cache["ids"])

    def __getitem__(self, idx):
        return {
            "belief":     self.cache["belief_content"][idx],
            "valid":      self.cache["valid_frames"][idx],
            "danger_pf":  self.cache["danger_pf"][idx],
            "hazard":     self.hazard[idx],
            "tick_action": int(self.cache["tick_action"][idx]),
        }


def collate(batch):
    return {
        "belief":      torch.stack([b["belief"] for b in batch]),
        "valid":       torch.stack([b["valid"] for b in batch]),
        "danger_pf":   torch.stack([b["danger_pf"] for b in batch]),
        "hazard":      torch.stack([b["hazard"] for b in batch]),
        "tick_action": torch.tensor([b["tick_action"] for b in batch],
                                     dtype=torch.long),
    }


@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    all_hazard_logits, all_hazard_t, all_alert_score, all_alert_t = [], [], [], []
    all_pf_logit, all_danger_pf, all_valid = [], [], []
    for b in loader:
        bc = b["belief"].to(device, dtype=torch.float32)
        v  = b["valid"].to(device)
        out = model(bc, valid_frames=v)
        all_hazard_logits.append(out["hazard_logits"].cpu().numpy())
        all_hazard_t.append(b["hazard"].numpy())
        all_alert_score.append(out["clip"].cpu().numpy())
        # alert ground-truth = (tick_action == 2)
        all_alert_t.append((b["tick_action"] == 2).numpy().astype(int))
        all_pf_logit.append(out["per_frame_logits"].cpu().numpy())
        all_danger_pf.append(b["danger_pf"].numpy())
        all_valid.append(v.cpu().numpy())

    hz_logits = np.concatenate(all_hazard_logits)
    hz_t = np.concatenate(all_hazard_t)
    hz_pred = hz_logits.argmax(axis=-1)
    a_s = np.concatenate(all_alert_score)
    a_t = np.concatenate(all_alert_t)

    metrics = {
        "hazard_acc":          float(accuracy_score(hz_t, hz_pred)),
        "hazard_balanced_acc": float(balanced_accuracy_score(hz_t, hz_pred)),
        "alert_AP":            float(average_precision_score(a_t, a_s)),
        "alert_AUROC":         float(roc_auc_score(a_t, a_s)),
    }
    return metrics


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--train_cache", type=Path,
                    default=ROOT / "data/belief_cache_v3/sft_x_v3__train_9k.pt")
    ap.add_argument("--val_cache", type=Path,
                    default=ROOT / "data/belief_cache_v3/sft_x_v3__multisrc_val.pt")
    ap.add_argument("--train_hazard", type=Path,
                    default=ROOT / "data/policy_labels/hazard_categories_train_9k.json")
    ap.add_argument("--val_hazard", type=Path,
                    default=ROOT / "data/policy_labels/hazard_categories_multisrc_val.json")
    ap.add_argument("--out_dir", type=Path,
                    default=ROOT / "checkpoints/danger_v3_hazard")
    ap.add_argument("--in_dim", type=int, default=10240)
    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("--lr", type=float, default=5e-4)
    ap.add_argument("--weight_decay", type=float, default=1e-4)
    ap.add_argument("--epochs", type=int, default=50)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--hazard_weight", type=float, default=0.3)
    ap.add_argument("--w_clip", type=float, default=0.5)
    ap.add_argument("--patience", type=int, default=15)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--max_samples", type=int, default=0,
                    help="if >0, truncate train+val for smoke testing")
    args = ap.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    train_ds = HazardDataset(args.train_cache, args.train_hazard)
    val_ds   = HazardDataset(args.val_cache, args.val_hazard)
    if args.max_samples > 0:
        train_ds.cache["ids"] = train_ds.cache["ids"][:args.max_samples]
        train_ds.hazard = train_ds.hazard[:args.max_samples]
        val_ds.cache["ids"] = val_ds.cache["ids"][:args.max_samples]
        val_ds.hazard = val_ds.hazard[:args.max_samples]
    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 * 2, shuffle=False,
                             num_workers=2, collate_fn=collate, pin_memory=True)

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

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

    log_records = []
    best_score = -1e9
    bad_epochs = 0
    for ep in range(args.epochs):
        model.train()
        run = {"loss": 0, "danger": 0, "hazard": 0}; n_b = 0
        pbar = tqdm(train_loader, ncols=80, desc=f"ep{ep}")
        for b in pbar:
            bc = b["belief"].to(device, dtype=torch.float32, non_blocking=True)
            v  = b["valid"].to(device, non_blocking=True)
            dpf = b["danger_pf"].to(device, non_blocking=True)
            hz = b["hazard"].to(device, non_blocking=True)
            out = model(bc, valid_frames=v)
            d_l = danger_loss(out, dpf, valid_frames=v, w_clip=args.w_clip)
            h_l = F.cross_entropy(out["hazard_logits"], hz)
            total = d_l["loss"] + args.hazard_weight * h_l
            total.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
            run["loss"]   += total.item()
            run["danger"] += d_l["loss"].item()
            run["hazard"] += h_l.item()
            n_b += 1
            pbar.set_postfix(loss=run["loss"]/n_b, hz=run["hazard"]/n_b)

        m = evaluate(model, val_loader, device)
        rec = {"ep": ep,
                "train_loss":   run["loss"]/max(1, n_b),
                "train_danger": run["danger"]/max(1, n_b),
                "train_hazard": run["hazard"]/max(1, n_b),
                "val": m}
        log_records.append(rec)
        logger.info(f"[ep{ep}] train={rec['train_loss']:.4f} "
                     f"haz={rec['train_hazard']:.4f}  | "
                     f"val: alert_AP={m['alert_AP']:.4f} "
                     f"alert_AUROC={m['alert_AUROC']:.4f} "
                     f"hazard_bal_acc={m['hazard_balanced_acc']:.4f}")

        # Composite: 0.5 alert_AP + 0.3 alert_AUROC + 0.2 hazard_bal_acc
        score = (0.5 * m["alert_AP"] + 0.3 * m["alert_AUROC"]
                 + 0.2 * m["hazard_balanced_acc"])
        if score > best_score:
            best_score = score; bad_epochs = 0
            save_dict = {
                "model": model.state_dict(),
                "in_dim": args.in_dim, "hidden": args.hidden,
                "k_queries": args.k_queries, "dropout": args.dropout,
                "n_hazards": N_HAZARDS,
                "val_metrics": m, "composite": score, "epoch": ep,
                "args": vars(args),
            }
            torch.save(save_dict, args.out_dir / "best.pt")
            logger.info(f"  [save best] composite={score:.4f}")
        else:
            bad_epochs += 1
            if bad_epochs >= args.patience:
                logger.info(f"  early stop @ ep{ep} (patience {args.patience})")
                break

    (args.out_dir / "training_log.json").write_text(
        json.dumps(log_records, indent=2, default=str))
    logger.info(f"\n[done] best composite={best_score:.4f}  saved to {args.out_dir}/best.pt")


if __name__ == "__main__":
    main()