Skip to content
/ mtp-go Public

Code for running the implementations proposed in: Westny, T., Oskarsson, J., Olofsson, B. and Frisk, E., "MTP-GO: Graph-Based Probabilistic Multi-Agent Trajectory Prediction with Neural ODEs", IEEE Transactions on Intelligent Vehicles

Notifications You must be signed in to change notification settings

westny/mtp-go

Repository files navigation

Graph-Based Probabilistic Multi-Agent Trajectory Prediction with Differentially Constrained Motion Models

Announcements

May 2024 📅

  • We have just released a toolbox designed for trajectory-prediction research on the Drone datasets!
  • It is made completely open-source here on GitHub. Make sure to check it out.

April 2023 📅

  • Update repository to include functionality to reproduce paper 2.
  • Migrate code to torch==2.0.0. Update requirements.

Description

mtp-go is a library containing the implementation for the papers:

  1. MTP-GO: Graph-Based Probabilistic Multi-Agent Trajectory Prediction with Neural ODEs (arXiv), published in IEEE Transactions on Intelligent Vehicles (TIV).
  2. Evaluation of Differentially Constrained Motion Models for Graph-Based Trajectory Prediction (arXiv), published in proceedings of IEEE 2023 Intelligent Vehicles Symposium (IV 2023).

Both papers are available in preprint format on ArXiv by the links above. All code is written using Python 3.11 using a combination of Pytorch logoPyTorch, PyG logoPyTorch Geometric, and Lightning logoPyTorch Lightning.

First page figure
If you found the content of this repository useful, please consider citing the papers in your work:
@article{westny2023mtp,
  title="{MTP-GO}: Graph-Based Probabilistic Multi-Agent Trajectory Prediction with Neural {ODEs}",
  author={Westny, Theodor and Oskarsson, Joel and Olofsson, Bj{\"o}rn and Frisk, Erik},
  journal={IEEE Transactions on Intelligent Vehicles},
  year={2023},
  volume={8},
  number={9},
  pages={4223-4236},
  doi={10.1109/TIV.2023.3282308}}
} 
@inproceedings{westny2023eval,
  title={Evaluation of Differentially Constrained Motion Models for Graph-Based Trajectory Prediction},
  author={Westny, Theodor and Oskarsson, Joel and Olofsson, Bj{\"o}rn and Frisk, Erik},
  booktitle={IEEE Intelligent Vehicles Symposium (IV)},
  pages={},
  year={2023},
  doi={10.1109/IV55152.2023.10186615}
}

Hardware requirements

The original implementation make use of a considerable amount of data (some gigabytes worth) for training and testing which can be demanding for some setups. For you reference all code has been tried and used on a computer with the following specs:

* Processor: Intel® Xeon(R) E-2144G CPU @ 3.60GHz x 8
* Memory: 32 GB
* GPU: NVIDIA Corporation TU102 [GeForce RTX 2080 Ti Rev. A]

Usage

Most of the necessary building blocks to implement MTP-GO is contained within the models folder. The main files of interest are:

In gru_gnn.py the complete encoder-decoder model implementation is contained. This includes a custom GRU cell implementation that make use of layers based on Graph Neural Networks.

In motion_models.py the implementations of the various motion models are contained, including the neural ODEs, used to learn road-user differential constraints. This is also where you will find functions used to perform the Jacobian calculations of the model.

In this work, Lightning logoPyTorch Lightning was used to implement the training and testing behavior. Since most of the functionality is still implemented using Pytorch logoPyTorch, you are not restricted to using lightning, but it is recommended given the additional functionality. In base_mdn.py the lightning-based abstraction of MTP-GO is contained. This module is used to implement batch-wise forward and backward passes as well as to specify training and testing behavior.

Assuming data is available, training a model based on MTP-GO is as easy as running train.py in an environment with the necessary libraries installed, e.g.:

python train.py --dataset rounD --motion-model neuralode --n-workers 8 --hidden-size 128

To learn more about the objective-scheduling algorithm described in the paper as well as the loss functions used, see losses.py.

Schematics of MTP-GO

Data sets

For model training and evaluation, the highD, rounD, and inD were used. The data sets contain recorded trajectories from different locations in Germany, including various highways, roundabouts, and intersections. The data includes several hours of naturalistic driving data recorded at 25 Hz of considerable quality. They are freely available for non-commercial use but do require applying for usage through the links above.

rounD.gif

Preprocessing

Assuming that you have been granted access to any of the above-mentioned data sets, proceed by moving the unzipped content (folder) into a folder named data_sets (you have to create this yourself) on the same level as this project. The contents may of course be placed in any accessible location of choice but do then require modifications of the preprocessing scripts (see the head of the .py files).

Methods of preprocessing are contained within Python scripts. Executing them may be done from a terminal or IDE of choice (from within this project folder), for example:

python rounD_preprocess.py

The output of the preprocessing scripts will be sent to a sub-folder with the name of the data set within the data folder in this project. Each data sample refers to a traffic sequence and is given a unique index used for easy access.

❗ A word of caution, by this approach, a lot of files are created that could be demanding for some systems.

To make sure unsaved work is not deleted, you will be prompted on your preferred cause of action should the folders already exist: either overwriting existing data or aborting the process.

🚩 The methods of preprocessing are by no means optimized for speed or computational efficiency. The process could take several hours depending on the data set and available hardware.

License

Creative Commons

Inquiries

Questions about the paper or the implementations found in this repository should be sent to theodor.westny [at] liu.se.

About

Code for running the implementations proposed in: Westny, T., Oskarsson, J., Olofsson, B. and Frisk, E., "MTP-GO: Graph-Based Probabilistic Multi-Agent Trajectory Prediction with Neural ODEs", IEEE Transactions on Intelligent Vehicles

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages