[[434376]]

想當(dāng)年，玩兒轉(zhuǎn)Dota 2的OpenAI Five訓(xùn)練起來竟然需要超過5萬核CPU。就在昨天，顏水成團(tuán)隊(duì)公開了最新的強(qiáng)化學(xué)習(xí)訓(xùn)練環(huán)境引擎，256核CPU的運(yùn)行速度直接達(dá)到1秒1百萬幀！就連筆記本上的i7-8750H也能有5萬幀每秒。

近年來，深度強(qiáng)化學(xué)習(xí)(Deep RL)取得了飛速的進(jìn)展，有很多影響深遠(yuǎn)的工作：從發(fā)表在Nature 2015的DQN，到后來打敗李世石、柯潔等世界冠軍的AlphaGo系列，再到復(fù)雜游戲：代表星際的AlphaStar和Dota 2的OpenAI Five。

除了算法上的進(jìn)步之外，最重要的是近年來對(duì)于深度強(qiáng)化學(xué)習(xí)智能體的訓(xùn)練速度及吞吐量的巨大改進(jìn)，當(dāng)年的DQN跑簡(jiǎn)單的Atari游戲都要花費(fèi)一周多的時(shí)間，而現(xiàn)在RL系統(tǒng)已經(jīng)能承載非常高的吞吐量，能在復(fù)雜的游戲和場(chǎng)景中訓(xùn)練起來。

世界上目前最大型的強(qiáng)化學(xué)習(xí)訓(xùn)練系統(tǒng)都是采取分布式訓(xùn)練方法，比如OpenAI Five，使用了超過5萬核CPU，以及上千個(gè)GPU來進(jìn)行訓(xùn)練。

OpenAI Five

類似的，AlphaStar里每個(gè)訓(xùn)練智能體都連接著同時(shí)跑的1萬多場(chǎng)星際游戲引擎。為了下游從游戲交互數(shù)據(jù)中迭代學(xué)習(xí)的訓(xùn)練速度，在這些訓(xùn)練系統(tǒng)中都不得不使用大量的CPU資源來跑RL環(huán)境游戲引擎。

AlphaStar

如此大的資源消耗令大型強(qiáng)化學(xué)習(xí)訓(xùn)練系統(tǒng)門檻非常高，學(xué)術(shù)界有限的資源難以在大型游戲或者復(fù)雜的RL環(huán)境進(jìn)行研究，不僅學(xué)術(shù)界，一個(gè)實(shí)現(xiàn)萬核CPU級(jí)別的用量在工業(yè)界研究院也是不小的負(fù)擔(dān)。

Sea AI Lab的研究人員注意到RL環(huán)境，包括游戲引擎，是整個(gè)RL訓(xùn)練系統(tǒng)里面最慢的部分，而且處于數(shù)據(jù)供應(yīng)端的位置，決定了整個(gè)系統(tǒng)吞吐量的上限。

然而這一部件并沒有得到研究人員的足夠重視，目前最常用的并行執(zhí)行RL環(huán)境的辦法是gym.vector_env，即使用Python多進(jìn)程來進(jìn)行簡(jiǎn)單的訓(xùn)練環(huán)境并行，使用起來接口能保持不變，可是由于Python的局限性，最后為了達(dá)到很高的吞吐量，只能使用更多的CPU資源來進(jìn)行環(huán)境的模擬。

為了提高RL環(huán)境的模擬性能及CPU利用效率，Sea AI Lab提供了一個(gè)高度并行的RL環(huán)境引擎解決方案EnvPool。

這個(gè)RL環(huán)境引擎底層使用C++線程池，通過異步的方式執(zhí)行多個(gè)RL環(huán)境實(shí)例，來大大加速并行的效率。EnvPool在經(jīng)典的RL環(huán)境模擬器Atari游戲上，利用起NVIDIA DGX A100單機(jī)上的256個(gè)CPU核，達(dá)到了驚人的一百萬幀每秒的執(zhí)行速度。

EnvPool系統(tǒng)概述

如此之高的吞吐量是研究人員最常用的執(zhí)行引擎gym.vector_env的近14倍。這意味著我們能用同樣的硬件資源達(dá)到一個(gè)數(shù)量級(jí)差別的資源利用效率，或者反過來，使用了EnvPool可以少用一個(gè)數(shù)量級(jí)的資源數(shù)量且能達(dá)到極高的吞吐量。

對(duì)于研究人員來說，EnvPool提供了方便易用的Python接口，如下圖所示，最簡(jiǎn)單的同步執(zhí)行模式的接口與Gym API完全一致，相比單環(huán)境的gym/dm_env, EnvPool僅僅將原本與單個(gè)環(huán)境交互的API拓展到批量交互。

批量獲取的state方便算法端直接將數(shù)據(jù)送往GPU上進(jìn)行inference，這樣下來算法端的實(shí)現(xiàn)難度也大大降低，從環(huán)境端出來的數(shù)據(jù)直接已經(jīng)是批量的適合GPU和TPU處理的格式，能更好的利用GPU/TPU的并行效率。

在異步模式下，EnvPool將原有的step函數(shù)拆分成send和recv函數(shù)，在調(diào)用send函數(shù)將action送往環(huán)境之后，線程無需等待當(dāng)前交互的環(huán)境返回新的state，而是交由EnvPool在線程池中執(zhí)行。

完成執(zhí)行的環(huán)境會(huì)將新的state放入隊(duì)列，通過recv函數(shù)來批量獲取。相比同步模式，異步模式需要提供環(huán)境總數(shù)，以及每次交互的批大小。

異步的方式下，算法和環(huán)境可以有效的在時(shí)間上重合，無需相互等待。

Sea AI Lab的研究員在兩種不同的機(jī)器型號(hào)上評(píng)測(cè)對(duì)比了EnvPool及別的RL環(huán)境執(zhí)行引擎的吞吐量，分別是TPU v3-8虛擬機(jī)， 以及NVIDIA DGX-A100機(jī)型。TPU虛擬機(jī)上有96 CPU cores，兩個(gè)NUMA節(jié)點(diǎn)，一臺(tái)NVIDIA DGX-A100機(jī)子有256 CPU核， 8 NUMA nodes.

與EnvPool對(duì)比的RL環(huán)境模擬系統(tǒng)包括簡(jiǎn)易的Python for-loop的多環(huán)境執(zhí)行，業(yè)界最常用的gym.vector_env，以及去年剛發(fā)布的據(jù)研究人員所知的之前最快的環(huán)境模擬系統(tǒng)Sample Factory。

DGX-100

TPU-VM

由結(jié)果可見，在不同的機(jī)型上，使用不同數(shù)量的worker，EnvPool都具有非常大的優(yōu)勢(shì)，特別是在NVIDIA DGX A100上，使用同樣的256個(gè)CPU核，EnvPool的性能達(dá)到了市面上最常用的gym.vector_env的13.3倍，并且達(dá)到了驚人的一百萬幀每秒的運(yùn)行速度。這個(gè)吞吐量下，產(chǎn)生10^9幀（10億）數(shù)據(jù)僅需要17分鐘。

在相對(duì)少資源的設(shè)置下，EnvPool表現(xiàn)也十分的優(yōu)異，在12核的情況下使用gym.vector_env運(yùn)行Atari僅能達(dá)到1.8萬幀每秒的速度，而EnvPool可以更有效地利用CPU起來，達(dá)到5萬幀每秒，是這個(gè)廣為使用的基準(zhǔn)的約3倍。

項(xiàng)目簡(jiǎn)介

研發(fā)團(tuán)隊(duì)Sea AI Lab (SAIL)隸屬于新加坡冬海（SEA）集團(tuán)，成立于2020年末，由顏水成掛帥，專注于前沿突破性基礎(chǔ)研究。

目前，EnvPool已經(jīng)在GitHub上開源。

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

這已經(jīng)是內(nèi)部迭代的第二版，對(duì)比第一版的內(nèi)部實(shí)現(xiàn)，開源的第二版著眼于簡(jiǎn)化開發(fā)者API，也就是更加方便社區(qū)開發(fā)者接入不同的游戲或者RL訓(xùn)練環(huán)境。

接入C++引擎時(shí)，開發(fā)者只需要定義好單個(gè)環(huán)境的執(zhí)行邏輯，EnvPool則負(fù)責(zé)分布式執(zhí)行并且提供批交互的API，這使得接入新的RL環(huán)境及游戲引擎并得到立刻的并行加速變得非常的簡(jiǎn)單。

目前EnvPool在進(jìn)行高度開發(fā)，接下來的核心更新包括接入更多的RL環(huán)境（包括連續(xù)動(dòng)作空間的近期宣布免費(fèi)使用的Mujoco），及會(huì)提供樣例讓用戶能很方便地利用EnvPool加速現(xiàn)有的開源RL訓(xùn)練庫，包括接入DeepMind的Acme，以及接入EnvPool同一第一作者的大受歡迎的RL訓(xùn)練庫天授Tianshou。

此外，EnvPool的成果也被邀請(qǐng)到NVIDIA GTC 2021大會(huì)上演講。