一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

在加强进修研究中,一个实验就要跑数天或数周,有没有更快的方法?近日,来自 SalesForce 的研究者提出了一种名为 WarpDrive(曲率引擎)的开源框架,它可以在一个 V100 GPU 上并行运转、训练数千个加强进修情况和上千个智能体。实验结果表明,与 CPU+GPU 的 RL 实行相比,WarpDrive 靠一个 GPU 实行的 RL 要快几个数量级。

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

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

为了提高加强进修系统的吞吐量,来自 Salesforce 的研究者提出了一种名为 WarpDrive 的开源框架,来建立极其快速的(多智能体)深度加强进修(MADRL)pipeline。

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

论文链接: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 这一名字的灵感来自科幻小说中的超光速宇宙飞船推进系统——曲率引擎,后者可以通过对时空本身的改造来驱动飞船。

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

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

用 WarpDrive 加速加强进修

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

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

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

CUDA 计算

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

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

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

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 非常高效。

实验结果

研究者应用 Tag 游戏演示了 WarpDrive 的工作流,并对其性能进行了测试。他们比较了以下两种情况的性能:

应用 CPU-simulations + GPU-agent 的模型;

在单个 Nvidia V100 GPU 上运转 WarpDrive。

测试结果表明,与应用 CPU-simulation + GPU-agent 的模型相比,WarpDrive 的吞吐量要高几个数量级。

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

一块V100运转上千个智能体、数千个情况,这个「曲率引擎」框架实行RL百倍提速

给TA打赏
共{{data.count}}人
人已打赏
AI

特斯拉8月在华销量首破4万台,或已达产能极限

2021-9-9 18:51:00

AI

撑持异构图、集成GraphGym,超好用的图神经网络库PyG更新2.0版本

2021-9-14 15:01:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
搜索