MoFlow/models/graph_interaction_nba_v4.py

"""
FutureInteractionGraphV4 — learnable sparse graph with spatio-temporal encoding.

Two-stage design:
  Stage 1 (cheap, all edges):
    Compute multi-scale geometric features [mean_rel, std_rel, min_dist] → [E, 5]
    Score each edge with a small MLP → [E] scalar importance.
    Select top-N neighbors per agent (hard, per-mode).

  Stage 2 (expensive, selected edges only):
    Run RelTrajEncoder (self-attention over T) only on the N*B*K*A selected edges.
    GNN runs on the resulting sparse graph.

Philosophy:
  - Per-mode: scoring and encoding both use per-mode predicted positions.
  - Timestep-specific: RelTrajEncoder retains full T-step temporal structure.
  - Speed: RelTrajEncoder runs on B*K*A*N edges instead of B*K*A*(A-1).
            With N=5, A=11 → ~2x fewer encoder calls; with N=3 → ~3x fewer.

Gradient flow:
  The scorer MLP receives gradients through the selected edges' GNN outputs.
  No straight-through needed — the scorer naturally learns to rank important
  neighbors because good selections lead to better trajectory predictions.
"""

import torch
import torch.nn as nn
from models.graph_interaction_nba import FutureInteractionGraph
from models.graph_interaction_nba_v3 import RelTrajEncoder


# ---------------------------------------------------------------------------
# V4 graph module
# ---------------------------------------------------------------------------

class FutureInteractionGraphV4(FutureInteractionGraph):
    """Sparse future interaction graph with learnable top-N neighbor selection.

    Extra constructor kwargs (beyond FutureInteractionGraph):
        top_n_neighbors (int, default 5): neighbors to keep per agent per mode.
        rel_traj_hidden (int, default 32): hidden dim in RelTrajEncoder.
    """

    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):
        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,
        )
        assert top_n_neighbors < num_agents, \
            f"top_n_neighbors ({top_n_neighbors}) must be < num_agents ({num_agents})"

        self.top_n = top_n_neighbors

        # Stage 1: cheap edge scorer
        # Features: mean_rel[2] + std_rel[2] + min_dist[1] = 5-dim
        # Learnable: MLP learns which geometric patterns indicate important interactions.
        self.edge_scorer = nn.Sequential(
            nn.Linear(5, 16),
            nn.ReLU(inplace=True),
            nn.Linear(16, 1),
        )

        # Stage 2: full spatio-temporal encoder (selected edges only)
        del self.rel_pos_proj   # replace V1's mean-position MLP
        self.rel_traj_encoder = RelTrajEncoder(
            out_dim  = embed_dim,
            T        = future_steps,
            D_hidden = rel_traj_hidden,
            num_heads = 4,
        )

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

    def forward(
        self,
        y_emb:       torch.Tensor,             # [B, K, A, D]
        y_abs:       torch.Tensor,             # [B, K, A, T, 2]
        t_emb:       torch.Tensor,             # [B, D]
        tau:         torch.Tensor,             # [B] ∈ [0, 1]
        sigma_agent: torch.Tensor = None,      # [B, K, A, T] or None
    ) -> torch.Tensor:                         # [B, K, A, D]
        B, K, A, D = y_emb.shape
        T  = y_abs.shape[3]
        E0 = self._E0                          # A*(A-1), edges per scene
        N  = self.top_n

        # ---- Full per-mode relative trajectories  [B*K*E0, T, 2] ----------
        pos_bk        = y_abs.reshape(B * K * A, T, 2)          # [B*K*A, T, 2]
        edge_index_bk = self._make_batched_edge_index(B * K)    # [2, B*K*E0]

        rel_pos_t = (pos_bk[edge_index_bk[0]] -
                     pos_bk[edge_index_bk[1]])                   # [B*K*E0, T, 2]

        # ---- Stage 1: cheap scorer on all edges  --------------------------
        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]
        score_feat = torch.cat([mean_rel, std_rel, min_dist], dim=-1)  # [E, 5]
        scores     = self.edge_scorer(score_feat).squeeze(-1)    # [B*K*E0]

        # ---- Top-N selection per target agent  ----------------------------
        # Edge ordering from _make_single_edge_index: outer loop is target i,
        # so each target has exactly (A-1) contiguous edges → safe to view.
        scores_grouped = scores.view(B * K * A, A - 1)          # [B*K*A, A-1]
        _, top_idx = scores_grouped.topk(N, dim=-1, sorted=False)# [B*K*A, N]
        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] bool

        # ---- Stage 2: RelTrajEncoder on selected edges only  --------------
        rel_pos_sparse = rel_pos_t[mask_flat]                    # [B*K*A*N, T, 2]

        if sigma_agent is not None:
            sigma_bka   = sigma_agent.reshape(B * K * A, T)     # [B*K*A, T]
            sigma_i_t   = sigma_bka[edge_index_bk[1][mask_flat]]# [B*K*A*N, T]
            sigma_j_t   = sigma_bka[edge_index_bk[0][mask_flat]]# [B*K*A*N, T]
            sigma_bias  = sigma_i_t - sigma_j_t                 # [B*K*A*N, T]
            tau_bka     = sigma_agent.mean(dim=-1).reshape(B * K * A)
        else:
            sigma_bias  = None
            tau_bka     = (tau
                           .unsqueeze(1).unsqueeze(2)
                           .expand(-1, K, A)
                           .reshape(B * K * A))                  # [B*K*A]

        edge_attr_sparse = self.rel_traj_encoder(
            rel_pos_sparse, sigma_bias
        )                                                        # [B*K*A*N, D]

        # ---- Sparse GNN pass  ---------------------------------------------
        edge_index_sparse = edge_index_bk[:, mask_flat]         # [2, B*K*A*N]

        temb_bka = (t_emb
                    .unsqueeze(1).unsqueeze(2)
                    .expand(-1, K, A, -1)
                    .reshape(B * K * A, D))                      # [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))
        out  = orig + gate * self.out_proj(nodes)

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