lkalert/models/danger_head.py

"""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(),
    }