MoFlow/models/graph_interaction_nba_v3.py

"""
FutureInteractionGraphV3 — spatio-temporal edge encoding via RelTrajEncoder.

Difference from V2:
  V2 encodes each directed edge as the mean relative future position
  (averaged over T timesteps) → [E, 2] → MLP → [E, D].

  V3 encodes the full relative trajectory sequence using a small
  transformer (RelTrajEncoder): [E, T, 2] → [E, D].
  Self-attention over T lets the model distinguish *when* and *how*
  agent pairs interact, not just their average spatial offset.

Edge features are still computed per-mode (like V2), so each of the K
modes receives interaction context from its own predicted trajectories.

Memory:
  Dominant intermediate: [B*K*E0, T, D_hidden] = [550K, 20, 32] ≈ 1.4 GB.
  Attention map:         [B*K*E0, H, T, T]     = [550K, 4, 20, 20] ≈ 880 MB.
  Total overhead ~3 GB beyond V2 — fits on an A6000.
"""

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


# ---------------------------------------------------------------------------
# Lightweight self-attention safe for large batch sizes (uses bmm not sdpa)
# ---------------------------------------------------------------------------

class _TemporalSelfAttn(nn.Module):
    """Pre-norm transformer block over a sequence dim, using torch.bmm.

    nn.TransformerEncoderLayer uses scaled_dot_product_attention (flash attn)
    which has a CUDA kernel batch-size limit (~2M).  With E=550K edges and
    H=4 heads the effective batch is 2.2M — exceeding the limit.
    torch.bmm has no such restriction.
    """

    def __init__(self, D: int, num_heads: int = 4):
        super().__init__()
        assert D % num_heads == 0
        self.H     = num_heads
        self.Dh    = D // num_heads
        self.scale = self.Dh ** -0.5
        self.qkv   = nn.Linear(D, 3 * D, bias=False)
        self.out   = nn.Linear(D, D)
        self.norm1 = nn.LayerNorm(D)
        self.norm2 = nn.LayerNorm(D)
        self.ff    = nn.Sequential(
            nn.Linear(D, D * 2), nn.ReLU(inplace=True), nn.Linear(D * 2, D),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # [E, T, D]
        E, T, D = x.shape
        H, Dh = self.H, self.Dh

        # self-attention
        res = x
        x   = self.norm1(x)
        qkv = self.qkv(x).reshape(E, T, 3, H, Dh)
        q   = qkv[:, :, 0].permute(0, 2, 1, 3).reshape(E * H, T, Dh)
        k   = qkv[:, :, 1].permute(0, 2, 1, 3).reshape(E * H, T, Dh)
        v   = qkv[:, :, 2].permute(0, 2, 1, 3).reshape(E * H, T, Dh)
        a   = torch.bmm(q, k.transpose(1, 2)).mul_(self.scale).softmax(dim=-1)
        out = torch.bmm(a, v).reshape(E, H, T, Dh).permute(0, 2, 1, 3).reshape(E, T, D)
        x   = self.out(out) + res

        # feed-forward
        x = x + self.ff(self.norm2(x))
        return x


# ---------------------------------------------------------------------------
# Spatio-temporal relative trajectory encoder
# ---------------------------------------------------------------------------

class RelTrajEncoder(nn.Module):
    """Encode a relative trajectory sequence [E, T, 2] → [E, out_dim].

    Architecture:
      1. Per-timestep linear projection: [E, T, 2]      → [E, T, D_hidden]
      2. Add learned temporal position encoding
      3. _TemporalSelfAttn (self-attention over T)       → [E, T, D_hidden]
      4. Attention-weighted pooling over T               → [E, D_hidden]
      5. Linear output projection                        → [E, out_dim]

    The self-attention over T captures patterns like:
      - when agents are closest (brief crossing vs sustained proximity)
      - convergence vs divergence
      - early vs late interaction in the prediction horizon
    """

    def __init__(self, out_dim: int, T: int = 20,
                 D_hidden: int = 32, num_heads: int = 4,
                 in_channels: int = 2):
        super().__init__()
        self.input_proj = nn.Linear(in_channels, D_hidden)
        self.t_pe       = nn.Embedding(T, D_hidden)        # learned temporal PE
        self.attn       = _TemporalSelfAttn(D_hidden, num_heads)
        self.pool_w     = nn.Linear(D_hidden, 1)           # attention pooling weights
        self.out_proj   = nn.Linear(D_hidden, out_dim)

    def forward(self, rel_pos_t: torch.Tensor,
                sigma_bias: torch.Tensor = None) -> torch.Tensor:
        """
        Args:
            rel_pos_t:   [E, T, 2]  relative position at each future timestep
            sigma_bias:  [E, T]     per-timestep uncertainty bias, or None.
                         Added to the pooling logits before softmax so that
                         timesteps where the target is uncertain and the source
                         is certain receive higher pooling weight.
        Returns:
            [E, out_dim]
        """
        E, T, _ = rel_pos_t.shape
        h = self.input_proj(rel_pos_t)                           # [E, T, D_hidden]
        h = h + self.t_pe(torch.arange(T, device=h.device))     # temporal PE
        h = self.attn(h)                                         # [E, T, D_hidden]
        w_logits = self.pool_w(h)                                # [E, T, 1]
        if sigma_bias is not None:
            w_logits = w_logits + sigma_bias.unsqueeze(-1)       # [E, T, 1]
        w = w_logits.softmax(dim=1)                              # [E, T, 1]
        pooled = (h * w).sum(dim=1)                              # [E, D_hidden]
        return self.out_proj(pooled)                             # [E, out_dim]


# ---------------------------------------------------------------------------
# V3 graph module
# ---------------------------------------------------------------------------

class FutureInteractionGraphV3(FutureInteractionGraph):
    """FutureInteractionGraph with spatio-temporal relative trajectory encoding.

    Replaces the mean-position MLP (rel_pos_proj inherited from V1) with
    RelTrajEncoder, which applies self-attention over the T future timesteps.
    Edge features are computed per-mode (same as V2).

    Extra constructor kwarg:
        rel_traj_hidden (int, default 32): hidden dim inside 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,
                 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,
        )

        # Replace the V1 mean-position projection with the spatio-temporal encoder
        del self.rel_pos_proj
        self.rel_traj_encoder = RelTrajEncoder(
            out_dim  = embed_dim,
            T        = future_steps,
            D_hidden = rel_traj_hidden,
            num_heads = 4,
        )

    # ------------------------------------------------------------------
    # Forward — per-mode edges + spatio-temporal encoding
    # ------------------------------------------------------------------

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

        # ---- Per-mode relative trajectory sequences  [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]

        # ---- Per-timestep uncertainty bias for RelTrajEncoder pooling ----
        # sigma_agent [B, K, A, T]: per-agent per-timestep uncertainty.
        # For edge (j→i): bias_t = σ_i_t − σ_j_t (uncertain target,
        # certain source → higher pooling weight at that timestep).
        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]]              # [E, T]  target
            sigma_j_t = sigma_bka[edge_index_bk[0]]              # [E, T]  source
            sigma_bias = sigma_i_t - sigma_j_t                   # [E, T]
            # For GNN node_tau: mean over T (scalar per agent)
            tau_bka = sigma_agent.mean(dim=-1).reshape(B * K * A)# [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_bk = self.rel_traj_encoder(rel_pos_t,
                                             sigma_bias)          # [B*K*E0, D]

        # ---- Per-agent time embeddings  ----------------------------------
        temb_bka = (t_emb
                    .unsqueeze(1).unsqueeze(2)
                    .expand(-1, K, A, -1)
                    .reshape(B * K * A, D))                      # [B*K*A, D]

        # ---- GNN pass  ---------------------------------------------------
        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=tau_bka)

        # ---- 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)