training/Policy/train_lkalert_mcb.py

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