""" E4 plug-in baselines — alternative interaction/refinement modules that honor the EXACT FutureInteractionGraphV6 forward contract, for a protocol-matched comparison vs SRA: forward(y_emb[B,K,A,D], y_abs[B,K,A,T,2], t_emb[B,D], tau[B], sigma_agent[B,K,A,T]|None, agent_mask[B,A]|None=None) -> [B,K,A,D] (gated residual, embedding->embedding, NO sigma out; robust to K in {1,10,20}, any A) Modules: * GameFormerInteraction [ref 20]: FAITHFUL GameFormer level-k InteractionDecoder — host denoiser = InitialDecoder (level 0); denoiser K samples = GameFormer K modes; FutureEncoder + score-weighted mode aggregation + interaction self-attention + own-future-masked cross-attention + L-level GMM refinement. No neighbor selection, no uncertainty (that is SRA). See the detailed faithfulness note above the class. * CoarseToFineRefine [ref 22]: Coarse-to-Fine (Jia et al. 2022) TEMPORAL refinement — a per-agent temporal module (GRU / 1D-CNN over the future horizon). No interaction. Orthogonal axis (temporal). Both accept the SAME __init__ kwargs as V6 (extras absorbed by **kwargs) and return the SAME gated-residual embedding, so they are true drop-ins at each host's V6 instantiation line. Selection is via env var SRA_MODULE in {sra, gameformer, c2f} (see build_interaction_module). """ import os import torch import torch.nn as nn class _GatedResidualHead(nn.Module): """V6's output convention: out = input + gate(cat[input,refined]) * out_proj(refined). Random-initialised (active at init), exactly like V6 — hosts that need no-op-at-init (MID/LED) supply their own external zero-init projection; MoFlow uses the output directly.""" def __init__(self, embed_dim): super().__init__() self.gate_proj = nn.Sequential(nn.Linear(2 * embed_dim, embed_dim), nn.Sigmoid()) self.out_proj = nn.Linear(embed_dim, embed_dim) nn.init.zeros_(self.out_proj.weight) # NO-OP AT INIT: module starts as identity nn.init.zeros_(self.out_proj.bias) # (stabilizes MoFlow, which uses output directly) def forward(self, orig, refined): # both [N, D] gate = self.gate_proj(torch.cat([orig, refined], dim=-1)) res = self.out_proj(refined) _cap = float(os.environ.get('MOFLOW_DAMP', 0) or 0) # cap residual norm -> MoFlow flow-field stability if _cap > 0: rn = res.norm(dim=-1, keepdim=True) res = res * (rn.clamp(max=_cap) / (rn + 1e-6)) return orig + gate * res # ============================================================================ # GameFormer [ref 20] -- FAITHFUL level-k InteractionDecoder as a denoiser plug-in # ============================================================================ # Faithful to Liu et al. ICCV'23 (github.com/MCZhi/GameFormer, model/modules.py): # * The host DENOISER plays the InitialDecoder (level 0): its current estimate # (y_emb, y_abs) IS the level-0 prediction, and the denoiser's K SAMPLES are # treated as GameFormer's K MODES (the modal set the level-k step reasons over). # * Each InteractionDecoder level implements the real mechanism: # - FutureEncoder: MLP on per-mode [x,y,heading,vx,vy], max-pooled over time # (modules.FutureEncoder; box-size channels dropped -- no size in NBA/sport). # - score-softmax-weighted aggregation of the K modes -> one future token/agent. # - interaction SelfTransformer over agents (game-theoretic "condition on all # other agents' level-(k-1) futures"). # - cross-attention of each mode's content (prev content + own future) to # [interaction ; scene-context], with the agent's OWN future token MASKED # (GameFormer's own-future masking). # - a GMM-mu trajectory decoder emits the level-k refinement; iterate L levels. # * The final level's content -> gated residual to the host embedding. # ONE gap vs the paper: per-level imitation SUPERVISION (GameFormer sums a GMM loss # over every level). A plug-in is trained only through the host denoising loss, so # the levels are learned end-to-end rather than level-wise-supervised. class _AgentSelfAttn(nn.Module): """SelfTransformer over agents (modules.SelfTransformer). x: [B, A, D].""" def __init__(self, dim, heads, dropout): super().__init__() self.heads = heads self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) self.n1 = nn.LayerNorm(dim); self.n2 = nn.LayerNorm(dim) self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim * 4, dim)) def forward(self, x, key_padding_mask=None): a, _ = self.attn(x, x, x, key_padding_mask=key_padding_mask, need_weights=False) x = self.n1(a + x) return self.n2(self.ffn(x) + x) class _CrossAttn(nn.Module): """CrossTransformer (modules.CrossTransformer). q:[N,Lq,D] k/v:[N,Lk,D].""" def __init__(self, dim, heads, dropout): super().__init__() self.heads = heads self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) self.n1 = nn.LayerNorm(dim); self.n2 = nn.LayerNorm(dim) self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim * 4, dim)) def forward(self, q, k, v, attn_mask=None): a, _ = self.attn(q, k, v, attn_mask=attn_mask, need_weights=False) a = self.n1(a) return self.n2(self.ffn(a) + a) class _FutureEncoder(nn.Module): """modules.FutureEncoder: per-mode future -> vector via max-pool over an MLP on [x, y, heading, vx, vy]. trajs [B,K,A,T,2] (K modes), cur_xy [B,K,A,2].""" def __init__(self, dim, dt=0.2): super().__init__() self.dt = dt self.mlp = nn.Sequential(nn.Linear(5, 64), nn.ReLU(inplace=True), nn.Linear(64, dim)) def forward(self, trajs, cur_xy): pos = trajs - cur_xy.unsqueeze(-2) # centered position xy = torch.cat([cur_xy.unsqueeze(-2), trajs], dim=-2) # [B,K,A,T+1,2] dxy = torch.diff(xy, dim=-2) v = dxy / self.dt theta = torch.atan2(dxy[..., 1], dxy[..., 0].clamp(min=1e-3)).unsqueeze(-1) state = torch.cat([pos, theta, v], dim=-1) # [B,K,A,T,5] h = self.mlp(state.detach()) # GameFormer DETACHES future feats return h.max(dim=-2).values # [B,K,A,D] class GameFormerInteraction(nn.Module): """FAITHFUL GameFormer level-k InteractionDecoder as a denoiser plug-in (see header).""" 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, edge_mode='full', neighbor_mode='rag', num_levels=None, **kwargs): super().__init__() self.embed_dim = embed_dim self.num_agents = num_agents # LED wrapper reads this L = int(os.environ.get('GF_LEVELS', num_levels if num_levels is not None else 3)) self.num_levels = L self.t_proj = nn.Linear(embed_dim, embed_dim) self.future_encoder = _FutureEncoder(embed_dim) # SHARED across levels self.score_head = nn.ModuleList([ # per-mode score -> softmax weight nn.Sequential(nn.Linear(embed_dim, 64), nn.ELU(), nn.Linear(64, 1)) for _ in range(L)]) self.interaction_enc = nn.ModuleList([_AgentSelfAttn(embed_dim, num_heads, dropout) for _ in range(L)]) self.query_enc = nn.ModuleList([_CrossAttn(embed_dim, num_heads, dropout) for _ in range(L)]) self.traj_decode = nn.ModuleList([nn.Linear(embed_dim, future_steps * 2) for _ in range(L)]) for dec in self.traj_decode: # each level starts as identity (stable) nn.init.zeros_(dec.weight); nn.init.zeros_(dec.bias) self.head = _GatedResidualHead(embed_dim) def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None, agent_mask=None): B, K, A, D = y_emb.shape T = y_abs.shape[3] if A <= 1: return y_emb dev = y_emb.device cur_xy = y_abs[..., 0, :] # [B,K,A,2] level-0 reference (1st future step) y_traj = y_abs # [B,K,A,T,2] current (level-0) future t = self.t_proj(t_emb).view(B, 1, 1, D) # denoiser timestep context agent_ctx = y_emb.mean(dim=1) + t.squeeze(1) # [B,A,D] fixed scene context (over modes) kp = (~agent_mask.bool()) if agent_mask is not None else None # [B,A] True=pad ar = torch.arange(A, device=dev) heads = self.query_enc[0].heads content = y_emb # level-0 content = host embedding for l in range(self.num_levels): multi_fut = self.future_encoder(y_traj, cur_xy) # [B,K,A,D] per-mode future w = self.score_head[l](multi_fut).softmax(dim=1) # [B,K,A,1] over K modes agg_fut = (multi_fut * w).sum(dim=1) # [B,A,D] aggregated future interaction = self.interaction_enc[l](agg_fut, kp) # [B,A,D] game-theoretic ctx = torch.cat([interaction, agent_ctx], dim=1) # [B,2A,D] q = (content + multi_fut + t).reshape(B, K * A, D) # prev content + own future + t am = torch.zeros(B, K, A, 2 * A, dtype=torch.bool, device=dev) am[:, :, ar, ar] = True # mask OWN future (interaction block) if agent_mask is not None: pad = (~agent_mask.bool())[:, None, None, :] # [B,1,1,A] am[..., :A] = am[..., :A] | pad am[..., A:] = am[..., A:] | pad am = am.reshape(B, K * A, 2 * A)[:, None].expand( B, heads, K * A, 2 * A).reshape(B * heads, K * A, 2 * A) content = self.query_enc[l](q, ctx, ctx, am).reshape(B, K, A, D) y_traj = y_traj + self.traj_decode[l](content).view(B, K, A, T, 2) # level-k refinement refined = content.reshape(B * K * A, D) orig = y_emb.reshape(B * K * A, D) return self.head(orig, refined).view(B, K, A, D) class CoarseToFineRefine(nn.Module): """Coarse-to-Fine temporal refiner: per-agent GRU/1D-CNN over the future horizon (no interaction).""" 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, edge_mode='full', neighbor_mode='rag', refine_type='gru', hidden=200, **kwargs): super().__init__() self.embed_dim = embed_dim self.num_agents = num_agents # LED wrapper reads this (skips variable-A rebuild when ==A) self.refine_type = os.environ.get('C2F_REFINE', refine_type) self.pos_emb = nn.Linear(2, hidden) if self.refine_type == 'cnn': # 1D-CNN temporal refiner (per-timestep output) self.temporal = nn.Sequential( nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True), nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True), nn.Conv1d(hidden, hidden, 3, padding=1), nn.ReLU(inplace=True)) else: # AUTOREGRESSIVE (unidirectional) GRU self.temporal = nn.GRU(hidden, hidden, num_layers=2, batch_first=True, dropout=dropout) self.delta_head = nn.Linear(hidden, 2) # per-timestep coarse->fine correction Δ_t self.traj_encode = nn.Sequential( # re-encode the FINE trajectory -> host feature nn.Linear(future_steps * 2, embed_dim), nn.ReLU(inplace=True), nn.Linear(embed_dim, embed_dim)) self.t_proj = nn.Linear(embed_dim, embed_dim) self.head = _GatedResidualHead(embed_dim) def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None, agent_mask=None): # Faithful coarse-to-fine: host prediction = COARSE trajectory; walk it temporally # (autoregressive GRU / 1D-CNN) emitting a per-timestep correction Δ_t -> FINE trajectory, # then re-encode the fine trajectory into the host residual. Per-agent (no interaction). B, K, A, D = y_emb.shape T = y_abs.shape[3] N = B * K * A y_coarse = (y_abs - y_abs[..., :1, :]).reshape(N, T, 2) # coarse trajectory (centered) seq = self.pos_emb(y_coarse) # [N, T, H] if self.refine_type == 'cnn': h = self.temporal(seq.transpose(1, 2)).transpose(1, 2) # [N, T, H] else: h, _ = self.temporal(seq) # [N, T, H] per-timestep (autoregressive) delta = self.delta_head(h) # [N, T, 2] coarse->fine correction y_fine = y_coarse + delta # refined (fine) trajectory refined = self.traj_encode(y_fine.reshape(N, T * 2)) # re-encode fine traj -> feature t = self.t_proj(t_emb).view(B, 1, 1, D).expand(B, K, A, D).reshape(N, D) refined = refined + t orig = y_emb.reshape(N, D) return self.head(orig, refined).view(B, K, A, D) def build_interaction_module(name=None, **kwargs): """Factory used at each host's V6 instantiation line. `name` defaults to env SRA_MODULE (then 'sra'). SRA branch imports the real V6 and forwards only kwargs it accepts.""" name = (name or os.environ.get('SRA_MODULE') or 'sra').lower() if name in ('sra', 'v6', 'graph'): import inspect from models.graph_interaction_nba_v6 import FutureInteractionGraphV6 allowed = set(inspect.signature(FutureInteractionGraphV6.__init__).parameters) v6kw = {k: v for k, v in kwargs.items() if k in allowed} return FutureInteractionGraphV6(**v6kw) if name in ('gameformer', 'gf', 'gameformer_int'): return GameFormerInteraction(**kwargs) if name in ('c2f', 'coarse2fine', 'coarsetofine'): return CoarseToFineRefine(**kwargs) raise ValueError(f"unknown SRA_MODULE / interaction_module: {name!r}")