lkalert/models/adaptive_window.py

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