training/Policy/train_focal_pomdp_v3.py

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