MoFlow/models/graph_interaction_nba_v6.py

"""
FutureInteractionGraphV6 — RAG-style neighbor retrieval with sparse RelTrajEncoder.

Replaces V5's per-edge MLP scorer with a RAG-style query-key dot product:

  Per-node (55K, computed once):
    q_i = W_q([y0_emb_i, σ_i])   — what agent i is looking for
    k_j = W_k([y0_emb_j, σ_j])   — what agent j offers

  Per-edge (550K, cheap):
    semantic_score = (q_i · k_j) / √D_s
    geo_bias       = geo_mlp([mean_rel, std_rel, min_dist, heading_diff_mean])
    score          = semantic_score + geo_bias

  Top-N selection → RelTrajEncoder([rel_pos, heading_diff]) on selected edges.

Design rationale:
  - Asymmetric W_q / W_k: "what to look for" ≠ "how to present oneself"
  - Uncertainty in query/key: uncertain agents learn to seek certain neighbors
    through training, not hard-coded scaling
  - Geometric features as additive re-ranking bias: cleanly separates
    semantic retrieval from spatial prior
  - Heading in geo_bias: converging agents are more likely to interact
  - Per-node query/key is ~10x cheaper than per-edge MLP scorer in V5

Document content (unchanged from V5):
  RelTrajEncoder([rel_pos(2), heading_diff(1)]) → [E_selected, D]
  → GNN message passing on sparse graph.
"""

import os
import torch
import torch.nn as nn
from models.graph_interaction_nba_v4 import FutureInteractionGraphV4
from models.graph_interaction_nba_v3 import RelTrajEncoder
from models.graph_interaction_nba_v5 import _heading_diff


class FutureInteractionGraphV6(FutureInteractionGraphV4):
    """RAG-style sparse interaction graph.

    Extra constructor kwargs (beyond V4):
        y0_score_dim (int, default 32): dim of query/key embedding space.
    """

    EDGE_MODES = {
        'full':          3,  # rel_pos(2) + heading_diff(1)
        'dist_only':     1,  # ||rel_pos||(1)
        'relpos_only':   2,  # rel_pos(2)
        'heading_only':  1,  # heading_diff(1)
        'full_relvel':   5,  # rel_pos(2) + heading_diff(1) + rel_vel(2)
        'vel_only':      2,  # rel_vel(2)
    }

    def __init__(self, embed_dim: int, future_steps: int, num_agents: int,
                 num_heads: int = 4, dropout: float = 0.1,
                 num_gnn_layers: int = 2, time_dim: int = 128,
                 top_n_neighbors: int = 5, rel_traj_hidden: int = 32,
                 y0_score_dim: int = 32, edge_mode: str = 'full',
                 neighbor_mode: str = 'rag'):
        super().__init__(
            embed_dim       = embed_dim,
            future_steps    = future_steps,
            num_agents      = num_agents,
            num_heads       = num_heads,
            dropout         = dropout,
            num_gnn_layers  = num_gnn_layers,
            time_dim        = time_dim,
            top_n_neighbors = top_n_neighbors,
            rel_traj_hidden = rel_traj_hidden,
        )
        assert edge_mode in self.EDGE_MODES, f"Unknown edge_mode: {edge_mode}"
        self.edge_mode = edge_mode
        in_ch = self.EDGE_MODES[edge_mode]
        assert neighbor_mode in ('rag', 'l2', 'semantic'), f"Unknown neighbor_mode: {neighbor_mode}"
        self.neighbor_mode = neighbor_mode

        self.y0_score_dim = y0_score_dim
        self.scale = y0_score_dim ** -0.5
        self._original_top_n = top_n_neighbors

        # Per-node: project y_0_hat trajectory to scoring space
        self.y0_score_proj = nn.Sequential(
            nn.Linear(future_steps * 2, y0_score_dim),
            nn.ReLU(inplace=True),
        )

        # Asymmetric query / key encoders (W_q ≠ W_k)
        self.W_q = nn.Linear(y0_score_dim + 1, y0_score_dim)
        self.W_k = nn.Linear(y0_score_dim + 1, y0_score_dim)

        # Per-edge geometric re-ranking bias
        # Input: mean_rel(2) + std_rel(2) + min_dist(1) + heading_diff_mean(1) = 6
        self.geo_mlp = nn.Sequential(
            nn.Linear(6, 16),
            nn.ReLU(inplace=True),
            nn.Linear(16, 1),
        )

        del self.edge_scorer

        self.rel_traj_encoder = RelTrajEncoder(
            out_dim     = embed_dim,
            T           = future_steps,
            D_hidden    = rel_traj_hidden,
            num_heads   = 4,
            in_channels = in_ch,
        )

        # ---- SRA_SOFT_START: identity-at-init for the graph branch ----------
        # By default V6 inherits a randomly-initialised out_proj and a sigmoid
        # gate that starts around 0.5, so a randomly-initialised graph output is
        # injected into the host from the very first step.  With SRA_EDGE_FIX=1
        # every node now receives its full neighbour set (previously half the
        # nodes were orphans), which makes that initial shock large enough to
        # destabilise one-step flow matching (MoFlow).
        #
        # MID does not suffer from this because it zero-inits its own
        # graph_out_proj and opens a learnable, clamped gate over a warmup
        # schedule.  SRA_SOFT_START ports that recipe to V6:
        #   * out_proj zero-init  -> graph contributes exactly 0 at step 0
        #   * gate bias -> large negative, so sigmoid(gate) starts near 0
        # Both stay LEARNABLE, so unlike a fixed SRA_GATE_SCALE the graph can
        # still grow to full strength during training.
        if os.environ.get('SRA_SOFT_START', '') not in ('', '0', 'false', 'False'):
            nn.init.zeros_(self.out_proj.weight)
            nn.init.zeros_(self.out_proj.bias)
            _gb = float(os.environ.get('SRA_GATE_BIAS', -4.0) or -4.0)
            for _m in self.gate_proj.modules():
                if isinstance(_m, nn.Linear):
                    nn.init.constant_(_m.bias, _gb)   # sigmoid(-4) ~ 0.018

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------

    def forward(
        self,
        y_emb:       torch.Tensor,             # [B, K, A, D]
        y_abs:       torch.Tensor,             # [B, K, A, T, 2]  ← y_0_hat unnorm
        t_emb:       torch.Tensor,             # [B, D]
        tau:         torch.Tensor,             # [B] ∈ [0, 1]
        sigma_agent: torch.Tensor = None,      # [B, K, A, T] or None
        agent_mask:  torch.Tensor = None,      # [B, A] bool, True=real; for padded batches
    ) -> torch.Tensor:                         # [B, K, A, D]
        B, K, A, D = y_emb.shape
        T  = y_abs.shape[3]

        # Degenerate case: a single-agent scene has no edges — the graph is a
        # no-op, so just return the input untouched.
        if A <= 1:
            return y_emb

        # Variable-A support (e.g. SDD, ETH/UCY): the model is instantiated with
        # a max num_agents (padding budget), but each batch may contain a smaller
        # real A. Rebuild the graph skeleton on the fly when A changes.
        if A != self.num_agents:
            self.num_agents = A
            self._E0 = A * (A - 1)
            self.top_n = max(1, min(self._original_top_n, A - 1))
            src, dst = [], []
            for i in range(A):
                for j in range(A):
                    if i != j:
                        src.append(j); dst.append(i)
            self._single_edge_index = torch.tensor(
                [src, dst], dtype=torch.long, device=y_emb.device
            )

        E0 = self._E0
        N  = self.top_n

        # ---- Per-node: y_0_hat trajectory embedding  ---------------------
        y0_flat = y_abs.reshape(B * K * A, T * 2)               # [B*K*A, T*2]
        y0_emb  = self.y0_score_proj(y0_flat)                   # [B*K*A, D_s]

        # ---- Per-node: uncertainty scalar  -------------------------------
        if sigma_agent is not None:
            sigma_mean = sigma_agent.mean(dim=-1).reshape(B * K * A, 1)  # [B*K*A, 1]
            tau_bka    = sigma_mean.squeeze(-1)
        else:
            sigma_mean = torch.zeros(B * K * A, 1, device=y_abs.device)
            tau_bka    = (tau
                          .unsqueeze(1).unsqueeze(2)
                          .expand(-1, K, A)
                          .reshape(B * K * A))

        # ---- Per-node: query and key  ------------------------------------
        node_feat = torch.cat([y0_emb, sigma_mean], dim=-1)     # [B*K*A, D_s+1]
        q_bka     = self.W_q(node_feat)                         # [B*K*A, D_s]
        k_bka     = self.W_k(node_feat)                         # [B*K*A, D_s]

        # ---- Per-edge: build positions and relative features  ------------
        pos_bk        = y_abs.reshape(B * K * A, T, 2)
        edge_index_bk = self._make_batched_edge_index(B * K)    # [2, B*K*E0]

        pos_i_t   = pos_bk[edge_index_bk[1]]                    # [E, T, 2]
        pos_j_t   = pos_bk[edge_index_bk[0]]                    # [E, T, 2]
        rel_pos_t = pos_j_t - pos_i_t                           # [E, T, 2]

        mean_rel  = rel_pos_t.mean(dim=1)                        # [E, 2]
        std_rel   = rel_pos_t.std(dim=1)                         # [E, 2]
        min_dist  = rel_pos_t.norm(dim=-1).min(dim=1).values.unsqueeze(-1)  # [E, 1]

        # Mean heading diff over T as geometric feature
        heading_full = _heading_diff(pos_i_t, pos_j_t)          # [E, T, 1]
        heading_mean = heading_full.mean(dim=1)                  # [E, 1]

        # ---- Per-edge scoring  -----------------------------------------------
        if self.neighbor_mode == 'l2':
            # L2 distance: select closest neighbors by avg future distance
            avg_dist = rel_pos_t.norm(dim=-1).mean(dim=-1)      # [E]
            scores = -avg_dist  # negate so topk picks smallest distance
        elif self.neighbor_mode == 'semantic':
            # Semantic-only: learned query-key dot product, no geo bias
            q_i = q_bka[edge_index_bk[1]]                       # [E, D_s]
            k_j = k_bka[edge_index_bk[0]]                       # [E, D_s]
            scores = (q_i * k_j).sum(dim=-1) * self.scale       # [E]
        else:
            # RAG: semantic dot product + geometric bias
            q_i = q_bka[edge_index_bk[1]]                       # [E, D_s]
            k_j = k_bka[edge_index_bk[0]]                       # [E, D_s]
            semantic_score = (q_i * k_j).sum(dim=-1) * self.scale  # [E]

            geo_feat  = torch.cat([mean_rel, std_rel,
                                   min_dist, heading_mean], dim=-1)  # [E, 6]
            geo_bias  = self.geo_mlp(geo_feat).squeeze(-1)       # [E]

            scores = semantic_score + geo_bias                   # [E]

        # ---- Mask padded agents so top-N never selects them  --------------
        if agent_mask is not None:
            # Build per-edge "both endpoints real" mask
            # agent_mask: [B, A] -> broadcast to [B, K, A] for node-ness
            node_real = agent_mask.unsqueeze(1).expand(B, K, A).reshape(B * K * A)  # [B*K*A]
            # edge_index_bk has shape [2, B*K*E0], rows 0=src=j, 1=dst=i
            src_real = node_real[edge_index_bk[0]]               # [E]
            dst_real = node_real[edge_index_bk[1]]               # [E]
            edge_real = src_real & dst_real                       # [E]
            scores = scores.masked_fill(~edge_real, float('-inf'))

        # ---- Top-N selection per target agent  ---------------------------
        scores_grouped = scores.view(B * K * A, A - 1)
        # Cap N so we never ask for more neighbors than rows.
        N_use = min(N, A - 1)
        _, top_idx     = scores_grouped.topk(N_use, dim=-1, sorted=False)
        mask           = torch.zeros(B * K * A, A - 1,
                                     device=scores.device, dtype=torch.bool)
        mask.scatter_(1, top_idx, True)
        mask_flat      = mask.view(-1)                           # [B*K*E0]

        # When padding is present, selected edges whose target agent is padded
        # produce zero messages anyway; additionally drop any edge whose
        # endpoint is padded so the GNN sees only real edges.
        if agent_mask is not None:
            mask_flat = mask_flat & edge_real                    # [B*K*E0]

        # ---- RelTrajEncoder on selected edges only  ----------------------
        rel_pos_sel    = rel_pos_t[mask_flat]                    # [E_sel, T, 2]
        heading_sel    = heading_full[mask_flat]                  # [E_sel, T, 1]

        if self.edge_mode == 'full':
            encoder_input = torch.cat([rel_pos_sel, heading_sel], dim=-1)  # [E_sel, T, 3]
        elif self.edge_mode == 'dist_only':
            encoder_input = rel_pos_sel.norm(dim=-1, keepdim=True)        # [E_sel, T, 1]
        elif self.edge_mode == 'relpos_only':
            encoder_input = rel_pos_sel                                    # [E_sel, T, 2]
        elif self.edge_mode == 'heading_only':
            encoder_input = heading_sel                                    # [E_sel, T, 1]
        elif self.edge_mode == 'full_relvel':
            rel_vel = rel_pos_sel[:, 1:] - rel_pos_sel[:, :-1]           # [E_sel, T-1, 2]
            rel_vel = torch.cat([rel_vel, rel_vel[:, -1:]], dim=1)        # [E_sel, T, 2]
            encoder_input = torch.cat([rel_pos_sel, heading_sel, rel_vel], dim=-1)  # [E_sel, T, 5]
        elif self.edge_mode == 'vel_only':
            rel_vel = rel_pos_sel[:, 1:] - rel_pos_sel[:, :-1]           # [E_sel, T-1, 2]
            rel_vel = torch.cat([rel_vel, rel_vel[:, -1:]], dim=1)        # [E_sel, T, 2]
            encoder_input = rel_vel

        if sigma_agent is not None:
            sigma_full  = sigma_agent.reshape(B * K * A, T)
            sigma_i_t   = sigma_full[edge_index_bk[1][mask_flat]]
            sigma_j_t   = sigma_full[edge_index_bk[0][mask_flat]]
            sigma_bias  = sigma_i_t - sigma_j_t                 # [E_sel, T]
        else:
            sigma_bias  = None

        edge_attr_sparse = self.rel_traj_encoder(
            encoder_input, sigma_bias
        )                                                        # [E_sel, D]

        # ---- Sparse GNN pass  --------------------------------------------
        edge_index_sparse = edge_index_bk[:, mask_flat]

        temb_bka = (t_emb
                    .unsqueeze(1).unsqueeze(2)
                    .expand(-1, K, A, -1)
                    .reshape(B * K * A, D))

        nodes = y_emb.reshape(B * K * A, D)
        for layer in self.gnn_layers:
            nodes = layer(nodes, edge_index_sparse, edge_attr_sparse,
                          temb_agent=temb_bka, tau=tau_bka)

        # ---- Gated residual  ---------------------------------------------
        orig = y_emb.reshape(B * K * A, D)
        gate = self.gate_proj(torch.cat([orig, nodes], dim=-1))
        # SRA_GATE_SCALE: global damping on the graph residual (default 1.0 = off).
        _gs  = float(os.environ.get('SRA_GATE_SCALE', 1.0) or 1.0)
        res  = _gs * gate * self.out_proj(nodes)                # [N, D] graph perturbation

        # SRA_RES_CAP: per-node residual-norm cap (default 0 = off).  With
        # SRA_EDGE_FIX=1 training is healthy (loss decreases monotonically) but
        # *sampling* diverges: the graph is applied at every one of MoFlow's flow
        # steps and its output feeds the next step, so any oversized per-node
        # perturbation compounds geometrically over the integration.  The old
        # scene-mixing bug hid this by leaving half the nodes orphaned (weaker
        # perturbation, less compounding).  Capping each node's residual NORM
        # bounds the per-step perturbation that drives the blow-up, while leaving
        # the (majority) small perturbations untouched — unlike a global scale it
        # only clips the outliers, so the graph keeps its normal expressive range.
        # 주의: 예전 구현 `res * (rn.clamp(max=cap) / (rn + 1e-6))` 은 두 가지가 틀렸다.
        #  (1) res 가 정확히 0 이면 스케일이 0/1e-6 = 0 이 되어 **gradient 도 0** 이다.
        #      SRA_SOFT_START(out_proj zero-init) 와 같이 켜면 out_proj 가 0 에 영구히
        #      갇혀 그래프가 학습되지 않는다(=사실상 host 단독). 실제로 그 조합으로
        #      돌린 실행들의 out_proj 는 58 epoch 뒤에도 정확히 0 이었다.
        #  (2) cap 미만인데도 rn/(rn+1e-6) 만큼 축소된다 (rn=1e-5 이면 0.909 배).
        # torch.where 로 바꾸면 cap 이하는 정확히 무연산(스케일 1)이고 res=0 에서도
        # gradient 가 흐른다.  clamp_min 은 미선택 분기의 Inf 를 막는다.
        _cap = float(os.environ.get('SRA_RES_CAP', 0.0) or 0.0)
        if _cap > 0:
            rn  = res.norm(dim=-1, keepdim=True)                # [N, 1]
            res = res * torch.where(rn > _cap, _cap / rn.clamp_min(1e-6),
                                    torch.ones_like(rn))

        # SRA_RES_CAP_REL: 노드별 상한을 **호스트 임베딩 norm 에 비례**해 정한다.
        # 절대 cap 은 임베딩 스케일에 의존해 호스트마다 의미가 달라진다(MoFlow 에서
        # 튜닝한 3.0 이 MID 에서는 사실상 무연산일 수 있다).  샘플링 발산은 결국
        # "스텝당 상대 섭동"이 누적되는 문제이므로, ‖res‖ ≤ ratio·‖orig‖ 로 두면
        # 스케일 무관하게 누적률을 직접 제한한다.  측정값 기준 무제한 시 비율은 ~0.39.
        _rel = float(os.environ.get('SRA_RES_CAP_REL', 0.0) or 0.0)
        if _rel > 0:
            lim = _rel * orig.norm(dim=-1, keepdim=True)        # [N, 1] 노드별 상한
            rn  = res.norm(dim=-1, keepdim=True)
            res = res * torch.where(rn > lim, lim / rn.clamp_min(1e-6),
                                    torch.ones_like(rn))
        out  = orig + res

        return out.view(B, K, A, D)