lkalert/models/policy_head_v2.py

"""VLAlert-X v2 Phase 4 — Policy Head with dual-stream + danger conditioning.

Inputs (per tick):
  • POLICY_POSITION[B, 8, 2560]  — decision-time register from cache
  • perception_summary[B, 4, 512] — from frozen DangerHead (PMA pooled)
  • danger_per_frame[B, 8]        — from frozen DangerHead (continuous)
  • prev_action[B] long           — previous tick's action (0/1/2 or BOS=3)

Architecture:
        POLICY_POSITION ──> GRU(2 layers, h=512) ──> last_state [B, 512]
                                                          │
        perception_summary ──> proj [B, 256] ─────────────┤
                                                          ▼
                  [last_state, percep, danger, prev_act] ── MLP ── [B, 3]

Loss: CE with class-balanced weights + label smoothing + entropy reg.
The frozen DangerHead provides perception_summary and danger_per_frame as
pre-computed features (just forward DangerHead once on cached
belief_content, then save). Policy Head's gradient does not flow into
DangerHead.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


class PolicyHeadV2(nn.Module):
    def __init__(self,
                 policy_dim: int = 2560,
                 perception_dim_per_query: int = 512,
                 k_queries: int = 4,
                 prev_act_emb: int = 16,
                 gru_hidden: int = 512,
                 n_classes: int = 3,
                 dropout: float = 0.2,
                 with_anticipation: bool = False):
        super().__init__()
        # Temporal GRU on POLICY_POSITION
        self.gru = nn.GRU(policy_dim, gru_hidden, num_layers=2,
                            batch_first=True, dropout=dropout)
        # Project perception summary (PMA flat) to a compact vector
        self.perception_proj = nn.Sequential(
            nn.Linear(perception_dim_per_query * k_queries, 256),
            nn.GELU(),
            nn.LayerNorm(256),
            nn.Dropout(dropout),
        )
        # Previous-action embedding (BOS index = n_classes)
        self.action_emb = nn.Embedding(n_classes + 1, prev_act_emb)

        # Fusion + classifier
        # input dim = gru_hidden + 256 + 8 (danger_pf) + prev_act_emb
        fuse_in = gru_hidden + 256 + 8 + prev_act_emb
        self.fuse_pre = nn.Sequential(
            nn.Linear(fuse_in, 256), nn.GELU(),
            nn.Dropout(dropout),
        )
        self.cls_head = nn.Linear(256, n_classes)

        # Optional anticipation aux head: predicts whether the NEXT tick is
        # ALERT-class (binary). OBSERVE samples whose next tick is ALERT should
        # have high anticipation score; this encourages OBSERVE-as-anticipation.
        self.with_anticipation = with_anticipation
        if with_anticipation:
            self.anticipation_head = nn.Linear(256, 1)

    # Backwards-compat alias so old code referencing `policy.fuse` keeps working.
    @property
    def fuse(self) -> nn.Module:
        return nn.Sequential(self.fuse_pre, self.cls_head)

    def forward(self,
                policy_position: torch.Tensor,        # [B, 8, 2560]
                perception_summary: torch.Tensor,     # [B, K, perc_dim]
                danger_per_frame: torch.Tensor,       # [B, 8]
                prev_action: torch.Tensor,            # [B] long
                valid_frames: torch.Tensor | None = None,
                return_aux: bool = False,
                ):
        # Zero out clamped / invalid timesteps before the GRU so the recurrent
        # hidden state isn't poisoned by duplicate-padded boundary frames. This
        # was the root cause of the streaming demo's all-SILENT collapse: at
        # tick_t < window_span, 5-6/8 frames are clamped to frame=0 and the GRU
        # was processing 6 duplicates as a real temporal sequence.
        if valid_frames is not None:
            mask = valid_frames.unsqueeze(-1).to(policy_position.dtype)
            policy_position = policy_position * mask
        gru_out, _ = self.gru(policy_position)         # [B, 8, gru_hidden]
        # Pick the *latest* valid timestep — `sum(valid) - 1` is only correct
        # when valid frames are contiguous at the start; in streaming, clamped
        # frames sit at the BEGINNING (e.g. valid=[F,F,T,T,T,T,T,T] at boundary
        # ticks), so we instead find the highest index where valid is True.
        if valid_frames is not None:
            T = valid_frames.shape[1]
            idx_t = torch.arange(T, device=valid_frames.device).expand_as(valid_frames)
            masked = torch.where(valid_frames, idx_t, torch.full_like(idx_t, -1))
            last_idx = masked.max(dim=1).values.clamp(min=0)
            last_state = gru_out[torch.arange(gru_out.size(0)), last_idx]
        else:
            last_state = gru_out[:, -1]
        percep = self.perception_proj(perception_summary.flatten(1))    # [B, 256]
        prev = self.action_emb(prev_action)                              # [B, emb]
        fused = torch.cat([last_state, percep, danger_per_frame, prev], dim=-1)
        h = self.fuse_pre(fused)                                          # [B, 256]
        logits = self.cls_head(h)                                         # [B, 3]
        if return_aux and self.with_anticipation:
            antic_logit = self.anticipation_head(h).squeeze(-1)           # [B]
            return logits, antic_logit
        return logits


def policy_loss(logits: torch.Tensor,
                  targets: torch.Tensor,
                  class_weights: torch.Tensor | None = None,
                  label_smoothing: float = 0.05,
                  entropy_reg: float = 0.02,
                  use_focal: bool = False,
                  focal_gamma: float = 2.0,
                  focal_alpha: torch.Tensor | None = None,
                  use_ordinal: bool = False,
                  ordinal_margin: float = 1.0,
                  ordinal_lax: float = 0.5,
                  ordinal_weight: float = 0.5,
                  antic_logit: torch.Tensor | None = None,
                  antic_target: torch.Tensor | None = None,
                  antic_weight: float = 0.3,
                  prev_p_alert: torch.Tensor | None = None,
                  cur_p_alert: torch.Tensor | None = None,
                  temporal_weight: float = 0.1) -> dict:
    """Composite loss for OBSERVE-encouraging supervised training.

    Components (each optional, controlled by flag):
      - Base CE (or Focal CE) with class weights + label smoothing
      - Entropy regulariser (keep policy soft for RL warm-start)
      - Ordinal margin: penalise "skip OBSERVE" predictions
      - Anticipation aux: BCE on "next tick is ALERT" logit
      - Temporal consistency: penalise negative P(ALERT) jumps in consecutive ticks

    Args:
      use_focal:       if True replace CE with focal-CE (γ=focal_gamma).
      focal_alpha:     per-class weight tensor [3]. SILENT/OBSERVE/ALERT
                       suggested (1.0, 2.5, 1.5).
      use_ordinal:     if True add ordinal-margin loss enforcing logit
                       SILENT < OBSERVE < ALERT ordering.
      ordinal_margin:  required gap between predicted class and the *correct*
                       neighbour (e.g. OBSERVE must beat SILENT by margin).
      ordinal_lax:     allowed slack for "non-correct neighbour" (e.g. OBSERVE
                       can be ≤ ALERT but not by more than `ordinal_lax`).
      antic_logit:     [B] anticipation head logits (None to skip).
      antic_target:    [B] {0,1} target: 1 if next-tick is ALERT-class.
      prev_p_alert:    [B] P(ALERT) of previous tick in the same video
                       (None to skip temporal consistency).
      cur_p_alert:     [B] P(ALERT) of current tick.
      temporal_weight: weight on temporal-consistency penalty.
    """
    log_p = F.log_softmax(logits, dim=-1)
    probs = log_p.exp()

    # ── base CE / focal CE ────────────────────────────────────────────────
    if use_focal:
        # focal: α_c · (1 - p_y)^γ · -log p_y   per sample
        p_y = probs.gather(1, targets.unsqueeze(1)).squeeze(1).clamp(min=1e-8)
        focal_w = (1.0 - p_y).pow(focal_gamma)
        log_p_y = log_p.gather(1, targets.unsqueeze(1)).squeeze(1)
        if focal_alpha is not None:
            a = focal_alpha.to(logits.device).gather(0, targets)
            ce_per = -a * focal_w * log_p_y
        else:
            ce_per = -focal_w * log_p_y
        # apply optional class_weights on top (acts like a sample weight)
        if class_weights is not None:
            cw = class_weights.to(logits.device).gather(0, targets)
            ce_per = ce_per * cw
        ce = ce_per.mean()
    else:
        ce = F.cross_entropy(logits, targets, weight=class_weights,
                              label_smoothing=label_smoothing)

    # ── ordinal margin ────────────────────────────────────────────────────
    # Enforce logit[SIL] < logit[OBS] < logit[ALR] near the target.
    ord_loss = logits.new_zeros(())
    if use_ordinal:
        l_sil = logits[:, 0]
        l_obs = logits[:, 1]
        l_alr = logits[:, 2]
        sil_mask = (targets == 0)
        obs_mask = (targets == 1)
        alr_mask = (targets == 2)

        # When GT=SILENT: require l_sil > l_obs by margin, l_obs > l_alr by lax
        if sil_mask.any():
            ord_loss = ord_loss + F.relu(
                (l_obs[sil_mask] - l_sil[sil_mask]) + ordinal_margin
            ).mean()
            ord_loss = ord_loss + F.relu(
                (l_alr[sil_mask] - l_obs[sil_mask]) + ordinal_lax
            ).mean() * 0.5
        # When GT=OBSERVE: require l_obs > l_sil by margin AND l_obs ≥ l_alr - lax
        if obs_mask.any():
            ord_loss = ord_loss + F.relu(
                (l_sil[obs_mask] - l_obs[obs_mask]) + ordinal_margin
            ).mean()
            ord_loss = ord_loss + F.relu(
                (l_alr[obs_mask] - l_obs[obs_mask]) - ordinal_lax  # allow slight ALR > OBS
            ).clamp(min=0).mean() * 0.5
        # When GT=ALERT: require l_alr > l_obs by margin, l_obs > l_sil by lax
        # (penalise SILENT→ALERT skip: l_sil ≥ l_obs is the skip pattern)
        if alr_mask.any():
            ord_loss = ord_loss + F.relu(
                (l_obs[alr_mask] - l_alr[alr_mask]) + ordinal_margin
            ).mean()
            ord_loss = ord_loss + F.relu(
                (l_sil[alr_mask] - l_obs[alr_mask]) + ordinal_lax
            ).mean()  # strong penalty: SILENT > OBSERVE under ALERT GT is the skip pattern

    # ── anticipation aux ──────────────────────────────────────────────────
    antic_loss = logits.new_zeros(())
    if antic_logit is not None and antic_target is not None:
        antic_loss = F.binary_cross_entropy_with_logits(
            antic_logit, antic_target.float()
        )

    # ── temporal consistency ──────────────────────────────────────────────
    temp_loss = logits.new_zeros(())
    if prev_p_alert is not None and cur_p_alert is not None:
        delta = cur_p_alert - prev_p_alert
        # penalise *negative* jumps (P(ALERT) dropping too fast = risk denial)
        # AND large positive jumps (SILENT→ALERT skip)
        temp_loss = (F.relu(-delta).pow(2).mean()
                     + F.relu(delta - 0.5).pow(2).mean())

    # ── entropy regulariser ───────────────────────────────────────────────
    entropy = -(probs * (probs + 1e-9).log()).sum(dim=-1).mean()

    total = (ce
             + (ordinal_weight if use_ordinal else 0.0) * ord_loss
             + (antic_weight if antic_logit is not None else 0.0) * antic_loss
             + temporal_weight * temp_loss
             - entropy_reg * entropy)

    return {
        "loss":      total,
        "ce":        ce.detach(),
        "ordinal":   ord_loss.detach(),
        "antic":     antic_loss.detach(),
        "temporal":  temp_loss.detach(),
        "entropy":   entropy.detach(),
    }


# Recommended per-class Focal α for the 9k legacy class distribution
# (SILENT 41% / OBSERVE 18% / ALERT 40%). Sets OBSERVE 2.5× stronger.
FOCAL_ALPHA_9K = torch.tensor([1.0, 2.5, 1.5], dtype=torch.float32)