training/Policy/policy_model.py

#!/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}")