File size: 8,852 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | #!/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}")
|