MoFlow/models/graph_interaction_nba.py

"""
Future Euclidean Interaction Graph Module for MoFlow-NBA.

Core idea (ported from srtp/models/denoiser_graph_v4.py::LocalEncoder_v4_3
and srtp/models/denoiser.py::graphDenoiser_uncertainty_v4_4_nba):

  During flow matching denoising, the current noisy trajectory y_t carries
  approximate future agent positions.  For each agent pair at each future
  timestep, the Euclidean relative position is computed from y_t, encoded
  as an edge feature, and propagated through graph attention — letting each
  agent's embedding be informed by where other agents are predicted to be.

Memory-efficient design for MoFlow's [B=250, K=20, A=11] NBA setting:
  The naive approach (B×K = 5000 scenes × 110 edges = 550 K edges) OOMs on
  the transformer FFN in the original EdgeTemporalEncoderCausal.
  Solution: compute edge features once from K-mode-averaged future positions
  (→ B = 250 scenes, 27.5 K edges), then expand to B×K by replicating the
  per-scene edge attributes before the GNN pass (55 K nodes, 550 K edges in
  one vectorised forward — each set of 128-D tensors is ~280 MB, well within
  the A6000's 48 GB).
"""

import os
import torch
import torch.nn as nn
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.typing import OptTensor
from torch_geometric.utils import softmax


# ---------------------------------------------------------------------------
# Graph attention layer (adapted from GlobalInteractorLayer_v2 in
# srtp/models/denoiser_graph_v4.py)
# ---------------------------------------------------------------------------

class FutureGeomAttnLayer(MessagePassing):
    """Time-conditioned graph attention with future-geometry edge features.

    Convention (PyG source_to_target, matching srtp's variable naming):
      edge_index[0] = source = neighbor j  (sends message)
      edge_index[1] = target = center  i   (receives message)
      x_i = center feature, x_j = neighbor feature
    """

    def __init__(self, embed_dim: int, num_heads: int = 4,
                 dropout: float = 0.1, time_dim: int = 128,
                 attn_gamma: float = 4.0, **kwargs):
        super().__init__(aggr='add', node_dim=0, **kwargs)
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim  = embed_dim // num_heads
        self.scale     = self.head_dim ** 0.5
        self.attn_gamma = attn_gamma

        # node-feature projections
        self.lin_q = nn.Linear(embed_dim, embed_dim)
        self.lin_k = nn.Linear(embed_dim, embed_dim)
        self.lin_v = nn.Linear(embed_dim, embed_dim)

        # time-conditioned bias into q / k / v
        self.t2qkv = nn.Linear(time_dim, 3 * embed_dim)

        # edge projections (carry the future geometry)
        self.lin_k_edge = nn.Linear(embed_dim, embed_dim)
        self.lin_v_edge = nn.Linear(embed_dim, embed_dim)

        self.out_proj  = nn.Linear(embed_dim, embed_dim)
        self.proj_drop = nn.Dropout(dropout)
        self.attn_drop = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(embed_dim * 4, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x, edge_index, edge_attr, temb_agent, tau, size=None):
        x = x + self._mha_block(self.norm1(x), edge_index, edge_attr,
                                 temb_agent, tau, size)
        x = x + self._ff_block(self.norm2(x))
        return x

    def message(self, x_i, x_j, edge_attr,
                temb_agent_i, temb_agent_j, tau_i, tau_j,
                index, ptr, size_i):
        bias_i = self.t2qkv(temb_agent_i)
        bias_j = self.t2qkv(temb_agent_j)
        tq_i = bias_i.chunk(3, dim=-1)[0]
        tk_j = bias_j.chunk(3, dim=-1)[1]
        tv_j = bias_j.chunk(3, dim=-1)[2]

        H, Dh = self.num_heads, self.head_dim
        q = (self.lin_q(x_i) + tq_i).view(-1, H, Dh)
        k = (self.lin_k(x_j) + tk_j + self.lin_k_edge(edge_attr)).view(-1, H, Dh)
        v = (self.lin_v(x_j) + tv_j + self.lin_v_edge(edge_attr)).view(-1, H, Dh)

        logits = (q * k).sum(dim=-1) / self.scale  # [E, H]

        # Directional bias: in flow matching all agents in a scene share the
        # same tau, so tau_i == tau_j and w_ij = sigmoid(0.5) ≈ 0.62.
        w_ij = torch.sigmoid(
            self.attn_gamma * (tau_i - tau_j) + 0.5
        ).unsqueeze(-1)                                    # [E, 1]

        alpha = softmax(logits, index, ptr, size_i)        # [E, H]
        alpha = self.attn_drop(alpha) * w_ij               # [E, H]
        return v * alpha.unsqueeze(-1)                     # [E, H, Dh]

    def update(self, inputs):
        return self.out_proj(self.proj_drop(inputs.view(-1, self.embed_dim)))

    def _mha_block(self, x, edge_index, edge_attr, temb_agent, tau, size):
        return self.propagate(edge_index=edge_index, x=x, edge_attr=edge_attr,
                              temb_agent=temb_agent, tau=tau, size=size)

    def _ff_block(self, x):
        return self.mlp(x)


# ---------------------------------------------------------------------------
# Main module
# ---------------------------------------------------------------------------

class FutureInteractionGraph(nn.Module):
    """Augment per-agent embeddings with future Euclidean interaction context.

    Memory-efficient forward pass:
      1. Average y_abs over K modes → y_abs_mean [B, A, T, 2].
         This gives one representative future trajectory per agent per scene.
      2. Build a B-scene fully-connected graph (E_b = B × A×(A-1) ≈ 27.5 K).
      3. Compute the mean relative future position for each edge:
           rel_pos_mean[e] = mean_t( pos[j, t] - pos[i, t] )   (shape [E_b, 2])
         Then project to [E_b, D].  (No transformer over T; avoids the FFN
         memory bottleneck that caused OOM with 550 K-edge batches.)
      4. Expand edge features to all K modes: edge_attr_bk [B*K*E0, D] by
         repeating each scene's edge features K times.
      5. Run GNN on B*K*A = 55 K nodes with B*K*E0 = 550 K edges in one
         vectorised forward (peak ~3 GB with gradients — fits on A6000).
      6. Gated residual back into y_emb.
    """

    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):
        super().__init__()
        self.future_steps = future_steps
        self.num_agents   = num_agents
        self.embed_dim    = embed_dim
        self.time_dim     = time_dim
        self._E0          = num_agents * (num_agents - 1)  # edges per scene

        # Project mean relative position [E, 2] → [E, D]
        self.rel_pos_proj = nn.Sequential(
            nn.Linear(2, embed_dim),
            nn.ReLU(inplace=True),
            nn.Linear(embed_dim, embed_dim),
        )

        # GNN layers (shared across all K modes / all scenes)
        self.gnn_layers = nn.ModuleList([
            FutureGeomAttnLayer(
                embed_dim=embed_dim, num_heads=num_heads,
                dropout=dropout, time_dim=time_dim,
            )
            for _ in range(num_gnn_layers)
        ])

        # Gated residual
        self.gate_proj = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),
            nn.Sigmoid(),
        )
        self.out_proj = nn.Linear(embed_dim, embed_dim)

        # Pre-built edge index for one fully-connected A-agent graph
        single_ei = self._make_single_edge_index(num_agents)
        self.register_buffer('_single_edge_index', single_ei)

    # ------------------------------------------------------------------
    # Graph construction helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _make_single_edge_index(A: int) -> torch.Tensor:
        """Fully-connected directed graph on A nodes: A*(A-1) edges.

        Convention:
          edge_index[0] = neighbor j  (source / sender)
          edge_index[1] = center  i   (target / receiver)
        → PyG x_i = center, x_j = neighbor (source_to_target flow).
        """
        src, dst = [], []
        for i in range(A):      # center / target
            for j in range(A):  # neighbor / source
                if i != j:
                    src.append(j)
                    dst.append(i)
        return torch.tensor([src, dst], dtype=torch.long)

    def _make_batched_edge_index(self, num_scenes: int) -> torch.Tensor:
        """Stack num_scenes copies of the A-node graph with correct offsets.

        NOTE (edge-index scene-mixing bug):
          `batched` is [S, 2, E0]; its contiguous layout is
          s0_src, s0_dst, s1_src, s1_dst, ...  A direct `.reshape(2, -1)` therefore
          packs *scene blocks* into each row instead of the src/dst rows, so row 0
          ends up holding s0_src followed by s0_dst, etc.  Consequences (measured,
          A=11): 0% of edges stay inside a scene, exactly half the nodes receive no
          incoming edge at all, and the other half receive 2x the intended degree.
          The correct behaviour is to move the src/dst axis first via permute.

          SRA_EDGE_FIX=1 selects the correct per-scene graph.  The default keeps the
          original (buggy) behaviour so that previously trained checkpoints and any
          in-flight runs remain reproducible.  See models/graph_interaction_nba_v6.py
          (line ~171), which is the call site used by MID / LED / MoFlow.
        """
        single  = self._single_edge_index              # [2, E0]
        A       = self.num_agents
        offsets = torch.arange(num_scenes, device=single.device) * A  # [S]
        batched = (single.unsqueeze(0)
                         .expand(num_scenes, -1, -1)   # [S, 2, E0]
                   + offsets.view(-1, 1, 1))
        if os.environ.get('SRA_EDGE_FIX', '') not in ('', '0', 'false', 'False'):
            # [S, 2, E0] -> [2, S, E0] -> [2, S*E0]: keeps src/dst rows intact
            return batched.permute(1, 0, 2).reshape(2, -1)
        return batched.reshape(2, -1)                  # [2, S*E0]  (legacy, scene-mixing)

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

    def forward(
        self,
        y_emb: torch.Tensor,             # [B, K, A, D]
        y_abs: torch.Tensor,             # [B, K, A, T, 2]  absolute-ish future pos
        t_emb: torch.Tensor,             # [B, D]  time embedding from backbone
        tau:   torch.Tensor,             # [B] ∈ [0, 1]
        sigma_agent: torch.Tensor = None,# [B, K, A] per-agent uncertainty, 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

        # ---- Step 1: mode-averaged future positions ----------------------
        # [B, A, T, 2]  — representative trajectory per agent per scene
        y_abs_mean = y_abs.mean(dim=1)

        # ---- Step 2: B-scene edge index  (E_b = B * E0 ≈ 27.5 K) -------
        edge_index_b = self._make_batched_edge_index(B)    # [2, B*E0]

        # ---- Step 3: mean relative position as edge feature  [E_b, D] ---
        pos_b = y_abs_mean.reshape(B * A, T, 2)            # [B*A, T, 2]
        # neighbor j = edge_index_b[0], center i = edge_index_b[1]
        rel_pos_mean = (pos_b[edge_index_b[0]] -
                        pos_b[edge_index_b[1]]).mean(dim=1) # [E_b, 2]
        edge_attr_b = self.rel_pos_proj(rel_pos_mean)       # [E_b, D]

        # ---- Step 4: expand edge features to B*K scenes  [B*K*E0, D] ---
        # Each scene b has K modes; all modes share the same edge features.
        edge_index_bk = self._make_batched_edge_index(B * K)  # [2, B*K*E0]

        # edge_attr_b [B*E0, D] → view [B, E0, D] → expand [B, K, E0, D]
        #  → reshape [B*K*E0, D]
        edge_attr_bk = (edge_attr_b
                        .view(B, E0, D)
                        .unsqueeze(1)
                        .expand(-1, K, -1, -1)
                        .reshape(B * K * E0, D))             # [B*K*E0, D]

        # ---- Step 5: per-agent time / uncertainty embeddings -------------
        # t_emb: [B, D] → expand to [B, K, A, D] → flatten [B*K*A, D]
        temb_bka = (t_emb
                    .unsqueeze(1).unsqueeze(2)
                    .expand(-1, K, A, -1)
                    .reshape(B * K * A, D))                   # [B*K*A, D]

        # Directional weight signal for graph attention:
        #   - if sigma_agent provided: collapse to scalar per agent (mean over T
        #     if [B,K,A,T], or use directly if [B,K,A]).
        #   - else: fall back to scalar denoising tau (same for all agents).
        if sigma_agent is not None:
            s = sigma_agent.mean(dim=-1) if sigma_agent.dim() == 4 else sigma_agent
            node_tau = s.reshape(B * K * A)                   # [B*K*A]
        else:
            node_tau = (tau
                        .unsqueeze(1).unsqueeze(2)
                        .expand(-1, K, A)
                        .reshape(B * K * A))                  # [B*K*A]

        # ---- Step 6: GNN pass on B*K*A nodes  ----------------------------
        nodes = y_emb.reshape(B * K * A, D)                  # [B*K*A, D]
        for layer in self.gnn_layers:
            nodes = layer(nodes, edge_index_bk, edge_attr_bk,
                          temb_agent=temb_bka, tau=node_tau)  # [B*K*A, D]

        # ---- Step 7: gated residual  -------------------------------------
        orig  = y_emb.reshape(B * K * A, D)
        gate  = self.gate_proj(torch.cat([orig, nodes], dim=-1))  # [N, D]
        out   = orig + gate * self.out_proj(nodes)                # [N, D]

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