This repository contains the official code for the paper:
Multi-agent Trajectory Prediction with Fuzzy Query Attention
Nitin Kamra, Hao Zhu, Dweep Trivedi, Ming Zhang and Yan Liu
Advances in Neural Information Processing Systems (NeurIPS), 2020
The code has been written for python3.5
and above. Apart from the python3
standard libraries, the following other packages are required:
- dill>=0.3.2
- numpy>=1.19.2
- torch>=1.4.0
- matplotlib==3.0.0
- pandas>=1.1.3
- torch-scatter==1.4.0 (https://github.com/rusty1s/pytorch_scatter)
The directory structure of the repository is as follows:
- config: This directory contains the configuration files for the provided models. These config files contain important code settings and model hyperparameters and are needed as arguments to
src/run.py
script which trains and evaluates all models. The directory contains one base config filecfg.py
which is extended by other model specific config files. - data: This folder contains the file
TrajectoryDataset.py
which provides theTrajectoryDataset
class to convert and load all processed datasets in the requiredtorch
format. It also contains compressed and processed versions of all datasets used in the paper. We do not provide the raw datasets and the pre-processing scripts since the datasets are huge in size and the pre-processing scripts only add additional clutter to the repository. However, the raw datasets can be downloaded from their respective online sources if needed.- While we provide our copy of
ETH-UCY
,NGsim
,Collisions
andCharges
, we do not provide theNBA
dataset since it requires special access and needs to be requested independently from the official sources if required. - Each of our processed datasets contains three folders:
train
,test
andvalidation
with their respective scene files. Each scene file contains the trajectories of all agents in the format:Time-step, AgentID, Normalized X-coord, Normalized Y-coord
. - Each dataset also contains a file
scale.txt
which contains the normalization scale which is later used to de-normalize model predictions before calculating final evaluation metrics. - Lastly each dataset contains a python file
<dataset>.py
which is used as an argument to thesrc/run.py
script to load the dataset.
- While we provide our copy of
- models: This folder contains some of the models used in the paper including our
FQA
architecture. From amongst the baselines, we only includeVanillaLSTM
and all the ablations ofFQA
(namelyInertiaModel
which comes inmodels
directory and others which can be called by setting appropriate flags inconfig/cfg_FQA.py
while training the FQA model). For the other baselines, we recommend contacting the original authors of the papers. Most of the common functionality and training procedure is provided in the abstractBaseModel
class in themodels/BaseModel.py
file. All other models inherit theBaseModel
class and add or override methods to provide additional functionality. - src: This folder contains all the runnable scripts:
run.py
is the major script which enables training and evaluation of all models.generate_eval_metrics.py
: Oncerun.py
has been called in both train and test modes on a dataset-model combination, this script de-normalizes the model's predictions on the dataset and produces final evaluation metrics.viz.py
produces the joint visualization from all models on a specific dataset.display_metrics.py
: This script collects scores from multiple training runs of all specified models on all specified datasets and aggregates them to produce mean, min, max and stdev for all scores.exps.sh
: This shell script contains commands to easily reproduce most of the results presented in the paper. The commands exemplify the calling format forrun.py
along with providing all required hyperparameters to be specified for each model/dataset combination.
- utils: This folder contains utility modules for plotting, model design, argument parsing etc.
- Install all dependencies mentioned above.
- Clone the repo and enter the main project directory:
git clone https://github.com/nitinkamra1992/FQA.git
cd FQA
- Unzip all datasets in the
data
folder:
cd data
unzip charged.zip
unzip collisions.zip
unzip ethucy.zip
unzip ngsim.zip
cd ..
- Open
src/exps.sh
and uncomment line 7 to set the value of run IDi
. We set it to1
by default but this can be set to any random integer since it controls the random seed for any run. The results in the paper have been produced by averaging over 5 runs withi = {1,2,3,4,5}
for reproducibility. Note that this integer also controls the name of the output directory for most experiments to keep results from different runs from over-writing each other. - Now uncomment the remaining experiment lines one-by-one and run the
src/exps.sh
script with each of them to run all experiments one-by-one. - Note that within most sections of
src/exps.sh
, there are commented subsections of the form:Train
,Eval_test
,Eval metrics
. For any model/dataset combination these three sub-sections must be run sequentially since they rely on the previous sub-section's results. But two different sections can be run independently in parallel. For instance, you can un-comment lines 49 and 57 simultaneously to launch two training experiments in parallel: one to trainVanillaLSTM
onCollisions
data and another to train it onETH-UCY
data. - You can run multiple experiments in parallel as long as your GPU memory permits. You can also run multiple experiments in parallel on different GPU devices (or CPU) by specifying the device in the arguments to
src/run.py
with--device
argument.