VBD - Versatile Behavior Diffusion

This repository contains the source code for the following paper:

Versatile Behavior Diffusion for Generalized Traffic Simulation

Zhiyu Huang^1,*, Zixu Zhang^2,*, Ameya Vaidya², Yuxiao Chen³, Chen Lv¹, and Jaime Fernández Fisac²

¹ Nanyang Technological University, ² Princeton University, ³ NVIDIA Research

Build Enviroment

To set up the required environment, execute the following commands:

conda env create -n vbd -f environment.yml
conda activate vbd 
# Install waymax from source
pip install git+https://github.com/waymo-research/waymax.git@main#egg=waymo-waymax
pip install -e .

Get Dataset

Download the Waymo Open Motion Dataset from Waymo Open Motion Dataset. Please use version V1.2 tf_example data to work with Waymax.

Data Preparation

Preprocess the dataset using the following command:

python script/extract_data.py \
    --data_dir /path/to/waymo_open_motion_dataset_dir \
    --save_dir /path/to/data_save_dir \
    --num_workers 16 \
    --save_raw # Extract Waymax Scenario for visualization

Make sure to process training and validation data separately.

Training

To train the VBD model, use the following command:

python script/train.py --cfg config/VBD.yaml --num_gpus 8

Update the VBD.yaml configuration file with appropriate parameters, such as the paths to the training and validation datasets.

Testing

Run the following command to test the model in closed-loop simulation:

python script/test.py --test_set /path/to/data --model_path ./train_log/VBD/model.pth --save_simulation

Ensure that both --test_set and --model_path parameters are provided. The test_set parameter should point to the raw tf_example data file.

Examples

Explore the following examples for different use cases:

• Unguided Scenario with extracted data
• Guided Scenario with extracted data

Citation

If you find our repo or our paper useful, please use the following citation:

@article{huang2024versatile,
  title={Versatile Behavior Diffusion for Generalized Traffic Agent Simulation},
  author={Huang, Zhiyu and Zhang, Zixu and Vaidya, Ameya and Chen, Yuxiao and Lv, Chen and Fisac, Jaime Fern{\'a}ndez},
  journal={arXiv preprint arXiv:2404.02524},
  year={2024}
}