training/Policy/train_head_kto.py

"""Head-RL KTO — train PolicyHeadV2 with KTO objective on labeled samples.

KTO (Kahneman-Tversky Optimization) does NOT require paired preferences.
Each sample has a binary `label` ∈ {desirable, undesirable}, and the loss
is asymmetric between desirable and undesirable cases.

Loss (Ethayarajh et al. 2024):
  For desirable y_w:    loss = λ_d · σ(-β · ( log π(y_w|x) − log π_ref(y_w|x) − z_ref ))
  For undesirable y_l:  loss = λ_u · σ( β · ( log π(y_l|x) − log π_ref(y_l|x) − z_ref ))

where z_ref is the KL anchor (avg over batch). β=0.1 typical.

Input: preference_kto.jsonl with {video_id, completion, label} per row.
       The `completion` is the action token (we map to {S,O,A} class index).

Usage:
    python -m training.Policy.train_head_kto \
        --pref_jsonl data/cot_corpus_v2/preference_kto.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_kto
"""
from __future__ import annotations

import argparse
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_kto")

ACTION_NAME_TO_IDX = {"SILENT": 0, "OBSERVE": 1, "ALERT": 2}


def extract_completion_action(completion: str) -> int:
    """The completion contains BELIEF blocks; find the LAST action token's class."""
    # Find last occurrence of one of the action tokens
    last_pos = -1
    last_act = "SILENT"
    for act_name, _ in ACTION_NAME_TO_IDX.items():
        tok = f"<|{act_name}|>"
        pos = completion.rfind(tok)
        if pos > last_pos:
            last_pos = pos
            last_act = act_name
    return ACTION_NAME_TO_IDX[last_act]


class KTODataset(Dataset):
    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.samples = []
        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]
                self.samples.append({
                    "cache_idx": ci,
                    "action_idx": extract_completion_action(obj["completion"]),
                    "label": bool(obj["label"]),
                    "tick_action": int(self.cache["tick_action"][ci]),
                })

        # OBSERVE oversample: duplicate samples whose action_idx == OBSERVE
        # AND label is True (so we reinforce desirable-OBSERVE more strongly)
        if observe_oversample > 1:
            base = list(self.samples)
            for s in base:
                if s["action_idx"] == 1 and s["label"]:
                    self.samples.extend([s] * (observe_oversample - 1))

        n_pos = sum(1 for s in self.samples if s["label"])
        n_obs = sum(1 for s in self.samples if s["action_idx"] == 1)
        logger.info(f"  loaded {len(self.samples)} samples "
                     f"(skipped {skipped}, pos={n_pos}, "
                     f"neg={len(self.samples)-n_pos}, action=OBS:{n_obs})")

    def __len__(self): return len(self.samples)
    def __getitem__(self, idx):
        s = self.samples[idx]
        ci = s["cache_idx"]
        return {
            "belief": self.cache["belief_content"][ci],
            "policy": self.cache["policy_position"][ci],
            "valid": self.cache["valid_frames"][ci],
            "action_idx": s["action_idx"],
            "label": s["label"],
            "tick_action": s["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]),
        "action_idx": torch.tensor([b["action_idx"] for b in batch], dtype=torch.long),
        "label": torch.tensor([b["label"] for b in batch], dtype=torch.bool),
        "tick_action": torch.tensor([b["tick_action"] for b in batch], dtype=torch.long),
    }


def kto_loss(logits, ref_logits, action_idx, label, beta=0.1,
              lambda_d=1.0, lambda_u=1.0):
    """KTO loss. logits/ref_logits: [B, 3]; action_idx: [B]; label: [B] bool."""
    log_p = F.log_softmax(logits, dim=-1)
    log_p_ref = F.log_softmax(ref_logits, dim=-1)
    B = logits.shape[0]
    idx = torch.arange(B, device=logits.device)

    log_p_y     = log_p[idx, action_idx]
    log_p_ref_y = log_p_ref[idx, action_idx]
    delta = log_p_y - log_p_ref_y

    # z_ref = KL(π || π_ref) batch average (detached)
    with torch.no_grad():
        kl = (log_p.exp() * (log_p - log_p_ref)).sum(dim=-1)
        z_ref = kl.mean()

    # σ(-β·(delta - z_ref)) for desirable; σ( β·(delta - z_ref)) for undesirable
    arg = beta * (delta - z_ref)
    pos_loss = lambda_d * torch.sigmoid(-arg)
    neg_loss = lambda_u * torch.sigmoid(arg)
    loss_per = torch.where(label, pos_loss, neg_loss)
    loss = loss_per.mean()

    with torch.no_grad():
        delta_mean = delta.mean().item()
        kl_mean = z_ref.item()
        pos_frac = label.float().mean().item()
    return loss, {"delta_mean": delta_mean, "kl_mean": kl_mean,
                   "pos_frac": pos_frac}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pref_jsonl", type=Path,
                    default=ROOT / "data/cot_corpus_v2/preference_kto.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)
    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)
    ap.add_argument("--alert_anchor_weight", type=float, default=1.0)
    ap.add_argument("--oversample_observe", type=int, default=3)
    ap.add_argument("--max_samples", type=int, default=0)
    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"

    ds = KTODataset(args.pref_jsonl, args.train_cache,
                     observe_oversample=args.oversample_observe)
    if args.max_samples > 0:
        ds.samples = ds.samples[:args.max_samples]
        logger.info(f"  truncated to {len(ds.samples)} samples (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)

    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)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(
        opt, T_max=args.epochs * len(loader))

    best_composite = -1e9
    log_records = []
    for ep in range(args.epochs):
        policy.train()
        run = {"loss": 0, "kto": 0, "anc": 0, "delta": 0, "kl": 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)
            ai = b["action_idx"].to(device, non_blocking=True)
            lbl = b["label"].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)

            kto_l, stats = kto_loss(logits, ref_logits, ai, lbl, beta=args.beta)

            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 * kto_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["kto"] += kto_l.item()
            run["anc"] += anchor_l.item()
            run["delta"] += stats["delta_mean"]
            run["kl"] += stats["kl_mean"]
            n_b += 1
            pbar.set_postfix(loss=run["loss"]/n_b, kto=run["kto"]/n_b,
                             anc=run["anc"]/n_b)

        rec = {"epoch": ep,
               "train_loss": run["loss"]/max(1,n_b),
               "train_kto": run["kto"]/max(1,n_b),
               "train_anchor": run["anc"]/max(1,n_b),
               "delta": run["delta"]/max(1,n_b),
               "kl_mean": run["kl"]/max(1,n_b)}

        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"kto 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:
            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"kto={rec['train_kto']:.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()