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