WarpDrive：在GPU上实现极快的端到端单代理或多代理深度强化学习

WarpDrive是一个灵活、轻量级且易于使用的开源强化学习(RL)框架，可在单个或多个GPU(图形处理单元)上实现端到端多代理RL。

利用GPU的极致并行能力，WarpDrive 能够实现比CPU模拟+GPU模型实现快几个数量级的RL。它非常高效，因为它避免了CPU和GPU之间的来回数据复制，并能跨多个代理和多个环境副本并行运行模拟。这些共同使用户能够运行数千甚至数百万个并发模拟，并在极大规模的经验批次上进行训练，实现比基于CPU的对应方案至少100倍的吞吐量。

下表直观展示了Warpdrive的主要特性和在各个维度上的可扩展性。

	支持	并发性	版本
环境	单个 ✅ 多个 ✅	每GPU 1至1000个	1.0+
代理	单个 ✅ 多个 ✅	每环境1至1024个	1.0+
代理	跨块多个 ✅	每块1024个	1.6+
离散动作	单个 ✅ 多个 ✅	-	1.0+
连续动作	单个 ✅ 多个 ✅	-	2.7+
在线策略梯度	A2C ✅ PPO ✅	-	1.0+
离线策略梯度	DDPG ✅	-	2.7+
自动缩放	✅	-	1.3+
分布式模拟	1 GPU ✅ 2-16 GPU节点 ✅	-	1.4+
环境后端	CUDA C ✅	-	1.0+
环境后端	CUDA C ✅ Numba ✅	-	2.0+
训练后端	Pytorch ✅	-	1.0+

环境

"捉迷藏"游戏：在"捉迷藏"游戏中，追捕者试图追赶并标记逃跑者。这些游戏相当复杂，适合用于基准测试和测试，涉及线程同步、共享内存、数千个交互代理的高维索引等。下面我们展示了使用WarpDrive为不同追捕者:逃跑者速度比训练的多代理RL策略。这些环境可以在单个GPU上以每秒数百万步的速度运行，并在仅仅几个小时内完成训练！

基于WarpDrive开发了复杂的两级多代理环境，如新冠疫情环境和气候变化环境，您可以在真实世界问题和合作中查看示例。

经典控制：我们包含了[gym.classic_control](https://github.com/salesforce/warp-drive/blob/master/ https://github.com/openai/gym/tree/master/gym/envs/classic_control)中的环境。在WarpDrive中，单代理是多代理环境的特殊情况。由于每个环境只有一个代理，可扩展性更高。

催化反应路径：我们包含了将量子密度泛函理论转换为强化学习表示的环境，使得能够从嘈杂的化学系统中自动搜索最优化学反应路径。您可以在真实世界问题和合作中查看示例。

吞吐量、可扩展性和收敛性

多代理

下面，我们比较了在N1 16-CPU节点和单个A100 GPU（使用WarpDrive）上的训练速度，环境为100名逃跑者和5名追捕者的捉迷藏游戏。在相同的环境配置和训练参数下，使用GPU上的WarpDrive比CPU快约10倍。两种情况都并行运行60个环境副本。在CPU节点上使用更多环境是不可行的，因为数据复制会变得过于昂贵。使用WarpDrive，可以将环境副本数量至少增加10倍，实现更快的训练。

单代理

下面，我们比较了在单个A100 GPU（使用WarpDrive）上的训练速度，分别针对（上）Cartpole-v1和（下）Acrobot-v1，并行运行10、100、1K和10K个环境副本，训练3000个周期（超参数相同）。您可以看到通过WarpDrive扩展的大量环境实现了惊人的收敛性和速度。

代码结构

WarpDrive提供了CUDA（或Numba）+ Python框架和便利工具，使您能够快速构建快速、灵活且大规模分布式的多代理RL系统。下图从底层到顶层说明了WarpDrive的设计和组件概览。用户只需在CUDA环境层编写CUDA或Numba步进函数，而其余部分是纯Python接口。我们提供了循序渐进的教程，帮助您掌握工作流程。

Python接口

WarpDrive提供了工具，只需几行代码就能快速构建和训练多代理RL系统。以下是一个简短示例，用于训练追捕者和逃跑者代理：

# 通过EnvWrapper创建一个包装的环境对象
# 确保env_backend设置为'pycuda'或'numba'（以便在GPU上运行）
env_wrapper = EnvWrapper(
    TagContinuous(**run_config["env"]),
    num_envs=run_config["trainer"]["num_envs"], 
    env_backend="pycuda"
)
# 代理可以共享策略模型：此字典将策略模型名称映射到代理ID。
policy_tag_to_agent_id_map = {
    "tagger": list(env_wrapper.env.taggers),
    "runner": list(env_wrapper.env.runners),
}

# 创建训练器对象
trainer = Trainer(
    env_wrapper=env_wrapper,
    config=run_config,
    policy_tag_to_agent_id_map=policy_tag_to_agent_id_map,
)

# 开始训练！
trainer.train()

论文和引用 WarpDrive

我们的论文发表在《机器学习研究杂志》(JMLR) https://jmlr.org/papers/v23/22-0185.html。您也可以在我们的白皮书中找到更多详细信息：https://arxiv.org/abs/2108.13976。

如果您在研究或应用中使用 WarpDrive，请使用以下 BibTeX 进行引用：

@article{JMLR:v23:22-0185,
  author  = {Tian Lan and Sunil Srinivasa and Huan Wang and Stephan Zheng},
  title   = {WarpDrive: Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU},
  journal = {Journal of Machine Learning Research},
  year    = {2022},
  volume  = {23},
  number  = {316},
  pages   = {1--6},
  url     = {http://jmlr.org/papers/v23/22-0185.html}
}

@misc{lan2021warpdrive,
  title={WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU}, 
  author={Tian Lan and Sunil Srinivasa and Huan Wang and Caiming Xiong and Silvio Savarese and Stephan Zheng},
  year={2021},
  eprint={2108.13976},
  archivePrefix={arXiv},
  primaryClass={cs.LG}
}

教程和快速入门

教程

通过在 Colab 或 NGC 容器上运行这些教程来熟悉 WarpDrive！

您也可以在本地运行这些教程，但您需要安装了 nvcc 编译器的 GPU 机器和兼容的 Nvidia GPU 驱动程序。您还需要安装 Jupyter。有关安装说明，请参阅 https://jupyter.readthedocs.io/en/latest/install.html

示例训练脚本

我们提供了一些示例脚本，供您快速开始端到端训练。例如，如果您想使用 2 个 GPU 和 CUDA C 后端训练 tag_continuous 环境（10 个标记器和 100 个跑步者）

python example_training_script_pycuda.py -e tag_continuous -n 2

或者切换到 JIT 编译的 Numba 后端，使用 1 个 GPU

python example_training_script_numba.py -e tag_continuous

您可以在这里找到完整的参考文档。

实际问题和合作

使用 WarpDrive 的 AI 经济学家 Covid 环境：我们使用 AI-Economist Foundation 训练两级多代理经济模拟，并使用 WarpDrive 进行训练。在这个例子中，我们特别考虑了 COVID-19 和经济模拟。
基于第一性原理的高通量强化学习：我们将量子密度泛函理论转换为强化学习表示，并实现了从嘈杂的化学系统中自动搜索最佳化学反应路径。有关更多详细信息，请查看我们发表在 Nature Communications 上的工作。
气候变化合作竞争与 Mila 合作。我们提供了 RICE（区域综合气候环境）模拟环境的基础版本。
使用 WarpDrive 的 Pytorch Lightning 训练器：我们提供了一个教程示例和一篇博客文章，介绍了使用 WarpDrive 和 Pytorch Lightning 的多代理强化学习训练循环。
NVIDIA NGC 目录和快速部署到 VertexAI：WarpDrive 镜像由 NGC 目录托管。NGC 目录"托管顶级 AI 和数据科学软件的容器，经过 NVIDIA 的调优、测试和优化"。我们的教程还支持快速部署到 NGC 支持的 VertexAI。

安装说明

要开始使用，您需要安装 Python 3.7+ 和 nvcc 编译器，并有兼容的 Nvidia GPU CUDA 驱动程序。

可以按照 Nvidia 的说明在此处安装 CUDA（包括 nvcc）：https://developer.nvidia.com/cuda-downloads。

Docker 镜像

V100 GPU：您可以参考示例 Dockerfile 来配置您的系统。 A100 GPU：我们的最新镜像由NVIDIA NGC发布和维护。我们建议您从NGC目录下载最新镜像。

如果您想构建自定义环境，我们建议您访问Nvidia Docker Hub下载与您系统兼容的CUDA和cuDNN镜像。您应该能够使用命令行工具来监控系统中的NVIDIA GPU设备：

nvidia-smi

并看到类似这样的输出

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   37C    P0    32W / 300W |      0MiB / 16160MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

在这个快照中，您可以看到我们正在使用Tesla V100 GPU和CUDA版本11.0。

使用Pip安装

您可以使用Python包管理器安装WarpDrive：

pip install rl_warp_drive

从源代码安装

将此仓库克隆到您的机器上：

git clone https://www.github.com/salesforce/warp-drive

可选，但建议首次尝试时使用： 创建一个新的conda环境（下面命名为"warp_drive"）并激活它：
```
conda create --name warp_drive python=3.7 --yes
conda activate warp_drive
```
作为可编辑的Python包安装：
```
cd warp_drive
pip install -e .
```

测试您的安装

您可以直接从Python命令调用来测试所有模块和端到端训练工作流。

python warp_drive/utils/unittests/run_unittests_pycuda.py
python warp_drive/utils/unittests/run_unittests_numba.py
python warp_drive/utils/unittests/run_trainer_tests.py