Home

Awesome

<div style="text-align: center;"> <img src="https://github.com/vita-epfl/UniTraj/blob/main/docs/assets/unitraj.gif" alt="Demo" width="300"> </div>

UniTraj: A Unified Framework for Scalable Vehicle Trajectory Prediction

PWC

Website | Paper

πŸ’‘UniTraj allows users to train and evaluate trajectory prediction models from real-world datasets like Waymo, nuPlan, nuScenes and Argoverse2 in a unified pipeline.

system

πŸ”₯Powered by Hydra, Pytorch-lightinig, and WandB, the framework is easy to configure, train and log.

system

πŸ“° News & Updates

Dec. 2024

Nov. 2024

Sep. 2024

Jul. 2024

Mar. 2024

πŸ›  Quick Start

  1. Create a new conda environment
conda create -n unitraj python=3.9
conda activate unitraj
  1. Install ScenarioNet: https://scenarionet.readthedocs.io/en/latest/install.html

  2. Install Unitraj:

git clone https://github.com/vita-epfl/UniTraj.git
cd unitraj
pip install -r requirements.txt
python setup.py develop

You can verify the installation of UniTraj via running the training script:

python train.py method=autobot

The model will be trained on several sample data.

Code Structure

There are three main components in UniTraj: dataset, model and config. The structure of the code is as follows:

unitraj
β”œβ”€β”€ configs
β”‚Β Β  β”œβ”€β”€ config.yaml
β”‚Β Β  β”œβ”€β”€ method
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ autobot.yaml
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ MTR.yaml
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ wayformer.yaml
β”œβ”€β”€ datasets
β”‚Β Β  β”œβ”€β”€ base_dataset.py
β”‚Β Β  β”œβ”€β”€ autobot_dataset.py
β”‚Β Β  β”œβ”€β”€ wayformer_dataset.py
β”‚Β Β  β”œβ”€β”€ MTR_dataset.py
β”œβ”€β”€ models
β”‚Β Β  β”œβ”€β”€ autobot
β”‚Β Β  β”œβ”€β”€ mtr
β”‚Β Β  β”œβ”€β”€ wayformer
β”‚Β Β  β”œβ”€β”€ base_model
β”œβ”€β”€ utils

There is a base config, dataset and model class, and each model has its own config, dataset and model class that inherit from the base class.

Pipeline

1. Data Preparation

UniTraj takes data from ScenarioNet as input. Process the data with ScenarioNet in advance.

2. Configuration

UniTraj uses Hydra to manage configuration files.

Universal configuration file is located in unitraj/config/config.yaml. Each model has its own configuration file in unitraj/config/method/, for example, unitraj/config/method/autobot.yaml.

The configuration file is organized in a hierarchical structure, and the configuration of the model is inherited from the universal configuration file.

Configuration Example

Please refer to config.yaml and method/autobot.yaml for more details.

2. Train

python train.py

3. Evaluation

  1. In config.yaml, set the ckpt_path to the path of the trained model and val_data_path to the validation data path.
  2. (Optional) In config.yaml, set eval_waymo or eval_nuscenes to True if you want to evaluate the model with Waymo or nuScenes official evaluation tool. (Install waymo-open-dataset and nuscenes-devkit first)
  3. Runpython evaluation.py

4. Dataset Analysis

python data_analysis.py

Contribute to UniTraj

Implement a new model

  1. Create a new config file in unitraj/config/ folder, for example, unitraj/config/new_model.yaml
  2. (Optional) Create a new dataset class in unitraj/datasets/ folder, for example, unitraj/datasets/new_dataset.py, and inherit unitraj/dataset/base_dataset.py
  3. Create a new model class in unitraj/model/ folder, for example, unitraj/model/lanegcn.py, and inherit from pl.LightningModule

Dataset Structure

Scenario Metadata

Object Trajectories

Centered Objects

Map Information

Future State Predictions

Ground Truth Data

Additional Metadata


For citation:

@article{feng2024unitraj,
  title={UniTraj: A Unified Framework for Scalable Vehicle Trajectory Prediction},
  author={Feng, Lan and Bahari, Mohammadhossein and Amor, Kaouther Messaoud Ben and Zablocki, {\'E}loi and Cord, Matthieu and Alahi, Alexandre},
  journal={arXiv preprint arXiv:2403.15098},
  year={2024}
}