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