# 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:

```python
# 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

```bash
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:

```bash
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`](sample_data/README.md):

```bash
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

```bash
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

```bash
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

```bash
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:

```bash
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:

```python
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 |