"""LKAlert-MCB head: gated multi-channel belief fusion. Day-11 baseline = 2 channels: Channel 1 (Qwen semantic): belief_seq [B, T, 2560] → POMDP trunk → 256 Channel 3 (V-JEPA dynamics): clip-level [B, 1024] → MLP → 256 Channel 2 (object motion) is NOT a learned input here — failed Day-10 gate. It can be re-introduced in Day-11.5 stretch via a teacher-trained critical_actor_selector + filtered features. Fusion modes (configurable): - "concat_mlp" [256+256] → MLP → 1 (default) - "gated_concat" per-channel gate g ∈ [0,1] then concat; the gate is learned from the joint state. Robust under `vjepa_mask=0` (V-JEPA missing). Output: a single binary collision logit `p_any`. Auxiliary slots (Day-11.5 stretch, controlled by `--with_teacher_aux`): - ego_relevance_logit (3-class CE) - path_conflict_logit (3-class CE) - risk_resolution_logit (3-class soft-label CE) - recommended_policy_logit (3-class CE) - tracking_assessment_logit (3-class CE) """ from __future__ import annotations from typing import Dict, Optional import torch import torch.nn as nn import torch.nn.functional as F class _QwenChannelTrunk(nn.Module): """Mirrors POMDPTemporalHead trunk: in_proj → GRU → masked attn pool. Returns the [B, gru_hidden] pooled state without the binary classifier.""" def __init__(self, in_dim: int = 2560, proj_dim: int = 512, gru_hidden: int = 256, dropout: float = 0.2): super().__init__() self.in_proj = nn.Sequential( nn.Linear(in_dim, proj_dim), nn.LayerNorm(proj_dim), nn.GELU(), nn.Dropout(dropout), ) self.text_proj = nn.Sequential( nn.Linear(in_dim, gru_hidden), nn.LayerNorm(gru_hidden), nn.Tanh(), ) self.gru = nn.GRU(proj_dim, gru_hidden, num_layers=1, batch_first=True) self.attn = nn.Linear(gru_hidden, 1) def forward(self, beliefs: torch.Tensor, valid: torch.Tensor, text: torch.Tensor) -> torch.Tensor: x = self.in_proj(beliefs) h0 = self.text_proj(text).unsqueeze(0).contiguous() out, _ = self.gru(x, h0) attn_logits = self.attn(out).squeeze(-1) attn_logits = attn_logits.masked_fill(~valid, float("-inf")) empty = (~valid).all(dim=1) if empty.any(): attn_logits[empty] = 0.0 w = F.softmax(attn_logits, dim=1).unsqueeze(-1) pooled = (out * w).sum(dim=1) return pooled # [B, gru_hidden] class _VJEPAChannel(nn.Module): """V-JEPA clip-level [B, 1024] → 256-D projection.""" def __init__(self, in_dim: int = 1024, out_dim: int = 256, dropout: float = 0.2): super().__init__() self.proj = nn.Sequential( nn.Linear(in_dim, 512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(dropout), nn.Linear(512, out_dim), nn.LayerNorm(out_dim), nn.GELU(), ) def forward(self, vjepa: torch.Tensor) -> torch.Tensor: return self.proj(vjepa) # [B, out_dim] class LKAlertMCB(nn.Module): """2-channel MCB head. Compatible with `multichannel_dataset` schema. Args: qwen_in_dim: Channel 1 belief feature dim (2560 for Qwen3-VL-4B). vjepa_in_dim: 1024 for V-JEPA frozen. use_vjepa: if False, the V-JEPA channel is replaced by zeros; used to ablate Channel 3 in the 8-row ablation matrix. use_qwen: if False, the Qwen channel is replaced by zeros; Day-11 ablation only — for Channel-3-only baseline. fusion: "concat_mlp" (default) or "gated_concat". with_teacher_aux: if True, adds 5 auxiliary slot heads (Day-11.5 stretch, gated on teacher pilot pass). """ def __init__(self, qwen_in_dim: int = 2560, proj_dim: int = 512, gru_hidden: int = 256, vjepa_in_dim: int = 1024, vjepa_out_dim: int = 256, dropout: float = 0.2, use_qwen: bool = True, use_vjepa: bool = True, fusion: str = "concat_mlp", with_teacher_aux: bool = False): super().__init__() assert fusion in ("concat_mlp", "gated_concat") self.use_qwen = use_qwen self.use_vjepa = use_vjepa self.fusion = fusion self.with_teacher_aux = with_teacher_aux self.qwen_trunk = _QwenChannelTrunk(in_dim=qwen_in_dim, proj_dim=proj_dim, gru_hidden=gru_hidden, dropout=dropout) self.vjepa_trunk = _VJEPAChannel(in_dim=vjepa_in_dim, out_dim=vjepa_out_dim, dropout=dropout) # gates (only used if fusion == "gated_concat") if fusion == "gated_concat": self.gate_qwen = nn.Linear(gru_hidden + vjepa_out_dim, 1) self.gate_vjepa = nn.Linear(gru_hidden + vjepa_out_dim, 1) clf_in = gru_hidden + vjepa_out_dim self.fuse_mlp = nn.Sequential( nn.Linear(clf_in, 128), nn.GELU(), nn.Dropout(dropout), ) self.head_p_any = nn.Linear(128, 1) # Day-11.5 stretch heads — present iff `with_teacher_aux=True` if with_teacher_aux: self.head_ego_relevance = nn.Linear(128, 3) # ego/non_ego/ambiguous self.head_path_conflict = nn.Linear(128, 3) # none/potential/active self.head_risk_resolution = nn.Linear(128, 3) # not/partial/resolved self.head_recommended_policy = nn.Linear(128, 3) # SILENT/OBSERVE/ALERT self.head_tracking_assessment = nn.Linear(128, 3) # yes/no/unclear # ────────────────────────────────────────────────────────────────────── def forward(self, beliefs: torch.Tensor, # [B, T, qwen_in_dim] valid: torch.Tensor, # [B, T] text: torch.Tensor, # [B, qwen_in_dim] vjepa: torch.Tensor, # [B, vjepa_in_dim] vjepa_mask: torch.Tensor, # [B] (1.0 if present) ) -> Dict[str, torch.Tensor]: B = beliefs.shape[0] # Channel 1 (Qwen) q_pool = self.qwen_trunk(beliefs, valid, text) # [B, H_q] if not self.use_qwen: q_pool = torch.zeros_like(q_pool) # Channel 3 (V-JEPA) v_pool = self.vjepa_trunk(vjepa) # [B, H_v] # mask out missing V-JEPA samples v_pool = v_pool * vjepa_mask.unsqueeze(-1) if not self.use_vjepa: v_pool = torch.zeros_like(v_pool) if self.fusion == "gated_concat": joint = torch.cat([q_pool, v_pool], dim=-1) g_q = torch.sigmoid(self.gate_qwen(joint)) g_v = torch.sigmoid(self.gate_vjepa(joint)) q_pool = q_pool * g_q v_pool = v_pool * g_v joint = torch.cat([q_pool, v_pool], dim=-1) # [B, H_q+H_v] h = self.fuse_mlp(joint) # [B, 128] out: Dict[str, torch.Tensor] = { "p_any": self.head_p_any(h).squeeze(-1), # [B] "fused": h, } if self.with_teacher_aux: out["ego_relevance_logits"] = self.head_ego_relevance(h) out["path_conflict_logits"] = self.head_path_conflict(h) out["risk_resolution_logits"] = self.head_risk_resolution(h) out["recommended_policy_logits"] = self.head_recommended_policy(h) out["tracking_assessment_logits"] = self.head_tracking_assessment(h) return out # ── warm-start from LKAlert-BD trunk ────────────────────────────────── def warm_start_qwen_trunk_from_bd(self, bd_state_dict: Dict[str, torch.Tensor]): """Copy Qwen trunk weights from a `lkalert_bd_best/best.pt` head_state.""" my_sd = self.qwen_trunk.state_dict() copied = [] for k in my_sd: full = f"qwen_trunk.{k}" # BD trunk parameters live under in_proj.* / text_proj.* / gru.* / attn.* # — same names as POMDPTemporalHead. if k in bd_state_dict and bd_state_dict[k].shape == my_sd[k].shape: my_sd[k] = bd_state_dict[k].clone() copied.append(k) self.qwen_trunk.load_state_dict(my_sd) return copied