SRA — Spatial Reasoning Adapter: How to Run

SRA is a modular future-interaction graph (FutureInteractionGraphV6) that plugs into three stochastic trajectory predictors at a single insertion point:

Host	Type	Entry point (NBA)	Entry point (soccer/football)
MID	DDPM diffusion	`MID/main_nba_mid_graphv6_v3.py`	`MID/main_{soccer,football}_mid_graphv5_sigma*.py`
LED	leapfrog-DDPM	`LED/main_led_nba_graph.py`	`LED/main_sport_led.py`
MoFlow	flow matching	`MoFlow/fm_nba_graph_v6.py`	`MoFlow/fm_sport_graph_v6.py`

Datasets: NBA (11 agents), Soccer / Football (23 agents). Metric: min-ADE₂₀ / min-FDE₂₀ @ 4.0 s.

0. CRITICAL: directory layout

MID/ and LED/ do not contain the SRA graph module. They import it from a sibling MoFlow/ directory at runtime:

# MID/main_nba_mid_graphv6_v3.py
MOFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'MoFlow'))
sys.path.insert(0, MOFLOW_ROOT)
from models.graph_interaction_nba_v6 import FutureInteractionGraphV6

So you must preserve this layout — do not move the three folders apart:

<root>/
├── MoFlow/          # owns the SRA graph + baseline modules
│   └── models/
│       ├── graph_interaction_nba_v6.py   # <-- SRA (the adapter itself)
│       ├── interaction_baselines.py      # <-- E4 baselines (GameFormer / C2F) + factory
│       └── backbone_graph_v6.py          # MoFlow's insertion point
├── MID/
└── LED/

1. Environment

conda create -n sra python=3.11 -y && conda activate sra
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
pip install -r MoFlow/requirements.txt
pip install -r LED/requirements.txt
pip install easydict pyyaml tensorboard tqdm scipy matplotlib gitpython

Verified on: PyTorch 2.4.0 + CUDA 12.1, NVIDIA A6000 (48 GB) and RTX 3090 (24 GB).

cuDNN gotcha (MID + LED). If a system CUDA shadows the env's cuDNN you get Could not load library libcudnn_cnn_train.so.8 ... undefined symbol at .backward(). Fix by prepending the env's bundled NVIDIA libs:

ENV=$CONDA_PREFIX
export LD_LIBRARY_PATH="$(ls -d $ENV/lib/python3.11/site-packages/nvidia/*/lib | tr '\n' ':')$LD_LIBRARY_PATH"

2. Data (not included in this repo)

Place the arrays as follows, then pass the paths via --data_dir:

MoFlow/data/nba/original/nba_train.npy      # 32500 scenes, 30 frames x 11 agents x 2
MoFlow/data/nba/original/nba_test.npy       # 12500 scenes
<raw_data>/soccer/{train.npy,val.npy}       # 23 agents
<raw_data>/football/{train.npy,val.npy}     # 23 agents

NBA past/future = 10/20 frames @ 5 Hz (4.0 s horizon). Sport uses the same convention.

--data_dir points at the PARENT of original/ for NBA — the loader appends original/ itself. So use --data_dir ./data/nba, not ./data/nba/original.

Several scripts still carry absolute --data_dir defaults from the development machine (/mnt/...). Always pass --data_dir explicitly — the CLI value overrides the default.

Smoke test with the bundled 100-scene sample

sample_data/nba/ contains a 100-scene NBA train/test subset (≈0.5 MB) so you can verify the pipeline before setting up the full data — see sample_data/README.md:

cd MoFlow
CUDA_VISIBLE_DEVICES=0 python fm_nba_graph_v6.py \
  --cfg cfg/nba/cor_fm.yml --exp smoke --data_dir ../sample_data/nba \
  --n_train 100 --n_test 100 --batch_size 8 --epochs 1 \
  --fm_in_scaling --tied_noise --top_n_neighbors 5 --uncertainty_weight 0.01

100 scenes cannot train a usable model — this only confirms the data path, model construction, training step and eval loop execute.

3. Run — host × dataset

All commands below are the exact recipes used to produce the reported numbers. Set the GPU with CUDA_VISIBLE_DEVICES=<id> (MoFlow) or --gpu <id> (MID/LED).

3.1 MoFlow

cd MoFlow

# NBA + SRA
CUDA_VISIBLE_DEVICES=0 python fm_nba_graph_v6.py \
  --cfg cfg/nba/cor_fm.yml --exp nba_sra \
  --batch_size 192 --epochs 150 --fm_in_scaling --tied_noise \
  --top_n_neighbors 5 --uncertainty_weight 0.01 \
  --data_dir ./data/nba

# Soccer / Football + SRA   (swap football <-> soccer)
CUDA_VISIBLE_DEVICES=0 python fm_sport_graph_v6.py \
  --cfg cfg/sport/football.yml --exp football_sra \
  --batch_size 64 --epochs 100 \
  --top_n_neighbors 5 --uncertainty_weight 0.01 \
  --data_dir <raw_data>/football

--resume <ckpt_name> (e.g. --resume checkpoint_epoch_20) restores model+optimizer+EMA+step from <cfg_dir>/models/. The run directory name is derived from the args, so resume only works if you pass the identical args as the original run.

3.2 MID

cd MID

# NBA + SRA
CUDA_VISIBLE_DEVICES=0 python main_nba_mid_graphv6_v3.py \
  --data_dir <root>/MoFlow/data/nba/original --exp_name mid_nba_sra \
  --epochs 100 --batch_size 32 --lr 1e-3 --eval_every 5 \
  --sampling ddim --sampling_step 10

# Football + SRA
python main_football_mid_graphv5_sigma.py \
  --data_dir <raw_data>/football --exp_name mid_football_sra --gpu 0 \
  --epochs 100 --batch_size 64 --lr 1e-3 --graph_lr_mult 1.0 --eval_every 1 \
  --top_n_neighbors 5 --uncertainty_weight 0.01 \
  --train_mode two_pass --sampling ddim --sampling_step 20

# Soccer + SRA  -> same flags, script main_soccer_mid_graphv5_sigma_output.py

⚠️ --graph_lr_mult matters. 3.0 diverges for the sparse (top_n=5) sport configuration (ADE stuck ≈0.46). Use 1.0 — that is what the reported results use.

3.3 LED

cd LED

# NBA + SRA
python main_led_nba_graph.py --cfg led_augment --gpu 0 --train 1 \
  --use_v6_graph --use_sigma --top_n 5 --uncertainty_weight 1.0 --residual_on eps

# Soccer / Football + SRA
python main_sport_led.py --cfg football --gpu 0 --train 1 \
  --use_v6_graph --residual_on eps

⚠️ LED-sport σ behaviour. main_sport_led.py has no --use_sigma flag: it passes the leapfrog initializer's variance_estimation to the graph automatically. To run a genuine no-σ ablation on sport you must set LED_NO_SIGMA=1 (see §5), otherwise the "no-σ" run silently uses σ and reproduces the full-SRA result.

4. E2 — cumulative ablation settings

The ablation adds one component at a time. A = number of agents (11 NBA / 23 sport).

Configuration	`top_n_neighbors`	uncertainty (σ)
host baseline	(graph disabled)	–
+ relational encoding (all neighbors, no σ)	`A-1` (10 NBA / 22 sport)	off
+ sparse neighbor selection (no σ)	`5`	off
+ uncertainty = full SRA	`5`	on

Turning σ off:

Host	How
MoFlow	`--uncertainty_weight 0.0` (code sets `use_sigma_gating = uncertainty_weight > 0`)
MID	`--uncertainty_weight 0.0`
LED (NBA)	omit `--use_sigma`
LED (sport)	`LED_NO_SIGMA=1` environment variable

Memory note. Dense sport (top_n_neighbors 22) does not fit at batch 64 on 48 GB — use --batch_size 32 (roughly doubles the wall-clock per epoch).

5. Environment variables

Variable	Applies to	Meaning
`SRA_MODULE`	all hosts	which module fills the insertion slot: `sra` (default) / `gameformer` / `c2f`
`GF_LEVELS`	GameFormer baseline	number of level-k reasoning levels (default 3)
`C2F_REFINE`	C2F baseline	refiner type: `gru` / `cnn`
`MOFLOW_DAMP`	MoFlow	cap on the residual norm (flow-field stability); `0` = off
`LED_TOP_N`	LED sport	neighbor budget override
`LED_NO_SIGMA`	LED sport	`1` = do not feed σ to the graph (required for no-σ ablation)

Example — swap SRA for the GameFormer baseline in the same slot:

SRA_MODULE=gameformer GF_LEVELS=2 CUDA_VISIBLE_DEVICES=0 python fm_nba_graph_v6.py \
  --cfg cfg/nba/cor_fm.yml --exp nba_gameformer \
  --batch_size 192 --epochs 150 --fm_in_scaling --tied_noise \
  --top_n_neighbors 5 --uncertainty_weight 0.01 --data_dir ./data/nba

6. The adapter contract

Any module in the slot must implement:

forward(y_emb,        # [B, K, A, D]      per-agent embedding of the current estimate
        y_abs,        # [B, K, A, T, 2]   current future-trajectory estimate
        t_emb,        # [B, D]            denoising-timestep embedding
        tau,          # [B]               noise level
        sigma_agent,  # [B, K, A, T] or None   per-step uncertainty
        agent_mask,   # [B, A] or None    validity mask (padded agents)
        ) -> [B, K, A, D]                 # gated residual, same shape as y_emb

It must tolerate K ∈ {1, 10, 20} and any A. MoFlow/models/interaction_baselines.py provides build_interaction_module(...), which honours SRA_MODULE and is the single place each host constructs its interaction module.

7. Known gotchas

Symptom	Cause / fix
`undefined symbol ... libcudnn_cnn_train.so.8` at backward	system CUDA shadows env cuDNN → set `LD_LIBRARY_PATH` (§1)
`ImportError: tensorboardX` (MID)	use `torch.utils.tensorboard` instead
`InvalidGitRepositoryError` (MoFlow)	`back_up_code_git` needs a git repo; already wrapped in try/except
MID sport ADE stuck ≈0.46	`--graph_lr_mult 3.0` → use `1.0`
Dense sport OOM at BS64	use `--batch_size 32`
"no-σ" LED-sport equals full SRA	missing `LED_NO_SIGMA=1`
Reported "best" keeps falling while eval worsens	`best ADE_MIN` is a monotone best-so-far tracker — inspect the per-eval `ADE_min(4.0s)` series to detect divergence