#!/usr/bin/env python3 """ PolicyModel — SFTModel with an attached trainable PolicyHead. Stage 1 architecture: Frozen : VLM + LoRA + BeliefAggregator (mean_pool) + HazardHead + TTAHead Trainable: PolicyHead only (~1.2M params) PolicyHead input (all from frozen SFT world model): belief [B, hidden_dim] — mean_pool of last hidden states tta_mean [B] — from frozen TTAHead (softplus, always positive) tta_var [B] — exp(tta_logvar), clamped for numerical safety prev_action [B] — constant SILENT (0) in Stage 1 (no temporal history) PolicyHead output: action_logits [B, 3] → SILENT=0 / OBSERVE=1 / ALERT=2 """ from __future__ import annotations import json import logging from pathlib import Path from typing import Any, Dict, List, Optional import torch import torch.nn as nn from torch.amp import autocast import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from training.SFT.trainer import SFTModel, load_sft_heads, _is_sft_ckpt_dir from lkalert.models.components import PolicyHead logger = logging.getLogger("Policy.model") SYSTEM = "You are a driving safety AI analyzing dashcam footage for collision risk." ACTION_NAMES = {0: "SILENT", 1: "OBSERVE", 2: "ALERT"} N_ACTIONS = 3 def _build_prompt(metadata: dict) -> str: """Build VLM prompt from window metadata. Identical to SFT/DPO trainers.""" parts = [] if metadata.get("weather"): parts.append(f"Weather: {metadata['weather']}") if metadata.get("road_type"): parts.append(f"Road: {metadata['road_type']}") if metadata.get("time_of_day"): parts.append(f"Time: {metadata['time_of_day']}") ctx = ", ".join(parts) or "Urban driving" return ( f"Analyze this driving sequence.\n" f"Context: {ctx}\n" f"Estimate the time to potential collision. Output a single number in seconds." ) class PolicyModel(nn.Module): """ Wraps SFTModel and attaches a trainable PolicyHead. All SFT modules are frozen. Only PolicyHead parameters receive gradients. BeliefAggregator strategy (mean_pool) is unchanged from SFT. """ def __init__( self, sft_checkpoint_dir: str, use_bf16: bool = True, ): super().__init__() ckpt = Path(sft_checkpoint_dir) if not _is_sft_ckpt_dir(ckpt): raise RuntimeError(f"Not a valid SFT checkpoint directory: {ckpt}") with open(ckpt / "config.json") as f: cfg = json.load(f) logger.info(f"Loading SFTModel from {ckpt} ...") self.sft = SFTModel( model_name = cfg["model_name"], pretrained_lora_path = str(ckpt / "vlm_lora"), belief_strategy = cfg.get("belief_strategy", "mean_pool"), tta_intermediate_dim = cfg.get("tta_intermediate_dim", 512), use_lora = True, use_bf16 = use_bf16, device = "auto", ) load_sft_heads(self.sft, ckpt) # ── freeze all SFT parameters ───────────────────────────────────────── for param in self.sft.parameters(): param.requires_grad = False logger.info(" SFT parameters frozen.") # ── attach trainable PolicyHead ──────────────────────────────────────── # hidden_dim + 2 (tta_mean, tta_var) + 16 (prev_action embedding) → 512 → 256 → 3 self.policy_head = PolicyHead( hidden_dim = self.sft.hidden_dim, num_actions = N_ACTIONS, ).to(self.sft.device, dtype=torch.float32) # PolicyHead runs in float32 for training stability even when SFT uses bf16 trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) total = sum(p.numel() for p in self.parameters()) logger.info( f"PolicyModel ready. " f"Trainable: {trainable:,} (PolicyHead) / Total: {total:,}" ) self.processor = self.sft.processor self.hidden_dim = self.sft.hidden_dim self._amp_dtype = torch.bfloat16 if use_bf16 else torch.float32 self._ckpt_dir = ckpt @property def device(self) -> torch.device: return self.sft.device # ── input builder ───────────────────────────────────────────────────────── def _build_inputs( self, images: List[List], # [B, n_frames] list of PIL images per sample metadata: List[dict], ) -> Dict[str, Any]: proc = self.processor apply_chat = ( proc.apply_chat_template if hasattr(proc, "apply_chat_template") else proc.tokenizer.apply_chat_template ) texts = [] for i in range(len(images)): frames = images[i] content = [{"type": "image"} for _ in range(len(frames))] content.append({"type": "text", "text": _build_prompt(metadata[i])}) msgs = [ {"role": "system", "content": SYSTEM}, {"role": "user", "content": content}, ] texts.append(apply_chat(msgs, tokenize=False, add_generation_prompt=False)) return proc( text=texts, images=images, return_tensors="pt", padding=True, truncation=True, ) # ── forward (image mode) ───────────────────────────────────────────────── def forward( self, images: List[List], # [B, n_frames] metadata: List[dict], ) -> torch.Tensor: """ Slow path: encodes images via frozen VLM, then runs PolicyHead. Used by make_belief_cache.py and evaluate_policy.py. Returns action_logits [B, 3] in float32. """ inputs = self._build_inputs(images, metadata) with torch.no_grad(): with autocast(device_type="cuda", dtype=self._amp_dtype, enabled=True): belief = self.sft.encode_observation(inputs) tta_mean, tta_logvar = self.sft.tta_head(belief) tta_var = torch.exp(tta_logvar.float().clamp(-20.0, 20.0)) tta_mean_f = tta_mean.float() B = belief.shape[0] prev_action = torch.zeros(B, dtype=torch.long, device=self.device) logits = self.policy_head( belief.detach().float(), tta_mean_f.detach(), tta_var.detach(), prev_action, ) return logits # [B, 3] # ── forward_cached (cache mode) ─────────────────────────────────────────── def forward_cached( self, beliefs: torch.Tensor, # [B, hidden_dim] pre-computed tta_means: torch.Tensor, # [B] tta_vars: torch.Tensor, # [B] ) -> torch.Tensor: """ Fast path: skips VLM, runs only PolicyHead on pre-computed beliefs. Used by warm_start_trainer.py when belief_cache is available. ~1000× faster than forward() for training. Returns action_logits [B, 3] in float32. """ dev = self.device B = beliefs.shape[0] prev_action = torch.zeros(B, dtype=torch.long, device=dev) logits = self.policy_head( beliefs.to(dev), tta_means.to(dev), tta_vars.to(dev), prev_action, ) return logits # [B, 3] # ── checkpointing ───────────────────────────────────────────────────────── def save_checkpoint(self, save_dir: str, meta: Optional[dict] = None): save_dir = Path(save_dir) save_dir.mkdir(parents=True, exist_ok=True) torch.save(self.policy_head.state_dict(), save_dir / "policy_head.pt") if meta is not None: with open(save_dir / "policy_meta.json", "w") as f: json.dump(meta, f, indent=2) logger.info(f" PolicyHead saved → {save_dir}") def load_policy_checkpoint(self, ckpt_dir: str): path = Path(ckpt_dir) / "policy_head.pt" if not path.exists(): raise FileNotFoundError(f"policy_head.pt not found in {ckpt_dir}") self.policy_head.load_state_dict( torch.load(path, map_location=self.device) ) logger.info(f" PolicyHead loaded from {path}")