编辑:好困
想当年,玩儿转Dota2的OpenAIFive训练起来竟然需要超过5万核CPU。就在昨天,颜水成团队公开了最新的强化学习训练环境引擎,核CPU的运行速度直接达到1秒1百万帧!就连笔记本上的i7-H也能有5万帧每秒。
近年来,深度强化学习(DeepRL)取得了飞速的进展,有很多影响深远的工作:从发表在Nature的DQN,到后来打败李世石、柯洁等世界冠军的AlphaGo系列,再到复杂游戏:代表星际的AlphaStar和Dota2的OpenAIFive。
除了算法上的进步之外,最重要的是近年来对于深度强化学习智能体的训练速度及吞吐量的巨大改进,当年的DQN跑简单的Atari游戏都要花费一周多的时间,而现在RL系统已经能承载非常高的吞吐量,能在复杂的游戏和场景中训练起来。
世界上目前最大型的强化学习训练系统都是采取分布式训练方法,比如OpenAIFive,使用了超过5万核CPU,以及上千个GPU来进行训练。
OpenAIFive
类似的,AlphaStar里每个训练智能体都连接着同时跑的1万多场星际游戏引擎。为了下游从游戏交互数据中迭代学习的训练速度,在这些训练系统中都不得不使用大量的CPU资源来跑RL环境游戏引擎。
AlphaStar
如此大的资源消耗令大型强化学习训练系统门槛非常高,学术界有限的资源难以在大型游戏或者复杂的RL环境进行研究,不仅学术界,一个实现万核CPU级别的用量在工业界研究院也是不小的负担。
SeaAILab的研究人员注意到RL环境,包括游戏引擎,是整个RL训练系统里面最慢的部分,而且处于数据供应端的位置,决定了整个系统吞吐量的上限。
然而这一部件并没有得到研究人员的足够重视,目前最常用的并行执行RL环境的办法是gym.vector_env,即使用Python多进程来进行简单的训练环境并行,使用起来接口能保持不变,可是由于Python的局限性,最后为了达到很高的吞吐量,只能使用更多的CPU资源来进行环境的模拟。
为了提高RL环境的模拟性能及CPU利用效率,SeaAILab提供了一个高度并行的RL环境引擎解决方案EnvPool。
这个RL环境引擎底层使用C++线程池,通过异步的方式执行多个RL环境实例,来大大加速并行的效率。EnvPool在经典的RL环境模拟器Atari游戏上,利用起NVIDIADGXA单机上的个CPU核,达到了惊人的一百万帧每秒的执行速度。
EnvPool系统概述
如此之高的吞吐量是研究人员最常用的执行引擎gym.vector_env的近14倍。这意味着我们能用同样的硬件资源达到一个数量级差别的资源利用效率,或者反过来,使用了EnvPool可以少用一个数量级的资源数量且能达到极高的吞吐量。
对于研究人员来说,EnvPool提供了方便易用的Python接口,如下图所示,最简单的同步执行模式的接口与GymAPI完全一致,相比单环境的gym/dm_env,EnvPool仅仅将原本与单个环境交互的API拓展到批量交互。
批量获取的state方便算法端直接将数据送往GPU上进行inference,这样下来算法端的实现难度也大大降低,从环境端出来的数据直接已经是批量的适合GPU和TPU处理的格式,能更好的利用GPU/TPU的并行效率。
在异步模式下,EnvPool将原有的step函数拆分成send和recv函数,在调用send函数将action送往环境之后,线程无需等待当前交互的环境返回新的state,而是交由EnvPool在线程池中执行。
完成执行的环境会将新的state放入队列,通过recv函数来批量获取。相比同步模式,异步模式需要提供环境总数,以及每次交互的批大小。
异步的方式下,算法和环境可以有效的在时间上重合,无需相互等待。
SeaAILab的研究员在两种不同的机器型号上评测对比了EnvPool及别的RL环境执行引擎的吞吐量,分别是TPUv3-8虚拟机,以及NVIDIADGX-A机型。TPU虚拟机上有96CPUcores,两个NUMA节点,一台NVIDIADGX-A机子有CPU核,8NUMAnodes.
与EnvPool对比的RL环境模拟系统包括简易的Pythonfor-loop的多环境执行,业界最常用的gym.vector_env,以及去年刚发布的据研究人员所知的之前最快的环境模拟系统SampleFactory。
DGX-
TPU-VM
由结果可见,在不同的机型上,使用不同数量的worker,EnvPool都具有非常大的优势,特别是在NVIDIADGXA上,使用同样的个CPU核,EnvPool的性能达到了市面上最常用的gym.vector_env的13.3倍,并且达到了惊人的一百万帧每秒的运行速度。这个吞吐量下,产生10^9帧(10亿)数据仅需要17分钟。
在相对少资源的设置下,EnvPool表现也十分的优异,在12核的情况下使用gym.vector_env运行Atari仅能达到1.8万帧每秒的速度,而EnvPool可以更有效地利用CPU起来,达到5万帧每秒,是这个广为使用的基准的约3倍。
项目简介
研发团队SeaAILab(SAIL)隶属于新加坡冬海(SEA)集团,成立于年末,由颜水成挂帅,专注于前沿突破性基础研究。
目前,EnvPool已经在GitHub上开源。