training/Policy/policy_model_v4.py

#!/usr/bin/env python3
"""
EvidentialPolicyModel — SFTModel + EvidentialPolicyHead (Dirichlet output).

Drop-in replacement for PolicyModel with evidential uncertainty output.
All SFT modules frozen; only EvidentialPolicyHead is trainable (~1.2M params).

Output: Dirichlet concentration α [B, 3] instead of logits [B, 3].
"""

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 EvidentialPolicyHead

logger = logging.getLogger("Policy.model_v4")

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:
    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 EvidentialPolicyModel(nn.Module):
    """
    Wraps frozen SFTModel and attaches a trainable EvidentialPolicyHead.
    Output is Dirichlet α [B, 3] instead of class logits.
    """

    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)

        for param in self.sft.parameters():
            param.requires_grad = False
        logger.info("  SFT parameters frozen.")

        self.policy_head = EvidentialPolicyHead(
            hidden_dim=self.sft.hidden_dim,
            num_actions=N_ACTIONS,
        ).to(self.sft.device, dtype=torch.float32)

        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"EvidentialPolicyModel ready.  "
            f"Trainable: {trainable:,} (EvidentialPolicyHead) / 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

    def _build_inputs(self, images: List[List], 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,
        )

    def forward(self, images: List[List], metadata: List[dict]) -> torch.Tensor:
        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)

        alpha = self.policy_head(
            belief.detach().float(),
            tta_mean_f.detach(),
            tta_var.detach(),
            prev_action,
        )
        return alpha  # [B, 3] Dirichlet concentration

    def forward_cached(
        self,
        beliefs: torch.Tensor,
        tta_means: torch.Tensor,
        tta_vars: torch.Tensor,
    ) -> torch.Tensor:
        dev = self.device
        B = beliefs.shape[0]
        prev_action = torch.zeros(B, dtype=torch.long, device=dev)

        alpha = self.policy_head(
            beliefs.to(dev),
            tta_means.to(dev),
            tta_vars.to(dev),
            prev_action,
        )
        return alpha  # [B, 3] Dirichlet concentration

    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"  EvidentialPolicyHead 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"  EvidentialPolicyHead loaded from {path}")