Promptable Closed-loop Traffic Simulation
Shuhan Tan1, Boris Ivanovic2, Yuxiao Chen2, Boyi Li2, Xinshuo Weng2, Yulong Cao2, Philipp Krähenbühl1, Marco Pavone2
1UT Austin, 2 NVIDIA
Conference on Robot Learning (CoRL), 2024
Webpage | Video | Paper | Model Colab | Data Colab | ProSim-Instruct-520k Dataset | Mini Dataset
21 Oct, 2024: Initial code and dataset release. Try out the online demo!
We have also released raw data of ProSim-Instruct-520k. Stay tuned for the training pipeline!
We provide an online demo in Colab. You can try it without any local installation. This demo includes:
- Loading and visualize ground-truth data.
- Use ProSim to conduct unconditional closed-loop simulation for all agents.
- Use ProSim to conduct promptable closed-loop simulation for all agents.
Please also refer to the code inside for demonstrations of basic model usage.
Note that this demo does not include text prompting, which requires Llama3-8B model which cannot be runned in standard colab environment.
Check the next section to see how to setup your own environment for text prompting.
# Clone the code to local
git clone https://github.com/Ariostgx/ProSim.git
cd ProSim
# Create virtual environment
conda create -n prosim python=3.8
conda activate prosim
# Install dependencies
bash install_local_env.shWe provide a demo dataset and a pretrained model for quick start for text prompting.
Please download the example model checkpoint prosim_demo_model.ckpt from here. Then put it into prosim_demo/ckpt folder.
Please download the Llama3-8B-Instruct model from Meta's model hub.
We provide a demo notebook prosim_demo/text_prompt_inference.ipynb for text prompting. Just replace the LLAMA_PATH with your local path to Llama3-8B-Instruct model.
We provide a colab demo for loading the ProSim-Instruct-520k dataset here. In this demo, we only load a subset of the dataset from here as a quick start. Check out this demo without any local installation or data preparation!
In the demo, we go through how to load the ProSim-Instruct-520k dataset paired with Waymo Open Dataset. Feel free to checkout the code in the demo notebook for details of data loading!
We release the raw data of ProSim-Instruct-520k here. It contains all the text prompts and motion tags for all scenarios in the Waymo Open Dataset training and validation sets. Feel free to download and preview the data.
Once the dataset is downloaded, unzip it and replace the PROSIM_INSTRUCT_520K_DATA_DIR parameter in prosim/config/path_cfg.py with the path to the unzipped folder.
We use Trajdata to process the Waymo Open Dataset. Please follow the instructions in the Trajdata repo to download the Waymo Open Dataset and process it with Trajdata.
Once the dataset is processed by Trajdata, set TRAJDATA_CACHE_DIR in prosim/config/path_cfg.py to the path of the Trajdata cache folder. Also set WAYMO_DATA_DIR to the path of the downloaded Waymo Open Dataset.
We provide a demo notebook prosim_demo/load_prosim_instruct_520k.ipynb for loading the ProSim-Instruct-520k dataset. Feel free to checkout the code in the demo notebook for details of data loading!
We use code in LCTGen and Trajdata for data processing and visualization.
- Release dataset labeling tools.
- Release training pipeline.
- Release dataset and data loading code.
- Initial repo & demo
@inproceedings{
tan2024promptable,
title={Promptable Closed-loop Traffic Simulation},
author={Tan, Shuhan and Ivanovic, Boris and Chen, Yuxiao and Li, Boyi and Weng, Xinshuo and Cao, Yulong and Kr{\"a}henb{\"u}hl, Philipp and Pavone, Marco},
booktitle={8th Annual Conference on Robot Learning},
year={2024},
}