100万帧数据仅1秒！AI大牛颜水成团队强化学习新作，代码已开源

编辑：好困

【新智元导读】想当年，玩儿转Dota 2的OpenAI Five训练起来竟然需要超过5万核CPU。就在昨天，颜水成团队公开了最新的强化学习训练环境引擎，256核CPU的运行速度直接达到1秒1百万帧！就连笔记本上的i7-8750H也能有5万帧每秒。

近年来，深度强化学习(Deep RL)取得了飞速的进展，有很多影响深远的工作：从发表在Nature 2015的DQN，到后来打败李世石、柯洁等世界冠军的AlphaGo系列，再到复杂游戏：代表星际的AlphaStar和Dota 2的OpenAI Five。

除了算法上的进步之外，最重要的是近年来对于深度强化学习智能体的训练速度及吞吐量的巨大改进，当年的DQN跑简单的Atari游戏都要花费一周多的时间，而现在RL系统已经能承载非常高的吞吐量，能在复杂的游戏和场景中训练起来。

世界上目前最大型的强化学习训练系统都是采取分布式训练方法，比如OpenAI Five，使用了超过5万核CPU，以及上千个GPU来进行训练。

OpenAI Five

类似的，AlphaStar里每个训练智能体都连接着同时跑的1万多场星际游戏引擎。为了下游从游戏交互数据中迭代学习的训练速度，在这些训练系统中都不得不使用大量的CPU资源来跑RL环境游戏引擎。

AlphaStar

如此大的资源消耗令大型强化学习训练系统门槛非常高，学术界有限的资源难以在大型游戏或者复杂的RL环境进行研究，不仅学术界，一个实现万核CPU级别的用量在工业界研究院也是不小的负担。

Sea AI Lab的研究人员注意到RL环境，包括游戏引擎，是整个RL训练系统里面最慢的部分，而且处于数据供应端的位置，决定了整个系统吞吐量的上限。

然而这一部件并没有得到研究人员的足够重视，目前最常用的并行执行RL环境的办法是gym.vector_env，即使用Python多进程来进行简单的训练环境并行，使用起来接口能保持不变，可是由于Python的局限性，最后为了达到很高的吞吐量，只能使用更多的CPU资源来进行环境的模拟。

为了提高RL环境的模拟性能及CPU利用效率，Sea AI Lab提供了一个高度并行的RL环境引擎解决方案EnvPool。

这个RL环境引擎底层使用C++线程池，通过异步的方式执行多个RL环境实例，来大大加速并行的效率。EnvPool在经典的RL环境模拟器Atari游戏上，利用起NVIDIA DGX A100单机上的256个CPU核，达到了惊人的一百万帧每秒的执行速度。

EnvPool系统概述

如此之高的吞吐量是研究人员最常用的执行引擎gym.vector_env的近14倍。这意味着我们能用同样的硬件资源达到一个数量级差别的资源利用效率，或者反过来，使用了EnvPool可以少用一个数量级的资源数量且能达到极高的吞吐量。

对于研究人员来说，EnvPool提供了方便易用的Python接口，如下图所示，最简单的同步执行模式的接口与Gym API完全一致，相比单环境的gym/dm_env, EnvPool仅仅将原本与单个环境交互的API拓展到批量交互。

批量获取的state方便算法端直接将数据送往GPU上进行inference，这样下来算法端的实现难度也大大降低，从环境端出来的数据直接已经是批量的适合GPU和TPU处理的格式，能更好的利用GPU/TPU的并行效率。

在异步模式下，EnvPool将原有的step函数拆分成send和recv函数，在调用send函数将action送往环境之后，线程无需等待当前交互的环境返回新的state，而是交由EnvPool在线程池中执行。

完成执行的环境会将新的state放入队列，通过recv函数来批量获取。相比同步模式，异步模式需要提供环境总数，以及每次交互的批大小。

异步的方式下，算法和环境可以有效的在时间上重合，无需相互等待。

Sea AI Lab的研究员在两种不同的机器型号上评测对比了EnvPool及别的RL环境执行引擎的吞吐量，分别是TPU v3-8虚拟机， 以及NVIDIA DGX-A100机型。TPU虚拟机上有96 CPU cores，两个NUMA节点，一台NVIDIA DGX-A100机子有256 CPU核， 8 NUMA nodes.

与EnvPool对比的RL环境模拟系统包括简易的Python for-loop的多环境执行，业界最常用的gym.vector_env，以及去年刚发布的据研究人员所知的之前最快的环境模拟系统Sample Factory。

DGX-100

TPU-VM

由结果可见，在不同的机型上，使用不同数量的worker，EnvPool都具有非常大的优势，特别是在NVIDIA DGX A100上，使用同样的256个CPU核，EnvPool的性能达到了市面上最常用的gym.vector_env的13.3倍，并且达到了惊人的一百万帧每秒的运行速度。这个吞吐量下，产生10^9帧（10亿）数据仅需要17分钟。

在相对少资源的设置下，EnvPool表现也十分的优异，在12核的情况下使用gym.vector_env运行Atari仅能达到1.8万帧每秒的速度，而EnvPool可以更有效地利用CPU起来，达到5万帧每秒，是这个广为使用的基准的约3倍。

项目简介

研发团队Sea AI Lab (SAIL)隶属于新加坡冬海（SEA）集团，成立于2020年末，由颜水成挂帅，专注于前沿突破性基础研究。

目前，EnvPool已经在GitHub上开源。

https://github.com/sail-sg/envpool

这已经是内部迭代的第二版，对比第一版的内部实现，开源的第二版着眼于简化开发者API，也就是更加方便社区开发者接入不同的游戏或者RL训练环境。

接入C++引擎时，开发者只需要定义好单个环境的执行逻辑，EnvPool则负责分布式执行并且提供批交互的API，这使得接入新的RL环境及游戏引擎并得到立刻的并行加速变得非常的简单。

目前EnvPool在进行高度开发，接下来的核心更新包括接入更多的RL环境（包括连续动作空间的近期宣布免费使用的Mujoco），及会提供样例让用户能很方便地利用EnvPool加速现有的开源RL训练库，包括接入DeepMind的Acme，以及接入EnvPool同一第一作者的大受欢迎的RL训练库天授Tianshou。

此外，EnvPool的成果也被邀请到NVIDIA GTC 2021大会上演讲。

参考资料：

https://github.com/sail-sg/envpool

GTC 2021演讲：

https://events.rainfocus.com/widget/nvidia/nvidiagtc/sessioncatalog/session/1630239583490001Z5dE