"""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)