""" FutureInteractionGraphV9 — V8 (kinematic) + receiver-adaptive temporal readout. Instead of fixed attention-pooling over T, the receiving node's embedding queries the temporal edge features via cross-attention. This makes the edge encoding adaptive: a defender cares about early timesteps (before screen), a help defender cares about late timesteps (after rotation). Combines V8's kinematic channels with V9's adaptive readout. """ import torch import torch.nn as nn from models.graph_interaction_nba_v8 import FutureInteractionGraphV8, _compute_kinematic_features from models.graph_interaction_nba_v3 import _TemporalSelfAttn from models.graph_interaction_nba_v5 import _heading_diff class AdaptiveRelTrajEncoder(nn.Module): """RelTrajEncoder with cross-attention readout conditioned on receiver node. Architecture: 1. Per-timestep projection: [E, T, in_ch] → [E, T, D_hidden] 2. Temporal PE + self-attention over T 3. Cross-attention: node_i queries temporal features → [E, D_hidden] 4. Output projection → [E, out_dim] """ def __init__(self, out_dim, T=20, D_hidden=32, num_heads=4, in_channels=8, node_dim=128): super().__init__() self.input_proj = nn.Linear(in_channels, D_hidden) self.t_pe = nn.Embedding(T, D_hidden) self.attn = _TemporalSelfAttn(D_hidden, num_heads) # Cross-attention: node_i (query) attends to temporal features (KV) self.q_proj = nn.Linear(node_dim, D_hidden) self.k_proj = nn.Linear(D_hidden, D_hidden) self.v_proj = nn.Linear(D_hidden, D_hidden) self.cross_scale = D_hidden ** -0.5 self.out_proj = nn.Linear(D_hidden, out_dim) def forward(self, rel_features, sigma_bias=None, node_i_emb=None): """ Args: rel_features: [E, T, in_ch] sigma_bias: [E, T] or None node_i_emb: [E, node_dim] — receiving node's embedding Returns: [E, out_dim] """ E, T, _ = rel_features.shape h = self.input_proj(rel_features) # [E, T, D_h] h = h + self.t_pe(torch.arange(T, device=h.device)) h = self.attn(h) # [E, T, D_h] if node_i_emb is not None: # Cross-attention: node_i queries temporal features q = self.q_proj(node_i_emb).unsqueeze(1) # [E, 1, D_h] k = self.k_proj(h) # [E, T, D_h] v = self.v_proj(h) # [E, T, D_h] attn_logits = (q * k).sum(dim=-1) * self.cross_scale # [E, T] if sigma_bias is not None: attn_logits = attn_logits + sigma_bias attn_w = attn_logits.softmax(dim=-1).unsqueeze(-1) # [E, T, 1] pooled = (v * attn_w).sum(dim=1) # [E, D_h] else: # Fallback: standard pooling (for backward compatibility) w_logits = (h * h.mean(dim=1, keepdim=True)).sum(dim=-1, keepdim=True) if sigma_bias is not None: w_logits = w_logits + sigma_bias.unsqueeze(-1) w = w_logits.softmax(dim=1) pooled = (h * w).sum(dim=1) return self.out_proj(pooled) class FutureInteractionGraphV9(FutureInteractionGraphV8): """V8 (kinematic) + receiver-adaptive temporal readout.""" def __init__(self, embed_dim, future_steps, num_agents, num_heads=4, dropout=0.1, num_gnn_layers=2, time_dim=128, top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=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, top_n_neighbors=top_n_neighbors, rel_traj_hidden=rel_traj_hidden, y0_score_dim=y0_score_dim) # Replace standard RelTrajEncoder with adaptive version self.rel_traj_encoder = AdaptiveRelTrajEncoder( out_dim=embed_dim, T=future_steps, D_hidden=rel_traj_hidden, num_heads=4, in_channels=8, node_dim=embed_dim) def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None): B, K, A, D = y_emb.shape T = y_abs.shape[3] E0 = self._E0 N = self.top_n # ---- Scoring (same as V6) ---- y0_flat = y_abs.reshape(B * K * A, T * 2) y0_emb = self.y0_score_proj(y0_flat) if sigma_agent is not None: sigma_mean = sigma_agent.mean(dim=-1).reshape(B * K * A, 1) tau_bka = sigma_mean.squeeze(-1) else: sigma_mean = torch.zeros(B * K * A, 1, device=y_abs.device) tau_bka = (tau.unsqueeze(1).unsqueeze(2) .expand(-1, K, A).reshape(B * K * A)) node_feat = torch.cat([y0_emb, sigma_mean], dim=-1) q_bka = self.W_q(node_feat) k_bka = self.W_k(node_feat) pos_bk = y_abs.reshape(B * K * A, T, 2) edge_index_bk = self._make_batched_edge_index(B * K) pos_i_t = pos_bk[edge_index_bk[1]] pos_j_t = pos_bk[edge_index_bk[0]] # Geometric features for scoring rel_pos_t = pos_j_t - pos_i_t mean_rel = rel_pos_t.mean(dim=1) std_rel = rel_pos_t.std(dim=1) min_dist = rel_pos_t.norm(dim=-1).min(dim=1).values.unsqueeze(-1) heading_full = _heading_diff(pos_i_t, pos_j_t) heading_mean = heading_full.mean(dim=1) q_i = q_bka[edge_index_bk[1]] k_j = k_bka[edge_index_bk[0]] semantic_score = (q_i * k_j).sum(dim=-1) * self.scale geo_feat = torch.cat([mean_rel, std_rel, min_dist, heading_mean], dim=-1) geo_bias = self.geo_mlp(geo_feat).squeeze(-1) scores = semantic_score + geo_bias # ---- Top-N selection ---- scores_grouped = scores.view(B * K * A, A - 1) _, top_idx = scores_grouped.topk(N, dim=-1, sorted=False) 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) # ---- Kinematic features + adaptive readout (NEW) ---- kinematic_sparse = _compute_kinematic_features( pos_i_t[mask_flat], pos_j_t[mask_flat]) if sigma_agent is not None: sigma_full = sigma_agent.reshape(B * K * A, T) sigma_i_t = sigma_full[edge_index_bk[1][mask_flat]] sigma_j_t = sigma_full[edge_index_bk[0][mask_flat]] sigma_bias = sigma_i_t - sigma_j_t else: sigma_bias = None # Get receiver node embeddings for cross-attention node_i_emb = y_emb.reshape(B * K * A, D)[edge_index_bk[1][mask_flat]] edge_attr_sparse = self.rel_traj_encoder( kinematic_sparse, sigma_bias, node_i_emb=node_i_emb) # ---- GNN + gated residual ---- edge_index_sparse = edge_index_bk[:, mask_flat] temb_bka = (t_emb.unsqueeze(1).unsqueeze(2) .expand(-1, K, A, -1).reshape(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) 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)