| """ |
| FutureInteractionGraphV9b — V6 + receiver-adaptive readout (WITHOUT kinematic channels). |
| |
| Tests the adaptive readout idea in isolation (3 channels like V6, not 8). |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from models.graph_interaction_nba_v6 import FutureInteractionGraphV6 |
| from models.graph_interaction_nba_v9 import AdaptiveRelTrajEncoder |
| from models.graph_interaction_nba_v5 import _heading_diff |
|
|
|
|
| class FutureInteractionGraphV9b(FutureInteractionGraphV6): |
| """V6 scoring + adaptive temporal readout (3 channels).""" |
|
|
| 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) |
|
|
| |
| self.rel_traj_encoder = AdaptiveRelTrajEncoder( |
| out_dim=embed_dim, T=future_steps, |
| D_hidden=rel_traj_hidden, num_heads=4, |
| in_channels=3, 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] |
| N = self.top_n |
|
|
| |
| 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]] |
| 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 |
|
|
| |
| 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) |
|
|
| |
| heading_sparse = heading_full[mask_flat] |
| rel_pos_sparse = torch.cat([rel_pos_t[mask_flat], heading_sparse], dim=-1) |
|
|
| 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 |
|
|
| |
| node_i_emb = y_emb.reshape(B * K * A, D)[edge_index_bk[1][mask_flat]] |
|
|
| edge_attr_sparse = self.rel_traj_encoder( |
| rel_pos_sparse, sigma_bias, node_i_emb=node_i_emb) |
|
|
| |
| 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) |
|
|
