#!/usr/bin/env python3 """LKAlert-MCB Day-11 trainer (2-channel: Qwen semantic + V-JEPA dynamics). **Channel 2 (object motion) is intentionally absent** — failed Red Line 4 gate on Day 10 (object+POMDP regresses by −5.5 pp). The object-motion cache is repurposed for teacher pilot input + qualitative analysis only. Ablation rows (Day 11 -- 8-row matrix becomes a 4-row matrix without Channel 2): | Variant | b_sem | b_vid | hysteresis | |---|---|---|---| | Qwen-only | ✓ | ✗ | ✗ | | Video-only | ✗ | ✓ | ✗ | | **mcb_no_aux** (headline) | ✓ | ✓ | ✗ | | Full MCB + hyst (Day 12) | ✓ | ✓ | ✓ | The 8-row matrix in §5 of the plan is reduced; the 4 dropped rows involving b_obj will be reported in the appendix Table 6 negative ablation. Training hyper-parameters mirror Day-3 LKAlert-BD trainer for direct comparability. Trunk warm-starts from `checkpoints/Policy/lkalert_bd_best/best.pt`. """ from __future__ import annotations import argparse import json import logging import random import sys from pathlib import Path from typing import Dict, List import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from training.Policy.multichannel_dataset import ( MultichannelDataset, collate as mc_collate, ) from lkalert.models.multichannel_belief import LKAlertMCB logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("lkalert_mcb") CACHE_DIR = Path("data/belief_cache_perframe_qwen3vl4b") DIAG_DIR = Path("data/policy_labels") # ─── eval ──────────────────────────────────────────────────────────────────── @torch.no_grad() def evaluate(model: LKAlertMCB, ds: MultichannelDataset, device: torch.device, batch_size: int = 64) -> Dict: from sklearn.metrics import average_precision_score, roc_auc_score model.eval() loader = DataLoader(ds, batch_size=batch_size, shuffle=False, collate_fn=mc_collate) probs, labels = [], [] for b in loader: out = model(b["belief"].to(device), b["valid"].to(device), b["text"].to(device), b["vjepa"].to(device), b["vjepa_mask"].to(device)) probs.append(torch.sigmoid(out["p_any"]).cpu().numpy()) labels.append(b["y_p_any"].cpu().numpy() if "y_p_any" in b else np.zeros(out["p_any"].shape[0])) p = np.concatenate(probs); y = np.concatenate(labels) if y.min() == y.max(): return {"ap": 0.0, "auc": 0.0, "n": int(y.size), "n_pos": int(y.sum())} return {"ap": float(average_precision_score(y, p)), "auc": float(roc_auc_score(y, p)), "n": int(y.size), "n_pos": int(y.sum())} # ─── training loop ────────────────────────────────────────────────────────── def train_one_seed(args) -> Dict: torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_ds = MultichannelDataset(args.train_cache, "train") val_caches: Dict[str, MultichannelDataset] = {} for vc in args.val_caches: try: val_caches[vc] = MultichannelDataset(vc, "val") except FileNotFoundError as e: logger.warning(f" skip val cache {vc}: {e}") # weighted sampler to balance pos/neg y = train_ds.y_any pos = (y == 1).sum(); neg = (y == 0).sum() weights = np.where(y == 1, 1.0 / max(pos, 1), 1.0 / max(neg, 1)) sampler = torch.utils.data.WeightedRandomSampler( weights=weights, num_samples=len(weights), replacement=True) train_loader = DataLoader(train_ds, batch_size=args.batch_size, sampler=sampler, collate_fn=mc_collate, num_workers=args.num_workers) in_dim = train_ds.bf.shape[-1] model = LKAlertMCB( qwen_in_dim = in_dim, proj_dim = args.proj_dim, gru_hidden = args.gru_hidden, vjepa_in_dim = 1024, vjepa_out_dim = args.vjepa_out_dim, dropout = args.dropout, use_qwen = args.use_qwen, use_vjepa = args.use_vjepa, fusion = args.fusion, with_teacher_aux = args.with_teacher_aux, ) if args.warm_start and Path(args.warm_start).exists(): ck = torch.load(args.warm_start, weights_only=False, map_location="cpu") copied = model.warm_start_qwen_trunk_from_bd(ck["head_state"]) logger.info(f"warm-started Qwen trunk: {len(copied)} params from " f"{args.warm_start}") model.to(device) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd) out_dir = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True) best = {"epoch": -1, "macro_ap": -1.0, "per_cache": {}} for epoch in range(args.epochs): model.train() ep_loss = 0.0; n_batches = 0 for b in train_loader: opt.zero_grad() out = model(b["belief"].to(device), b["valid"].to(device), b["text"].to(device), b["vjepa"].to(device), b["vjepa_mask"].to(device)) y = b["y_p_any"].to(device) loss = F.binary_cross_entropy_with_logits(out["p_any"], y) loss.backward() opt.step() ep_loss += float(loss.detach()); n_batches += 1 # eval per_cache: Dict[str, Dict] = {} macro = 0.0 for vc, ds in val_caches.items(): m = evaluate(model, ds, device, args.batch_size) per_cache[vc] = m macro += m.get("ap", 0.0) macro /= max(1, len(val_caches)) logger.info(f"ep {epoch:02d} loss={ep_loss / max(1, n_batches):.4f} " f"macro AP={macro:.4f}") for vc, m in per_cache.items(): logger.info(f" {vc}: AP={m['ap']:.4f} AUC={m['auc']:.4f} " f"n_pos={m['n_pos']}/{m['n']}") if macro > best["macro_ap"]: best = {"epoch": epoch, "macro_ap": float(macro), "per_cache": per_cache, "head_state": model.state_dict(), "args": vars(args)} torch.save(best, out_dir / "best.pt") logger.info(f" -> saved best.pt @ macro_ap={macro:.4f}") return best def main(): ap = argparse.ArgumentParser() ap.add_argument("--train_cache", default="nexar_train_diag") ap.add_argument("--val_caches", nargs="+", default=["nexar_val", "dota_val", "dad_test", "dada_test"]) ap.add_argument("--out_dir", default="checkpoints/Policy/lkalert_mcb_seed0") ap.add_argument("--epochs", type=int, default=30) ap.add_argument("--batch_size", type=int, default=64) ap.add_argument("--num_workers", type=int, default=2) ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--wd", type=float, default=1e-4) ap.add_argument("--proj_dim", type=int, default=512) ap.add_argument("--gru_hidden", type=int, default=256) ap.add_argument("--vjepa_out_dim", type=int, default=256) ap.add_argument("--dropout", type=float, default=0.2) ap.add_argument("--use_qwen", action="store_true", default=True) ap.add_argument("--no_qwen", dest="use_qwen", action="store_false") ap.add_argument("--use_vjepa", action="store_true", default=True) ap.add_argument("--no_vjepa", dest="use_vjepa", action="store_false") ap.add_argument("--fusion", default="concat_mlp", choices=["concat_mlp", "gated_concat"]) ap.add_argument("--with_teacher_aux", action="store_true", help="Day-11.5 stretch — adds 5 aux slot heads") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--warm_start", default="checkpoints/Policy/lkalert_bd_best/best.pt", help="LKAlert-BD best.pt for trunk warm-start") args = ap.parse_args() train_one_seed(args) if __name__ == "__main__": main()