""" MotionTransformerGraph: extends MotionTransformer with a FutureInteractionGraph module inserted after the per-agent self-attention. Two-pass forward design: Pass 1 (torch.no_grad): run the full model with y_t as the graph edge source and no sigma weighting → produces y_0_hat and logvar_. Pass 2 (grad enabled): run the identical model again, now using y_0_hat from pass 1 as the graph edge source, and sigma derived from logvar_ to soft-weight the pairwise agent interaction (certain → uncertain). Uncertainty: A logvar_head (same MLP structure as reg_head) predicts per-agent per-timestep log-variance [B, K, A, T*2]. From this, a per-agent uncertainty scalar sigma = sqrt(exp(logvar).mean()) is computed and passed to the graph as the directional weight: w_ij = sigmoid(γ * (σ_i − σ_j) + 0.5) so uncertain agents receive more information from certain neighbors. The NLL loss on logvar is applied in flow_matching.py. """ import math import torch import torch.nn as nn from einops import rearrange, repeat from models.backbone import MotionTransformer from models.graph_interaction_nba_v6 import FutureInteractionGraphV6 as FutureInteractionGraph from models.utils.common_layers import build_mlps from utils.normalization import unnormalize_min_max, unnormalize_sqrt class MotionTransformerGraph(MotionTransformer): """MotionTransformer augmented with a future-interaction graph module and a per-agent uncertainty head. Extra constructor kwargs: graph_num_gnn_layers (int, default 2) graph_dropout (float, default 0.1) """ def __init__(self, model_config, logger, config, graph_num_gnn_layers: int = 2, graph_dropout: float = 0.1): super().__init__(model_config, logger, config) self.T_future = config.future_frames # 20 self.A = config.agents # 11 self.data_norm = config.get('data_norm', 'min_max') D = self.dim # 128 time_dim = D self.future_graph = FutureInteractionGraph( embed_dim = D, future_steps = self.T_future, num_agents = self.A, num_heads = 4, dropout = graph_dropout, num_gnn_layers = graph_num_gnn_layers, time_dim = time_dim, ) # Uncertainty head — same MLP structure as reg_head. # Predicts log-variance [B, K, A, T*2] from readout_token. self.logvar_head = build_mlps( c_in = self.dim, mlp_channels = self.model_cfg.REGRESSION_MLPS, ret_before_act = True, without_norm = True, ) params_graph = sum(p.numel() for p in self.future_graph.parameters()) params_logvar = sum(p.numel() for p in self.logvar_head.parameters()) logger.info("FutureInteractionGraph parameters: {:,}".format(params_graph)) logger.info("LogvarHead parameters: {:,}".format(params_logvar)) # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ def _unnormalize_y(self, y_norm: torch.Tensor) -> torch.Tensor: """Unnormalize [B, K, A, T, 2] from training norm back to metres.""" if self.data_norm == 'min_max': return unnormalize_min_max( y_norm, self.config.fut_traj_min, self.config.fut_traj_max, -1, 1, ) elif self.data_norm == 'sqrt': sqrt_a_ = torch.tensor( [self.config.sqrt_x_a, self.config.sqrt_y_a], device=y_norm.device, ) sqrt_b_ = torch.tensor( [self.config.sqrt_x_b, self.config.sqrt_y_b], device=y_norm.device, ) return unnormalize_sqrt(y_norm, sqrt_a_, sqrt_b_) else: return y_norm # ------------------------------------------------------------------ # Shared single-pass implementation # ------------------------------------------------------------------ def _forward_impl(self, y, time, x_data, y_0_for_graph=None, sigma_for_graph=None, skip_graph=False): """Single forward pass. y_0_for_graph: [B, K, A, T*2] normalized, or None. None → graph edge features built from y_t (noisy). Given → graph edge features built from this cleaner prediction. sigma_for_graph: [B, K, A] per-agent uncertainty scalar, or None. None → graph uses scalar denoising tau for directional weight. Given → graph uses sigma to promote certain→uncertain flow. skip_graph: if True, skip the future interaction graph entirely. Returns: (denoiser_x [B,K,A,T*2], denoiser_cls [B,K,A], logvar [B,K,A,T*2]) """ # ---- Shape normalisation (variable-A compatible) ----------------- # Accept [B, K, A, T_future, 2] or [B, K, A, T_future*2] with A taken # from the input, not from self.A (supports SDD padded batches). if y.dim() == 5 and y.size(-1) == 2 and y.size(-2) == self.T_future: B, K, A = y.size(0), y.size(1), y.size(2) y = y.reshape(B, K, A, self.T_future * 2) else: assert y.size(-1) == self.T_future * 2, \ f"Unexpected y shape: {y.shape}" device = y.device B, K, A, _ = y.shape # ---- Context encoder (past trajectories) ------------------------- agent_type = self.config.get('agent_type', 'sport') agent_mask_ctx = x_data.get('agent_mask', None) if isinstance(x_data, dict) else None encoder_out = self.context_encoder( x_data['past_traj_original_scale'], agent_type=agent_type, agent_mask=agent_mask_ctx, ) # [B, A, D] encoder_out_batch = repeat( encoder_out, 'b a d -> b k a d', k=K, a=A ) # [B, K, A, D] # ---- Noisy-y embedding ------------------------------------------- y_emb = self.noisy_y_mlp(y) # [B, K, A, D] # ---- Time embedding (keep time_ ∈ [0,1] for tau) ---------------- time_ = time if self.config.denoising_method == 'fm': time = time * 1000.0 t_emb = self.time_mlp(time) # [B, D] t_emb_batch = repeat(t_emb, 'b d -> b k a d', b=B, k=K, a=A) # [B, K, A, D] # ---- Positional encodings ---------------------------------------- k_pe = self.motion_query_embedding( torch.arange(self.model_cfg.NUM_PROPOSED_QUERY, device=device) ) k_pe_batch = repeat(k_pe, 'k d -> b k a d', b=B, a=A) # Use actual A from input (supports variable-A SDD). a_pe = self.agent_order_embedding(torch.arange(A, device=device)) a_pe_batch = repeat(a_pe, 'a d -> b k a d', b=B, k=K) # ---- K-level self-attention -------------------------------------- y_emb_k = rearrange( self.apply_PE(y_emb, k_pe_batch, a_pe_batch), 'b k a d -> (b a) k d', ) y_emb_k = self.noisy_y_attn_k(y_emb_k) y_emb = rearrange(y_emb_k, '(b a) k d -> b k a d', b=B, a=A) # ---- Agent-level self-attention (with padding mask for SDD) ----- y_emb_a = rearrange(y_emb, 'b k a d -> (b k) a d') agent_mask_bka = x_data.get('agent_mask', None) if isinstance(x_data, dict) else None if agent_mask_bka is not None: kp_mask_a = ~agent_mask_bka.unsqueeze(1).expand(-1, K, -1).reshape(B * K, A) y_emb_a = self.noisy_y_attn_a(y_emb_a, src_key_padding_mask=kp_mask_a) else: y_emb_a = self.noisy_y_attn_a(y_emb_a) y_emb = rearrange(y_emb_a, '(b k) a d -> b k a d', b=B, k=K) # ---- Embedding dropout (training only) --------------------------- if self.training and self.config.get('drop_method', None) == 'emb': m, k_drop = self.config.drop_logi_m, self.config.drop_logi_k p_m = 1 / (1 + torch.exp(-k_drop * (time_ - m))) p_m = p_m[:, None, None, None] y_emb = y_emb.masked_fill(torch.rand_like(p_m) < p_m, 0.) # ================================================================== # >>> Future Euclidean Interaction Graph <<< # ================================================================== if not skip_graph: # Edge source: y_0_hat from pass 1 (cleaner), or y_t (fallback). y_graph_src = (y_0_for_graph.view(B, K, A, self.T_future, 2) if y_0_for_graph is not None else y.view(B, K, A, self.T_future, 2)) y_graph_unnorm = self._unnormalize_y(y_graph_src) init_pos = x_data['past_traj_original_scale'][:, :, -1, :2] # [B, A, 2] y_abs = y_graph_unnorm + init_pos.unsqueeze(1).unsqueeze(3) tau = time_ # [B] ∈ [0, 1] agent_mask = x_data.get('agent_mask', None) if isinstance(x_data, dict) else None y_emb_graph = self.future_graph( y_emb, y_abs, t_emb, tau, sigma_agent=sigma_for_graph, # None in pass 1; [B,K,A] in pass 2 agent_mask=agent_mask, # [B,A] for padded SDD batches ) y_emb = y_emb_graph # ================================================================== # ---- Context fusion + motion decoder ---------------------------- emb_fusion = self.init_emb_fusion_mlp( torch.cat((encoder_out_batch, y_emb, t_emb_batch), dim=-1) ) query_token = self.post_pe_cat_mlp( self.apply_PE(emb_fusion, k_pe_batch, a_pe_batch) ) readout_token = self.motion_decoder(query_token, t_emb) # [B, K, A, D] # ---- Readout heads ---------------------------------------------- denoiser_x = self.reg_head(readout_token) # [B, K, A, T*2] denoiser_cls = self.cls_head(readout_token).squeeze(-1) # [B, K, A] logvar = self.logvar_head(readout_token) # [B, K, A, T*2] return denoiser_x, denoiser_cls, logvar # ------------------------------------------------------------------ # Two-pass forward (overrides MotionTransformer.forward) # ------------------------------------------------------------------ def forward(self, y, time, x_data, y_0_prev=None): """Two-pass forward — identical architecture both passes. Pass 1 (no_grad): - Training: graph uses GT future trajectory for edge features. - Inference: graph uses y_0_prev from previous sampling step, or skips graph if y_0_prev is None (first step). - Produces y_0_hat (clean prediction) and logvar_ (uncertainty). Pass 2 (grad): - Graph uses y_0_hat for cleaner edge features. - Graph uses sigma derived from logvar_ to weight agent messages: w_ij = sigmoid(γ * (σ_i − σ_j) + 0.5) promoting information flow from certain → uncertain agents. Returns: (denoiser_x, denoiser_cls, logvar) logvar is used by flow_matching.p_losses for the NLL uncertainty loss. """ # Pass 1 — no gradient, get y_0_hat and preliminary logvar if self.training: # Use GT future as graph edge source during training. # A comes from input (variable for SDD), not from self.A. K = self.model_cfg.NUM_PROPOSED_QUERY gt_fut = x_data['fut_traj'] # [B, A, T, 2] B_gt, A_gt = gt_fut.shape[0], gt_fut.shape[1] y_0_gt = gt_fut.unsqueeze(1).expand(-1, K, -1, -1, -1) # [B, K, A, T, 2] y_0_for_pass1 = y_0_gt.reshape(B_gt, K, A_gt, -1) # [B, K, A, T*2] else: y_0_for_pass1 = y_0_prev # None at first step → skip graph with torch.no_grad(): y_0_hat, _, logvar_ = self._forward_impl( y, time, x_data, y_0_for_graph=y_0_for_pass1, sigma_for_graph=None, skip_graph=(y_0_for_pass1 is None), ) # Derive per-agent per-timestep uncertainty from logvar_: # logvar_ [B, K, A, T*2] → exp → mean over xy → sqrt → [B, K, A, T] B = y_0_hat.shape[0] K = self.model_cfg.NUM_PROPOSED_QUERY A = y_0_hat.shape[2] # variable A from input (not self.A) T = self.T_future sigma_ = (logvar_.detach() .view(B, K, A, T, 2) .clamp(-10, 10) .exp() .mean(dim=-1) # mean over xy only → [B, K, A, T] .sqrt()) # [B, K, A, T] # Pass 2 — full gradient, graph uses y_0_hat edges + sigma weighting use_sigma = self.config.get('use_sigma_gating', True) return self._forward_impl( y, time, x_data, y_0_for_graph=y_0_hat.detach(), sigma_for_graph=sigma_ if use_sigma else None, )