【WarpDrive】GPU加速RL

在强化学习研究中，一个实验就要跑数天或数周，有没有更快的方法？近日，来自 SalesForce 的研究者提出了一种名为 WarpDrive（曲率引擎）的开源框架，它可以在一个 V100 GPU 上并行运行、训练数千个强化学习环境和上千个智能体。实验结果表明，与 CPU+GPU 的 RL 实现相比，WarpDrive 靠一个 GPU 实现的 RL 要快几个数量级。

深度强化学习是一种训练 AI 智能体的强大方法。然而，如何将强化学习应用于多个智能体或高维状态、观察结果、行动空间等复杂场景仍然是一个具有挑战性的问题。多智能体系统是强化学习研究和应用的前沿，尤其是那些存在多个智能体交互的系统。它们对于解决经济、自动驾驶、对话智能体、机器人学以及其他很多领域的工程和科学问题具有重要意义。

在这一问题中，训练智能体的时间是一大挑战。如果训练的智能体比较多，一个强化学习实验往往要花费数天或数周才能完成。主要原因在于：在在线训练中，强化学习需要反复运行模拟、训练智能体模型。此外，学习所需的重复次数会随着问题的复杂度呈指数级增长。这在无模型设置中是最突出的，在这种设置中，强化学习智能体在训练时通常对当前的模拟或任务没有任何初始知识。由于深度强化学习实现通常将基于 CPU 的模拟和 GPU 神经网络模型结合在一起，现实世界中的训练时间可能会非常长。

为了提高强化学习系统的吞吐量，来自 Salesforce 的研究者提出了一种名为 WarpDrive 的开源框架，来构建极其快速的（多智能体）深度强化学习（MADRL）pipeline。

论文链接：https://arxiv.org/pdf/2108.13976.pdf
项目地址：https://github.com/salesforce/warp-drive

WarpDrive 在一个 GPU 上端到端地运行整个 MADRL 工作流，使用单个数据存储进行模拟 roll-out、推理和训练。这将最小化通信和复制的成本，显著提升模拟采样和学习速率。WarpDrive 还利用 GPU 的并行能力，并行地运行模拟和每次模拟中的智能体。综合来看，这些设计选择使得运行数千个并发模拟成为可能，每个模拟包含数千个智能体，在非常大的批次上进行训练。

论文中的基准测试表明，与混合 CPU 模拟和 GPU 模型的一般实现相比，WarpDrive 实现的 RL 要快几个数量级。比如，在有 2000 个模拟、1000 个智能体的 Tag 模拟中，WarpDrive 的吞吐量实现了至少 100 倍的性能提升。

此外，研究者还将 WarpDrive 构建得非常灵活、易用、易扩展，并在现有库的基础上进行构建和补充，如 PyCUDA 和 PyTorch。在这个初始版本中，WarpDrive 使用在 CUDA C 中实现的 step 和 reset 方法的模拟，使用 OpenAI gym-API 风格。这使得用户在实现模拟的时候能够使用 CUDA 编程的全部特征集。WarpDrive 在 PyCUDA 的基础上构建，可以方便地访问用 CUDA C 编写的 GPU 数据和模拟方法。它还提供了一个轻量级的 Python 包装器来自动构建 gym 风格的环境对象，实现模拟 API。因此，WarpDrive 可以让用户轻松地创建和扩展自定义 RL pipeline，以最大化 GPU 的用途。

作者表示，WarpDrive 这一名字的灵感来自科幻小说中的超光速宇宙飞船推进系统——曲率引擎，后者可以通过对时空本身的改造来驱动飞船。

在这篇论文中，研究者概述了 WarpDrive 的设计，并展示了它在基准模拟中的吞吐量和可扩展性。最后，他们总结了未来的发展方向。

用 WarpDrive 加速强化学习

WarpDrive 提供了一个框架和一些高质量的工具，帮助用户快速构建快捷而灵活的多智能体强化学习系统。图 3 显示了 WarpDrive 的一个基本用例。

接下来，研究者自下而上地介绍了 WarpDrive 的设计和组件，以及它们在加速强化学习方面的作用。

CUDA 计算

WarpDrive 专注于 CUDA（Compute Unified Device Architecture），这是一个流行的平台和编程模型，允许用户在（CUDAenabled）GPU 硬件上运行程序。CUDA 程序也被称为计算 kernel。CUDA API 可以直接访问 GPU 的虚拟指令集和并行计算元素。

GPU 的一个关键特性是它们可以并行运行许多计算线程。线程组成线程块，多个线程块组成一个网格结构。一个 CUDA kernel 可以访问并定义这些线程的并行计算，见下图 2。

DataManager 和 FunctionManager

图 2 显示了 WarpDrive 框架的一些细节。所有在 GPU 上运行的 CUDA 程序都需要通过 CPU 触发。按照 CUDA 的规则，CPU 被称为主机（host），GPU 被称为设备。任何 CUDA 程序的执行都包括三个主要步骤：

主机到设备的传输（push）：将输入数据从主机内存复制到设备内存，例如在训练开始时。
加载 CUDA 函数，在 GPU 上执行并缓存数据以提高性能。
设备到主机的传输（pull）：将数据从设备内存复制回主机，例如训练结束的时候。

按照这个模式，WarpDrive 实现了两个重要的 Python 类（驻留在 CPU 上）——DataManager 和 FunctionManager，以方便主机和设备之间所有与 RL 相关的通信和交互。DataManager 提供了一个简易 API 来处理与 RL 相关的主机和设备之间的所有数据传输（pull 和 push）。FunctionManager 允许用户从 CPU 调用 CUDA kernel，并在 GPU 上执行它们。

WarpDrive RL 工作流

使用 WarpDrive，在 GPU 上收集 rollout 和训练模型的一个典型 RL 工作流如下所示：

1. 初始化并重置环境对象后，将主机上的所有数据复制到设备上。DataManager 提供 API 方法来执行此 push 操作。从主机复制到设备的数据可能包括环境配置参数、在第一次重置结束时创建的数据数组，以及观察、动作、奖励和「完成」标志的占位符。DataManager 还帮助维护变量的副本，这些变量需要在每次重置时重新初始化。在这点之后，就不再从主机向设备推送数据了。

2. FunctionManager 提供 API 方法来初始化和调用 CUDA C kernel 函数，这些函数用于从主机节点执行环境步骤。这些函数只在设备上执行，所有的数据数组都被就地修改。出于可视化或分析的目的，主机可能会不时地拉取（pulled）数据，但所有数据本质上只能在训练期间驻留在 GPU 上。

3.WarpDrive 还包含一个 Sampler 类，用于对动作进行采样，以便逐步遍历环境。使用策略模型计算出的概率对动作进行抽样。WarpDrive 的采样器在每个智能体线程上并行运行，运行速度比等效的 PyTorch 实现快 2 倍。详见实验部分。

4. 在 GPU 中，研究者并行执行多个环境副本。每个环境运行在一个单独的线程块上。由于一个典型的 GPU 有数千个块，在一个 GPU 上并行执行数千个环境也是可行的。

5. 任何环境都可能达到终端状态并被「完成」。WarpDrive 提供了一个 EnvironmentReset 类，用于自动识别和重置那些已完成的环境。此时，那些环境也被重置并给出（新的）初始数据。例如，它们可能使用在初始重置时复制的初始数据数组。

6. 在每个环境中，每个智能体还可以在智能体独占（agent-exclusive）的线程上执行它自己的逻辑。图 2 显示了一个示例，其中智能体 i、j 和 k 分别在单个线程 i、j 和 k 上并行操作。这在多智能体 RL 中非常有用，因为可以完全将智能体操作并行化，因此模拟时间复杂度保持不变，即使智能体数量增加（达到可用线程的数量）。

7. 一旦从几个环境和智能体中收集到的数据被送到训练数据 batch 中，也可以用 WarpDrive 执行端到端训练。这次推出的 WarpDrive 包括一个示例训练脚本和 Trainer 类，它目前实现了 Advantage Actor Critic（A2C）、近端策略优化（PPO）RL 算法和一个全连接神经网络策略模型。Trainer 基于 PyTorch 构建，并从 CPU 主机调用所有 CUDA kernel 和 PyTorch 操作。然而，PyTorch 可以直接访问和解释以 torch 形式存储在设备上的 WarpDrive 数据 batch（状态、动作、奖励和完成标志）。这使得用户可以计算损失并修改模型参数，同时消除数据复制。考虑到 WarpDrive 的模块化特性，这一框架很容易扩展。

这些设计选择使得 WarpDrive 非常高效。