GPUDrive

一款用C++编写的极速、数据驱动的自动驾驶仿真器。

亮点

• ⚡️ 通过Madrona引擎实现每秒100万帧的高速仿真，适用于智能体开发与评估。
• 🐍 提供Python绑定及基于torch和jax的gymnasium封装。
• 🏃‍➡️ 兼容Waymo开放运动数据集，包含超过10万个由人类驾驶演示的场景。
• 📜 可通过SB3和CleanRL / Pufferlib轻松获取现成的PPO实现。
• 👀 轻松配置仿真器和智能体视角。
• 🎨 多样化的智能体类型：车辆、骑行者和行人。

仿真器状态	智能体观测

更多详情请参阅我们的论文以及入门教程，其中将引导您完成基本使用。

安装

要构建GPUDrive，请确保已安装所有必需的依赖项，这些依赖项列于此处，包括CMake、Python和CUDA工具包。具体步骤如下。

安装完必要依赖后，克隆仓库（别忘了加上--recursive标志！）：

git clone --recursive https://github.com/Emerge-Lab/gpudrive.git
cd gpudrive

随后有两种方式来构建仿真器：

mkdir build
cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j # 核心数，例如32
cd ..

uv sync --frozen

pyenv virtualenv 3.11 gpudrive
pyenv activate gpudrive

pyenv local gpudrive

conda env create -f ./environment.yml
conda activate gpudrive

# macOS和Linux。
pip install -e .

# Windows。
pip install -e . -Cpackages.madrona_escape_room.ext-out-dir=<PATH_TO_YOUR_BUILD_DIR on Windows>

DOCKER_BUILDKIT=1 docker build --build-arg USE_CUDA=true --tag gpudrive:latest --progress=plain .

docker run --gpus all -it --rm --shm-size=20G -v ${PWD}:/workspace gpudrive:latest /bin/bash

通过导入仿真器来测试安装是否成功：

import madrona_gpudrive

为避免每次都在GPU模式下编译，可以设置以下环境变量，并指定自定义路径。例如，您可以将编译后的程序存储在名为gpudrive_cache的缓存中：

export MADRONA_MWGPU_KERNEL_CACHE=./gpudrive_cache

请注意，如果对C++代码进行任何更改，都需要删除缓存并重新编译。

  sudo apt install libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev mesa-common-dev libc++1

集成

入门指南

要开始使用，请参考以下入门资源：

• 我们的入门教程。这些教程大约需要30–60分钟完成，将引导您了解数据集、模拟器以及如何在模拟器中添加不同类型的智能体。
• 环境文档 提供了关于环境设置和支持功能的详细信息。

预训练策略

通过 🤗 huggingface_hub 中的 PyTorchModelHubMixin 类，提供了若干预训练策略。

• 最佳策略（10,000 场景）。来自论文《通过扩展自我博弈构建可靠的自动驾驶代理》的最佳策略在此处提供：链接。该策略基于 WOMD 训练数据集中随机采样的 10,000 个场景进行训练。

• 备选策略（1,000 场景）。一个基于 1,000 个场景训练的策略可在此处找到：链接

注意：这些模型是在 examples/experimental/config/reliable_agents_params.yaml 中定义的环境配置下训练的，更改环境或观测配置将影响性能。

使用方法

要加载预训练策略，可以使用以下代码：

from gpudrive.networks.late_fusion import NeuralNet

# 通过 huggingface_hub 加载预训练模型
agent = NeuralNet.from_pretrained("daphne-cornelisse/policy_S10_000_02_27")

更多详细信息请参阅教程 04。

数据集

下载数据集

• 数据集有两个版本：一个包含 1,000 个训练文件和 300 个测试/验证文件的迷你版，以及一个包含 10 万个唯一场景的大型数据集。
• 如果您希望下载完整数据集，请将下方路径中的“GPUDrive_mini”替换为“GPUDrive”。

pip install huggingface_hub

>>> from huggingface_hub import snapshot_download
>>> snapshot_download(repo_id="EMERGE-lab/GPUDrive_mini", repo_type="dataset", local_dir="data/processed")

huggingface-cli login

huggingface-cli download EMERGE-lab/GPUDrive_mini --local-dir data/processed --repo-type "dataset"

python data_utils/post_processing.py # 如果您使用了自定义下载路径，请使用 --help 参数

重新构建数据集

如果您希望手动生成数据集，GPUDrive 兼容完整的 Waymo 开放运动数据集，该数据集包含超过 10 万个场景。要下载新文件并为模拟器创建场景，请按照以下步骤操作。

waymo_open_dataset_motion_v_1_2_1/
│
├── uncompressed/
│   ├── lidar_and_camera/
│   ├── scenario/
│   │   ├── testing_interactive/
│   │   ├── testing/
│   │   ├── training_20s/
│   │   ├── training/
│   │   ├── validation_interactive/
│   │   └── validation/
│   └── tf_example/

gcloud auth login

gsutil -m cp -r gs://waymo_open_dataset_motion_v_1_2_1/uncompressed/scenario/validation/ data/raw

python data_utils/process_waymo_files.py '<raw-data-path>' '<storage-path>' '<dataset>'

pip install waymo-open-dataset-tf-2-12-0 trimesh[easy] python-fcl

python data_utils/process_waymo_files.py 'data/raw/' 'data/processed/' 'validation'
>>>
Processing Waymo files: 100%|████████████████████████████████████████████████████████████████| 150/150 [00:05<00:00, 28.18it/s]
INFO:root:Done!

后处理

• 运行 python data_utils/postprocessing.py 可以过滤掉损坏的文件，并取消 Hugging Face 的目录分组。

📜 引用 GPUDrive

如果您在研究中使用了 GPUDrive，请引用我们即将发表于 ICLR 2025 的论文：

@inproceedings{kazemkhani2025gpudrive,
      title={GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS},
      author={Saman Kazemkhani and Aarav Pandya and Daphne Cornelisse and Brennan Shacklett and Eugene Vinitsky},
      booktitle={Proceedings of the International Conference on Learning Representations (ICLR)},
      year={2025},
      url={https://arxiv.org/abs/2408.01584},
      eprint={2408.01584},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
}

贡献

如果您遇到 bug、发现缺少功能，或希望参与贡献，请随时创建 issue 或联系我们！我们非常期待您的加入！

GPUDrive 快速上手指南

GPUDrive 是一个基于 C++ 开发的超高速数据驱动驾驶模拟器，利用 Madrona 引擎可实现百万级 FPS 的仿真速度，支持 PyTorch 和 JAX，并兼容 Waymo Open Motion Dataset。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐), macOS, 或 Windows
• Python: >= 3.11
• CMake: >= 3.24
• CUDA Toolkit: >= 12.2 且 <= 12.4 (注意: 暂不支持 12.5 及以上版本)
  • 验证版本命令：nvcc --version
• 编译工具:
  • Linux: GCC/G++
  • macOS: Xcode Command Line Tools
  • Windows: Visual Studio C++ 工具

提示: 国内用户若遇到依赖下载缓慢，可配置 pip 使用清华或阿里镜像源。

安装步骤

1. 克隆仓库

务必使用 --recursive 参数以拉取子模块：

git clone --recursive https://github.com/Emerge-Lab/gpudrive.git
cd gpudrive

2. 构建模拟器 (二选一)

方案 A：手动编译 (推荐 Linux/macOS)

mkdir build
cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j # 根据 CPU 核心数调整，例如 make -j32
cd ..

Windows 用户: 请使用 Visual Studio 打开克隆的目录，利用内置的 CMake 功能进行构建。

方案 B：Docker 快速启动

确保已安装 Docker 和 NVIDIA Container Toolkit。

构建镜像:

DOCKER_BUILDKIT=1 docker build --build-arg USE_CUDA=true --tag gpudrive:latest --progress=plain .

运行容器:

docker run --gpus all -it --rm --shm-size=20G -v ${PWD}:/workspace gpudrive:latest /bin/bash

3. 配置 Python 环境

推荐使用 uv (最快)，也可选择 conda 或 pyenv。

使用 uv (推荐):

uv sync --frozen

使用 conda:

conda env create -f ./environment.yml
conda activate gpudrive

4. 安装 Python 包

如果是通过 uv 安装的可跳过此步（Linux/macOS），其他情况需执行：

Linux/macOS:

pip install -e .

Windows:

pip install -e . -Cpackages.madrona_escape_room.ext-out-dir=<PATH_TO_YOUR_BUILD_DIR>

(请将 <PATH_TO_YOUR_BUILD_DIR> 替换为实际的 build 目录路径)

5. 验证安装

尝试导入模块，若无报错则安装成功：

import madrona_gpudrive

优化建议: 为避免每次重复编译 GPU 内核，可设置缓存路径：

export MADRONA_MWGPU_KERNEL_CACHE=./gpudrive_cache

注意：若修改了 C++ 代码，需删除缓存并重新编译。

基本使用

加载预训练模型

GPUDrive 提供了基于 Hugging Face 的预训练策略。以下是最简单的加载示例：

from gpudrive.networks.late_fusion import NeuralNet

# 加载在 10,000 个场景上训练的最佳策略
agent = NeuralNet.from_pretrained("daphne-cornelisse/policy_S10_000_02_27")

获取数据集

模拟器兼容 Waymo Open Motion Dataset。您可以下载精简版（Mini）或完整版。

下载精简版数据集 (约 1000 个训练文件):

1. 安装依赖：

   pip install huggingface_hub
   

2. 使用 Python 下载：

   from huggingface_hub import snapshot_download
   
   # 下载到本地 data/processed 目录
   snapshot_download(repo_id="EMERGE-lab/GPUDrive_mini", repo_type="dataset", local_dir="data/processed")
   

下一步学习

• 入门教程: 查看 examples/tutorials 目录，包含约 30-60 分钟的完整引导，涵盖数据集处理、模拟器配置及智能体设置。
• 环境文档: 参考 gpudrive/env 了解详细的环境参数配置。
• 基线训练: 运行内置的 PPO 基线脚本开始训练：

  python baselines/ppo/ppo_sb3.py