File size: 7,618 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """
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)
|