File size: 9,099 Bytes
1e05592 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """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
|