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