lkalert/models/multichannel_belief.py

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