training/Policy/dynamic_features.py

"""Dynamic-residual features for LKAlert-BD.

Day-1 diagnostic showed mean adjacent cosine distance ≈ 0.03 across all
caches — the per-frame Qwen3-VL belief is dynamically smooth. The whole
"belief is too invariant" observation is concrete evidence that the GRU
head can't recover motion residual on its own. This module builds explicit
hand-crafted features from the belief sequence so a small MLP can decide
how much motion residual is recoverable from the existing belief cache.

Features (all differentiable, all derivable from `beliefs_frame [B, T, D]`
and `valid_frames [B, T]`):

  belief-pool channels (length D):
    b_last         : last valid belief
    b_first        : first valid belief
    b_mean         : valid-mean
    b_max          : valid-max-pool (per-dim)
    delta_last     : b_last - b_first      (motion direction over the clip)

  scalar dynamics (length 1 each):
    mean_adj_cos_dist   : mean cosine distance between adjacent valid frames
    p95_adj_cos_dist    : 95-percentile of same
    max_norm_jump       : max ||b_t - b_{t-1}|| / max_t ||b_t||
    mean_norm_slope     : (||b_last|| - ||b_first||) / max(1, T-1)
    n_valid             : number of valid frames

  optional TTA channels (length 2 each, if `tta_means`/`tta_vars` provided):
    tta_mean_last, tta_var_last
    tta_mean_max,  tta_var_max
    tta_mean_first
    tta_mean_slope (last - first) / valid steps

This module exposes only `build_features(...)`. It returns a single
dict and never makes architectural decisions for the caller.
"""
from __future__ import annotations

from typing import Dict, Optional

import torch
import torch.nn.functional as F


@torch.no_grad()
def _adjacent_cos_distance(b: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
    """[B,T,D] → [B,T-1] adjacent (1-cos), 0 where either side invalid."""
    eps = 1e-6
    bn = b / b.norm(dim=-1, keepdim=True).clamp(min=eps)
    cos = (bn[:, 1:] * bn[:, :-1]).sum(dim=-1)         # [B, T-1]
    pair = (valid[:, 1:] & valid[:, :-1]).float()      # [B, T-1]
    return (1.0 - cos) * pair                          # 0 where invalid


def build_features(
    beliefs:     torch.Tensor,                 # [B, T, D]
    valid:       torch.Tensor,                 # [B, T] bool
    tta_means:   Optional[torch.Tensor] = None, # [B, T]
    tta_vars:    Optional[torch.Tensor] = None, # [B, T]
) -> Dict[str, torch.Tensor]:
    """Returns a dict of named feature tensors. All keys are length-axis B.

    `pooled` is a single concatenated [B, F] tensor for downstream MLPs.
    """
    B, T, D = beliefs.shape
    valid_f = valid.float().unsqueeze(-1)                       # [B, T, 1]
    n_valid = valid_f.sum(dim=1).squeeze(-1).clamp(min=1.0)     # [B]

    # last valid index per row (fallback to T-1 if all invalid)
    pos = torch.arange(T, device=beliefs.device).unsqueeze(0).expand(B, T)
    last_idx = (pos * valid.long()).max(dim=1).values           # [B]
    first_idx = (pos.masked_fill(~valid, T) ).min(dim=1).values # [B]
    first_idx = first_idx.clamp(max=T - 1)

    bidx = torch.arange(B, device=beliefs.device)
    b_last  = beliefs[bidx, last_idx]                           # [B, D]
    b_first = beliefs[bidx, first_idx]                          # [B, D]
    b_mean  = (beliefs * valid_f).sum(dim=1) / n_valid.unsqueeze(-1)   # [B, D]
    # masked max
    masked = beliefs.masked_fill(~valid.unsqueeze(-1), float("-inf"))
    b_max   = masked.max(dim=1).values                          # [B, D]
    # if a row had no valid frames the max collapses to -inf — recover with mean
    b_max   = torch.where(b_max == float("-inf"), b_mean, b_max)
    delta_last = b_last - b_first                               # [B, D]

    # scalar dynamics
    adj = _adjacent_cos_distance(beliefs, valid)                # [B, T-1]
    pair_count = (valid[:, 1:] & valid[:, :-1]).float().sum(dim=1).clamp(min=1.0)
    mean_adj = adj.sum(dim=1) / pair_count                      # [B]
    p95_adj  = torch.quantile(adj, q=0.95, dim=1)               # [B]

    norm_t = beliefs.norm(dim=-1)                                # [B, T]
    max_norm = norm_t.max(dim=1).values.clamp(min=1e-6)         # [B]
    diffs = (beliefs[:, 1:] - beliefs[:, :-1]).norm(dim=-1)      # [B, T-1]
    pair_mask = (valid[:, 1:] & valid[:, :-1]).float()
    diffs = diffs * pair_mask
    max_norm_jump = diffs.max(dim=1).values / max_norm           # [B]

    norm_last  = beliefs[bidx, last_idx].norm(dim=-1)
    norm_first = beliefs[bidx, first_idx].norm(dim=-1)
    mean_norm_slope = (norm_last - norm_first) / n_valid.clamp(min=1.0)   # [B]

    out: Dict[str, torch.Tensor] = {
        "b_last":  b_last,
        "b_first": b_first,
        "b_mean":  b_mean,
        "b_max":   b_max,
        "delta_last": delta_last,
        "mean_adj_cos_dist": mean_adj,
        "p95_adj_cos_dist":  p95_adj,
        "max_norm_jump":     max_norm_jump,
        "mean_norm_slope":   mean_norm_slope,
        "n_valid":           n_valid,
    }

    if tta_means is not None and tta_vars is not None:
        # The qwen3vl4b cache stores tta as a clip-level scalar [B], not [B,T].
        # Older caches store [B,T]. Handle both transparently.
        if tta_means.dim() == 1:
            out.update({
                "tta_mean_last":  tta_means,
                "tta_var_last":   tta_vars,
                "tta_mean_first": tta_means,
                "tta_mean_max":   tta_means,
                "tta_var_max":    tta_vars,
                "tta_mean_slope": torch.zeros_like(tta_means),
            })
        else:
            tm = tta_means * valid.float()
            tv = tta_vars  * valid.float()
            out.update({
                "tta_mean_last":  tta_means[bidx, last_idx],
                "tta_var_last":   tta_vars [bidx, last_idx],
                "tta_mean_first": tta_means[bidx, first_idx],
                "tta_mean_max":   tm.max(dim=1).values,
                "tta_var_max":    tv.max(dim=1).values,
                "tta_mean_slope": (tta_means[bidx, last_idx]
                                    - tta_means[bidx, first_idx]) / n_valid,
            })

    # convenience: a single concatenated pooled tensor
    pieces = [
        out["b_last"], out["b_mean"], out["b_max"], out["delta_last"],
        out["mean_adj_cos_dist"].unsqueeze(-1),
        out["p95_adj_cos_dist"].unsqueeze(-1),
        out["max_norm_jump"].unsqueeze(-1),
        out["mean_norm_slope"].unsqueeze(-1),
    ]
    if "tta_mean_last" in out:
        pieces += [
            out["tta_mean_last"].unsqueeze(-1),
            out["tta_var_last"].unsqueeze(-1),
            out["tta_mean_first"].unsqueeze(-1),
            out["tta_mean_max"].unsqueeze(-1),
            out["tta_var_max"].unsqueeze(-1),
            out["tta_mean_slope"].unsqueeze(-1),
        ]
    out["pooled"] = torch.cat(pieces, dim=-1)                   # [B, F]
    return out


def feature_dim(belief_dim: int, with_tta: bool = True) -> int:
    """Returns F = D*4 + 4 (+6 if with_tta)."""
    return belief_dim * 4 + 4 + (6 if with_tta else 0)