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