#!/usr/bin/env python3 """3-class POMDP head trainer for Qwen3-VL-4B per-frame belief cache. Generalization of train_pomdp_head_v2.py from binary sigmoid to 3-class softmax (SILENT/OBSERVE/ALERT) with focal CE + label smoothing + manifest ce_weight + EMA. Key differences from v2 (binary): - Output dim: 1 (sigmoid) → 3 (softmax) - Loss: BCE → focal cross-entropy (γ=2, α-vec for class balance) - Label: action_label as-is (0/1/2), not (>0).int() - ce_weight: per-sample weight from manifest used (non_ego = 0.4) - EMA: teacher EMA weights for eval (decay=0.999) - Selection: policy_score (PS_v3) on val instead of binary_AP Reuses POMDPTemporalHead skeleton but with last linear changed to 3 outputs. Usage: python -m training.Policy.train_focal_pomdp_v3 \ --train_cache data/belief_cache_perframe_qwen3vl4b/multisrc_train.pt \ --val_cache data/belief_cache_perframe_qwen3vl4b/multisrc_val.pt \ --label_dir data/policy_labels \ --output_dir checkpoints/Policy \ --experiment_name focal_pomdp_qwen3vl4b_seed0 \ --seed 0 """ from __future__ import annotations import argparse import copy import json import logging import math import random import sys import time from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sklearn.metrics import average_precision_score from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler ROOT = Path(__file__).resolve().parents[2] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) logger = logging.getLogger("focal_pomdp_v3") logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s") # ─── 3-class POMDP head ──────────────────────────────────────────────── class POMDPTemporalHead3Class(nn.Module): """3-class softmax variant of POMDPTemporalHead.""" def __init__(self, in_dim: int = 2560, proj_dim: int = 512, gru_hidden: int = 256, dropout: float = 0.2, n_actions: int = 3): super().__init__() self.in_proj = nn.Sequential( nn.Linear(in_dim, proj_dim), nn.LayerNorm(proj_dim), nn.GELU(), nn.Dropout(dropout), ) self.text_proj = nn.Sequential( nn.Linear(in_dim, gru_hidden), nn.LayerNorm(gru_hidden), nn.Tanh(), ) self.gru = nn.GRU(proj_dim, gru_hidden, num_layers=1, batch_first=True) self.attn = nn.Linear(gru_hidden, 1) self.cls = nn.Sequential( nn.Linear(gru_hidden, 128), nn.GELU(), nn.Dropout(dropout), nn.Linear(128, n_actions), ) def forward(self, beliefs, valid, text): x = self.in_proj(beliefs) h0 = self.text_proj(text).unsqueeze(0).contiguous() out, _ = self.gru(x, h0) attn_logits = self.attn(out).squeeze(-1) attn_logits = attn_logits.masked_fill(~valid, float("-inf")) empty = (~valid).all(dim=1) if empty.any(): attn_logits[empty] = 0.0 w = F.softmax(attn_logits, dim=1).unsqueeze(-1) pooled = (out * w).sum(dim=1) return self.cls(pooled) # [B, 3] # ─── focal cross-entropy ─────────────────────────────────────────────── def focal_cross_entropy(logits, target, alpha=None, gamma=2.0, label_smoothing=0.0, sample_weight=None): """3-class focal CE with label smoothing. logits: [B, C], target: [B] long, alpha: [C] tensor or None, sample_weight: [B] or None.""" log_probs = F.log_softmax(logits, dim=-1) probs = log_probs.exp() n_classes = logits.size(-1) # one-hot with label smoothing with torch.no_grad(): true_dist = torch.zeros_like(log_probs) true_dist.fill_(label_smoothing / (n_classes - 1)) true_dist.scatter_(1, target.unsqueeze(1), 1.0 - label_smoothing) # focal weight: (1 - pt)^gamma where pt = prob of true class pt = probs.gather(1, target.unsqueeze(1)).squeeze(1).clamp_min(1e-8) focal_w = (1.0 - pt).pow(gamma) # alpha (per-class weight vector) if alpha is not None: alpha_t = alpha[target] focal_w = focal_w * alpha_t # CE per sample per_sample = -(true_dist * log_probs).sum(dim=-1) per_sample = per_sample * focal_w # per-sample weight from manifest (non_ego ce_weight, etc.) if sample_weight is not None: per_sample = per_sample * sample_weight return per_sample.mean() # ─── Cache dataset ───────────────────────────────────────────────────── class CacheDataset3Class(Dataset): def __init__(self, cache_path: Path, label_path: Path): logger.info(f"loading cache {cache_path}") self.cache = torch.load(cache_path, weights_only=False, map_location="cpu") self.bf = self.cache["beliefs_frame"] # [N, T, D] self.vf = self.cache["valid_frames"] # [N, T] self.bt = self.cache["beliefs_text"] # [N, D] self.meta = self.cache.get("meta", {}) # Load labels manifest aligned by sample id manifest = json.loads(label_path.read_text()) samples = manifest.get("samples", []) id_to_meta = {s["video_id"]: s for s in samples} ids = self.meta.get("ids", []) self.action_labels = [] self.ce_weights = [] self.categories = [] for i, vid in enumerate(ids): m = id_to_meta.get(vid, {}) self.action_labels.append(int(m.get("action_label", -1))) self.ce_weights.append(float(m.get("ce_weight", 1.0))) self.categories.append(m.get("category", "unknown")) self.action_labels = torch.tensor(self.action_labels, dtype=torch.long) self.ce_weights = torch.tensor(self.ce_weights, dtype=torch.float32) n_dropped = (self.action_labels < 0).sum().item() logger.info(f" N={len(ids)}, T={self.bf.shape[1]}, D={self.bf.shape[2]}, " f"dropped(label<0)={n_dropped}") # filter out invalid labels keep = (self.action_labels >= 0) self.indices = torch.nonzero(keep).squeeze(1).tolist() # log class distribution labs = self.action_labels[keep] for c in (0, 1, 2): n = (labs == c).sum().item() logger.info(f" class={c}: n={n}") def __len__(self): return len(self.indices) def __getitem__(self, i): idx = self.indices[i] return { "belief": self.bf[idx].float(), # [T, D] "valid": self.vf[idx], # [T] "text": self.bt[idx].float(), # [D] "label": int(self.action_labels[idx]), "weight": float(self.ce_weights[idx]), } def collate_fn(batch): return { "belief": torch.stack([b["belief"] for b in batch]), "valid": torch.stack([b["valid"] for b in batch]), "text": torch.stack([b["text"] for b in batch]), "label": torch.tensor([b["label"] for b in batch], dtype=torch.long), "weight": torch.tensor([b["weight"] for b in batch], dtype=torch.float32), } # ─── EMA wrapper ─────────────────────────────────────────────────────── class EMA: def __init__(self, model, decay=0.999): self.decay = decay self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()} def update(self, model): for k, v in model.state_dict().items(): if v.dtype.is_floating_point: self.shadow[k].mul_(self.decay).add_(v.detach(), alpha=1 - self.decay) else: self.shadow[k] = v.detach().clone() def apply(self, model): self.backup = {k: v.detach().clone() for k, v in model.state_dict().items()} model.load_state_dict(self.shadow) def restore(self, model): model.load_state_dict(self.backup) del self.backup # ─── eval ────────────────────────────────────────────────────────────── @torch.no_grad() def evaluate(model, val_loader, device): model.eval() all_logits, all_labels, all_cats = [], [], [] for batch in val_loader: b = batch["belief"].to(device) v = batch["valid"].to(device) t = batch["text"].to(device) logits = model(b, v, t) all_logits.append(logits.cpu()) all_labels.append(batch["label"]) logits = torch.cat(all_logits, dim=0) labels = torch.cat(all_labels, dim=0).numpy() probs = F.softmax(logits, dim=-1).numpy() p_alert = probs[:, 2] # binary AP (ego = label==2 vs others) bin_target = (labels == 2).astype(int) binary_ap = float(average_precision_score(bin_target, p_alert)) # PS_v3 on argmax decisions preds = probs.argmax(axis=1) # categories from val_loader.dataset cats = np.array(val_loader.dataset.categories)[ val_loader.dataset.indices] ego_mask = (cats == "ego_positive") safe_mask = (cats == "safe_neg") ego_recall = float(((preds == 2) & ego_mask & (labels == 2)).sum() / max((ego_mask & (labels == 2)).sum(), 1)) safe_silent = float(((preds == 0) & safe_mask).sum() / max(safe_mask.sum(), 1)) safe_alert = float(((preds == 2) & safe_mask).sum() / max(safe_mask.sum(), 1)) ps_v3 = 0.65 * ego_recall + 0.25 * safe_silent - 0.15 * safe_alert return { "binary_ap": binary_ap, "ps_v3": ps_v3, "ego_alert_recall": ego_recall, "safe_neg_silent": safe_silent, "safe_neg_alert_leak": safe_alert, } # ─── train ───────────────────────────────────────────────────────────── def train(args): out = Path(args.output_dir) / args.experiment_name (out / "best").mkdir(parents=True, exist_ok=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_ds = CacheDataset3Class( Path(args.train_cache), Path(args.label_dir) / "train.json", ) val_ds = CacheDataset3Class( Path(args.val_cache), Path(args.label_dir) / "val.json", ) # balanced sampler (inverse class freq) if args.use_balanced_sampler: labs = train_ds.action_labels[train_ds.indices].numpy() counts = np.array([(labs == c).sum() for c in range(3)]) weights_per_class = 1.0 / np.maximum(counts, 1) sample_weights = weights_per_class[labs] sampler = WeightedRandomSampler(sample_weights, num_samples=len(labs), replacement=True) train_loader = DataLoader( train_ds, batch_size=args.batch_size, sampler=sampler, collate_fn=collate_fn, num_workers=args.num_workers, pin_memory=True, ) else: train_loader = DataLoader( train_ds, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, num_workers=args.num_workers, pin_memory=True, ) val_loader = DataLoader( val_ds, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn, num_workers=args.num_workers, pin_memory=True, ) in_dim = train_ds.bf.shape[-1] logger.info(f"in_dim={in_dim}") model = POMDPTemporalHead3Class( in_dim=in_dim, proj_dim=args.proj_dim, gru_hidden=args.gru_hidden, dropout=args.dropout, ).to(device) logger.info(f" n_params = {sum(p.numel() for p in model.parameters())}") # focal alpha = inverse class freq (for class balance) labs = train_ds.action_labels[train_ds.indices].numpy() counts = np.array([(labs == c).sum() for c in range(3)], dtype=np.float64) alpha_vec = (counts.sum() / (3 * np.maximum(counts, 1))).astype(np.float32) alpha = torch.tensor(alpha_vec, device=device) logger.info(f"focal alpha (inv-freq) = {alpha_vec}") opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) n_steps = args.num_epochs * len(train_loader) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( opt, T_max=n_steps, eta_min=args.lr * 0.003) ema = EMA(model, decay=args.ema_decay) if args.use_ema else None best_ps = -1.0 best_meta = {} step = 0 for epoch in range(args.num_epochs): model.train() t_epoch = time.time() for batch in train_loader: b = batch["belief"].to(device) v = batch["valid"].to(device) t = batch["text"].to(device) y = batch["label"].to(device) w = batch["weight"].to(device) logits = model(b, v, t) loss = focal_cross_entropy( logits, y, alpha=alpha, gamma=args.focal_gamma, label_smoothing=args.label_smoothing, sample_weight=w, ) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) opt.step() scheduler.step() if ema is not None: ema.update(model) step += 1 if step % args.log_every == 0: logger.info(f" ep{epoch} step{step:>5d} loss={loss.item():.4f} " f"lr={scheduler.get_last_lr()[0]:.2e}") if step % args.val_every_n_steps == 0: if ema is not None: ema.apply(model) metrics = evaluate(model, val_loader, device) if ema is not None: ema.restore(model) ps = metrics["ps_v3"] logger.info(f" [val ep{epoch} step{step}] " f"PS_v3={ps:.4f} AP={metrics['binary_ap']:.4f} " f"ego_recall={metrics['ego_alert_recall']:.3f} " f"safe_silent={metrics['safe_neg_silent']:.3f} " f"fa_leak={metrics['safe_neg_alert_leak']:.3f}") if ps > best_ps: best_ps = ps best_meta = {**metrics, "epoch": epoch, "step": step, "experiment": args.experiment_name} if ema is not None: ema.apply(model) torch.save({ "head_state": model.state_dict(), "args": vars(args), "metrics": metrics, }, out / "best" / "head.pt") if ema is not None: ema.restore(model) logger.info(f" ✓ saved new best PS_v3={ps:.4f}") model.train() logger.info(f"epoch {epoch} done in {time.time()-t_epoch:.1f}s") (out / "best" / "best_meta.json").write_text(json.dumps(best_meta, indent=2)) logger.info(f"\nbest PS_v3 = {best_ps:.4f}") logger.info(f" meta: {best_meta}") logger.info(f" ckpt: {out / 'best' / 'head.pt'}") def main(): p = argparse.ArgumentParser("focal_pomdp_v3") p.add_argument("--train_cache", required=True) p.add_argument("--val_cache", required=True) p.add_argument("--label_dir", default="data/policy_labels") p.add_argument("--output_dir", default="checkpoints/Policy") p.add_argument("--experiment_name", default="focal_pomdp_qwen3vl4b_seed0") p.add_argument("--proj_dim", type=int, default=512) p.add_argument("--gru_hidden", type=int, default=256) p.add_argument("--dropout", type=float, default=0.2) p.add_argument("--num_epochs", type=int, default=6) p.add_argument("--batch_size", type=int, default=128) p.add_argument("--num_workers", type=int, default=4) p.add_argument("--lr", type=float, default=2e-4) p.add_argument("--weight_decay", type=float, default=1e-4) p.add_argument("--grad_clip", type=float, default=1.0) p.add_argument("--focal_gamma", type=float, default=2.0) p.add_argument("--label_smoothing", type=float, default=0.05) p.add_argument("--use_balanced_sampler", action="store_true", default=True) p.add_argument("--use_ema", action="store_true", default=True) p.add_argument("--ema_decay", type=float, default=0.999) p.add_argument("--log_every", type=int, default=100) p.add_argument("--val_every_n_steps", type=int, default=200) p.add_argument("--seed", type=int, default=0) args = p.parse_args() random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.seed) train(args) if __name__ == "__main__": main()