Add GAMEFORMER_SRA.md (negative result) + standalone GameFormer/C2F/SRA-GameFormer scripts

GAMEFORMER_SRA.md

@@ -0,0 +1,147 @@
+# GameFormer + SRA — a negative result
+
+**Question.** SRA improves three *iterative denoisers* (MID, LED, MoFlow). Do its ideas also help
+a **feedforward, non-diffusion** predictor? If so, SRA generalizes beyond denoising.
+
+**Answer. No.** Injecting SRA's three components into GameFormer is **neutral-to-slightly-negative**.
+
+| Model (NBA, min-ADE₂₀ / min-FDE₂₀ @ 4.0 s) | ADE₄ | FDE₄ | params |
+|---|---|---|---|
+| Vanilla GameFormer (control) | **0.8496** | **0.9557** | 11.85 M |
+| GameFormer **+ SRA** (`GF_SRA=1`) | 0.8554 | 0.9564 | 11.90 M |
+| **Δ (SRA − vanilla)** | **+0.0058 (worse)** | +0.0007 (≈neutral) | +0.05 M |
+
+Both peaked at epoch 48 of 50. Same file, same seed (3407), same schedule — **only the three flags differ**.
+
+---
+
+## 1. Experimental design
+
+The comparison is a **within-file ablation**, not a cross-codebase one:
+
+* One script, `gameformer_sra_nba_standalone.py`, contains both arms.
+* Three env-gated components; **with all flags off the code path is exactly vanilla GameFormer.**
+* Verified: flags-off reports **11.85 M params** — identical to the separate vanilla
+  implementation, which scores 0.848/0.951. The refactor perturbs nothing.
+* SRA adds only **+0.05 M params (+0.4 %)**, so any effect is attributable to the *mechanism*,
+  not to capacity.
+
+Config: level-k = 3, K = 20 modes, 6 encoder layers, AdamW lr 1e-4, wd 1e-4, grad-clip 5,
+MultiStepLR ×0.5, batch 128, 50 epochs, seed 3407. Standalone predictor — history only,
+**no denoiser and no external future signal**.
+
+---
+
+## 2. What was injected
+
+All three live inside GameFormer's `InteractionDecoder` (one per level-k stage).
+
+### 2.1 Uncertainty-gated message passing — `GF_SRA_SIGMA`
+
+GameFormer's GMM head already *emits* a per-step log-σ but never **uses** it for interaction.
+SRA turns it into a message gate: uncertain agents send weaker interaction messages.
+
+```python
+sig       = torch.exp(last_traj[..., 2:].clamp(0, 5)).mean(-1)   # [B,K,A,T] from prev level's GMM
+sig_agent = (sig * w[..., None]).sum(2).mean(-1)                 # [B,A] score-weighted, mean over T
+...
+futures = futures * self.sigma_gate(sig_agent[..., None])        # MLP(1→32→1)+Sigmoid ⇒ gate ∈ (0,1)
+```
+
+### 2.2 Sparse RAG neighbor selection — `GF_SRA_TOPN`
+
+GameFormer attends **densely** over all agents. SRA restricts each agent to its top-N (=5)
+neighbours, chosen by *future* proximity and biased toward confident agents.
+
+```python
+prox = (rep[:, :, None] - rep[:, None]).norm(dim=-1).mean(-1)    # [B,A,A] mean future distance
+if sig_agent is not None:
+    prox = prox + SIG_SEL_W * sig_agent[:, None, :]              # prefer confident neighbours
+prox    = prox.masked_fill(eye[None], float('inf'))              # exclude self
+nbr_idx = prox.topk(N, largest=False, dim=-1).indices
+```
+
+The resulting mask is applied **both** to the interaction self-attention (`attn_mask`) and to the
+cross-attention key mask, so sparsity is enforced everywhere, not just in one place.
+
+### 2.3 Explicit relational future features — `GF_SRA_RELFUT`
+
+GameFormer's `FutureEncoder` encodes each agent's future **individually**. SRA adds explicit
+**pairwise-relative** future geometry (Δposition, Δvelocity), max-pooled over time.
+
+```python
+dpos = rep[:, :, None] - rep[:, None]                            # [B,A,A,T,2]
+dvel = F.pad(torch.diff(dpos, dim=-2), (0, 0, 1, 0))
+relf = self.rel_mlp(torch.cat([dpos, dvel], -1)).max(dim=-2).values   # MLP(4→64→D)
+futures = futures + relf.sum(2) / denom                          # neighbour-mean
+```
+
+---
+
+## 3. Reproducing
+
+```bash
+# control (vanilla GameFormer)
+python gameformer_sra_nba_standalone.py --gpu 0 --exp van  --level 3 --epochs 50 --batch_size 128
+
+# + SRA (all three components)
+GF_SRA=1 python gameformer_sra_nba_standalone.py --gpu 0 --exp sra --level 3 --epochs 50 --batch_size 128
+```
+
+| Variable | Effect |
+|---|---|
+| `GF_SRA=1` | enable all three components |
+| `GF_SRA_SIGMA` / `GF_SRA_TOPN` / `GF_SRA_RELFUT` | enable individually |
+| `GF_TOPN` | neighbour budget (default 5) |
+| `GF_SIG_SEL_W` | weight of σ in neighbour selection (default 0.5) |
+
+The script prints its configuration so every run is self-documenting:
+
+```
+[SRA-GF] SRA components: neighbor-topN=True(N=5) sigma-gate=True relational-future=True
+```
+
+---
+
+## 4. Why it does not transfer
+
+1. **The future-interaction idea is already native to GameFormer.** Its level-k reasoning
+   *is* "condition each agent on the others' predicted futures". SRA's relational features
+   are therefore largely **redundant** here, adding parameters and noise rather than new signal.
+   In MID/LED/MoFlow nothing plays this role, which is exactly why SRA helps there.
+
+2. **The uncertainty signal is weaker.** In a denoiser, σ is *principled* — a noise schedule or
+   the leapfrog initializer's variance estimate, tied to how much the model should trust its
+   current iterate. GameFormer's log-σ is just a learned regression head; gating on it does not
+   carry comparable information.
+
+3. **Sparsity costs more than it saves at A = 11.** Dense attention over 11 agents is cheap and
+   fully informative; discarding all but 5 neighbours removes usable context. SRA's sparse
+   selection pays off on 23-agent sport scenes and on hosts where dense interaction is unstable.
+
+4. **One-shot vs iterative (the core reason).** SRA is built to *refine an evolving estimate*: it
+   is applied at every denoising step, so its gated residual is corrected repeatedly and errors
+   are damped. A feedforward decoder gets **one** pass — no iterative correction, no damping.
+   The regime SRA was designed for is absent.
+
+---
+
+## 5. How to read this result
+
+This is a **scope finding, not a failure**, and it is consistent with the paper's framing:
+SRA is a *denoiser* adapter. Its gains come from repeatedly reshaping an iteratively refined
+future estimate — a mechanism a single-pass predictor does not have.
+
+Stated honestly:
+
+> SRA's gains are specific to the iterative-denoising setting. Applied to a strong feedforward
+> predictor (GameFormer), it is neutral-to-slightly-negative, which supports rather than
+> contradicts the paper's positioning.
+
+It mirrors the MoFlow observation from the other direction: the *host's* update mechanism
+determines what an interaction module can contribute.
+
+**Caveats.** Single seed; the ADE gap (0.006, 0.7 %) is small enough that multi-seed runs could
+place it within noise — i.e. "neutral" rather than "worse". It is clearly **not** an improvement
+either way. No per-component ablation was run, so it remains open whether one component alone
+transfers even though the combination does not.

README.md

@@ -18,6 +18,12 @@ It covers the required directory layout, environment setup, the exact training c
 host × dataset, the E2 ablation settings, environment-variable switches, the adapter contract,
 and known gotchas.
 
+## Documents
+
+- [`RUNNING.md`](RUNNING.md) — how to run every host × dataset
+- [`GAMEFORMER_SRA.md`](GAMEFORMER_SRA.md) — GameFormer+SRA negative result (does SRA generalize to feedforward models?)
+- [`sample_data/README.md`](sample_data/README.md) — bundled 100-scene NBA smoke-test subset
+
 ## Contents
 
 ```

standalone/c2f_nba_standalone.py

@@ -0,0 +1,434 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+c2f_nba_standalone.py
+================================================================================
+FAITHFUL standalone re-implementation of the **Coarse-to-Fine** trajectory
+predictor of ref [22] ("Towards Capturing the Temporal Dynamics for Trajectory
+Prediction: A Coarse-to-Fine Approach") as a STANDALONE predictor for the NBA
+basketball dataset.
+
+It is a STANDALONE predictor: it consumes ONLY agent HISTORY (past trajectories)
+and predicts K multi-modal futures via a native COARSE stage followed by an
+autoregressive temporal FINE-refinement stage. It does NOT consume any external
+denoiser / diffusion prediction (unlike the plug-in `CoarseToFineRefine` module,
+which was handed the host's intermediate future estimate -> that borrowed SRA's
+future-interaction signal and was therefore unfair; this standalone version
+generates its own coarse trajectory, so it is faithful to the native method).
+
+--------------------------------------------------------------------------------
+WHAT THE NATIVE METHOD IS (and what we replicate)
+--------------------------------------------------------------------------------
+Coarse-to-fine = two-stage decoding that *captures temporal dynamics* by first
+predicting a rough full trajectory and then refining it step-by-step:
+
+  STAGE 0  (COARSE):  a multimodal decoder emits K rough full-horizon
+    trajectories per agent from the social-encoded context (mode + agent query
+    cross-attended to all agents' context, then an MLP trajectory head). This is
+    the model's OWN coarse prediction -- it is NOT received from a host.
+
+  STAGE 1..S (FINE):  the coarse trajectory is walked TEMPORALLY by a
+    unidirectional (autoregressive) GRU -- the mechanism the paper uses to
+    capture temporal dynamics -- conditioned on the mode's context, emitting a
+    per-timestep residual correction  delta_t ; y_fine = y_coarse + delta.
+    Repeated S times (progressive coarse -> fine -> finer). delta head is
+    zero-initialised so training first fits the coarse stage, then the refiner
+    engages (stable).
+
+  Mode scores: a per-mode classification head; winner-takes-all training.
+
+ENCODER (shared, IDENTICAL to the GameFormer standalone so the E4 comparison
+isolates the DECODER mechanism, coarse-to-fine vs level-k, not the backbone):
+  * AgentHistoryEncoder = 2-layer LSTM(6->256) over agent history + learned
+    player/ball type embedding.
+  * FusionEncoder = nn.TransformerEncoder (d=256, heads=8, ff=1024, gelu),
+    `encoder_layers` deep -> agent<->agent social self-attention (no map: NBA).
+
+LOSS (native coarse-to-fine supervision):
+  * coarse WTA-L2 (variety loss over K modes)  +  fine WTA-L2  +  mode
+    cross-entropy (label smoothing 0.2) on the fine-stage winning mode.
+  * endpoint-emphasised distance (mean_t + checkpoint-step sum) for mode
+    selection, marginal PER AGENT (NBA Table-1 metric is marginal min-ADE_20).
+
+--------------------------------------------------------------------------------
+NBA I/O + METRIC (matched to MoFlow / Table 1 -- identical to GameFormer standalone)
+--------------------------------------------------------------------------------
+  * Data: MoFlow's data/dataloader_nba.py::NBADatasetMinMax, same .npy, split,
+    scaling (traj_scale=94/28, traj_mean=[14,7.5]). Past=10, Future=20 (4.0 s).
+  * Predict in centered-abs frame (pos/scale - mean); cur_xy = last past step;
+    gt_center = fut_traj_original_scale (displacement) + cur_xy.
+  * Metric (identical to eval_perscene_moflow.py):
+      d = ||pred_disp - gt_disp|| ; ADE4 = d[:,:20].mean_t.min_K ; FDE4 =
+      d[:,19].min_K ; mean over 11 agents & scenes  (marginal min-of-K=20 @ 4s).
+
+Prints:  [C2F-NBA] epoch N ADE4=.. FDE4=..
+"""
+
+import os
+import sys
+import argparse
+import time
+
+# ------------------------------------------------------------------ GPU FIRST
+def _early_gpu():
+    for i, a in enumerate(sys.argv):
+        if a == '--gpu' and i + 1 < len(sys.argv):
+            return sys.argv[i + 1]
+        if a.startswith('--gpu='):
+            return a.split('=', 1)[1]
+    return None
+
+
+_g = _early_gpu()
+if _g is not None:
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(_g)
+os.environ.setdefault('MPLBACKEND', 'Agg')
+
+# ------------------------------------------------ reuse MoFlow's NBA pipeline
+MOFLOW_ROOT = '/mnt/jaewoo4tb/srtp/MoFlow'
+if MOFLOW_ROOT not in sys.path:
+    sys.path.insert(0, MOFLOW_ROOT)
+
+import types as _types
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+try:
+    os.chdir(MOFLOW_ROOT)
+except Exception:
+    pass
+from data.dataloader_nba import NBADatasetMinMax, seq_collate_nba
+
+
+# ============================================================================
+#  PRIMITIVES  (shared with the GameFormer standalone, verbatim)
+# ============================================================================
+D_MODEL = 256
+N_HEADS = 8
+DROPOUT = 0.1
+
+
+class CrossTransformer(nn.Module):
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, dropout=DROPOUT):
+        super().__init__()
+        self.cross_attention = nn.MultiheadAttention(dim, heads, dropout, batch_first=True)
+        self.norm_1 = nn.LayerNorm(dim)
+        self.norm_2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout))
+
+    def forward(self, query, key, value, mask=None):
+        attn, _ = self.cross_attention(query, key, value, key_padding_mask=mask)
+        attn = self.norm_1(attn)
+        return self.norm_2(self.ffn(attn) + attn)
+
+
+class AgentHistoryEncoder(nn.Module):
+    """2-layer LSTM(6->256) over agent history + player/ball type embedding."""
+
+    def __init__(self, in_dim=6, dim=D_MODEL, n_types=2):
+        super().__init__()
+        self.motion = nn.LSTM(in_dim, dim, 2, batch_first=True)
+        self.type_emb = nn.Embedding(n_types, dim)
+
+    def forward(self, hist, types):
+        B, A, T, C = hist.shape
+        traj, _ = self.motion(hist.reshape(B * A, T, C))
+        out = traj[:, -1].reshape(B, A, -1)
+        out = out + self.type_emb(types)[None]
+        return out
+
+
+class FusionEncoder(nn.Module):
+    """Agent<->agent social self-attention (no map for NBA)."""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, layers=6, dropout=DROPOUT):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(
+            d_model=dim, nhead=heads, dim_feedforward=dim * 4,
+            activation=F.gelu, dropout=dropout, batch_first=True)
+        self.encoder = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
+
+    def forward(self, tokens, mask=None):
+        return self.encoder(tokens, src_key_padding_mask=mask)
+
+
+# ============================================================================
+#  COARSE-TO-FINE DECODER
+# ============================================================================
+class CoarseDecoder(nn.Module):
+    """Stage-0: K multimodal rough full-horizon trajectories per agent.
+    Mode + agent query added to the agent's context token, cross-attended to the
+    full agent context (social), then an MLP trajectory head + a mode-score head."""
+
+    def __init__(self, modalities, n_agents, future_len, dim=D_MODEL):
+        super().__init__()
+        self.M = modalities
+        self.multi_modal_query_embedding = nn.Embedding(modalities, dim)
+        self.agent_query_embedding = nn.Embedding(n_agents, dim)
+        self.query_encoder = CrossTransformer(dim)
+        self.traj_head = nn.Sequential(
+            nn.Linear(dim, 512), nn.ELU(), nn.Dropout(0.1),
+            nn.Linear(512, future_len * 2))
+        self.score_head = nn.Sequential(
+            nn.Linear(dim, 64), nn.ELU(), nn.Dropout(0.1), nn.Linear(64, 1))
+        self.future_len = future_len
+        self.register_buffer('modal', torch.arange(modalities).long())
+        self.register_buffer('agent', torch.arange(n_agents).long())
+
+    def forward(self, encoding, cur_xy, mask=None):
+        B, A, D = encoding.shape
+        M, T = self.M, self.future_len
+        mm = self.multi_modal_query_embedding(self.modal)          # [M, D]
+        ag = self.agent_query_embedding(self.agent)                # [A, D]
+        query = encoding[:, :, None, :] + mm[None, None] + ag[None, :, None]  # [B,A,M,D]
+
+        q = query.reshape(B * A, M, D)
+        kv = encoding[:, None, :, :].expand(B, A, A, D).reshape(B * A, A, D)
+        km = None
+        if mask is not None:
+            km = mask[:, None, :].expand(B, A, A).reshape(B * A, A)
+        content = self.query_encoder(q, kv, kv, km)                # [B*A, M, D]
+
+        coarse = self.traj_head(content).view(B, A, M, T, 2)
+        coarse = coarse + cur_xy[:, :, None, None, :]              # centered-abs
+        score = self.score_head(content).view(B, A, M)
+        return content.view(B, A, M, D), coarse, score
+
+
+class FineRefiner(nn.Module):
+    """Stage-1..S: autoregressive temporal refinement of the coarse trajectory.
+    A unidirectional GRU walks the (centered) coarse trajectory, conditioned on
+    the mode context, emitting a per-timestep residual delta_t. Repeated S times."""
+
+    def __init__(self, dim=D_MODEL, hidden=256, n_stages=2, dropout=DROPOUT):
+        super().__init__()
+        self.n_stages = n_stages
+        self.pos_emb = nn.Linear(2, hidden)
+        self.ctx_proj = nn.Linear(dim, hidden)
+        self.gru = nn.GRU(hidden, hidden, 2, batch_first=True, dropout=dropout)
+        self.delta = nn.Linear(hidden, 2)
+        nn.init.zeros_(self.delta.weight)      # start as identity: fine == coarse at init
+        nn.init.zeros_(self.delta.bias)
+
+    def forward(self, coarse, content, cur_xy):
+        # coarse:[B,A,M,T,2] content:[B,A,M,D] cur_xy:[B,A,2]
+        B, A, M, T, _ = coarse.shape
+        ctx = self.ctx_proj(content).reshape(B * A * M, 1, -1)     # [N,1,H]
+        y = coarse
+        for _ in range(self.n_stages):
+            yc = (y - cur_xy[:, :, None, None, :]).reshape(B * A * M, T, 2)   # center
+            seq = self.pos_emb(yc) + ctx                          # broadcast ctx over T
+            h, _ = self.gru(seq)                                  # [N,T,H] autoregressive
+            d = self.delta(h).view(B, A, M, T, 2)
+            y = y + d
+        return y
+
+
+class CoarseToFineNBA(nn.Module):
+    """Standalone coarse-to-fine predictor (NBA, map dropped)."""
+
+    def __init__(self, n_agents=11, past_dim=6, future_len=20, modalities=20,
+                 n_stages=2, dim=D_MODEL, heads=N_HEADS, enc_layers=6, hidden=256,
+                 n_types=2, ball_idx=10):
+        super().__init__()
+        self.history_encoder = AgentHistoryEncoder(past_dim, dim, n_types)
+        self.fusion_encoder = FusionEncoder(dim, heads, enc_layers)
+        self.coarse_decoder = CoarseDecoder(modalities, n_agents, future_len, dim)
+        self.fine_refiner = FineRefiner(dim, hidden, n_stages)
+        types = torch.zeros(n_agents, dtype=torch.long)
+        if 0 <= ball_idx < n_agents and n_types > 1:
+            types[ball_idx] = 1
+        self.register_buffer('agent_types', types)
+
+    def forward(self, feats, cur_xy):
+        enc = self.history_encoder(feats, self.agent_types)       # [B,A,D]
+        enc = self.fusion_encoder(enc, None)                      # [B,A,D] social
+        content, coarse, score = self.coarse_decoder(enc, cur_xy, None)
+        fine = self.fine_refiner(coarse, content, cur_xy)         # [B,A,M,T,2]
+        return {'coarse': coarse, 'fine': fine, 'scores': score}
+
+
+# ============================================================================
+#  LOSS  (coarse WTA + fine WTA + mode CE; marginal per agent)
+# ============================================================================
+def wta_l2(pred, gt, metric_idx):
+    """Winner-takes-all L2 (variety loss). pred:[B,A,M,T,2] gt:[B,A,T,2].
+    Mode selection uses endpoint-emphasised distance; returns best mode's mean
+    displacement (reconstruction loss) and the winning mode index [B,A]."""
+    B, A, M, T, _ = pred.shape
+    d = torch.norm(pred - gt[:, :, None], dim=-1)                 # [B,A,M,T]
+    sel = d.mean(-1) + d[..., metric_idx].sum(-1)                 # [B,A,M]
+    best = sel.argmin(-1)                                         # [B,A]
+    gi = best[..., None, None, None].expand(B, A, 1, T, 2)
+    best_d = torch.norm(torch.gather(pred, 2, gi).squeeze(2) - gt, dim=-1)  # [B,A,T]
+    reg = best_d.mean(-1) + best_d[..., metric_idx].sum(-1)       # [B,A] endpoint emphasis
+    return reg.mean(), best
+
+
+def c2f_loss(out, gt_center, metric_idx):
+    coarse_reg, _ = wta_l2(out['coarse'], gt_center, metric_idx)
+    fine_reg, best = wta_l2(out['fine'], gt_center, metric_idx)
+    B, A, M = out['scores'].shape
+    cls = F.cross_entropy(out['scores'].reshape(B * A, M), best.reshape(B * A),
+                          label_smoothing=0.2)
+    return coarse_reg + fine_reg + 2.0 * cls
+
+
+# ============================================================================
+#  DATA  (identical to the GameFormer standalone)
+# ============================================================================
+def build_loaders(args):
+    cfg = _types.SimpleNamespace(
+        traj_mean=[14, 7.5], data_norm='min_max',
+        past_frames=args.past_len, future_frames=args.future_len, agents=args.agents)
+    train_set = NBADatasetMinMax(
+        obs_len=args.past_len, pred_len=args.future_len, training=True,
+        num_scenes=args.n_train, cfg=cfg, data_dir=args.data_dir,
+        rotate=False, data_norm='min_max')
+    test_set = NBADatasetMinMax(
+        obs_len=args.past_len, pred_len=args.future_len, training=False,
+        test_scenes=args.n_test, cfg=cfg, data_dir=args.data_dir,
+        rotate=False, data_norm='min_max')
+    train_loader = DataLoader(
+        train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
+        collate_fn=seq_collate_nba, pin_memory=True, drop_last=True)
+    test_loader = DataLoader(
+        test_set, batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
+        collate_fn=seq_collate_nba, pin_memory=True)
+    return train_loader, test_loader
+
+
+def unpack(data, device):
+    feats = data['past_traj'].to(device)                       # [B,A,T,6] normalized
+    past_orig = data['past_traj_original_scale'].to(device)    # [B,A,T,6] metric
+    gt_disp = data['fut_traj_original_scale'].to(device)       # [B,A,Tf,2] displacement
+    cur_xy = past_orig[:, :, -1, 0:2]
+    gt_center = gt_disp + cur_xy[:, :, None, :]
+    return feats, cur_xy, gt_disp, gt_center
+
+
+# ============================================================================
+#  EVAL  (marginal min-of-K=20 at 4.0 s; identical to eval_perscene_moflow.py)
+# ============================================================================
+@torch.no_grad()
+def evaluate(model, loader, device, end):
+    model.eval()
+    ade_sum = fde_sum = 0.0
+    n = 0
+    for data in loader:
+        feats, cur_xy, gt_disp, _ = unpack(data, device)
+        out = model(feats, cur_xy)
+        pred = out['fine'][..., :2]                               # [B,A,M,T,2] fine stage
+        pred_disp = pred - cur_xy[:, :, None, None, :]
+        d = torch.norm(pred_disp - gt_disp[:, :, None], dim=-1)   # [B,A,M,T]
+        ade = d[..., :end].mean(-1).min(dim=-1).values            # [B,A]
+        fde = d[..., end - 1].min(dim=-1).values                  # [B,A]
+        ade_sum += ade.sum().item()
+        fde_sum += fde.sum().item()
+        n += ade.numel()
+    return ade_sum / max(n, 1), fde_sum / max(n, 1)
+
+
+# ============================================================================
+#  TRAIN
+# ============================================================================
+def parse_args():
+    p = argparse.ArgumentParser('Coarse-to-Fine standalone predictor for NBA')
+    p.add_argument('--data_dir', default='/mnt/jaewoo4tb/srtp/MoFlow/data/nba', type=str)
+    p.add_argument('--gpu', default='0', type=str)
+    p.add_argument('--exp', default='c2f_nba', type=str)
+    p.add_argument('--epochs', default=50, type=int)
+    p.add_argument('--batch_size', default=128, type=int)
+    p.add_argument('--test_batch', default=500, type=int)
+    p.add_argument('--workers', default=4, type=int)
+    p.add_argument('--seed', default=3407, type=int)
+
+    p.add_argument('--modalities', default=20, type=int, help='K modes (NBA min-ADE_20)')
+    p.add_argument('--stages', default=2, type=int, help='coarse->fine refinement stages')
+    p.add_argument('--encoder_layers', default=6, type=int)
+    p.add_argument('--hidden', default=256, type=int, help='fine-refiner GRU hidden')
+    p.add_argument('--agents', default=11, type=int)
+    p.add_argument('--past_len', default=10, type=int)
+    p.add_argument('--future_len', default=20, type=int)
+    p.add_argument('--ball_idx', default=10, type=int, help='-1 to disable type emb')
+
+    p.add_argument('--lr', default=1e-4, type=float)
+    p.add_argument('--weight_decay', default=1e-4, type=float)
+    p.add_argument('--grad_clip', default=5.0, type=float)
+    p.add_argument('--n_train', default=32500, type=int)
+    p.add_argument('--n_test', default=12500, type=int)
+    p.add_argument('--eval_every', default=2, type=int)
+    p.add_argument('--save_dir', default=os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), 'c2f_nba_ckpt'), type=str)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    train_loader, test_loader = build_loaders(args)
+
+    n_types = 2 if (0 <= args.ball_idx < args.agents) else 1
+    model = CoarseToFineNBA(
+        n_agents=args.agents, past_dim=6, future_len=args.future_len,
+        modalities=args.modalities, n_stages=args.stages, dim=D_MODEL, heads=N_HEADS,
+        enc_layers=args.encoder_layers, hidden=args.hidden, n_types=n_types,
+        ball_idx=args.ball_idx).to(device)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f'[C2F-NBA] model params: {n_params/1e6:.2f} M | stages={args.stages} '
+          f'K={args.modalities} enc_layers={args.encoder_layers} device={device}', flush=True)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    milestones = sorted(set(int(args.epochs * f) for f in (0.5, 0.65, 0.8, 0.9)))
+    milestones = [m for m in milestones if 0 < m < args.epochs]
+    sched = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=milestones, gamma=0.5)
+
+    end = args.future_len
+    metric_idx = sorted(set([max(0, args.future_len // 2 - 1), args.future_len - 1]))
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    best_ade = float('inf')
+
+    for epoch in range(args.epochs):
+        model.train()
+        t0 = time.time()
+        running = 0.0
+        nb = 0
+        for data in train_loader:
+            feats, cur_xy, _, gt_center = unpack(data, device)
+            out = model(feats, cur_xy)
+            loss = c2f_loss(out, gt_center, metric_idx)
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            opt.step()
+            running += float(loss.item())
+            nb += 1
+        sched.step()
+        print(f'[C2F-NBA] epoch {epoch+1}/{args.epochs} loss={running/max(nb,1):.4f} '
+              f'lr={opt.param_groups[0]["lr"]:.2e} ({time.time()-t0:.1f}s)', flush=True)
+
+        if (epoch + 1) % args.eval_every == 0 or epoch == args.epochs - 1:
+            ade, fde = evaluate(model, test_loader, device, end)
+            print(f'[C2F-NBA] epoch {epoch+1} ADE4={ade:.4f} FDE4={fde:.4f}', flush=True)
+            if ade < best_ade:
+                best_ade = ade
+                torch.save({'model': model.state_dict(), 'epoch': epoch + 1,
+                            'ade4': ade, 'fde4': fde, 'args': vars(args)},
+                           os.path.join(args.save_dir, f'{args.exp}_best.pt'))
+                print(f'[C2F-NBA] saved best (ADE4={ade:.4f}) -> '
+                      f'{os.path.join(args.save_dir, args.exp + "_best.pt")}', flush=True)
+
+    print(f'[C2F-NBA] done. best ADE4={best_ade:.4f}', flush=True)
+
+
+if __name__ == '__main__':
+    main()

standalone/gameformer_nba_standalone.py

@@ -0,0 +1,614 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+gameformer_nba_standalone.py
+================================================================================
+FAITHFUL standalone re-implementation of **GameFormer** (Liu et al., ICCV 2023,
+"GameFormer: Game-theoretic Modeling and Learning of Transformer-based
+Interactive Prediction", github.com/MCZhi/GameFormer) as a trajectory predictor
+for the NBA basketball dataset.
+
+It is a STANDALONE predictor: it consumes ONLY agent HISTORY (past trajectories)
+and predicts K multi-modal futures via the level-k game-theoretic decoder. It
+does NOT consume any external denoiser / diffusion prediction.
+
+--------------------------------------------------------------------------------
+WHAT WAS REPLICATED (evidence = MCZhi/GameFormer @ main)
+--------------------------------------------------------------------------------
+Reference files studied verbatim:
+  * model/GameFormer.py        (Encoder, Decoder, GameFormer)
+  * model/modules.py           (AgentEncoder, FutureEncoder, GMMPredictor,
+                                SelfTransformer, CrossTransformer,
+                                InitialDecoder, InteractionDecoder)
+  * utils/inter_pred_utils.py  (level_k_loss, gmm_loss)
+  * interaction_prediction/train.py (AdamW 1e-4, MultiStepLR, grad-clip 5,
+                                     level=3, modalities=6, encoder_layers=6)
+
+ENCODER  (GameFormer.Encoder + modules.AgentEncoder)
+  * Per-agent temporal encoder = 2-layer LSTM, hidden 256  (nn.LSTM(8,256,2)).
+    We adapt the input feature dim 8 -> 6 (NBA has [abs xy, rel xy, vel xy], no
+    heading/size channels). A learned agent-type embedding (player vs. ball) is
+    added to the LSTM last-step output, exactly as GameFormer adds `type_emb`.
+  * Interaction / self-attention over agents = the GameFormer `fusion_encoder`,
+    a nn.TransformerEncoder of TransformerEncoderLayer(d_model=256, nhead=8,
+    dim_feedforward=1024, activation=gelu, dropout=0.1, batch_first=True),
+    `encoder_layers` deep. In GameFormer this fuses agent tokens WITH per-agent
+    map (lane/crosswalk) tokens; NBA has NO map, so the map encoders
+    (LaneEncoder / CrosswalkEncoder / segment_map) are DROPPED and the fusion is
+    pure agent<->agent self-attention over the 11 tokens (identical for every
+    target agent when there is no per-agent map, so it is computed once).
+
+LEVEL-k GAME-THEORETIC DECODER  (GameFormer.Decoder)
+  * Level 0 = InitialDecoder: for each agent, a set of `modalities` (=K) learned
+    modal query embeddings + a per-agent query embedding are added to the agent's
+    own context token, then cross-attended (CrossTransformer) to the full agent
+    context. A GMMPredictor head outputs K trajectories of (mu_x,mu_y,
+    log_sig_x,log_sig_y) per step + K mode scores.  predictions[...,:2] += the
+    agent's current (last observed) position.
+  * Levels 1..L = InteractionDecoder (one per level, sharing ONE FutureEncoder):
+      - FutureEncoder encodes EVERY agent's K level-(k-1) modal futures into a
+        vector (max-pool over time of an MLP on [x,y,heading,vx,vy]); GameFormer
+        also appends box size (w,l,h) -> we drop those (no size in NBA), 8->5.
+      - The per-mode future encodings are score-softmax-weighted and mean-pooled
+        to one vector per agent, then passed through a SelfTransformer to produce
+        the interaction encoding over all agents (this is the game-theoretic
+        "condition on all other agents' level-(k-1) predictions" step).
+      - The interaction tokens are concatenated with the context tokens; each
+        agent's OWN future token is masked out; the previous-level content plus
+        the agent's own future encoding form the query, cross-attended to that
+        combined context, and a fresh GMMPredictor emits the level-k K futures.
+  * Output dict: level_0_interactions/scores ... level_L_interactions/scores,
+    exactly as GameFormer.Decoder.forward.
+
+LOSS  (utils.inter_pred_utils.level_k_loss + gmm_loss, gmm=True)
+  * Summed over ALL levels k=0..L (level-wise imitation, GameFormer's key term).
+  * Per level, GMM winner-takes-all: best mode = argmin over K of
+    (mean_t ||traj-gt|| + sum over checkpoint steps), then the Gaussian NLL
+    reg = (log_sig_x+log_sig_y) + 0.5*(dx^2/sx^2 + dy^2/sy^2) with log_sig
+    clipped to [0,5], reduced as mean_t + sum over checkpoints; plus mode
+    classification F.cross_entropy(scores, best_mode, label_smoothing=0.2)
+    weighted x2.  We MARGINALIZE per agent (GameFormer's interaction_prediction
+    task predicts 2 agents JOINTLY via distance.mean(1)/scores.sum(1); the NBA
+    Table-1 metric is MARGINAL min-ADE_20, so the best mode / scores are taken
+    per agent instead of jointly -- the only loss adaptation).
+
+--------------------------------------------------------------------------------
+NBA I/O + METRIC (matched to MoFlow / Table 1)
+--------------------------------------------------------------------------------
+  * Data: reuses MoFlow's exact loader data/dataloader_nba.py::NBADatasetMinMax
+    (the graph loader dataloader_nba_graph.py merely re-exports it). Same .npy
+    files (data/nba/original/nba_{train,test}.npy, 32500/12500 scenes,
+    30 frames x 11 agents x 2D), same split (trajs[:n_train] / trajs[:n_test]),
+    same coordinate scaling traj_scale_total=94/28 and traj_mean=[14,7.5].
+  * Past = 10 frames, Future = 20 frames = 4.0 s @ 5 Hz, 11 agents, 2D.
+  * Encoder input = data['past_traj'] (6-dim min-max normalized to [-1,1]).
+  * Coordinate frame for the GMM/interaction = data['past_traj_original_scale']
+    (centered-abs = pos/scale - traj_mean); cur_xy = its last past step.
+    Ground-truth future displacement = data['fut_traj_original_scale']
+    (= future - last_observed, in /scale units) -- the SAME tensor the MoFlow
+    evaluator (eval_perscene_moflow.py) compares against.
+  * Metric (identical to eval_perscene_moflow.py, verified vs Table 1):
+      d = ||pred_disp - gt_disp||                       # [B,A,K,T]
+      ADE4 = d[..., :20].mean(over T).min(over K), mean over 11 agents & scenes
+      FDE4 = d[..., 19]        .min(over K),          mean over 11 agents & scenes
+    i.e. marginal min-of-K=20 at the 4.0 s horizon (frame 20). Table-1 baselines
+    are ~0.70-1.04 ADE in this space.
+
+Prints:  [GF-NBA] epoch N ADE4=.. FDE4=..
+
+DO NOT run on GPU here (GPUs busy) -- this file is only `python -m py_compile`d.
+"""
+
+import os
+import sys
+import argparse
+import time
+
+# ------------------------------------------------------------------ GPU FIRST
+# CUDA context is fixed at torch import; select the device BEFORE importing torch
+def _early_gpu():
+    for i, a in enumerate(sys.argv):
+        if a == '--gpu' and i + 1 < len(sys.argv):
+            return sys.argv[i + 1]
+        if a.startswith('--gpu='):
+            return a.split('=', 1)[1]
+    return None
+
+
+_g = _early_gpu()
+if _g is not None:
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(_g)
+os.environ.setdefault('MPLBACKEND', 'Agg')  # dataloader imports matplotlib
+
+# ------------------------------------------------ reuse MoFlow's NBA pipeline
+MOFLOW_ROOT = '/mnt/jaewoo4tb/srtp/MoFlow'
+if MOFLOW_ROOT not in sys.path:
+    sys.path.insert(0, MOFLOW_ROOT)
+
+import math
+import types as _types
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+# so that `from utils.normalization import ...` inside the loader resolves
+try:
+    os.chdir(MOFLOW_ROOT)
+except Exception:
+    pass
+from data.dataloader_nba import NBADatasetMinMax, seq_collate_nba
+
+
+# ============================================================================
+#  MODEL  (faithful GameFormer, NBA-adapted: no map)
+# ============================================================================
+D_MODEL = 256
+N_HEADS = 8
+DROPOUT = 0.1
+
+
+# ---------------------------------------------------------------- primitives
+class SelfTransformer(nn.Module):
+    """modules.py::SelfTransformer (verbatim)."""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, dropout=DROPOUT):
+        super().__init__()
+        self.self_attention = nn.MultiheadAttention(dim, heads, dropout, batch_first=True)
+        self.norm_1 = nn.LayerNorm(dim)
+        self.norm_2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout))
+
+    def forward(self, inputs, mask=None):
+        attn, _ = self.self_attention(inputs, inputs, inputs, key_padding_mask=mask)
+        attn = self.norm_1(attn + inputs)
+        return self.norm_2(self.ffn(attn) + attn)
+
+
+class CrossTransformer(nn.Module):
+    """modules.py::CrossTransformer (verbatim)."""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, dropout=DROPOUT):
+        super().__init__()
+        self.cross_attention = nn.MultiheadAttention(dim, heads, dropout, batch_first=True)
+        self.norm_1 = nn.LayerNorm(dim)
+        self.norm_2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout))
+
+    def forward(self, query, key, value, mask=None):
+        attn, _ = self.cross_attention(query, key, value, key_padding_mask=mask)
+        attn = self.norm_1(attn)
+        return self.norm_2(self.ffn(attn) + attn)
+
+
+class GMMPredictor(nn.Module):
+    """modules.py::GMMPredictor (verbatim). Out: (mu_x,mu_y,log_sig_x,log_sig_y)."""
+
+    def __init__(self, future_len, dim=D_MODEL):
+        super().__init__()
+        self._future_len = future_len
+        self.gaussian = nn.Sequential(
+            nn.Linear(dim, 512), nn.ELU(), nn.Dropout(0.1),
+            nn.Linear(512, future_len * 4))
+        self.score = nn.Sequential(
+            nn.Linear(dim, 64), nn.ELU(), nn.Dropout(0.1), nn.Linear(64, 1))
+
+    def forward(self, x):
+        # x: [B, M, D]
+        B, M, _ = x.shape
+        res = self.gaussian(x).view(B, M, self._future_len, 4)
+        score = self.score(x).squeeze(-1)  # [B, M]
+        return res, score
+
+
+# ------------------------------------------------------------------ encoder
+class AgentHistoryEncoder(nn.Module):
+    """modules.py::AgentEncoder -- 2-layer LSTM(->256) over agent history + type
+    embedding.  GameFormer uses nn.LSTM(8,256,2); NBA feature dim is 6."""
+
+    def __init__(self, in_dim=6, dim=D_MODEL, n_types=2):
+        super().__init__()
+        self.motion = nn.LSTM(in_dim, dim, 2, batch_first=True)
+        self.type_emb = nn.Embedding(n_types, dim)
+
+    def forward(self, hist, types):
+        # hist: [B, A, T, in_dim]  ;  types: [A] long
+        B, A, T, C = hist.shape
+        traj, _ = self.motion(hist.reshape(B * A, T, C))
+        out = traj[:, -1].reshape(B, A, -1)          # [B, A, D]
+        out = out + self.type_emb(types)[None]       # + per-agent type emb
+        return out
+
+
+class FusionEncoder(nn.Module):
+    """GameFormer.Encoder.fusion_encoder -- self-attention over agent tokens.
+    (No map tokens for NBA.)"""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, layers=6, dropout=DROPOUT):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(
+            d_model=dim, nhead=heads, dim_feedforward=dim * 4,
+            activation=F.gelu, dropout=dropout, batch_first=True)
+        self.encoder = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
+
+    def forward(self, tokens, mask=None):
+        return self.encoder(tokens, src_key_padding_mask=mask)
+
+
+# ------------------------------------------------------------- future encoder
+class FutureEncoder(nn.Module):
+    """modules.py::FutureEncoder -- encodes each agent's K modal futures into a
+    per-mode vector via max-pool over an MLP on [x,y,heading,vx,vy].  GameFormer
+    appends box size (w,l,h) -> dropped for NBA (8 -> 5 input dims)."""
+
+    def __init__(self, dim=D_MODEL, n_types=2, dt=0.2):
+        super().__init__()
+        self.dt = dt
+        self.mlp = nn.Sequential(nn.Linear(5, 64), nn.ReLU(), nn.Linear(64, dim))
+        self.type_emb = nn.Embedding(n_types, dim)
+
+    def state_process(self, trajs, cur_xy):
+        # trajs: [B, A, M, T, 2] ; cur_xy: [B, A, 2]
+        B, A, M, T, _ = trajs.shape
+        cur = cur_xy[:, :, None, None, :].expand(B, A, M, 1, 2)
+        xy = torch.cat([cur, trajs], dim=-2)                     # [B,A,M,T+1,2]
+        dxy = torch.diff(xy, dim=-2)                             # [B,A,M,T,2]
+        v = dxy / self.dt
+        theta = torch.atan2(dxy[..., 1], dxy[..., 0].clamp(min=1e-3)).unsqueeze(-1)
+        return torch.cat([trajs, theta, v], dim=-1)             # [B,A,M,T,5]
+
+    def forward(self, trajs, cur_xy, types):
+        feat = self.state_process(trajs, cur_xy)
+        h = self.mlp(feat.detach())                             # GameFormer detaches
+        out = torch.max(h, dim=-2).values                       # [B,A,M,D]
+        out = out + self.type_emb(types)[None, :, None, :]      # + type emb
+        return out
+
+
+# ----------------------------------------------------------- level-0 decoder
+class InitialDecoder(nn.Module):
+    """modules.py::InitialDecoder -- vectorized over agents."""
+
+    def __init__(self, modalities, n_agents, future_len, dim=D_MODEL):
+        super().__init__()
+        self.M = modalities
+        self.A = n_agents
+        self.multi_modal_query_embedding = nn.Embedding(modalities, dim)
+        self.agent_query_embedding = nn.Embedding(n_agents, dim)
+        self.query_encoder = CrossTransformer(dim)
+        self.predictor = GMMPredictor(future_len, dim)
+        self.register_buffer('modal', torch.arange(modalities).long())
+        self.register_buffer('agent', torch.arange(n_agents).long())
+
+    def forward(self, encoding, cur_xy, mask=None):
+        # encoding: [B, A, D]  (full agent context, shared as keys)
+        B, A, D = encoding.shape
+        M = self.M
+        multi_modal = self.multi_modal_query_embedding(self.modal)   # [M, D]
+        agent = self.agent_query_embedding(self.agent)               # [A, D]
+        mm_agent_query = multi_modal[None, :, :] + agent[:, None, :]  # [A, M, D]
+        query = encoding[:, :, None, :] + mm_agent_query[None]       # [B, A, M, D]
+
+        q = query.reshape(B * A, M, D)
+        kv = encoding[:, None, :, :].expand(B, A, A, D).reshape(B * A, A, D)
+        km = None
+        if mask is not None:
+            km = mask[:, None, :].expand(B, A, A).reshape(B * A, A)
+        content = self.query_encoder(q, kv, kv, km)                 # [B*A, M, D]
+        preds, scores = self.predictor(content)                     # [B*A,M,T,4],[B*A,M]
+
+        content = content.reshape(B, A, M, D)
+        T = preds.shape[-2]
+        preds = preds.reshape(B, A, M, T, 4)
+        scores = scores.reshape(B, A, M)
+        preds = preds.clone()
+        preds[..., :2] = preds[..., :2] + cur_xy[:, :, None, None, :]
+        return content, preds, scores
+
+
+# ------------------------------------------------------- level-k decoder
+class InteractionDecoder(nn.Module):
+    """modules.py::InteractionDecoder -- one game-theoretic reasoning level,
+    vectorized over agents.  Conditions each agent on ALL agents' level-(k-1)
+    predictions (interaction self-attention), masking the agent's own future."""
+
+    def __init__(self, future_encoder, future_len, n_agents, dim=D_MODEL):
+        super().__init__()
+        self.A = n_agents
+        self.future_encoder = future_encoder            # SHARED across levels
+        self.interaction_encoder = SelfTransformer(dim)
+        self.query_encoder = CrossTransformer(dim)
+        self.decoder = GMMPredictor(future_len, dim)
+
+    def forward(self, cur_xy, last_traj, last_scores, last_content, encoding, types, mask=None):
+        # cur_xy:[B,A,2] last_traj:[B,A,M,T,4] last_scores:[B,A,M]
+        # last_content:[B,A,M,D] encoding:[B,A,D]
+        B, A, M, T, _ = last_traj.shape
+        D = encoding.shape[-1]
+
+        # encode all agents' level-(k-1) modal futures
+        multi_futures = self.future_encoder(last_traj[..., :2], cur_xy, types)  # [B,A,M,D]
+        futures = (multi_futures * last_scores.softmax(-1).unsqueeze(-1)).mean(dim=2)  # [B,A,D]
+
+        # interaction over agents (self-attention)
+        interaction = self.interaction_encoder(futures, None)      # [B,A,D]
+
+        # combined context = [interaction tokens ; context tokens]
+        ctx = torch.cat([interaction, encoding], dim=1)            # [B, 2A, D]
+        # per-agent key mask: mask each agent's OWN future token in the 1st block
+        km = torch.zeros(B, A, 2 * A, dtype=torch.bool, device=encoding.device)
+        idx = torch.arange(A, device=encoding.device)
+        km[:, idx, idx] = True
+        if mask is not None:
+            # second block inherits context validity mask (all-valid for NBA)
+            km[:, :, A:] = km[:, :, A:] | mask[:, None, :]
+            km[:, :, :A] = km[:, :, :A] | mask[:, None, :]
+        km = km.reshape(B * A, 2 * A)
+        kv = ctx[:, None, :, :].expand(B, A, 2 * A, D).reshape(B * A, 2 * A, D)
+
+        query = (last_content + multi_futures).reshape(B * A, M, D)  # prev content + own future
+        content = self.query_encoder(query, kv, kv, km)             # [B*A, M, D]
+        traj, scores = self.decoder(content)                        # [B*A,M,T,4],[B*A,M]
+
+        content = content.reshape(B, A, M, D)
+        traj = traj.reshape(B, A, M, T, 4)
+        scores = scores.reshape(B, A, M)
+        traj = traj.clone()
+        traj[..., :2] = traj[..., :2] + cur_xy[:, :, None, None, :]
+        return content, traj, scores
+
+
+# -------------------------------------------------------------- full decoder
+class GameFormerDecoder(nn.Module):
+    """GameFormer.Decoder -- level 0 + L interaction levels (shared FutureEncoder)."""
+
+    def __init__(self, modalities, n_agents, future_len, levels, dim=D_MODEL,
+                 n_types=2, dt=0.2):
+        super().__init__()
+        self._levels = levels
+        self.future_encoder = FutureEncoder(dim, n_types, dt)
+        self.initial_stage = InitialDecoder(modalities, n_agents, future_len, dim)
+        self.interaction_stage = nn.ModuleList(
+            [InteractionDecoder(self.future_encoder, future_len, n_agents, dim)
+             for _ in range(levels)])
+
+    def forward(self, encoding, cur_xy, types, mask=None):
+        out = {}
+        content, traj, scores = self.initial_stage(encoding, cur_xy, mask)
+        out['level_0_interactions'] = traj
+        out['level_0_scores'] = scores
+        for k in range(1, self._levels + 1):
+            content, traj, scores = self.interaction_stage[k - 1](
+                cur_xy, traj, scores, content, encoding, types, mask)
+            out[f'level_{k}_interactions'] = traj
+            out[f'level_{k}_scores'] = scores
+        return out
+
+
+# ---------------------------------------------------------------- top model
+class GameFormerNBA(nn.Module):
+    """GameFormer.GameFormer (NBA, map dropped)."""
+
+    def __init__(self, n_agents=11, past_dim=6, future_len=20, modalities=20,
+                 levels=2, dim=D_MODEL, heads=N_HEADS, enc_layers=6,
+                 n_types=2, dt=0.2, ball_idx=10):
+        super().__init__()
+        self.levels = levels
+        self.history_encoder = AgentHistoryEncoder(past_dim, dim, n_types)
+        self.fusion_encoder = FusionEncoder(dim, heads, enc_layers)
+        self.decoder = GameFormerDecoder(modalities, n_agents, future_len, levels,
+                                         dim, n_types, dt)
+        types = torch.zeros(n_agents, dtype=torch.long)
+        if 0 <= ball_idx < n_agents and n_types > 1:
+            types[ball_idx] = 1
+        self.register_buffer('agent_types', types)
+
+    def forward(self, feats, cur_xy):
+        # feats: [B,A,T,past_dim] (encoder input) ; cur_xy: [B,A,2] (metric frame)
+        enc = self.history_encoder(feats, self.agent_types)   # [B,A,D]
+        enc = self.fusion_encoder(enc, None)                  # [B,A,D]
+        return self.decoder(enc, cur_xy, self.agent_types, None)
+
+
+# ============================================================================
+#  LOSS  (level_k_loss + gmm_loss, gmm=True; MARGINAL per agent for NBA)
+# ============================================================================
+def gmm_loss_marginal(traj, conv, scores, gt, metric_idx):
+    """Per-agent GMM winner-takes-all (marginal version of gmm_loss).
+      traj:  [B,A,M,T,2]   conv: [B,A,M,T,2]   scores: [B,A,M]   gt: [B,A,T,2]
+    """
+    B, A, M, T, _ = traj.shape
+    dist = torch.norm(traj - gt[:, :, None], dim=-1)                  # [B,A,M,T]
+    ndist = dist.mean(-1) + dist[..., metric_idx].sum(-1)            # [B,A,M]
+    best = ndist.argmin(-1)                                          # [B,A]
+
+    gi = best[..., None, None, None].expand(B, A, 1, T, 2)
+    best_traj = torch.gather(traj, 2, gi).squeeze(2)                 # [B,A,T,2]
+    best_conv = torch.gather(conv, 2, gi).squeeze(2)                 # [B,A,T,2]
+
+    dx = best_traj[..., 0] - gt[..., 0]
+    dy = best_traj[..., 1] - gt[..., 1]
+    log_std_x = torch.clip(best_conv[..., 0], 0, 5)
+    log_std_y = torch.clip(best_conv[..., 1], 0, 5)
+    std_x, std_y = torch.exp(log_std_x), torch.exp(log_std_y)
+
+    reg = (log_std_x + log_std_y) + 0.5 * ((dx ** 2) / (std_x ** 2) + (dy ** 2) / (std_y ** 2))
+    reg = reg.mean(-1) + reg[..., metric_idx].sum(-1)               # [B,A]
+
+    prob_loss = F.cross_entropy(scores.reshape(B * A, M), best.reshape(B * A),
+                                label_smoothing=0.2)
+    return reg.mean() + 2.0 * prob_loss, best
+
+
+def level_k_loss_marginal(outputs, gt_center, levels, metric_idx):
+    """Sum GMM loss across all levels 0..L (GameFormer's level-wise imitation)."""
+    total = 0.0
+    for k in range(levels + 1):
+        pred = outputs[f'level_{k}_interactions']       # [B,A,M,T,4]
+        scores = outputs[f'level_{k}_scores']           # [B,A,M]
+        traj, conv = pred[..., :2], pred[..., 2:]
+        l, _ = gmm_loss_marginal(traj, conv, scores, gt_center, metric_idx)
+        total = total + l
+    return total
+
+
+# ============================================================================
+#  DATA
+# ============================================================================
+def build_loaders(args):
+    cfg = _types.SimpleNamespace(
+        traj_mean=[14, 7.5], data_norm='min_max',
+        past_frames=args.past_len, future_frames=args.future_len, agents=args.agents)
+    # Train first so cfg min/max scalars are populated for the test set.
+    train_set = NBADatasetMinMax(
+        obs_len=args.past_len, pred_len=args.future_len, training=True,
+        num_scenes=args.n_train, cfg=cfg, data_dir=args.data_dir,
+        rotate=False, data_norm='min_max')
+    test_set = NBADatasetMinMax(
+        obs_len=args.past_len, pred_len=args.future_len, training=False,
+        test_scenes=args.n_test, cfg=cfg, data_dir=args.data_dir,
+        rotate=False, data_norm='min_max')
+    train_loader = DataLoader(
+        train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
+        collate_fn=seq_collate_nba, pin_memory=True, drop_last=True)
+    test_loader = DataLoader(
+        test_set, batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
+        collate_fn=seq_collate_nba, pin_memory=True)
+    return train_loader, test_loader
+
+
+def unpack(data, device):
+    feats = data['past_traj'].to(device)                       # [B,A,T,6] normalized
+    past_orig = data['past_traj_original_scale'].to(device)    # [B,A,T,6] metric
+    gt_disp = data['fut_traj_original_scale'].to(device)       # [B,A,Tf,2] displacement
+    cur_xy = past_orig[:, :, -1, 0:2]                          # centered-abs last pos
+    gt_center = gt_disp + cur_xy[:, :, None, :]                # centered-abs future
+    return feats, cur_xy, gt_disp, gt_center
+
+
+# ============================================================================
+#  EVAL  (marginal min-of-K=20 at 4.0 s; identical to eval_perscene_moflow.py)
+# ============================================================================
+@torch.no_grad()
+def evaluate(model, loader, device, levels, end):
+    model.eval()
+    ade_sum = fde_sum = 0.0
+    n = 0
+    for data in loader:
+        feats, cur_xy, gt_disp, _ = unpack(data, device)
+        out = model(feats, cur_xy)
+        pred = out[f'level_{levels}_interactions'][..., :2]        # [B,A,M,T,2]
+        pred_disp = pred - cur_xy[:, :, None, None, :]             # -> displacement
+        d = torch.norm(pred_disp - gt_disp[:, :, None], dim=-1)    # [B,A,M,T]
+        ade = d[..., :end].mean(-1).min(dim=-1).values            # [B,A] best-of-K
+        fde = d[..., end - 1].min(dim=-1).values                  # [B,A]
+        ade_sum += ade.sum().item()
+        fde_sum += fde.sum().item()
+        n += ade.numel()
+    return ade_sum / max(n, 1), fde_sum / max(n, 1)
+
+
+# ============================================================================
+#  TRAIN
+# ============================================================================
+def parse_args():
+    p = argparse.ArgumentParser('GameFormer standalone predictor for NBA')
+    p.add_argument('--data_dir', default='/mnt/jaewoo4tb/srtp/MoFlow/data/nba', type=str)
+    p.add_argument('--gpu', default='0', type=str)
+    p.add_argument('--exp', default='gf_nba', type=str)
+    p.add_argument('--epochs', default=50, type=int)
+    p.add_argument('--batch_size', default=128, type=int)
+    p.add_argument('--test_batch', default=500, type=int)
+    p.add_argument('--workers', default=4, type=int)
+    p.add_argument('--seed', default=3407, type=int)          # GameFormer default
+
+    # --- architecture (faithful GameFormer knobs) ---
+    p.add_argument('--modalities', default=20, type=int, help='K modes (NBA min-ADE_20)')
+    p.add_argument('--level', default=2, type=int, help='game-theoretic levels (paper: 3)')
+    p.add_argument('--encoder_layers', default=6, type=int)
+    p.add_argument('--agents', default=11, type=int)
+    p.add_argument('--past_len', default=10, type=int)
+    p.add_argument('--future_len', default=20, type=int)      # 4.0 s @ 5 Hz
+    p.add_argument('--ball_idx', default=10, type=int, help='-1 to disable type emb')
+
+    # --- optimization (GameFormer: AdamW 1e-4, MultiStepLR x0.5, clip 5) ---
+    p.add_argument('--lr', default=1e-4, type=float)
+    p.add_argument('--weight_decay', default=1e-4, type=float)
+    p.add_argument('--grad_clip', default=5.0, type=float)
+    p.add_argument('--n_train', default=32500, type=int)
+    p.add_argument('--n_test', default=12500, type=int)
+    p.add_argument('--eval_every', default=2, type=int)
+    p.add_argument('--save_dir', default=os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), 'gf_nba_ckpt'), type=str)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    train_loader, test_loader = build_loaders(args)
+
+    n_types = 2 if (0 <= args.ball_idx < args.agents) else 1
+    model = GameFormerNBA(
+        n_agents=args.agents, past_dim=6, future_len=args.future_len,
+        modalities=args.modalities, levels=args.level, dim=D_MODEL, heads=N_HEADS,
+        enc_layers=args.encoder_layers, n_types=n_types, dt=0.2,
+        ball_idx=args.ball_idx).to(device)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f'[GF-NBA] model params: {n_params/1e6:.2f} M | levels={args.level} '
+          f'K={args.modalities} enc_layers={args.encoder_layers} device={device}')
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    milestones = sorted(set(int(args.epochs * f) for f in (0.5, 0.65, 0.8, 0.9)))
+    milestones = [m for m in milestones if 0 < m < args.epochs]
+    sched = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=milestones, gamma=0.5)
+
+    end = args.future_len                                   # 4.0 s = frame 20
+    metric_idx = sorted(set([max(0, args.future_len // 2 - 1), args.future_len - 1]))
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    best_ade = float('inf')
+
+    for epoch in range(args.epochs):
+        model.train()
+        t0 = time.time()
+        running = 0.0
+        nb = 0
+        for data in train_loader:
+            feats, cur_xy, _, gt_center = unpack(data, device)
+            out = model(feats, cur_xy)
+            loss = level_k_loss_marginal(out, gt_center, args.level, metric_idx)
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            opt.step()
+            running += float(loss.item())
+            nb += 1
+        sched.step()
+        print(f'[GF-NBA] epoch {epoch+1}/{args.epochs} loss={running/max(nb,1):.4f} '
+              f'lr={opt.param_groups[0]["lr"]:.2e} ({time.time()-t0:.1f}s)', flush=True)
+
+        if (epoch + 1) % args.eval_every == 0 or epoch == args.epochs - 1:
+            ade, fde = evaluate(model, test_loader, device, args.level, end)
+            print(f'[GF-NBA] epoch {epoch+1} ADE4={ade:.4f} FDE4={fde:.4f}', flush=True)
+            if ade < best_ade:
+                best_ade = ade
+                torch.save({'model': model.state_dict(), 'epoch': epoch + 1,
+                            'ade4': ade, 'fde4': fde, 'args': vars(args)},
+                           os.path.join(args.save_dir, f'{args.exp}_best.pt'))
+                print(f'[GF-NBA] saved best (ADE4={ade:.4f}) -> '
+                      f'{os.path.join(args.save_dir, args.exp + "_best.pt")}', flush=True)
+
+    print(f'[GF-NBA] done. best ADE4={best_ade:.4f}', flush=True)
+
+
+if __name__ == '__main__':
+    main()

standalone/gameformer_sra_nba_standalone.py

@@ -0,0 +1,712 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+gameformer_sra_nba_standalone.py
+================================================================================
+SRA-AUGMENTED standalone GameFormer for NBA. Base = the faithful standalone
+GameFormer (Liu et al., ICCV 2023); this file INJECTS SRA's three philosophies
+into GameFormer's level-k InteractionDecoder to test whether the SRA idea
+generalizes to a NON-diffusion (feedforward) predictor.
+
+  SRA components injected (env-gated; GF_SRA=1 enables all):
+   1. NEIGHBOR SELECTION  (GF_SRA_TOPN) -- sparse RAG top-N (=GF_TOPN, default 5)
+      neighbor selection from future proximity; replaces GameFormer's DENSE
+      agent<->agent interaction with a sparse one (self-attn attn_mask + cross-
+      attn key mask restricted to each agent's top-N neighbors).
+   2. UNCERTAINTY GATING  (GF_SRA_SIGMA) -- turns the per-step GMM log-sigma
+      (which GameFormer emits but never USES for interaction) into a learned
+      message gate in (0,1): uncertain agents send weaker interaction messages
+      (uncertainty-weighted message passing). Also biases neighbor selection
+      toward confident agents (GF_SIG_SEL_W).
+   3. RELATIONAL FUTURE FEATURES  (GF_SRA_RELFUT) -- augments GameFormer's
+      per-agent future encoding with explicit PAIRWISE-RELATIVE future features
+      (MLP over relative position+velocity, max-pooled over time, neighbor-mean).
+
+  ALL FLAGS OFF  ==  vanilla GameFormer (identical to gameformer_nba_standalone,
+  ADE4/FDE4 = 0.848/0.951). So (vanilla) vs (GF_SRA=1) is a clean isolated
+  ablation of "does SRA's philosophy help a feedforward predictor too?".
+
+It is a STANDALONE predictor: it consumes ONLY agent HISTORY (past trajectories)
+and predicts K multi-modal futures via the level-k game-theoretic decoder. It
+does NOT consume any external denoiser / diffusion prediction.
+
+--------------------------------------------------------------------------------
+WHAT WAS REPLICATED (evidence = MCZhi/GameFormer @ main)
+--------------------------------------------------------------------------------
+Reference files studied verbatim:
+  * model/GameFormer.py        (Encoder, Decoder, GameFormer)
+  * model/modules.py           (AgentEncoder, FutureEncoder, GMMPredictor,
+                                SelfTransformer, CrossTransformer,
+                                InitialDecoder, InteractionDecoder)
+  * utils/inter_pred_utils.py  (level_k_loss, gmm_loss)
+  * interaction_prediction/train.py (AdamW 1e-4, MultiStepLR, grad-clip 5,
+                                     level=3, modalities=6, encoder_layers=6)
+
+ENCODER  (GameFormer.Encoder + modules.AgentEncoder)
+  * Per-agent temporal encoder = 2-layer LSTM, hidden 256  (nn.LSTM(8,256,2)).
+    We adapt the input feature dim 8 -> 6 (NBA has [abs xy, rel xy, vel xy], no
+    heading/size channels). A learned agent-type embedding (player vs. ball) is
+    added to the LSTM last-step output, exactly as GameFormer adds `type_emb`.
+  * Interaction / self-attention over agents = the GameFormer `fusion_encoder`,
+    a nn.TransformerEncoder of TransformerEncoderLayer(d_model=256, nhead=8,
+    dim_feedforward=1024, activation=gelu, dropout=0.1, batch_first=True),
+    `encoder_layers` deep. In GameFormer this fuses agent tokens WITH per-agent
+    map (lane/crosswalk) tokens; NBA has NO map, so the map encoders
+    (LaneEncoder / CrosswalkEncoder / segment_map) are DROPPED and the fusion is
+    pure agent<->agent self-attention over the 11 tokens (identical for every
+    target agent when there is no per-agent map, so it is computed once).
+
+LEVEL-k GAME-THEORETIC DECODER  (GameFormer.Decoder)
+  * Level 0 = InitialDecoder: for each agent, a set of `modalities` (=K) learned
+    modal query embeddings + a per-agent query embedding are added to the agent's
+    own context token, then cross-attended (CrossTransformer) to the full agent
+    context. A GMMPredictor head outputs K trajectories of (mu_x,mu_y,
+    log_sig_x,log_sig_y) per step + K mode scores.  predictions[...,:2] += the
+    agent's current (last observed) position.
+  * Levels 1..L = InteractionDecoder (one per level, sharing ONE FutureEncoder):
+      - FutureEncoder encodes EVERY agent's K level-(k-1) modal futures into a
+        vector (max-pool over time of an MLP on [x,y,heading,vx,vy]); GameFormer
+        also appends box size (w,l,h) -> we drop those (no size in NBA), 8->5.
+      - The per-mode future encodings are score-softmax-weighted and mean-pooled
+        to one vector per agent, then passed through a SelfTransformer to produce
+        the interaction encoding over all agents (this is the game-theoretic
+        "condition on all other agents' level-(k-1) predictions" step).
+      - The interaction tokens are concatenated with the context tokens; each
+        agent's OWN future token is masked out; the previous-level content plus
+        the agent's own future encoding form the query, cross-attended to that
+        combined context, and a fresh GMMPredictor emits the level-k K futures.
+  * Output dict: level_0_interactions/scores ... level_L_interactions/scores,
+    exactly as GameFormer.Decoder.forward.
+
+LOSS  (utils.inter_pred_utils.level_k_loss + gmm_loss, gmm=True)
+  * Summed over ALL levels k=0..L (level-wise imitation, GameFormer's key term).
+  * Per level, GMM winner-takes-all: best mode = argmin over K of
+    (mean_t ||traj-gt|| + sum over checkpoint steps), then the Gaussian NLL
+    reg = (log_sig_x+log_sig_y) + 0.5*(dx^2/sx^2 + dy^2/sy^2) with log_sig
+    clipped to [0,5], reduced as mean_t + sum over checkpoints; plus mode
+    classification F.cross_entropy(scores, best_mode, label_smoothing=0.2)
+    weighted x2.  We MARGINALIZE per agent (GameFormer's interaction_prediction
+    task predicts 2 agents JOINTLY via distance.mean(1)/scores.sum(1); the NBA
+    Table-1 metric is MARGINAL min-ADE_20, so the best mode / scores are taken
+    per agent instead of jointly -- the only loss adaptation).
+
+--------------------------------------------------------------------------------
+NBA I/O + METRIC (matched to MoFlow / Table 1)
+--------------------------------------------------------------------------------
+  * Data: reuses MoFlow's exact loader data/dataloader_nba.py::NBADatasetMinMax
+    (the graph loader dataloader_nba_graph.py merely re-exports it). Same .npy
+    files (data/nba/original/nba_{train,test}.npy, 32500/12500 scenes,
+    30 frames x 11 agents x 2D), same split (trajs[:n_train] / trajs[:n_test]),
+    same coordinate scaling traj_scale_total=94/28 and traj_mean=[14,7.5].
+  * Past = 10 frames, Future = 20 frames = 4.0 s @ 5 Hz, 11 agents, 2D.
+  * Encoder input = data['past_traj'] (6-dim min-max normalized to [-1,1]).
+  * Coordinate frame for the GMM/interaction = data['past_traj_original_scale']
+    (centered-abs = pos/scale - traj_mean); cur_xy = its last past step.
+    Ground-truth future displacement = data['fut_traj_original_scale']
+    (= future - last_observed, in /scale units) -- the SAME tensor the MoFlow
+    evaluator (eval_perscene_moflow.py) compares against.
+  * Metric (identical to eval_perscene_moflow.py, verified vs Table 1):
+      d = ||pred_disp - gt_disp||                       # [B,A,K,T]
+      ADE4 = d[..., :20].mean(over T).min(over K), mean over 11 agents & scenes
+      FDE4 = d[..., 19]        .min(over K),          mean over 11 agents & scenes
+    i.e. marginal min-of-K=20 at the 4.0 s horizon (frame 20). Table-1 baselines
+    are ~0.70-1.04 ADE in this space.
+
+Prints:  [SRA-GF] epoch N ADE4=.. FDE4=..
+
+DO NOT run on GPU here (GPUs busy) -- this file is only `python -m py_compile`d.
+"""
+
+import os
+import sys
+import argparse
+import time
+
+# ------------------------------------------------------------------ GPU FIRST
+# CUDA context is fixed at torch import; select the device BEFORE importing torch
+def _early_gpu():
+    for i, a in enumerate(sys.argv):
+        if a == '--gpu' and i + 1 < len(sys.argv):
+            return sys.argv[i + 1]
+        if a.startswith('--gpu='):
+            return a.split('=', 1)[1]
+    return None
+
+
+_g = _early_gpu()
+if _g is not None:
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(_g)
+os.environ.setdefault('MPLBACKEND', 'Agg')  # dataloader imports matplotlib
+
+# ------------------------------------------------ reuse MoFlow's NBA pipeline
+MOFLOW_ROOT = '/mnt/jaewoo4tb/srtp/MoFlow'
+if MOFLOW_ROOT not in sys.path:
+    sys.path.insert(0, MOFLOW_ROOT)
+
+import math
+import types as _types
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+# so that `from utils.normalization import ...` inside the loader resolves
+try:
+    os.chdir(MOFLOW_ROOT)
+except Exception:
+    pass
+from data.dataloader_nba import NBADatasetMinMax, seq_collate_nba
+
+
+# ============================================================================
+#  MODEL  (faithful GameFormer, NBA-adapted: no map)
+# ============================================================================
+D_MODEL = 256
+N_HEADS = 8
+DROPOUT = 0.1
+
+
+# ---------------------------------------------------------------- SRA toggles
+# Inject SRA's three philosophies into GameFormer's level-k InteractionDecoder.
+# Master flag GF_SRA=1 enables all; or enable individually. All-off == vanilla
+# GameFormer (identical to gameformer_nba_standalone.py, 0.848/0.951).
+def _sra_on(name):
+    if os.environ.get('GF_SRA', '') not in ('', '0', 'false', 'False'):
+        return True
+    return os.environ.get(name, '') not in ('', '0', 'false', 'False')
+
+
+USE_TOPN = _sra_on('GF_SRA_TOPN')       # sparse RAG top-N neighbor selection
+USE_SIGMA = _sra_on('GF_SRA_SIGMA')     # uncertainty-gated message passing
+USE_RELFUT = _sra_on('GF_SRA_RELFUT')   # explicit relational future features
+TOPN_K = int(os.environ.get('GF_TOPN', '5'))   # SRA neighbor budget (top_n=5)
+SIG_SEL_W = float(os.environ.get('GF_SIG_SEL_W', '0.5'))  # sigma weight in nbr selection
+
+
+# ---------------------------------------------------------------- primitives
+class SelfTransformer(nn.Module):
+    """modules.py::SelfTransformer. `attn_mask` added for SRA sparse-neighbor
+    selection (per-query [B*heads,A,A] boolean; True = block non-neighbor)."""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, dropout=DROPOUT):
+        super().__init__()
+        self.heads = heads
+        self.self_attention = nn.MultiheadAttention(dim, heads, dropout, batch_first=True)
+        self.norm_1 = nn.LayerNorm(dim)
+        self.norm_2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout))
+
+    def forward(self, inputs, mask=None, attn_mask=None):
+        attn, _ = self.self_attention(inputs, inputs, inputs, key_padding_mask=mask,
+                                      attn_mask=attn_mask)
+        attn = self.norm_1(attn + inputs)
+        return self.norm_2(self.ffn(attn) + attn)
+
+
+class CrossTransformer(nn.Module):
+    """modules.py::CrossTransformer (verbatim)."""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, dropout=DROPOUT):
+        super().__init__()
+        self.cross_attention = nn.MultiheadAttention(dim, heads, dropout, batch_first=True)
+        self.norm_1 = nn.LayerNorm(dim)
+        self.norm_2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout))
+
+    def forward(self, query, key, value, mask=None):
+        attn, _ = self.cross_attention(query, key, value, key_padding_mask=mask)
+        attn = self.norm_1(attn)
+        return self.norm_2(self.ffn(attn) + attn)
+
+
+class GMMPredictor(nn.Module):
+    """modules.py::GMMPredictor (verbatim). Out: (mu_x,mu_y,log_sig_x,log_sig_y)."""
+
+    def __init__(self, future_len, dim=D_MODEL):
+        super().__init__()
+        self._future_len = future_len
+        self.gaussian = nn.Sequential(
+            nn.Linear(dim, 512), nn.ELU(), nn.Dropout(0.1),
+            nn.Linear(512, future_len * 4))
+        self.score = nn.Sequential(
+            nn.Linear(dim, 64), nn.ELU(), nn.Dropout(0.1), nn.Linear(64, 1))
+
+    def forward(self, x):
+        # x: [B, M, D]
+        B, M, _ = x.shape
+        res = self.gaussian(x).view(B, M, self._future_len, 4)
+        score = self.score(x).squeeze(-1)  # [B, M]
+        return res, score
+
+
+# ------------------------------------------------------------------ encoder
+class AgentHistoryEncoder(nn.Module):
+    """modules.py::AgentEncoder -- 2-layer LSTM(->256) over agent history + type
+    embedding.  GameFormer uses nn.LSTM(8,256,2); NBA feature dim is 6."""
+
+    def __init__(self, in_dim=6, dim=D_MODEL, n_types=2):
+        super().__init__()
+        self.motion = nn.LSTM(in_dim, dim, 2, batch_first=True)
+        self.type_emb = nn.Embedding(n_types, dim)
+
+    def forward(self, hist, types):
+        # hist: [B, A, T, in_dim]  ;  types: [A] long
+        B, A, T, C = hist.shape
+        traj, _ = self.motion(hist.reshape(B * A, T, C))
+        out = traj[:, -1].reshape(B, A, -1)          # [B, A, D]
+        out = out + self.type_emb(types)[None]       # + per-agent type emb
+        return out
+
+
+class FusionEncoder(nn.Module):
+    """GameFormer.Encoder.fusion_encoder -- self-attention over agent tokens.
+    (No map tokens for NBA.)"""
+
+    def __init__(self, dim=D_MODEL, heads=N_HEADS, layers=6, dropout=DROPOUT):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(
+            d_model=dim, nhead=heads, dim_feedforward=dim * 4,
+            activation=F.gelu, dropout=dropout, batch_first=True)
+        self.encoder = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False)
+
+    def forward(self, tokens, mask=None):
+        return self.encoder(tokens, src_key_padding_mask=mask)
+
+
+# ------------------------------------------------------------- future encoder
+class FutureEncoder(nn.Module):
+    """modules.py::FutureEncoder -- encodes each agent's K modal futures into a
+    per-mode vector via max-pool over an MLP on [x,y,heading,vx,vy].  GameFormer
+    appends box size (w,l,h) -> dropped for NBA (8 -> 5 input dims)."""
+
+    def __init__(self, dim=D_MODEL, n_types=2, dt=0.2):
+        super().__init__()
+        self.dt = dt
+        self.mlp = nn.Sequential(nn.Linear(5, 64), nn.ReLU(), nn.Linear(64, dim))
+        self.type_emb = nn.Embedding(n_types, dim)
+
+    def state_process(self, trajs, cur_xy):
+        # trajs: [B, A, M, T, 2] ; cur_xy: [B, A, 2]
+        B, A, M, T, _ = trajs.shape
+        cur = cur_xy[:, :, None, None, :].expand(B, A, M, 1, 2)
+        xy = torch.cat([cur, trajs], dim=-2)                     # [B,A,M,T+1,2]
+        dxy = torch.diff(xy, dim=-2)                             # [B,A,M,T,2]
+        v = dxy / self.dt
+        theta = torch.atan2(dxy[..., 1], dxy[..., 0].clamp(min=1e-3)).unsqueeze(-1)
+        return torch.cat([trajs, theta, v], dim=-1)             # [B,A,M,T,5]
+
+    def forward(self, trajs, cur_xy, types):
+        feat = self.state_process(trajs, cur_xy)
+        h = self.mlp(feat.detach())                             # GameFormer detaches
+        out = torch.max(h, dim=-2).values                       # [B,A,M,D]
+        out = out + self.type_emb(types)[None, :, None, :]      # + type emb
+        return out
+
+
+# ----------------------------------------------------------- level-0 decoder
+class InitialDecoder(nn.Module):
+    """modules.py::InitialDecoder -- vectorized over agents."""
+
+    def __init__(self, modalities, n_agents, future_len, dim=D_MODEL):
+        super().__init__()
+        self.M = modalities
+        self.A = n_agents
+        self.multi_modal_query_embedding = nn.Embedding(modalities, dim)
+        self.agent_query_embedding = nn.Embedding(n_agents, dim)
+        self.query_encoder = CrossTransformer(dim)
+        self.predictor = GMMPredictor(future_len, dim)
+        self.register_buffer('modal', torch.arange(modalities).long())
+        self.register_buffer('agent', torch.arange(n_agents).long())
+
+    def forward(self, encoding, cur_xy, mask=None):
+        # encoding: [B, A, D]  (full agent context, shared as keys)
+        B, A, D = encoding.shape
+        M = self.M
+        multi_modal = self.multi_modal_query_embedding(self.modal)   # [M, D]
+        agent = self.agent_query_embedding(self.agent)               # [A, D]
+        mm_agent_query = multi_modal[None, :, :] + agent[:, None, :]  # [A, M, D]
+        query = encoding[:, :, None, :] + mm_agent_query[None]       # [B, A, M, D]
+
+        q = query.reshape(B * A, M, D)
+        kv = encoding[:, None, :, :].expand(B, A, A, D).reshape(B * A, A, D)
+        km = None
+        if mask is not None:
+            km = mask[:, None, :].expand(B, A, A).reshape(B * A, A)
+        content = self.query_encoder(q, kv, kv, km)                 # [B*A, M, D]
+        preds, scores = self.predictor(content)                     # [B*A,M,T,4],[B*A,M]
+
+        content = content.reshape(B, A, M, D)
+        T = preds.shape[-2]
+        preds = preds.reshape(B, A, M, T, 4)
+        scores = scores.reshape(B, A, M)
+        preds = preds.clone()
+        preds[..., :2] = preds[..., :2] + cur_xy[:, :, None, None, :]
+        return content, preds, scores
+
+
+# ------------------------------------------------------- level-k decoder
+class InteractionDecoder(nn.Module):
+    """modules.py::InteractionDecoder -- one game-theoretic reasoning level,
+    vectorized over agents.  Conditions each agent on ALL agents' level-(k-1)
+    predictions (interaction self-attention), masking the agent's own future."""
+
+    def __init__(self, future_encoder, future_len, n_agents, dim=D_MODEL,
+                 use_topn=None, use_sigma=None, use_relfut=None, top_n=None):
+        super().__init__()
+        self.A = n_agents
+        self.future_encoder = future_encoder            # SHARED across levels
+        self.interaction_encoder = SelfTransformer(dim)
+        self.query_encoder = CrossTransformer(dim)
+        self.decoder = GMMPredictor(future_len, dim)
+        # --- SRA augmentations (read module-level toggles unless overridden) ---
+        self.use_topn = USE_TOPN if use_topn is None else use_topn
+        self.use_sigma = USE_SIGMA if use_sigma is None else use_sigma
+        self.use_relfut = USE_RELFUT if use_relfut is None else use_relfut
+        self.top_n = TOPN_K if top_n is None else top_n
+        if self.use_sigma:                              # uncertainty -> message gate in (0,1)
+            self.sigma_gate = nn.Sequential(
+                nn.Linear(1, 32), nn.ReLU(), nn.Linear(32, 1), nn.Sigmoid())
+        if self.use_relfut:                             # relational future edge encoder
+            self.rel_mlp = nn.Sequential(nn.Linear(4, 64), nn.ReLU(), nn.Linear(64, dim))
+
+    def forward(self, cur_xy, last_traj, last_scores, last_content, encoding, types, mask=None):
+        # cur_xy:[B,A,2] last_traj:[B,A,M,T,4] last_scores:[B,A,M]
+        # last_content:[B,A,M,D] encoding:[B,A,D]
+        B, A, M, T, _ = last_traj.shape
+        D = encoding.shape[-1]
+        dev = encoding.device
+        w = last_scores.softmax(-1)                                # [B,A,M] mode weights
+
+        # encode all agents' level-(k-1) modal futures
+        multi_futures = self.future_encoder(last_traj[..., :2], cur_xy, types)  # [B,A,M,D]
+        futures = (multi_futures * w.unsqueeze(-1)).mean(dim=2)     # [B,A,D]
+
+        # ---- SRA: per-agent uncertainty from prev-level GMM log-sigma ----------
+        sig_agent = None
+        if self.use_sigma:
+            sig = torch.exp(last_traj[..., 2:].clamp(0, 5)).mean(-1)   # [B,A,M,T]
+            sig_agent = (sig * w[..., None]).sum(2).mean(-1)          # [B,A] score-wtd, mean-T
+
+        # ---- SRA: sparse RAG neighbor selection from future proximity ----------
+        nbr_block = None                                             # [B,A,A] True = block
+        rep = (last_traj[..., :2] * w[..., None, None]).sum(2)        # [B,A,T,2] rep future
+        eye = torch.eye(A, dtype=torch.bool, device=dev)
+        if self.use_topn and A > 1:
+            prox = (rep[:, :, None] - rep[:, None]).norm(dim=-1).mean(-1)  # [B,A,A]
+            if sig_agent is not None:                               # prefer confident neighbors
+                prox = prox + SIG_SEL_W * sig_agent[:, None, :]
+            prox = prox.masked_fill(eye[None], float('inf'))         # exclude self from top-N
+            N = min(max(self.top_n - 1, 1), A - 1)
+            nbr_idx = prox.topk(N, largest=False, dim=-1).indices    # [B,A,N] nearest others
+            nbr_block = torch.ones(B, A, A, dtype=torch.bool, device=dev)
+            nbr_block[:, torch.arange(A, device=dev), torch.arange(A, device=dev)] = False  # self ok
+            nbr_block.scatter_(2, nbr_idx, False)                    # neighbors ok
+
+        # ---- SRA: uncertainty-gated message passing ----------------------------
+        if self.use_sigma:
+            futures = futures * self.sigma_gate(sig_agent[..., None])  # weaken uncertain senders
+
+        # ---- SRA: explicit relational future features --------------------------
+        if self.use_relfut and A > 1:
+            dpos = rep[:, :, None] - rep[:, None]                    # [B,A,A,T,2]
+            dvel = F.pad(torch.diff(dpos, dim=-2), (0, 0, 1, 0))     # [B,A,A,T,2]
+            relf = self.rel_mlp(torch.cat([dpos, dvel], dim=-1)).max(dim=-2).values  # [B,A,A,D]
+            if nbr_block is not None:
+                relf = relf.masked_fill(nbr_block[..., None], 0.0)
+                denom = (~nbr_block).sum(-1, keepdim=True).clamp(min=1)
+                futures = futures + relf.sum(2) / denom             # neighbor-mean relational future
+            else:
+                futures = futures + relf.mean(2)
+
+        # interaction over agents (self-attention; SRA restricts to neighbors)
+        attn_mask = None
+        if nbr_block is not None:
+            h = self.interaction_encoder.heads
+            attn_mask = nbr_block[:, None].expand(B, h, A, A).reshape(B * h, A, A)
+        interaction = self.interaction_encoder(futures, None, attn_mask)  # [B,A,D]
+
+        # combined context = [interaction tokens ; context tokens]
+        ctx = torch.cat([interaction, encoding], dim=1)            # [B, 2A, D]
+        # per-agent key mask: mask each agent's OWN future token in the 1st block
+        km = torch.zeros(B, A, 2 * A, dtype=torch.bool, device=dev)
+        idx = torch.arange(A, device=dev)
+        km[:, idx, idx] = True
+        if mask is not None:
+            # second block inherits context validity mask (all-valid for NBA)
+            km[:, :, A:] = km[:, :, A:] | mask[:, None, :]
+            km[:, :, :A] = km[:, :, :A] | mask[:, None, :]
+        if nbr_block is not None:                                   # SRA sparsity on both blocks
+            km[:, :, :A] = km[:, :, :A] | nbr_block                 # interaction block: neighbors only
+            km[:, :, A:] = km[:, :, A:] | nbr_block                 # context block: neighbors(+self) only
+        km = km.reshape(B * A, 2 * A)
+        kv = ctx[:, None, :, :].expand(B, A, 2 * A, D).reshape(B * A, 2 * A, D)
+
+        query = (last_content + multi_futures).reshape(B * A, M, D)  # prev content + own future
+        content = self.query_encoder(query, kv, kv, km)             # [B*A, M, D]
+        traj, scores = self.decoder(content)                        # [B*A,M,T,4],[B*A,M]
+
+        content = content.reshape(B, A, M, D)
+        traj = traj.reshape(B, A, M, T, 4)
+        scores = scores.reshape(B, A, M)
+        traj = traj.clone()
+        traj[..., :2] = traj[..., :2] + cur_xy[:, :, None, None, :]
+        return content, traj, scores
+
+
+# -------------------------------------------------------------- full decoder
+class GameFormerDecoder(nn.Module):
+    """GameFormer.Decoder -- level 0 + L interaction levels (shared FutureEncoder)."""
+
+    def __init__(self, modalities, n_agents, future_len, levels, dim=D_MODEL,
+                 n_types=2, dt=0.2):
+        super().__init__()
+        self._levels = levels
+        self.future_encoder = FutureEncoder(dim, n_types, dt)
+        self.initial_stage = InitialDecoder(modalities, n_agents, future_len, dim)
+        self.interaction_stage = nn.ModuleList(
+            [InteractionDecoder(self.future_encoder, future_len, n_agents, dim)
+             for _ in range(levels)])
+
+    def forward(self, encoding, cur_xy, types, mask=None):
+        out = {}
+        content, traj, scores = self.initial_stage(encoding, cur_xy, mask)
+        out['level_0_interactions'] = traj
+        out['level_0_scores'] = scores
+        for k in range(1, self._levels + 1):
+            content, traj, scores = self.interaction_stage[k - 1](
+                cur_xy, traj, scores, content, encoding, types, mask)
+            out[f'level_{k}_interactions'] = traj
+            out[f'level_{k}_scores'] = scores
+        return out
+
+
+# ---------------------------------------------------------------- top model
+class GameFormerNBA(nn.Module):
+    """GameFormer.GameFormer (NBA, map dropped)."""
+
+    def __init__(self, n_agents=11, past_dim=6, future_len=20, modalities=20,
+                 levels=2, dim=D_MODEL, heads=N_HEADS, enc_layers=6,
+                 n_types=2, dt=0.2, ball_idx=10):
+        super().__init__()
+        self.levels = levels
+        self.history_encoder = AgentHistoryEncoder(past_dim, dim, n_types)
+        self.fusion_encoder = FusionEncoder(dim, heads, enc_layers)
+        self.decoder = GameFormerDecoder(modalities, n_agents, future_len, levels,
+                                         dim, n_types, dt)
+        types = torch.zeros(n_agents, dtype=torch.long)
+        if 0 <= ball_idx < n_agents and n_types > 1:
+            types[ball_idx] = 1
+        self.register_buffer('agent_types', types)
+
+    def forward(self, feats, cur_xy):
+        # feats: [B,A,T,past_dim] (encoder input) ; cur_xy: [B,A,2] (metric frame)
+        enc = self.history_encoder(feats, self.agent_types)   # [B,A,D]
+        enc = self.fusion_encoder(enc, None)                  # [B,A,D]
+        return self.decoder(enc, cur_xy, self.agent_types, None)
+
+
+# ============================================================================
+#  LOSS  (level_k_loss + gmm_loss, gmm=True; MARGINAL per agent for NBA)
+# ============================================================================
+def gmm_loss_marginal(traj, conv, scores, gt, metric_idx):
+    """Per-agent GMM winner-takes-all (marginal version of gmm_loss).
+      traj:  [B,A,M,T,2]   conv: [B,A,M,T,2]   scores: [B,A,M]   gt: [B,A,T,2]
+    """
+    B, A, M, T, _ = traj.shape
+    dist = torch.norm(traj - gt[:, :, None], dim=-1)                  # [B,A,M,T]
+    ndist = dist.mean(-1) + dist[..., metric_idx].sum(-1)            # [B,A,M]
+    best = ndist.argmin(-1)                                          # [B,A]
+
+    gi = best[..., None, None, None].expand(B, A, 1, T, 2)
+    best_traj = torch.gather(traj, 2, gi).squeeze(2)                 # [B,A,T,2]
+    best_conv = torch.gather(conv, 2, gi).squeeze(2)                 # [B,A,T,2]
+
+    dx = best_traj[..., 0] - gt[..., 0]
+    dy = best_traj[..., 1] - gt[..., 1]
+    log_std_x = torch.clip(best_conv[..., 0], 0, 5)
+    log_std_y = torch.clip(best_conv[..., 1], 0, 5)
+    std_x, std_y = torch.exp(log_std_x), torch.exp(log_std_y)
+
+    reg = (log_std_x + log_std_y) + 0.5 * ((dx ** 2) / (std_x ** 2) + (dy ** 2) / (std_y ** 2))
+    reg = reg.mean(-1) + reg[..., metric_idx].sum(-1)               # [B,A]
+
+    prob_loss = F.cross_entropy(scores.reshape(B * A, M), best.reshape(B * A),
+                                label_smoothing=0.2)
+    return reg.mean() + 2.0 * prob_loss, best
+
+
+def level_k_loss_marginal(outputs, gt_center, levels, metric_idx):
+    """Sum GMM loss across all levels 0..L (GameFormer's level-wise imitation)."""
+    total = 0.0
+    for k in range(levels + 1):
+        pred = outputs[f'level_{k}_interactions']       # [B,A,M,T,4]
+        scores = outputs[f'level_{k}_scores']           # [B,A,M]
+        traj, conv = pred[..., :2], pred[..., 2:]
+        l, _ = gmm_loss_marginal(traj, conv, scores, gt_center, metric_idx)
+        total = total + l
+    return total
+
+
+# ============================================================================
+#  DATA
+# ============================================================================
+def build_loaders(args):
+    cfg = _types.SimpleNamespace(
+        traj_mean=[14, 7.5], data_norm='min_max',
+        past_frames=args.past_len, future_frames=args.future_len, agents=args.agents)
+    # Train first so cfg min/max scalars are populated for the test set.
+    train_set = NBADatasetMinMax(
+        obs_len=args.past_len, pred_len=args.future_len, training=True,
+        num_scenes=args.n_train, cfg=cfg, data_dir=args.data_dir,
+        rotate=False, data_norm='min_max')
+    test_set = NBADatasetMinMax(
+        obs_len=args.past_len, pred_len=args.future_len, training=False,
+        test_scenes=args.n_test, cfg=cfg, data_dir=args.data_dir,
+        rotate=False, data_norm='min_max')
+    train_loader = DataLoader(
+        train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
+        collate_fn=seq_collate_nba, pin_memory=True, drop_last=True)
+    test_loader = DataLoader(
+        test_set, batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
+        collate_fn=seq_collate_nba, pin_memory=True)
+    return train_loader, test_loader
+
+
+def unpack(data, device):
+    feats = data['past_traj'].to(device)                       # [B,A,T,6] normalized
+    past_orig = data['past_traj_original_scale'].to(device)    # [B,A,T,6] metric
+    gt_disp = data['fut_traj_original_scale'].to(device)       # [B,A,Tf,2] displacement
+    cur_xy = past_orig[:, :, -1, 0:2]                          # centered-abs last pos
+    gt_center = gt_disp + cur_xy[:, :, None, :]                # centered-abs future
+    return feats, cur_xy, gt_disp, gt_center
+
+
+# ============================================================================
+#  EVAL  (marginal min-of-K=20 at 4.0 s; identical to eval_perscene_moflow.py)
+# ============================================================================
+@torch.no_grad()
+def evaluate(model, loader, device, levels, end):
+    model.eval()
+    ade_sum = fde_sum = 0.0
+    n = 0
+    for data in loader:
+        feats, cur_xy, gt_disp, _ = unpack(data, device)
+        out = model(feats, cur_xy)
+        pred = out[f'level_{levels}_interactions'][..., :2]        # [B,A,M,T,2]
+        pred_disp = pred - cur_xy[:, :, None, None, :]             # -> displacement
+        d = torch.norm(pred_disp - gt_disp[:, :, None], dim=-1)    # [B,A,M,T]
+        ade = d[..., :end].mean(-1).min(dim=-1).values            # [B,A] best-of-K
+        fde = d[..., end - 1].min(dim=-1).values                  # [B,A]
+        ade_sum += ade.sum().item()
+        fde_sum += fde.sum().item()
+        n += ade.numel()
+    return ade_sum / max(n, 1), fde_sum / max(n, 1)
+
+
+# ============================================================================
+#  TRAIN
+# ============================================================================
+def parse_args():
+    p = argparse.ArgumentParser('GameFormer standalone predictor for NBA')
+    p.add_argument('--data_dir', default='/mnt/jaewoo4tb/srtp/MoFlow/data/nba', type=str)
+    p.add_argument('--gpu', default='0', type=str)
+    p.add_argument('--exp', default='sragf_nba', type=str)
+    p.add_argument('--epochs', default=50, type=int)
+    p.add_argument('--batch_size', default=128, type=int)
+    p.add_argument('--test_batch', default=500, type=int)
+    p.add_argument('--workers', default=4, type=int)
+    p.add_argument('--seed', default=3407, type=int)          # GameFormer default
+
+    # --- architecture (faithful GameFormer knobs) ---
+    p.add_argument('--modalities', default=20, type=int, help='K modes (NBA min-ADE_20)')
+    p.add_argument('--level', default=2, type=int, help='game-theoretic levels (paper: 3)')
+    p.add_argument('--encoder_layers', default=6, type=int)
+    p.add_argument('--agents', default=11, type=int)
+    p.add_argument('--past_len', default=10, type=int)
+    p.add_argument('--future_len', default=20, type=int)      # 4.0 s @ 5 Hz
+    p.add_argument('--ball_idx', default=10, type=int, help='-1 to disable type emb')
+
+    # --- optimization (GameFormer: AdamW 1e-4, MultiStepLR x0.5, clip 5) ---
+    p.add_argument('--lr', default=1e-4, type=float)
+    p.add_argument('--weight_decay', default=1e-4, type=float)
+    p.add_argument('--grad_clip', default=5.0, type=float)
+    p.add_argument('--n_train', default=32500, type=int)
+    p.add_argument('--n_test', default=12500, type=int)
+    p.add_argument('--eval_every', default=2, type=int)
+    p.add_argument('--save_dir', default=os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), 'sragf_nba_ckpt'), type=str)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    train_loader, test_loader = build_loaders(args)
+
+    n_types = 2 if (0 <= args.ball_idx < args.agents) else 1
+    model = GameFormerNBA(
+        n_agents=args.agents, past_dim=6, future_len=args.future_len,
+        modalities=args.modalities, levels=args.level, dim=D_MODEL, heads=N_HEADS,
+        enc_layers=args.encoder_layers, n_types=n_types, dt=0.2,
+        ball_idx=args.ball_idx).to(device)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f'[SRA-GF] model params: {n_params/1e6:.2f} M | levels={args.level} '
+          f'K={args.modalities} enc_layers={args.encoder_layers} device={device}')
+    print(f'[SRA-GF] SRA components: neighbor-topN={USE_TOPN}(N={TOPN_K}) '
+          f'sigma-gate={USE_SIGMA} relational-future={USE_RELFUT}  '
+          f'(all False == vanilla GameFormer)', flush=True)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    milestones = sorted(set(int(args.epochs * f) for f in (0.5, 0.65, 0.8, 0.9)))
+    milestones = [m for m in milestones if 0 < m < args.epochs]
+    sched = torch.optim.lr_scheduler.MultiStepLR(opt, milestones=milestones, gamma=0.5)
+
+    end = args.future_len                                   # 4.0 s = frame 20
+    metric_idx = sorted(set([max(0, args.future_len // 2 - 1), args.future_len - 1]))
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    best_ade = float('inf')
+
+    for epoch in range(args.epochs):
+        model.train()
+        t0 = time.time()
+        running = 0.0
+        nb = 0
+        for data in train_loader:
+            feats, cur_xy, _, gt_center = unpack(data, device)
+            out = model(feats, cur_xy)
+            loss = level_k_loss_marginal(out, gt_center, args.level, metric_idx)
+            opt.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            opt.step()
+            running += float(loss.item())
+            nb += 1
+        sched.step()
+        print(f'[SRA-GF] epoch {epoch+1}/{args.epochs} loss={running/max(nb,1):.4f} '
+              f'lr={opt.param_groups[0]["lr"]:.2e} ({time.time()-t0:.1f}s)', flush=True)
+
+        if (epoch + 1) % args.eval_every == 0 or epoch == args.epochs - 1:
+            ade, fde = evaluate(model, test_loader, device, args.level, end)
+            print(f'[SRA-GF] epoch {epoch+1} ADE4={ade:.4f} FDE4={fde:.4f}', flush=True)
+            if ade < best_ade:
+                best_ade = ade
+                torch.save({'model': model.state_dict(), 'epoch': epoch + 1,
+                            'ade4': ade, 'fde4': fde, 'args': vars(args)},
+                           os.path.join(args.save_dir, f'{args.exp}_best.pt'))
+                print(f'[SRA-GF] saved best (ADE4={ade:.4f}) -> '
+                      f'{os.path.join(args.save_dir, args.exp + "_best.pt")}', flush=True)
+
+    print(f'[SRA-GF] done. best ADE4={best_ade:.4f}', flush=True)
+
+
+if __name__ == '__main__':
+    main()