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