training/Policy/train_head_dpo.py

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