COMPARING GRAPH-BASED AND NON-GRAPH METHODS FOR MULTI-AGENT NBA TRAJECTORY FORECASTING

Adrian Fluder, Marine Mailland, Julien Ropert, Edgar Desnos - EPFL

EE-452: Network Machine Learning - Spring 2025/2026

This repository contains our implementation for the NBA trajectory forecasting project. Given a context sequence of C = 8 observed time steps, the goal is to predict the future trajectories of all 11 entities (10 players + 1 ball) over a horizon of H = 12 steps, for a total window of 20 frames.

The evaluation metric is the ADE (Average Discrepancy Error), defined as the average Mean Square Error over the predicted horizon H:

ADE = (1 / H*N) * sum_{t=C+1}^{C+H} sum_n MSE(z_hat_t^n, z_t^n)

where z_hat_t^n = (x_hat, y_hat) is the predicted position of entity n at time t, and z_t^n the corresponding ground truth.

---

Method

Our approach centers on a modular intent-target refinement architecture in which trajectory prediction is decomposed into two stages: (1) predicting a spatial intent target for each entity at the end of the horizon, and (2) refining the full trajectory as a correction over a straight-line base path from the last observed position to that target.

Feature extraction. For each entity at each context step we construct a feature vector combining absolute position, position relative to the last observed frame, velocity, acceleration, speed, acceleration magnitude, static entity features (is-player, team), team one-hot encoding, ball-relative position, and a temporal index τ ∈ [−1, 0]. These are projected into a shared D-dimensional token space and enriched with learned time and team embeddings.

Context encoder. A standard Transformer encoder operates over all C × N tokens jointly, allowing every entity at every timestep to attend to all others. A learned temporal attention pool then collapses the context axis per entity, producing one D-dimensional state per entity. A global fusion step further enriches each entity state with scene-level context.

Target encoder (swappable). The per-entity states are passed through a target encoder whose output drives the intent-target prediction head. This is the main experimental variable of the project and can be instantiated as one of:

• Transformer (target_encoder_mode = "transformer") - a standard multi-head self-attention encoder with FFN blocks, treating all N entities symmetrically. This is our primary non-graph baseline.
• Graph-based encoders ("graph_transformer", "gat", "gatv2") - PyTorch Geometric convolutional layers with adjacency built from pairwise entity distances. These were explored but did not consistently outperform the transformer baseline.
• EGNN (egnn_refinement model) - an E(n)-equivariant graph neural network that maintains and updates geometric coordinate channels alongside node features, giving the model an explicit equivariant geometric inductive bias. The EGNN model is implemented as a separate architecture (EGNNRefinement) sharing the same context encoder and correction decoder backbone as the transformer baseline.

All target encoder variants share a common interface and can be swapped without any change to the surrounding pipeline.

Base trajectory and correction. The predicted intent target defines a straight-line base trajectory interpolating from the last observed position to the target over H steps. A Transformer decoder with cross-attention over the context memory then predicts per-step corrections, accumulated via cumulative sum, which are added to the base to produce the final trajectory.

Training. The main loss is a horizon-weighted MSE over normalized positions with an auxiliary velocity term. An optional auxiliary loss directly supervises intermediate target encoder states, providing gradient signal to all layers and stabilising training.

---

Dependencies

torch
torchvision
numpy
pandas
matplotlib
torch_geometric
ipython

Two requirements files are provided:

• requirements-conda.txt — base packages (Python, PyTorch, NumPy, etc.) installed via conda
• requirements-pip.txt — PyG and companion packages installed via pip after the conda environment is active

Step-by-step setup

# 1. Create and activate the conda environment
conda create --name nba_traj --file requirements-conda.txt -c pytorch -c conda-forge
conda activate nba_traj

# 2. Install PyG
pip install -r requirements-pip.txt

# 3. Verify the install
python -c "import torch; import torch_geometric; print('OK')"

Apple Silicon (M1/M2/M3) note: PyTorch Geometric uses scatter_reduce ops that are not yet implemented for MPS. If you hit a NotImplementedError on MPS, set this environment variable before running any script:

export PYTORCH_ENABLE_MPS_FALLBACK=1

Add it to your ~/.zshrc to make it permanent. The affected ops fall back to CPU automatically; everything else continues to run on MPS.

---

Dataset

The dataset is the NBA highlight dataset from SocialVAE (Xu et al., 2022), as provided by the course (see NML_Project_Proposal.pdf for the full description). It consists of more than 6K highlight plays from the 2015–2016 NBA season.

Each sequence is stored as a PyTorch tensor (.pt) of shape L x N x D, where:

• L - sequence length, which varies per play, usually between 20 and 100+ frames
• N = 11 - number of entities, corresponding to 10 players and 1 ball
• D = 4 - feature vector [x, y, isplayer, team]

The team field encodes membership: -1 = Team A, 0 = Ball, 1 = Team B. Coordinates are in feet, centered at the court midpoint (x in [-47.5, 47.5], y in [-25, 25]).

---

Dataset preparation

The dataset can be prepared automatically by running:

python data_creation.py

This script downloads the Kaggle competition files using kagglehub, then creates the following local folder structure:

data/
├── train/
│   └── train/          <- training sequences used for full training
├── train_subset/
│   └── train_subset/   <- 500 training sequences for quick experiments
├── val/
│   └── val/            <- 50 validation sequences
└── test/
    └── test/           <- official Kaggle test sequences

The split is created from the official Kaggle train/train/ folder as follows:

• the training files are shuffled with seed 42;
• the first 50 files are copied to data/val/val/;
• the next 500 files are copied to data/train_subset/train_subset/;
• all files except the validation files are copied to data/train/train/;
• the official Kaggle test files are copied from test/test/ to data/test/test/.

The script also copies the following CSV files into the data/ folder when they are available:

data/train.csv
data/test.csv
data/sample_submission.csv

Before running the script, make sure that your Kaggle API token is correctly set up and that you have accepted the Kaggle competition rules.

---

Manual setup

If data_creation.py does not work, the dataset can be prepared manually.

First, download the competition files from Kaggle:

kaggle competitions download -c nml-2026

Alternatively, download the files directly from the Kaggle competition page and extract the archive manually.

After extraction, the downloaded dataset should have the following structure:

nml-2026/
├── train/
│   └── train/          <- official training sequences (.pt)
├── test/
│   └── test/           <- official test sequences (.pt)
├── train.csv
├── test.csv
└── sample_submission.csv

Then create the following folder structure inside the project:

data/
├── train/
│   └── train/
├── train_subset/
│   └── train_subset/
├── val/
│   └── val/
└── test/
    └── test/

The files must then be copied manually according to the same logic as the automatic script.

1. Validation set

Randomly select 50 .pt files from:

nml-2026/train/train/

and copy them into:

data/val/val/

These files are used as the local validation set.

2. Training subset

Randomly select 500 additional .pt files from:

nml-2026/train/train/

and copy them into:

data/train_subset/train_subset/

This folder is only used for quick experiments and debugging.

3. Full training set

Copy all training .pt files except the 50 validation files into:

data/train/train/

This folder is used for full training runs.

4. Test set

Copy all official Kaggle test .pt files from:

nml-2026/test/test/

into:

data/test/test/

These files correspond to the held-out Kaggle test set.

5. CSV files

Finally, copy the CSV files into the data/ folder:

nml-2026/train.csv              -> data/train.csv
nml-2026/test.csv               -> data/test.csv
nml-2026/sample_submission.csv  -> data/sample_submission.csv

The final local structure should look like this:

data/
├── train/
│   └── train/
│       └── *.pt
├── train_subset/
│   └── train_subset/
│       └── *.pt
├── val/
│   └── val/
│       └── *.pt
├── test/
│   └── test/
│       └── *.pt
├── train.csv
├── test.csv
└── sample_submission.csv


The automatic script uses a fixed random seed (`42`) to make the split reproducible.
For exact reproducibility, it is recommended to use `data_creation.py` rather than
creating the split manually.

The default paths in config/train_config.py point to the full training set (data/train/train). Switch to the subset for quick iteration by passing --train data/train_subset/train_subset or editing train_path in config/train_config.py.

---

Code Usage

Train a model from scratch

# Default (full training data)
python main.py --model target_refinement_baseline

# Specify all paths explicitly
python main.py --model target_refinement_baseline \
               --train data/train/train \
               --val   data/val/val \
               --test  data/test/test \
               --ckpt  results

Results are always written to results/<model_name>/ (or results/<model_name>_<tag>/ if --tag is used). Always pass --ckpt results — the subdirectory is created automatically from the model name and tag.

Use --tag to run multiple experiments without overwriting each other:

# Two runs of the same model with different settings, kept separate
python main.py --model target_refinement_baseline --tag transformer_full --ckpt results
python main.py --model target_refinement_baseline --tag egnn_run1       --ckpt results

Checkpoints, training history CSV, training curves, trajectory plots, and attention heatmaps (for transformer and EGNN models) are written to the --ckpt directory (or results/<model>/ by default) automatically.

Evaluate an existing checkpoint (no retraining)

python main.py --model target_refinement_baseline \
               --test  data/test/test \
               --ckpt  results \
               --no-train

Available models

--model	Description
target_refinement_baseline	Intent-target refinement with Transformer target encoder (non-graph baseline)
egnn_refinement	Intent-target refinement with EGNN target encoder (graph-based)
nba_rnn	GRU sequence model (simple baseline)

See model_registry.py for the full registry and instructions on adding new models.

---

Project Structure

├── main.py                                         <- entry point (train / eval / submit)
├── metrics.py                                      <- loss functions and ADE metric
├── plots.py                                        <- figure generation (curves, trajectories,
│                                                      attention heatmaps)
├── model_registry.py                               <- model name → (class, config) registry
├── requirements.txt
│
├── config/
│   ├── train_config.py                             <- shared training hyper-parameters
│   ├── target_refinement_baseline_config.py        <- Transformer model config
│   ├── egnn_refinement_config.py                   <- EGNN model config
│   └── nba_rnn_config.py                           <- RNN baseline config
│
├── data/
│   ├── data_config.py                              <- data augmentation settings
│   ├── train/train/
│   ├── train_subset/train_subset/
│   ├── val/val/
│   └── test/test/
│
├── models_src/
│   ├── target_refinement_baseline_src.py           <- Transformer architecture
│   ├── egnn_refinement_src.py                      <- EGNN architecture
│   └── nba_rnn_src.py                              <- GRU baseline architecture
│
└── results/
    └── <model_name>/
        ├── <model_name>_ckpt.pth
        ├── training_history.csv
        ├── solution.csv
        ├── <model_name>_training_curves.pdf
        ├── <model_name>_best_2_trajectories.pdf
        ├── <model_name>_worst_2_trajectories.pdf
        └── <model_name>_attn_sample*_layer*.pdf    <- attention heatmaps

---

Adding a New Model

1. Architecture - create models_src/<model_name>_src.py exposing:

   def build_model(cfg, court_min_norm, court_max_norm) -> nn.Module

2. Config - create config/<model_name>_config.py exposing a MODEL_CONFIG dataclass with at minimum context_size, horizon_size, loss_name, loss_kwargs.
3. Register - add one entry to REGISTRY in model_registry.py.
4. Run - pass --model <model_name> to main.py.

---

Config File Reference

File	Controls
config/train_config.py	lr, weight_decay, epochs, batch_size_*, grad_clip, scheduler params, window sampling counts
config/<model>_config.py	Hidden dim, layers, heads, dropout, correction_scale, loss name and loss kwargs, auxiliary loss settings
data/data_config.py	Flip probability, team-swap probability, entity permutation probability, noise std, geometric augmentation settings

---

Kaggle Competition

Predictions for the held-out test set are written to results/<model_name>/solution.csv in the format expected by the course Kaggle leaderboard. The file is generated automatically at the end of every training run, or on demand with --no-train (requires an existing checkpoint).