This repository contains the official implementation of SMART: Scalable Multi-agent Real-time Motion Generation via Next-token Prediction. SMART is a novel autonomous driving motion generation paradigm that models vectorized map and agent trajectory data into discrete sequence tokens.
To set up the environment, you can use conda to create and activate a new environment with the necessary dependencies:
conda env create -f environment.yml
conda activate SMART
pip install -r requirements.txt
If you encounter issues while installing pyg dependencies, execute the following script:
bash install_pyg.sh
Alternatively, you can configure the environment in your preferred way. Installing the latest versions of PyTorch, PyG, and PyTorch Lightning should suffice.
Step 1: Download the Dataset
Download the Waymo Open Motion Dataset (scenario protocol
format) and organize the data as follows:
SMART
├── data
│ ├── waymo
│ │ ├── scenario
│ │ │ ├──training
│ │ │ ├──validation
│ │ │ ├──testing
├── model
├── tools
Step 2: Install the Waymo Open Dataset API
Follow the instructions here to install the Waymo Open Dataset API.
Step 3: Preprocess the Dataset
Preprocess the dataset by running:
python data_preprocess.py --input_dir ./data/waymo/scenario/training --output_dir ./data/waymo_processed/training
The first path is the raw data path, and the second is the output data path.
The processed data will be saved to the data/waymo_processed/
directory as follows:
SMART
├── data
│ ├── waymo_processed
│ │ ├── training
│ │ ├── validation
│ │ ├──testing
├── model
├── utils
To train the model, run the following command:
python train.py --config ${config_path}
The default config path is configs/train/train_scalable.yaml
. Ensure you have downloaded and prepared the Waymo data for training.
To evaluate the model, run:
python eval.py --config ${config_path} --pretrain_ckpt ${ckpt_path}
This will evaluate the model using the configuration and checkpoint provided.
To comply with the WOMD participation agreement, we will release the model parameters of a medium-sized model not trained on Waymo data. Users can fine-tune this model with Waymo data as needed.
Our model achieves the following performance on the Waymo Open Motion Dataset Sim Agents Challenge:
Model name | Metirc Score |
---|---|
SMART-tiny | 0.7591 |
SMART-large | 0.7614 |
SMART-zeroshot | 0.7210 |
If you find this repository useful, please consider citing our work and giving us a star:
@article{wu2024smart,
title={SMART: Scalable Multi-agent Real-time Simulation via Next-token Prediction},
author={Wu, Wei and Feng, Xiaoxin and Gao, Ziyan and Kan, Yuheng},
journal={arXiv preprint arXiv:2405.15677},
year={2024}
}
Special thanks to the QCNET repository for providing valuable reference code that significantly influenced this work.
All code in this repository is licensed under the Apache License 2.0.