training/Policy/train_policy_head_ablation.py

"""E3 — Module ablation trainer.

Wrapper around train_policy_head_v2 that supports ablation flags:

  --no_danger        : feed zeros for perception_summary + danger_per_frame
  --no_policy_pos    : feed zeros for policy_position (forces head to rely on
                        perception_summary only)
  --no_prev_action   : always pass BOS=3 for prev_action (same as current default
                        but explicitly logged; this is the "no temporal context"
                        baseline)
  --no_class_weight  : disable class_weights_from()
  --pool_mean        : replace PMA aggregator with mean pooling (need to bypass
                        DangerHead; uses mean over the 8 frames of belief_content
                        as perception_summary)

One seed per ablation, 15 epochs, ~5 min each.
"""
from __future__ import annotations

import argparse, json, logging, math, random, sys, gc
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 sklearn.metrics import accuracy_score, f1_score, confusion_matrix
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from lkalert.models.danger_head import DangerHead
from lkalert.models.policy_head_v2 import PolicyHeadV2, policy_loss

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("train_ablation")


def set_seed(s: int):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)


@torch.no_grad()
def precompute(cache_path, danger_ckpt, device, no_danger=False, pool_mean=False):
    d = torch.load(cache_path, weights_only=False, map_location="cpu")
    belief = d["belief_content"]
    valid = d["valid_frames"]
    N, T, D = belief.shape
    logger.info(f"[precompute] {cache_path.name}: belief={tuple(belief.shape)}")

    if no_danger:
        # zeros perception + danger
        perc = torch.zeros(N, 4, 512, dtype=torch.float32)
        dang = torch.zeros(N, 8, dtype=torch.float32)
    elif pool_mean:
        # mean over belief frames → 4-way replicate to match [N, K=4, 512]
        # Use first 512 dims as a "summary"
        mean_b = belief.float().mean(dim=1)  # [N, D=10240]
        # Project to 512: just take first 512 dims (cheap baseline)
        perc = mean_b[:, :512].unsqueeze(1).repeat(1, 4, 1)
        dang = torch.zeros(N, 8, dtype=torch.float32)
    else:
        ck = torch.load(danger_ckpt, weights_only=False, map_location="cpu")
        model = DangerHead(in_dim=ck["in_dim"]).to(device)
        model.load_state_dict(ck["model"]); model.eval()
        bs = 64
        all_p, all_d = [], []
        for i in tqdm(range(0, N, bs), desc="danger_precompute", ncols=80):
            bc = belief[i:i+bs].to(device, dtype=torch.float32)
            v = valid[i:i+bs].to(device)
            o = model(bc, valid_frames=v)
            all_p.append(o["perception_summary"].cpu())
            all_d.append(o["per_frame"].cpu())
        perc = torch.cat(all_p, 0)
        dang = torch.cat(all_d, 0)
        del model

    out = {
        "policy_position": d["policy_position"],
        "perception_summary": perc,
        "danger_per_frame": dang,
        "valid_frames": d["valid_frames"],
        "tick_action": d["tick_action"].long(),
    }
    del belief, d; gc.collect()
    return out


class AblationDataset(Dataset):
    def __init__(self, feats, no_policy_pos=False):
        self.pp = feats["policy_position"]
        self.perc = feats["perception_summary"]
        self.dang = feats["danger_per_frame"]
        self.v = feats["valid_frames"]
        self.y = feats["tick_action"]
        self.no_policy_pos = no_policy_pos
        self.n = self.pp.shape[0]
        self.prev_action = torch.full((self.n,), 3, dtype=torch.long)

    def __len__(self): return self.n

    def __getitem__(self, i):
        pp = self.pp[i]
        if self.no_policy_pos:
            pp = torch.zeros_like(pp)
        return {"policy_position": pp,
                "perception_summary": self.perc[i],
                "danger_per_frame": self.dang[i],
                "valid_frames": self.v[i],
                "tick_action": self.y[i],
                "prev_action": self.prev_action[i]}


def collate(b): return {k: torch.stack([x[k] for x in b]) for k in b[0]}


def train(args):
    set_seed(args.seed)
    args.out_dir.mkdir(parents=True, exist_ok=True)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    train_feats = precompute(args.train_cache, args.danger_ckpt, device,
                              no_danger=args.no_danger, pool_mean=args.pool_mean)
    val_feats = precompute(args.val_cache, args.danger_ckpt, device,
                            no_danger=args.no_danger, pool_mean=args.pool_mean)
    train_ds = AblationDataset(train_feats, no_policy_pos=args.no_policy_pos)
    val_ds = AblationDataset(val_feats, no_policy_pos=args.no_policy_pos)
    train_loader = DataLoader(train_ds, batch_size=64, shuffle=True,
                                num_workers=2, collate_fn=collate, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=64, shuffle=False,
                              num_workers=2, collate_fn=collate, pin_memory=True)

    in_dim = int(train_feats["policy_position"].shape[-1])
    perc_dim = int(train_feats["perception_summary"].shape[2])
    K = int(train_feats["perception_summary"].shape[1])
    model = PolicyHeadV2(policy_dim=in_dim,
                          perception_dim_per_query=perc_dim,
                          k_queries=K).to(device)

    cw = None
    if not args.no_class_weight:
        counts = torch.bincount(train_feats["tick_action"], minlength=3).float()
        inv = 1.0 / counts.clamp(min=1.0)
        cw = (inv * (counts.sum() / inv.sum())).to(device)
    logger.info(f"class_weights = {cw.tolist() if cw is not None else 'None'}")

    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
    n_steps = math.ceil(len(train_loader) * args.epochs)
    warmup = max(1, int(n_steps * 0.10))
    def lrlam(s):
        if s < warmup: return s / warmup
        p = (s - warmup) / max(1, n_steps - warmup)
        return 0.5 * (1 + math.cos(math.pi * p))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lrlam)

    best = -1
    for ep in range(args.epochs):
        model.train()
        for b in tqdm(train_loader, ncols=80, desc=f"ep{ep}"):
            pp = b["policy_position"].to(device, dtype=torch.float32, non_blocking=True)
            perc = b["perception_summary"].to(device)
            dang = b["danger_per_frame"].to(device)
            prev = b["prev_action"].to(device)
            v = b["valid_frames"].to(device)
            y = b["tick_action"].to(device)
            logits = model(pp, perc, dang, prev, valid_frames=v)
            losses = policy_loss(logits, y, class_weights=cw,
                                  label_smoothing=0.05, entropy_reg=0.02)
            losses["loss"].backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); sched.step(); opt.zero_grad(set_to_none=True)

        # eval
        model.eval()
        preds, targs = [], []
        with torch.no_grad():
            for b in val_loader:
                pp = b["policy_position"].to(device, dtype=torch.float32)
                logits = model(pp, b["perception_summary"].to(device),
                                b["danger_per_frame"].to(device),
                                b["prev_action"].to(device),
                                valid_frames=b["valid_frames"].to(device))
                preds.append(logits.argmax(-1).cpu())
                targs.append(b["tick_action"])
        p = torch.cat(preds).numpy(); t = torch.cat(targs).numpy()
        cm = confusion_matrix(t, p, labels=[0,1,2])
        per_class = cm.diagonal() / cm.sum(axis=1).clip(min=1)
        bal = float(per_class.mean())
        f1 = f1_score(t, p, average="macro")
        logger.info(f"ep{ep}  val_bal={bal:.4f}  f1={f1:.4f}  per_class={per_class.tolist()}")
        if bal > best:
            best = bal
            torch.save({"model": model.state_dict(),
                        "val_bal": bal, "val_f1": float(f1),
                        "policy_dim": in_dim, "perception_dim_per_query": perc_dim,
                        "k_queries": K, "in_dim": in_dim, "epoch": ep,
                        "val_metrics": {"balanced_acc": bal, "macro_f1": float(f1),
                                         "per_class_recall": {f"cls_{c}": float(per_class[c]) for c in range(3)}},
                        }, args.out_dir / "best.pt")
    logger.info(f"DONE best val_bal={best:.4f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train_cache", type=Path, required=True)
    ap.add_argument("--val_cache", type=Path, required=True)
    ap.add_argument("--danger_ckpt", type=Path, required=True)
    ap.add_argument("--out_dir", type=Path, required=True)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--epochs", type=int, default=15)
    # Ablation flags
    ap.add_argument("--no_danger", action="store_true",
                    help="zero out perception_summary + danger_per_frame")
    ap.add_argument("--no_policy_pos", action="store_true",
                    help="zero out POLICY_POSITION (force reliance on perception)")
    ap.add_argument("--no_prev_action", action="store_true",
                    help="(currently default; flag is informational)")
    ap.add_argument("--no_class_weight", action="store_true")
    ap.add_argument("--pool_mean", action="store_true",
                    help="replace DangerHead with simple mean-pool over belief frames")
    args = ap.parse_args()
    train(args)


if __name__ == "__main__":
    main()