/ECCO

Primary LanguagePython

Trajectory Prediction using Equivariant Continuous Convolution (ECCO)

This is the codebase for the ICLR 2021 paper Trajectory Prediction using Equivariant Continuous Convolution, by Robin Walters, Jinxi Li and Rose Yu.

Thumbnail

Installation

This codebase is trained on Python 3.6.6+. For the usage of argoverse dataset, argoverse-api is required. We recommand the reader to follow their guide to install the complete api and datasets. Other requirements include:

  • numpy>=1.18.3
  • pandas>=1.0.3
  • torch==1.5.0

Dependency can be installed using the following command:

pip install -r requirements.txt

Data Preparation

Original data could be downloaded from argoverse. To generate the training and validation data

  1. Set the path to argoverse_forecasting in datasets/preprocess_data.py scripts.
  2. Run the script
python preprocess_data.py

The data will be stored in path/to/argoverse_forecasting/train(val)/lane_data.

Data Download

If you want to skip the data generation part, the link to preprocessed data will be provided soon.

Model Training

Here are commands for training the model.

For -ECCO, run the following command

python train.py --dataset_path /path/to/argoverse_forecasting/ --rho1 --model_name rho_1_ecco

For -ECCO, run the following command

python train.py --dataset_path /path/to/argoverse_forecasting/ --rho-reg --model_name rho_1_ecco

Citation

If you find this repository useful in your research, please cite our paper:

@article{Walters2021ECCO,
  title={Trajectory Prediction using Equivariant Continuous Convolution},
  author={Robin Walters and Jinxi Li and Rose Yu},
  journal={International Conference on Learning Representations},
  year={2021},
}