"""Head-RL DPO — train PolicyHeadV2 with DPO objective on preference pairs. Frozen: SFT Qwen3-VL backbone + BELIEF cache features. Trainable: PolicyHeadV2 (~7M params). Reference policy: a frozen COPY of the supervised PolicyHeadV2 (`policy_v3_strong`). DPO objective on 3-class softmax: loss = -log σ(β · ( log π_θ(c|x) − log π_θ(r|x) − log π_ref(c|x) + log π_ref(r|x) )) where c=chosen_action_idx, r=rejected_action_idx, π_θ is the 3-class softmax output of PolicyHeadV2(x), π_ref is the same architecture with frozen weights. Pair structure (from preference_pairs.jsonl): Each pair has (video_id, frame_indices, chosen_action ∈ {S,O,A}, rejected_action ∈ {S,O,A}). We use the CACHED BELIEF features for that video's tick (looked up by video_id). PolicyHead predicts a single tick-level action; DPO loss applies on that tick's 3-class softmax with chosen / rejected as the preference target. Usage: python -m training.Policy.train_head_dpo \ --pref_jsonl data/cot_corpus_v2/preference_pairs.jsonl \ --train_cache data/belief_cache_v3/sft_x_v3__train_9k.pt \ --policy_warm checkpoints/policy_v3_strong/best.pt \ --out_dir checkpoints/policy_v3_head_dpo """ from __future__ import annotations import argparse import copy import json import logging import sys from pathlib import Path import torch import torch.nn.functional as F 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 from training.Policy._balance_eval import evaluate_policy_on_val, format_gate_row logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("head_dpo") ACTION_NAME_TO_IDX = {"SILENT": 0, "OBSERVE": 1, "ALERT": 2} class PreferenceDataset(Dataset): """For each preference pair, retrieve the cached BELIEF + POLICY features. `cache` is the full v3 train cache (9440 ticks). We look up each pair's `video_id` (or `id` minus tick suffix) and pick the matching cache row. If `observe_oversample > 1`, pairs whose chosen_action == OBSERVE are repeated `observe_oversample` times in the index (extra samples are duplicates, not novel pairs). """ def __init__(self, pref_jsonl: Path, cache_path: Path, observe_oversample: int = 1): self.cache = torch.load(cache_path, weights_only=False, map_location="cpu") self.id_to_idx = {iid: i for i, iid in enumerate(self.cache["ids"])} self.pairs = [] skipped = 0 with pref_jsonl.open() as f: for ln in f: ln = ln.strip() if not ln: continue obj = json.loads(ln) vid = obj.get("video_id") if vid not in self.id_to_idx: skipped += 1 continue ci = self.id_to_idx[vid] pair = { "cache_idx": ci, "chosen": ACTION_NAME_TO_IDX[obj["chosen_action"]], "rejected": ACTION_NAME_TO_IDX[obj["rejected_action"]], "pair_type": obj.get("pair_type", "?"), "tick_action": int(self.cache["tick_action"][ci]), } self.pairs.append(pair) # OBSERVE oversample: duplicate pairs where chosen_action == OBSERVE if observe_oversample > 1: base_pairs = list(self.pairs) for p in base_pairs: if p["chosen"] == ACTION_NAME_TO_IDX["OBSERVE"]: self.pairs.extend([p] * (observe_oversample - 1)) n_obs_chosen = sum(1 for p in self.pairs if p["chosen"] == 1) n_alr_chosen = sum(1 for p in self.pairs if p["chosen"] == 2) n_sil_chosen = sum(1 for p in self.pairs if p["chosen"] == 0) logger.info(f" loaded {len(self.pairs)} pairs (skipped {skipped} unmatched; " f"chosen SIL/OBS/ALR = {n_sil_chosen}/{n_obs_chosen}/{n_alr_chosen})") def __len__(self): return len(self.pairs) def __getitem__(self, idx): p = self.pairs[idx] ci = p["cache_idx"] return { "belief": self.cache["belief_content"][ci], "policy": self.cache["policy_position"][ci], "valid": self.cache["valid_frames"][ci], "chosen": p["chosen"], "rejected": p["rejected"], "tick_action": p["tick_action"], } def collate(batch): return { "belief": torch.stack([b["belief"] for b in batch]), "policy": torch.stack([b["policy"] for b in batch]), "valid": torch.stack([b["valid"] for b in batch]), "chosen": torch.tensor([b["chosen"] for b in batch], dtype=torch.long), "rejected": torch.tensor([b["rejected"] for b in batch], dtype=torch.long), "tick_action": torch.tensor([b["tick_action"] for b in batch], dtype=torch.long), } def dpo_loss(logits, ref_logits, chosen, rejected, beta=0.1): """3-class DPO loss. logits, ref_logits: [B, 3] chosen, rejected: [B] long ids into {0,1,2} """ log_p = F.log_softmax(logits, dim=-1) # [B, 3] log_p_ref = F.log_softmax(ref_logits, dim=-1) # [B, 3] B = logits.shape[0] idx = torch.arange(B, device=logits.device) log_p_chosen = log_p[idx, chosen] log_p_rejected = log_p[idx, rejected] log_p_ref_chosen = log_p_ref[idx, chosen] log_p_ref_rejected = log_p_ref[idx, rejected] # DPO advantage delta = beta * ((log_p_chosen - log_p_rejected) - (log_p_ref_chosen - log_p_ref_rejected)) loss = -F.logsigmoid(delta).mean() # Logging stats with torch.no_grad(): chosen_minus_rejected = (log_p_chosen - log_p_rejected).mean().item() prefers_chosen_rate = ((log_p_chosen > log_p_rejected).float().mean().item()) return loss, {"delta_mean": chosen_minus_rejected, "prefers_chosen_rate": prefers_chosen_rate} def main(): ap = argparse.ArgumentParser() ap.add_argument("--pref_jsonl", type=Path, default=ROOT / "data/cot_corpus_v2/preference_pairs.jsonl") ap.add_argument("--train_cache", type=Path, default=ROOT / "data/belief_cache_v3/sft_x_v3__train_9k.pt") ap.add_argument("--val_cache", type=Path, default=ROOT / "data/belief_cache_v3/sft_x_v3__multisrc_val.pt") ap.add_argument("--policy_warm", type=Path, default=ROOT / "checkpoints/policy_v3_strong/best.pt") ap.add_argument("--danger_ckpt", type=Path, default=ROOT / "checkpoints/danger_v2/seed2/best.pt") ap.add_argument("--out_dir", type=Path, required=True) ap.add_argument("--beta", type=float, default=0.05, help="DPO temperature; lower preserves supervised more") ap.add_argument("--lr", type=float, default=1e-5) ap.add_argument("--epochs", type=int, default=5) ap.add_argument("--batch_size", type=int, default=64) ap.add_argument("--rl_weight", type=float, default=0.7, help="α: scales DPO term in mixed loss (α·L_RL + (1-α)·L_anchor)") ap.add_argument("--alert_anchor_weight", type=float, default=1.0, help="Multiplicative weight applied to anchor CE term before " "(1-α) scaling; used for tuning anchor strength") ap.add_argument("--oversample_observe", type=int, default=3, help="Duplicate OBSERVE-chosen pairs by this factor in train loader") ap.add_argument("--max_samples", type=int, default=0, help="If >0, truncate the dataset to this many pairs (smoke testing)") ap.add_argument("--seed", type=int, default=0) args = ap.parse_args() args.out_dir.mkdir(parents=True, exist_ok=True) torch.manual_seed(args.seed) device = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"[load] pref={args.pref_jsonl}") ds = PreferenceDataset(args.pref_jsonl, args.train_cache, observe_oversample=args.oversample_observe) if args.max_samples > 0: ds.pairs = ds.pairs[:args.max_samples] logger.info(f" truncated to {len(ds.pairs)} pairs (smoke)") loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True, num_workers=2, collate_fn=collate, pin_memory=True) ck_d = torch.load(args.danger_ckpt, weights_only=False, map_location="cpu") dh = DangerHead(in_dim=ck_d["in_dim"]).to(device) dh.load_state_dict(ck_d["model"]); dh.eval() for p in dh.parameters(): p.requires_grad_(False) ck_p = torch.load(args.policy_warm, weights_only=False, map_location="cpu") ph_kwargs = dict( policy_dim=ck_p.get("policy_dim", 2560), perception_dim_per_query=ck_p.get("perception_dim_per_query", 512), k_queries=ck_p.get("k_queries", 4), ) policy = PolicyHeadV2(**ph_kwargs).to(device) policy.load_state_dict(ck_p["model"]) ref_policy = PolicyHeadV2(**ph_kwargs).to(device) ref_policy.load_state_dict(ck_p["model"]) ref_policy.eval() for p in ref_policy.parameters(): p.requires_grad_(False) logger.info(f" PolicyHead params: {sum(p.numel() for p in policy.parameters())/1e6:.2f} M (trainable)") # Load val cache once into CPU memory (5.6 GB but only forward-passed per epoch) val_cache = None if args.val_cache.exists() and args.epochs >= 1 and args.max_samples == 0: logger.info(f"[load] val_cache={args.val_cache}") val_cache = torch.load(args.val_cache, weights_only=False, map_location="cpu") logger.info(f" val N={len(val_cache['ids'])}") opt = torch.optim.AdamW(policy.parameters(), lr=args.lr, weight_decay=1e-5) n_steps = args.epochs * len(loader) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=n_steps) best_composite = -1e9 log_records = [] for ep in range(args.epochs): policy.train() run_loss = 0; run_dpo = 0; run_anc = 0; run_delta = 0; run_pref = 0; n_b = 0 pbar = tqdm(loader, ncols=80, desc=f"ep{ep}") for b in pbar: bc = b["belief"].to(device, dtype=torch.float32, non_blocking=True) pp = b["policy"].to(device, dtype=torch.float32, non_blocking=True) v = b["valid"].to(device, non_blocking=True) chosen = b["chosen"].to(device, non_blocking=True) rejected = b["rejected"].to(device, non_blocking=True) ta = b["tick_action"].to(device, non_blocking=True) prev = torch.full((bc.shape[0],), 3, dtype=torch.long, device=device) with torch.no_grad(): dh_out = dh(bc, valid_frames=v) logits = policy(pp, dh_out["perception_summary"], dh_out["per_frame"], prev, valid_frames=v) with torch.no_grad(): ref_logits = ref_policy(pp, dh_out["perception_summary"], dh_out["per_frame"], prev, valid_frames=v) dpo_l, stats = dpo_loss(logits, ref_logits, chosen, rejected, beta=args.beta) # ALERT-anchor CE loss: applied only on samples where the *true* # tick_action == 2 (real ALERT). This prevents DPO from drifting # the policy away from supervised behaviour on real-ALERT samples. anchor_mask = (ta == 2) if anchor_mask.any(): anchor_l = F.cross_entropy(logits[anchor_mask], ta[anchor_mask]) else: anchor_l = torch.zeros((), device=device) total = (args.rl_weight * dpo_l + (1 - args.rl_weight) * args.alert_anchor_weight * anchor_l) total.backward() torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0) opt.step(); sched.step(); opt.zero_grad(set_to_none=True) run_loss += total.item() run_dpo += dpo_l.item() run_anc += anchor_l.item() run_delta += stats["delta_mean"] run_pref += stats["prefers_chosen_rate"] n_b += 1 pbar.set_postfix(loss=run_loss/n_b, dpo=run_dpo/n_b, anc=run_anc/n_b) rec = { "epoch": ep, "train_loss": run_loss / max(1, n_b), "train_dpo": run_dpo / max(1, n_b), "train_anchor": run_anc / max(1, n_b), "delta_chosen_minus_rejected": run_delta / max(1, n_b), "prefers_chosen_rate": run_pref / max(1, n_b), } # Validation (universal balance gate metric) if val_cache is not None: val_m = evaluate_policy_on_val(policy, dh, val_cache, device, batch_size=256) rec["val"] = val_m logger.info(format_gate_row(val_m, tag=f"dpo ep{ep}")) composite = val_m["composite"] if composite > best_composite: best_composite = composite save_dict = { "model": policy.state_dict(), "policy_dim": ph_kwargs["policy_dim"], "perception_dim_per_query": ph_kwargs["perception_dim_per_query"], "k_queries": ph_kwargs["k_queries"], "args": vars(args), "epoch": ep, "val_metrics": val_m, "composite": composite, } torch.save(save_dict, args.out_dir / "best.pt") logger.info(f" [save best] composite={composite:.4f}") else: # No val cache (smoke) → save last save_dict = { "model": policy.state_dict(), "policy_dim": ph_kwargs["policy_dim"], "perception_dim_per_query": ph_kwargs["perception_dim_per_query"], "k_queries": ph_kwargs["k_queries"], "args": vars(args), "epoch": ep, } torch.save(save_dict, args.out_dir / "best.pt") log_records.append(rec) logger.info(f"[ep{ep}] train={rec['train_loss']:.4f} " f"dpo={rec['train_dpo']:.4f} anc={rec['train_anchor']:.4f}") (args.out_dir / "training_log.json").write_text( json.dumps(log_records, indent=2, default=str)) logger.info(f"\n[done] best composite={best_composite:.4f} saved to {args.out_dir}/best.pt") if __name__ == "__main__": main()