"""VLAlert-X v2 Phase 3 — Danger Head. Continuous per-frame and clip-level risk regressor on BELIEF_CONTENT features (the perception/risk-cue register from Phase 2 cache). Supervision: TTA-derived continuous danger ∈ [0, 1] danger[f] = sigmoid(4 * (L_alert - tta_f) / L_alert) for tta in (0, 5] danger[f] = 0.05 (floor) for SILENT clips danger[f] = 1.0 for post-event frames This is an interpretable, threshold-free risk score that the downstream Policy Head (Phase 4) consumes as an input feature. It also exposes a clip-level scalar useful as a fallback alert score (e.g., for ablations where Policy Head is removed). Architecture: BELIEF_CONTENT [B, 8, 10240] │ ├──> per-frame MLP ──> [B, 8] sigmoid (per-frame danger) │ └──> MultiQueryPMA (K=4) ──> [B, 4, 512] (perception_summary) │ └──> clip MLP ──> [B] sigmoid (clip danger) The `perception_summary` is returned alongside heads so the Policy Head (Phase 4) can re-use it without re-running the PMA aggregator. """ from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F class MultiQueryPMAAggregator(nn.Module): """Multi-query Pooling by Multi-head Attention (PMA, Lee et al. 2019). K learnable query vectors attend to the per-frame tokens to produce K summary vectors. Simpler and more parameter-efficient than a full Transformer encoder for fixed-length pooling. """ def __init__(self, in_dim: int, k_queries: int = 4, out_dim: int = 512, dropout: float = 0.1): super().__init__() self.k = k_queries self.out_dim = out_dim # Project input → out_dim self.in_proj = nn.Linear(in_dim, out_dim) # K learnable query vectors self.queries = nn.Parameter(torch.randn(k_queries, out_dim) * 0.02) self.attn = nn.MultiheadAttention(out_dim, num_heads=4, dropout=dropout, batch_first=True) self.norm = nn.LayerNorm(out_dim) self.ffn = nn.Sequential( nn.Linear(out_dim, out_dim * 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(out_dim * 2, out_dim)) self.norm2 = nn.LayerNorm(out_dim) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: """ x: [B, T, in_dim] — per-frame features mask: [B, T] — True = valid frame returns: [B, K, out_dim] """ B = x.size(0) h = self.in_proj(x) # [B, T, D] q = self.queries.unsqueeze(0).expand(B, -1, -1) # [B, K, D] key_padding_mask = None if mask is not None: key_padding_mask = ~mask # True = pad attn_out, _ = self.attn(q, h, h, key_padding_mask=key_padding_mask) h2 = self.norm(q + attn_out) h3 = self.norm2(h2 + self.ffn(h2)) return h3 # [B, K, D] class DangerHead(nn.Module): """Continuous risk regressor on BELIEF_CONTENT features. Args: in_dim: hidden dim of BELIEF_CONTENT (default 10240 for L4 concat) hidden: internal width k_queries: number of PMA queries dropout: dropout rate n_hazards: if > 0, also emit a k-way hazard classification logit over the AdaptiveWindow 8-way taxonomy (Phase G.0). New tensor in output dict: 'hazard_logits' [B, n_hazards]. Backward-compatible: defaults to 0 → no hazard head. """ def __init__(self, in_dim: int = 10240, hidden: int = 512, k_queries: int = 4, dropout: float = 0.2, n_hazards: int = 0): super().__init__() self.n_hazards = n_hazards # Per-frame head (no aggregation — independent per frame) self.frame_proj = nn.Sequential( nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, hidden // 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden // 2, 1)) # logit # Cross-frame perception summary (PMA) self.pma = MultiQueryPMAAggregator( in_dim=in_dim, k_queries=k_queries, out_dim=hidden, dropout=dropout) # Clip-level head consumes flattened PMA output self.clip_mlp = nn.Sequential( nn.Linear(hidden * k_queries, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, 1)) # logit # Phase G.0: optional 8-way hazard classification head if n_hazards > 0: self.hazard_head = nn.Sequential( nn.Linear(hidden * k_queries, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, n_hazards)) # logits def forward(self, belief_content: torch.Tensor, valid_frames: torch.Tensor | None = None) -> dict: """ belief_content: [B, 8, in_dim] valid_frames: [B, 8] bool (True = valid) Returns: { "per_frame": [B, 8] sigmoid prob "per_frame_logits": [B, 8] "clip": [B] sigmoid prob "clip_logit": [B] "perception_summary": [B, K, hidden] for downstream re-use "hazard_logits": [B, n_hazards] (only if n_hazards > 0) } """ # per-frame: apply MLP independently per_frame_logits = self.frame_proj(belief_content).squeeze(-1) # [B, 8] per_frame = torch.sigmoid(per_frame_logits) # perception summary via PMA pooled = self.pma(belief_content, mask=valid_frames) # [B, K, H] clip_logit = self.clip_mlp(pooled.flatten(1)).squeeze(-1) # [B] clip = torch.sigmoid(clip_logit) out = { "per_frame": per_frame, "per_frame_logits": per_frame_logits, "clip": clip, "clip_logit": clip_logit, "perception_summary": pooled, } if self.n_hazards > 0: out["hazard_logits"] = self.hazard_head(pooled.flatten(1)) return out def danger_loss(out: dict, danger_per_frame: torch.Tensor, valid_frames: torch.Tensor | None = None, w_clip: float = 0.5) -> dict: """BCE on per-frame + BCE on clip-level (clip target = max over frames). out: output dict of DangerHead.forward danger_per_frame: [B, 8] continuous targets in [0, 1] valid_frames: [B, 8] bool Returns dict with 'loss', 'frame_loss', 'clip_loss'. """ pf = out["per_frame_logits"] if valid_frames is not None: frame_target = danger_per_frame.clamp(0.0, 1.0) # mask invalid frames to zero contribution loss_per = F.binary_cross_entropy_with_logits( pf, frame_target, reduction="none") loss_per = loss_per * valid_frames.float() denom = valid_frames.float().sum().clamp(min=1.0) frame_loss = loss_per.sum() / denom else: frame_loss = F.binary_cross_entropy_with_logits( pf, danger_per_frame.clamp(0.0, 1.0)) clip_target = danger_per_frame.max(dim=1).values.clamp(0.0, 1.0) clip_loss = F.binary_cross_entropy_with_logits(out["clip_logit"], clip_target) return { "loss": frame_loss + w_clip * clip_loss, "frame_loss": frame_loss.detach(), "clip_loss": clip_loss.detach(), }