""" FutureInteractionGraphV6 — RAG-style neighbor retrieval with sparse RelTrajEncoder. Replaces V5's per-edge MLP scorer with a RAG-style query-key dot product: Per-node (55K, computed once): q_i = W_q([y0_emb_i, σ_i]) — what agent i is looking for k_j = W_k([y0_emb_j, σ_j]) — what agent j offers Per-edge (550K, cheap): semantic_score = (q_i · k_j) / √D_s geo_bias = geo_mlp([mean_rel, std_rel, min_dist, heading_diff_mean]) score = semantic_score + geo_bias Top-N selection → RelTrajEncoder([rel_pos, heading_diff]) on selected edges. Design rationale: - Asymmetric W_q / W_k: "what to look for" ≠ "how to present oneself" - Uncertainty in query/key: uncertain agents learn to seek certain neighbors through training, not hard-coded scaling - Geometric features as additive re-ranking bias: cleanly separates semantic retrieval from spatial prior - Heading in geo_bias: converging agents are more likely to interact - Per-node query/key is ~10x cheaper than per-edge MLP scorer in V5 Document content (unchanged from V5): RelTrajEncoder([rel_pos(2), heading_diff(1)]) → [E_selected, D] → GNN message passing on sparse graph. """ import os import torch import torch.nn as nn from models.graph_interaction_nba_v4 import FutureInteractionGraphV4 from models.graph_interaction_nba_v3 import RelTrajEncoder from models.graph_interaction_nba_v5 import _heading_diff class FutureInteractionGraphV6(FutureInteractionGraphV4): """RAG-style sparse interaction graph. Extra constructor kwargs (beyond V4): y0_score_dim (int, default 32): dim of query/key embedding space. """ EDGE_MODES = { 'full': 3, # rel_pos(2) + heading_diff(1) 'dist_only': 1, # ||rel_pos||(1) 'relpos_only': 2, # rel_pos(2) 'heading_only': 1, # heading_diff(1) 'full_relvel': 5, # rel_pos(2) + heading_diff(1) + rel_vel(2) 'vel_only': 2, # rel_vel(2) } 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, y0_score_dim: int = 32, edge_mode: str = 'full', neighbor_mode: str = 'rag'): 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, ) assert edge_mode in self.EDGE_MODES, f"Unknown edge_mode: {edge_mode}" self.edge_mode = edge_mode in_ch = self.EDGE_MODES[edge_mode] assert neighbor_mode in ('rag', 'l2', 'semantic'), f"Unknown neighbor_mode: {neighbor_mode}" self.neighbor_mode = neighbor_mode self.y0_score_dim = y0_score_dim self.scale = y0_score_dim ** -0.5 self._original_top_n = top_n_neighbors # Per-node: project y_0_hat trajectory to scoring space self.y0_score_proj = nn.Sequential( nn.Linear(future_steps * 2, y0_score_dim), nn.ReLU(inplace=True), ) # Asymmetric query / key encoders (W_q ≠ W_k) self.W_q = nn.Linear(y0_score_dim + 1, y0_score_dim) self.W_k = nn.Linear(y0_score_dim + 1, y0_score_dim) # Per-edge geometric re-ranking bias # Input: mean_rel(2) + std_rel(2) + min_dist(1) + heading_diff_mean(1) = 6 self.geo_mlp = nn.Sequential( nn.Linear(6, 16), nn.ReLU(inplace=True), nn.Linear(16, 1), ) del self.edge_scorer self.rel_traj_encoder = RelTrajEncoder( out_dim = embed_dim, T = future_steps, D_hidden = rel_traj_hidden, num_heads = 4, in_channels = in_ch, ) # ---- SRA_SOFT_START: identity-at-init for the graph branch ---------- # By default V6 inherits a randomly-initialised out_proj and a sigmoid # gate that starts around 0.5, so a randomly-initialised graph output is # injected into the host from the very first step. With SRA_EDGE_FIX=1 # every node now receives its full neighbour set (previously half the # nodes were orphans), which makes that initial shock large enough to # destabilise one-step flow matching (MoFlow). # # MID does not suffer from this because it zero-inits its own # graph_out_proj and opens a learnable, clamped gate over a warmup # schedule. SRA_SOFT_START ports that recipe to V6: # * out_proj zero-init -> graph contributes exactly 0 at step 0 # * gate bias -> large negative, so sigmoid(gate) starts near 0 # Both stay LEARNABLE, so unlike a fixed SRA_GATE_SCALE the graph can # still grow to full strength during training. if os.environ.get('SRA_SOFT_START', '') not in ('', '0', 'false', 'False'): nn.init.zeros_(self.out_proj.weight) nn.init.zeros_(self.out_proj.bias) _gb = float(os.environ.get('SRA_GATE_BIAS', -4.0) or -4.0) for _m in self.gate_proj.modules(): if isinstance(_m, nn.Linear): nn.init.constant_(_m.bias, _gb) # sigmoid(-4) ~ 0.018 # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ def forward( self, y_emb: torch.Tensor, # [B, K, A, D] y_abs: torch.Tensor, # [B, K, A, T, 2] ← y_0_hat unnorm t_emb: torch.Tensor, # [B, D] tau: torch.Tensor, # [B] ∈ [0, 1] sigma_agent: torch.Tensor = None, # [B, K, A, T] or None agent_mask: torch.Tensor = None, # [B, A] bool, True=real; for padded batches ) -> torch.Tensor: # [B, K, A, D] B, K, A, D = y_emb.shape T = y_abs.shape[3] # Degenerate case: a single-agent scene has no edges — the graph is a # no-op, so just return the input untouched. if A <= 1: return y_emb # Variable-A support (e.g. SDD, ETH/UCY): the model is instantiated with # a max num_agents (padding budget), but each batch may contain a smaller # real A. Rebuild the graph skeleton on the fly when A changes. if A != self.num_agents: self.num_agents = A self._E0 = A * (A - 1) self.top_n = max(1, min(self._original_top_n, A - 1)) src, dst = [], [] for i in range(A): for j in range(A): if i != j: src.append(j); dst.append(i) self._single_edge_index = torch.tensor( [src, dst], dtype=torch.long, device=y_emb.device ) E0 = self._E0 N = self.top_n # ---- Per-node: y_0_hat trajectory embedding --------------------- y0_flat = y_abs.reshape(B * K * A, T * 2) # [B*K*A, T*2] y0_emb = self.y0_score_proj(y0_flat) # [B*K*A, D_s] # ---- Per-node: uncertainty scalar ------------------------------- if sigma_agent is not None: sigma_mean = sigma_agent.mean(dim=-1).reshape(B * K * A, 1) # [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)) # ---- Per-node: query and key ------------------------------------ node_feat = torch.cat([y0_emb, sigma_mean], dim=-1) # [B*K*A, D_s+1] q_bka = self.W_q(node_feat) # [B*K*A, D_s] k_bka = self.W_k(node_feat) # [B*K*A, D_s] # ---- Per-edge: build positions and relative features ------------ pos_bk = y_abs.reshape(B * K * A, T, 2) edge_index_bk = self._make_batched_edge_index(B * K) # [2, B*K*E0] pos_i_t = pos_bk[edge_index_bk[1]] # [E, T, 2] pos_j_t = pos_bk[edge_index_bk[0]] # [E, T, 2] rel_pos_t = pos_j_t - pos_i_t # [E, T, 2] 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] # Mean heading diff over T as geometric feature heading_full = _heading_diff(pos_i_t, pos_j_t) # [E, T, 1] heading_mean = heading_full.mean(dim=1) # [E, 1] # ---- Per-edge scoring ----------------------------------------------- if self.neighbor_mode == 'l2': # L2 distance: select closest neighbors by avg future distance avg_dist = rel_pos_t.norm(dim=-1).mean(dim=-1) # [E] scores = -avg_dist # negate so topk picks smallest distance elif self.neighbor_mode == 'semantic': # Semantic-only: learned query-key dot product, no geo bias q_i = q_bka[edge_index_bk[1]] # [E, D_s] k_j = k_bka[edge_index_bk[0]] # [E, D_s] scores = (q_i * k_j).sum(dim=-1) * self.scale # [E] else: # RAG: semantic dot product + geometric bias q_i = q_bka[edge_index_bk[1]] # [E, D_s] k_j = k_bka[edge_index_bk[0]] # [E, D_s] semantic_score = (q_i * k_j).sum(dim=-1) * self.scale # [E] geo_feat = torch.cat([mean_rel, std_rel, min_dist, heading_mean], dim=-1) # [E, 6] geo_bias = self.geo_mlp(geo_feat).squeeze(-1) # [E] scores = semantic_score + geo_bias # [E] # ---- Mask padded agents so top-N never selects them -------------- if agent_mask is not None: # Build per-edge "both endpoints real" mask # agent_mask: [B, A] -> broadcast to [B, K, A] for node-ness node_real = agent_mask.unsqueeze(1).expand(B, K, A).reshape(B * K * A) # [B*K*A] # edge_index_bk has shape [2, B*K*E0], rows 0=src=j, 1=dst=i src_real = node_real[edge_index_bk[0]] # [E] dst_real = node_real[edge_index_bk[1]] # [E] edge_real = src_real & dst_real # [E] scores = scores.masked_fill(~edge_real, float('-inf')) # ---- Top-N selection per target agent --------------------------- scores_grouped = scores.view(B * K * A, A - 1) # Cap N so we never ask for more neighbors than rows. N_use = min(N, A - 1) _, top_idx = scores_grouped.topk(N_use, 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) # [B*K*E0] # When padding is present, selected edges whose target agent is padded # produce zero messages anyway; additionally drop any edge whose # endpoint is padded so the GNN sees only real edges. if agent_mask is not None: mask_flat = mask_flat & edge_real # [B*K*E0] # ---- RelTrajEncoder on selected edges only ---------------------- rel_pos_sel = rel_pos_t[mask_flat] # [E_sel, T, 2] heading_sel = heading_full[mask_flat] # [E_sel, T, 1] if self.edge_mode == 'full': encoder_input = torch.cat([rel_pos_sel, heading_sel], dim=-1) # [E_sel, T, 3] elif self.edge_mode == 'dist_only': encoder_input = rel_pos_sel.norm(dim=-1, keepdim=True) # [E_sel, T, 1] elif self.edge_mode == 'relpos_only': encoder_input = rel_pos_sel # [E_sel, T, 2] elif self.edge_mode == 'heading_only': encoder_input = heading_sel # [E_sel, T, 1] elif self.edge_mode == 'full_relvel': rel_vel = rel_pos_sel[:, 1:] - rel_pos_sel[:, :-1] # [E_sel, T-1, 2] rel_vel = torch.cat([rel_vel, rel_vel[:, -1:]], dim=1) # [E_sel, T, 2] encoder_input = torch.cat([rel_pos_sel, heading_sel, rel_vel], dim=-1) # [E_sel, T, 5] elif self.edge_mode == 'vel_only': rel_vel = rel_pos_sel[:, 1:] - rel_pos_sel[:, :-1] # [E_sel, T-1, 2] rel_vel = torch.cat([rel_vel, rel_vel[:, -1:]], dim=1) # [E_sel, T, 2] encoder_input = rel_vel 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 # [E_sel, T] else: sigma_bias = None edge_attr_sparse = self.rel_traj_encoder( encoder_input, sigma_bias ) # [E_sel, D] # ---- Sparse GNN pass -------------------------------------------- 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) # ---- Gated residual --------------------------------------------- orig = y_emb.reshape(B * K * A, D) gate = self.gate_proj(torch.cat([orig, nodes], dim=-1)) # SRA_GATE_SCALE: global damping on the graph residual (default 1.0 = off). _gs = float(os.environ.get('SRA_GATE_SCALE', 1.0) or 1.0) res = _gs * gate * self.out_proj(nodes) # [N, D] graph perturbation # SRA_RES_CAP: per-node residual-norm cap (default 0 = off). With # SRA_EDGE_FIX=1 training is healthy (loss decreases monotonically) but # *sampling* diverges: the graph is applied at every one of MoFlow's flow # steps and its output feeds the next step, so any oversized per-node # perturbation compounds geometrically over the integration. The old # scene-mixing bug hid this by leaving half the nodes orphaned (weaker # perturbation, less compounding). Capping each node's residual NORM # bounds the per-step perturbation that drives the blow-up, while leaving # the (majority) small perturbations untouched — unlike a global scale it # only clips the outliers, so the graph keeps its normal expressive range. # 주의: 예전 구현 `res * (rn.clamp(max=cap) / (rn + 1e-6))` 은 두 가지가 틀렸다. # (1) res 가 정확히 0 이면 스케일이 0/1e-6 = 0 이 되어 **gradient 도 0** 이다. # SRA_SOFT_START(out_proj zero-init) 와 같이 켜면 out_proj 가 0 에 영구히 # 갇혀 그래프가 학습되지 않는다(=사실상 host 단독). 실제로 그 조합으로 # 돌린 실행들의 out_proj 는 58 epoch 뒤에도 정확히 0 이었다. # (2) cap 미만인데도 rn/(rn+1e-6) 만큼 축소된다 (rn=1e-5 이면 0.909 배). # torch.where 로 바꾸면 cap 이하는 정확히 무연산(스케일 1)이고 res=0 에서도 # gradient 가 흐른다. clamp_min 은 미선택 분기의 Inf 를 막는다. _cap = float(os.environ.get('SRA_RES_CAP', 0.0) or 0.0) if _cap > 0: rn = res.norm(dim=-1, keepdim=True) # [N, 1] res = res * torch.where(rn > _cap, _cap / rn.clamp_min(1e-6), torch.ones_like(rn)) # SRA_RES_CAP_REL: 노드별 상한을 **호스트 임베딩 norm 에 비례**해 정한다. # 절대 cap 은 임베딩 스케일에 의존해 호스트마다 의미가 달라진다(MoFlow 에서 # 튜닝한 3.0 이 MID 에서는 사실상 무연산일 수 있다). 샘플링 발산은 결국 # "스텝당 상대 섭동"이 누적되는 문제이므로, ‖res‖ ≤ ratio·‖orig‖ 로 두면 # 스케일 무관하게 누적률을 직접 제한한다. 측정값 기준 무제한 시 비율은 ~0.39. _rel = float(os.environ.get('SRA_RES_CAP_REL', 0.0) or 0.0) if _rel > 0: lim = _rel * orig.norm(dim=-1, keepdim=True) # [N, 1] 노드별 상한 rn = res.norm(dim=-1, keepdim=True) res = res * torch.where(rn > lim, lim / rn.clamp_min(1e-6), torch.ones_like(rn)) out = orig + res return out.view(B, K, A, D)