一块V100运行上千个智能体、数千个环境，这个「曲率引擎」框架实现RL百倍提速
发布时间：2021-09-13 09:55:45来源：机器之心
机器之心报道
编辑：张倩
在强化学习研究中，一个实验就要跑数天或数周，有没有更快的方法？近日，来自SalesForce的研究者提出了一种名为WarpDrive（曲率引擎）的开源框架，它可以在一个V100GPU上并行运行、训练数千个强化学习环境和上千个智能体。实验结果表明，与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使用在CUDAC中实现的step和reset方法的模拟，使用OpenAIgym-API风格。这使得用户在实现模拟的时候能够使用CUDA编程的全部特征集。WarpDrive在PyCUDA的基础上构建，可以方便地访问用CUDAC编写的GPU数据和模拟方法。它还提供了一个轻量级的Python包装器来自动构建gym风格的环境对象，实现模拟API。因此，WarpDrive可以让用户轻松地创建和扩展自定义RLpipeline，以最大化GPU的用途。
作者表示，WarpDrive这一名字的灵感来自科幻小说中的超光速宇宙飞船推进系统——曲率引擎，后者可以通过对时空本身的改造来驱动飞船。
在这篇论文中，研究者概述了WarpDrive的设计，并展示了它在基准模拟中的吞吐量和可扩展性。最后，他们总结了未来的发展方向。
用WarpDrive加速强化学习
WarpDrive提供了一个框架和一些高质量的工具，帮助用户快速构建快捷而灵活的多智能体强化学习系统。图3显示了WarpDrive的一个基本用例。
接下来，研究者自下而上地介绍了WarpDrive的设计和组件，以及它们在加速强化学习方面的作用。
CUDA计算
WarpDrive专注于CUDA（ComputeUnifiedDeviceArchitecture），这是一个流行的平台和编程模型，允许用户在（CUDAenabled）GPU硬件上运行程序。CUDA程序也被称为计算kernel。CUDAAPI可以直接访问GPU的虚拟指令集和并行计算元素。
GPU的一个关键特性是它们可以并行运行许多计算线程。线程组成线程块，多个线程块组成一个网格结构。一个CUDAkernel可以访问并定义这些线程的并行计算，见下图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调用CUDAkernel，并在GPU上执行它们。
WarpDriveRL工作流
使用WarpDrive，在GPU上收集rollout和训练模型的一个典型RL工作流如下所示：
1.初始化并重置环境对象后，将主机上的所有数据复制到设备上。DataManager提供API方法来执行此push操作。从主机复制到设备的数据可能包括环境配置参数、在第一次重置结束时创建的数据数组，以及观察、动作、奖励和「完成」标志的占位符。DataManager还帮助维护变量的副本，这些变量需要在每次重置时重新初始化。在这点之后，就不再从主机向设备推送数据了。
2.FunctionManager提供API方法来初始化和调用CUDACkernel函数，这些函数用于从主机节点执行环境步骤。这些函数只在设备上执行，所有的数据数组都被就地修改。出于可视化或分析的目的，主机可能会不时地拉取（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类，它目前实现了AdvantageActorCritic（A2C）、近端策略优化（PPO）RL算法和一个全连接神经网络策略模型。Trainer基于PyTorch构建，并从CPU主机调用所有CUDAkernel和PyTorch操作。然而，PyTorch可以直接访问和解释以torch形式存储在设备上的WarpDrive数据batch（状态、动作、奖励和完成标志）。这使得用户可以计算损失并修改模型参数，同时消除数据复制。考虑到WarpDrive的模块化特性，这一框架很容易扩展。
这些设计选择使得WarpDrive非常高效。
实验结果
研究者使用Tag游戏演示了WarpDrive的工作流，并对其性能进行了测试。他们比较了以下两种情况的性能：
使用CPU-simulations+GPU-agent的模型；
在单个NvidiaV100GPU上运行WarpDrive。
测试结果表明，与使用CPU-simulation+GPU-agent的模型相比，WarpDrive的吞吐量要高几个数量级。
选择NVIDIA®DGX™系统五大理由
AI时代已来，人工智能几乎随处可见，所有AI应用表现出的「智能」背后都是数学运算的结果。而支撑运算需要三个关键因素：数据、算法和算力。数据在爆发，而算法越来越复杂，为了与前两者相匹配，算力自然也要不断提升，如何提高算力呢？
在这个问题上，GPU的出现使得AI发展得到了前所未有的加速。随着AI已经广泛应用于各行各业中，越来越多企业要投入构建硬件平台，而目前GPU专业平台很多，为什么要购买NVIDIADGX系统？请观看选择NVIDIADGX系统五大理由。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：content@jiqizhixin.com
