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