当前位置: 首页 > 科技观察

100万帧数据仅需1秒!AI大神严水城团队研发新的强化学习作品,代码已经开源

时间:2023-03-12 04:45:57 科技观察

想当年玩Dota2的OpenAIFive需要5万多个CPU核心来训练。就在昨天,闫水成团队发布了最新的强化学习训练环境引擎。256核CPU的运行速度直接可以达到100万帧/秒!即使是笔记本电脑上的i7-8750H也能达到每秒5万帧。近年来,深度强化学习(DeepRL)进展迅速,影响深远的著作很多:从2015年发表在Nature上的DQN,到后来击败李世石、柯洁等世界冠军的AlphaGo系列,到复杂游戏:星际穿越的AlphaStar和Dota2的OpenAIFive。除了算法的进步,最重要的是近年来深度强化学习代理的训练速度和吞吐量的巨大提升。DQN运行一个简单的Atari游戏需要一周多的时间,但现在RL系统已经能够承载非常高的吞吐量,并且可以在复杂的游戏和场景中进行训练。目前全球最大的强化学习训练系统采用分布式训练方式,如OpenAIFive,使用5万多个核心CPU和数千个GPU进行训练。与OpenAIFive类似,AlphaStar中的每个训练代理都连接到超过10,000个同时运行的星际游戏引擎。为了从游戏交互数据中进行下游迭代学习的训练速度,这些训练系统不得不使用大量的CPU资源来运行RL环境游戏引擎。AlphaStar如此大的资源消耗,使得大规模强化学习训练系统的门槛非常高。学术界有限的资源很难在大型游戏或复杂的RL环境中进行研究。不仅在学术界,在产业研究机构也是不小的负担。SeaAILab的研究人员注意到,包括游戏引擎在内的RL环境是整个RL训练系统中最慢的部分,处于数据供给端,决定了整个系统吞吐量的上限。然而,该组件并未得到研究人员的足够重视。目前最常用的RL环境并行执行的方法是gym.vector_env,它使用Python多进程在简单的环境下进行并行训练。使用时界面可以保持不变,但是由于Python的限制,为了达到高吞吐量,只能使用更多的CPU资源来模拟环境。为了提高RL环境的仿真性能和CPU利用效率,SeaAILab提供了高度并行的RL环境引擎解决方案EnvPool。该强化学习环境引擎底层采用C++线程池异步执行多个强化学习环境实例,大大加快并行效率。在经典的RL环境模拟器Atari游戏上,EnvPool在单台NVIDIADGXA100机器上利用256个CPU核心,实现了每秒百万帧的惊人执行速度。EnvPool系统概述如此高的吞吐量是研究人员最常用的执行引擎gym.vector_env的近14倍。这意味着我们可以使用相同的硬件资源来实现资源利用效率的数量级差异,或者反过来说,使用EnvPool可以少使用一个数量级的资源并获得极高的吞吐量。对于研究人员,EnvPool提供了一个方便易用的Python接口。如下图所示,最简单的同步执行方式的接口与GymAPI完全一致。与单环境gym/dm_env相比,EnvPool只是集成了原有环境交互的API,扩展为批量交互。批量获取的状态方便算法端直接将数据发送给GPU进行推理,使得算法端的实现难度大大降低,环境端的数据直接以适合GPU和TPU的批量格式处理,可以更好的利用GPU/TPU的并行效率。在异步模式下,EnvPool将原来的step函数拆分为send和recv函数。调用send函数将action发送到环境后,线程不需要等待当前交互环境返回一个新的状态,而是到池中执行的线程中的EnvPool。执行完成的环境会将新的状态放入队列中,通过recv函数批量获取。与同步模式相比,异步模式需要提供环境总数,以及每次交互的批大小。以异步的方式,算法和环境可以在时间上有效重叠,无需相互等待。SeaAILab研究人员评估并比较了EnvPool和其他RL环境执行引擎在两种不同机器模型(即TPUv3-8虚拟机和NVIDIADGX-A100模型)上的吞吐量。TPU虚拟机上有96个CPU核心和两个NUMA节点,一台NVIDIADGX-A100机器有256个CPU核心和8个NUMA节点。对比EnvPool的RL环境模拟系统包括简单的Pythonfor-loop多环境执行,业界最常用的gym.vector_env,以及去年刚刚发布的研究人员已知最快的环境模拟系统SampleFactory。从结果可以看出DGX-100TPU-VM。在不同的模型上,EnvPool在使用不同数量的worker时具有很大的优势。尤其是在NVIDIADGXA100上,同样使用256个CPU核心,EnvPool的性能达到了市面上最常用的gym.vector_env的13.3倍,达到了每秒百万帧的惊人速度。在这个吞吐量下,生成10^9帧(10亿)数据只需要17分钟。在资源相对较少的设置下,EnvPool的表现也非常出色。在12核的情况下,使用gym.vector_env运行Atari只能达到18000帧/秒的速度,而EnvPool可以更有效地利用CPU,达到510000帧/秒,比这个广泛使用的benchmark快3倍左右.项目介绍研发团队SeaAILab(SAIL)隶属于新加坡东海(SEA)集团。成立于2020年底,由颜水成牵头,专注于前沿突破性基础研究。目前,EnvPool已经在GitHub上开源。https://github.com/sail-sg/envpool这已经是内部迭代的第二个版本了。与第一版内部实现相比,第二版开源着重于简化开发者API,更方便社区开发者接入不同的游戏或RL训练环境。开发者接入C++引擎时,只需定义单一环境的执行逻辑,EnvPool负责分布式执行,提供批量交互API,更方便接入新的RL环境和游戏引擎,获得即时并行加速很简单。EnvPool目前正处于高度开发中,下一次核心更新将包括访问更多的RL环境(包括最近宣布在连续动作空间中免费使用Mujoco),并将提供示例以便用户可以轻松地使用EnvPool来加速现有的Open来源RL训练库,包括连接到DeepMind的Acme,以及连接到EnvPool的同一第一作者的流行的RL训练库Tianshou。此外,EnvPool的成果还受邀在NVIDIAGTC2021大会上发言。