File size: 8,806 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 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | """AdaptiveWindowModule — VLAlert-X core architectural innovation.
Maps (current policy distribution + hazard logits + belief summary) to a
window choice for the *next* tick:
w_{t+1} = AdaptiveWindow(pi_t, hazard_logits_t, belief_summary_t)
The next tick's belief vector is then extracted from frames sampled
according to w_{t+1} ∈ {narrow, mid, wide}. This closes the
"OBSERVE-as-action" loop: when the policy commits to OBSERVE, the
window narrows on the *next* tick, providing tighter temporal evidence
for the subsequent action decision.
Window index convention (matches build_adaptive_trajectories.py):
0 = narrow (1 s span, 8 frames at ~0.125 s stride)
1 = mid (2 s span, 8 frames at ~0.25 s stride) -- legacy default
2 = wide (4 s span, 8 frames at ~0.5 s stride)
Training protocol — 3-stage curriculum (see plan §3.2 of vlalert-x-upgrade.md):
Stage 1 (epoch 1-2): 100 % oracle window (deterministic from action)
Stage 2 (epoch 3-4): 50/50 oracle / student-predicted window
Stage 3 (epoch 5-6): 100 % student-predicted window (with
straight-through gradient on the discrete choice)
Hazard-conditional bias: at inference, the window logits are biased by
a learned per-hazard correction. The bias maps each of the 8 hazard
categories to a 3-D tilt over windows. Defaults (initialised from
empirical priors):
pedestrian / vrurider -> +1.0 bias on dim 0 (narrow)
vehicle_cross / oncoming -> +0.5 bias on dim 0 (narrow)
vehicle_lead -> +0.3 bias on dim 1 (mid)
weather / infrastructure -> +0.5 bias on dim 1 (mid)
none -> +1.0 bias on dim 2 (wide)
"""
from __future__ import annotations
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
# Window-index convention
WINDOW_NARROW = 0
WINDOW_MID = 1
WINDOW_WIDE = 2
# Hazard categories (matches Phase 1.1 GPT-5 schema)
HAZARD_PEDESTRIAN = 0
HAZARD_VRURIDER = 1
HAZARD_VEHICLE_CROSS = 2
HAZARD_VEHICLE_ONCOMING = 3
HAZARD_VEHICLE_LEAD = 4
HAZARD_WEATHER = 5
HAZARD_INFRASTRUCTURE = 6
HAZARD_NONE = 7
N_HAZARDS = 8
# Empirical hazard→window prior (used to initialise hazard_bias)
HAZARD_BIAS_INIT = torch.tensor([
# narrow, mid, wide
[ 1.0, 0.0, 0.0], # pedestrian
[ 1.0, 0.0, 0.0], # vrurider
[ 0.5, 0.5, 0.0], # vehicle_cross
[ 0.5, 0.5, 0.0], # vehicle_oncoming
[ 0.0, 0.5, 0.0], # vehicle_lead
[ 0.0, 0.5, 0.0], # weather
[ 0.0, 0.5, 0.0], # infrastructure
[ 0.0, 0.0, 1.0], # none
], dtype=torch.float32)
class AdaptiveWindowModule(nn.Module):
"""Lightweight MLP head that emits a 3-window choice.
Inputs:
pi_t : [B, 3] current-tick policy distribution (softmax)
hazard_logits: [B, 8] hazard-category logits from the SFT'd VLM
belief_summary: [B, D] mean-pooled belief at current tick (D=2560 for Qwen3-VL-4B)
Output:
window_logits: [B, 3] logits over {narrow, mid, wide}
"""
def __init__(self,
belief_dim: int = 2560,
hidden: int = 128,
dropout: float = 0.1,
use_hazard_bias: bool = True,
hazard_bias_lr_mult: float = 0.5):
super().__init__()
# Belief summariser (compresses 2560-D belief to 256-D)
self.belief_proj = nn.Sequential(
nn.Linear(belief_dim, 256),
nn.GELU(),
nn.LayerNorm(256),
)
# Main classifier: pi_t (3) + hazard_logits (8) + belief_proj (256) -> 3 windows
in_dim = 3 + N_HAZARDS + 256
self.mlp = nn.Sequential(
nn.Linear(in_dim, hidden),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden, 3),
)
# Hazard-conditional bias on window logits, initialised from empirical prior.
# Uses a smaller LR multiplier so the prior survives early epochs.
self.use_hazard_bias = use_hazard_bias
if use_hazard_bias:
self.hazard_bias = nn.Parameter(HAZARD_BIAS_INIT.clone())
self.hazard_bias_lr_mult = hazard_bias_lr_mult
def forward(self,
pi_t: torch.Tensor,
hazard_logits: torch.Tensor,
belief_summary: torch.Tensor) -> torch.Tensor:
"""Returns raw window logits [B, 3]."""
b_proj = self.belief_proj(belief_summary)
z = torch.cat([pi_t, hazard_logits, b_proj], dim=-1)
logits = self.mlp(z)
if self.use_hazard_bias:
# Soft hazard mixture: bias = hazard_softmax · HAZARD_BIAS_INIT [B, 3]
hazard_probs = F.softmax(hazard_logits, dim=-1) # [B, 8]
bias = hazard_probs @ self.hazard_bias # [B, 3]
logits = logits + bias
return logits
@torch.no_grad()
def predict_window(self,
pi_t: torch.Tensor,
hazard_logits: torch.Tensor,
belief_summary: torch.Tensor,
temperature: float = 1.0,
sample: bool = False) -> torch.Tensor:
"""Inference-time window choice as integer in {0,1,2}.
Args:
sample: if True, sample from softmax (Stage 2/3 of training-loop
with stochastic sampling); if False, take argmax (deployment).
"""
logits = self.forward(pi_t, hazard_logits, belief_summary) / max(temperature, 1e-3)
if sample:
probs = F.softmax(logits, dim=-1)
choice = torch.multinomial(probs, num_samples=1).squeeze(-1)
else:
choice = logits.argmax(dim=-1)
return choice
def param_groups(self, base_lr: float):
"""Yield optimiser param groups, applying lr-mult to hazard_bias."""
bias_params, other_params = [], []
for n, p in self.named_parameters():
if n.endswith("hazard_bias"):
bias_params.append(p)
else:
other_params.append(p)
groups = [{"params": other_params, "lr": base_lr}]
if bias_params:
groups.append({"params": bias_params,
"lr": base_lr * self.hazard_bias_lr_mult})
return groups
# ───────────────────────────── helpers ──────────────────────────────────
def oracle_window_from_action(action: torch.Tensor) -> torch.Tensor:
"""Map per-tick action label {0=SILENT, 1=OBSERVE, 2=ALERT} to window.
SILENT → wide (window_idx 2)
OBSERVE → mid (window_idx 1)
ALERT → narrow (window_idx 0)
"""
table = torch.tensor([WINDOW_WIDE, WINDOW_MID, WINDOW_NARROW],
dtype=torch.long, device=action.device)
return table[action.clamp(min=0, max=2)]
def scheduled_sampling_window(stage: int,
oracle_window: torch.Tensor,
student_window: torch.Tensor,
rng: Optional[torch.Generator] = None,
p_oracle_stage2: float = 0.5
) -> torch.Tensor:
"""Pick window per-tick according to curriculum stage.
Stage 1: 100 % oracle.
Stage 2: per-tick coin flip (p_oracle_stage2) between oracle / student.
Stage 3: 100 % student.
"""
if stage == 1:
return oracle_window
if stage == 3:
return student_window
# Stage 2: mixed
p = torch.rand(oracle_window.shape, generator=rng,
device=oracle_window.device)
return torch.where(p < p_oracle_stage2, oracle_window, student_window)
def straight_through_window_select(window_logits: torch.Tensor,
belief_per_window: torch.Tensor) -> torch.Tensor:
"""Differentiable window-conditioned belief lookup with straight-through.
Args:
window_logits : [B, 3]
belief_per_window : [B, 3, F, D] pre-computed beliefs for all 3 windows
Returns:
belief : [B, F, D] the chosen window's belief, with straight-through
gradient flowing back into window_logits.
"""
probs = F.softmax(window_logits, dim=-1) # [B, 3]
onehot = F.one_hot(window_logits.argmax(dim=-1), 3).float() # [B, 3]
# straight-through: forward = onehot, backward = softmax probs
soft = onehot + (probs - probs.detach())
soft = soft.unsqueeze(-1).unsqueeze(-1) # [B, 3, 1, 1]
belief = (belief_per_window * soft).sum(dim=1) # [B, F, D]
return belief
|