如何做到史上最快的基于CPU的通用強(qiáng)化學(xué)習(xí)環(huán)境并行模擬器？
點(diǎn)擊上方“程序員大白”，選擇“星標(biāo)”公眾號(hào)
重磅干貨，第一時(shí)間送達(dá)
導(dǎo)讀
?一個(gè)基于C++的、高效的、通用的RL并行環(huán)境模擬器EnvPool，并且能兼容已有的gym/dm_env API。根據(jù)現(xiàn)有測(cè)試結(jié)果，能在正常的筆記本上比python subprocess快個(gè)3倍，在256核的CPU上能全吃滿，跑出一百萬FPS。?
這不是標(biāo)題黨謝謝，這是我暑假做的項(xiàng)目）
TL; DR: EnvPool是一個(gè)基于C++的、高效的、通用的RL并行環(huán)境（俗稱vector env）模擬器，并且能兼容已有的gym/dm_env API。根據(jù)現(xiàn)有測(cè)試結(jié)果，能在正常的筆記本上比python subprocess快個(gè)3倍，并且在服務(wù)器上加速比也很高（核越多越高），在256核的CPU上能全吃滿，跑出一百萬FPS，比subprocess快十幾倍，比之前最快的異步模擬器 Sample Factory（https://github.com/alex-petrenko/sample-factory/）還快個(gè)一倍多兩倍（有的時(shí)候甚至同步版EnvPool跑的都比Sample Factory快）！
現(xiàn)在已經(jīng)在GitHub開源了！歡迎批評(píng)以及各種評(píng)測(cè)（雖然還有很多東西沒從內(nèi)部版checkout出來……在做了在做了……咕咕咕）
https://github.com/sail-sg/envpool
使用的話……原來SubprocVecEnv怎么用，EnvPool也可以一樣用，當(dāng)然還有一些高端用法后面會(huì)慢慢說
為啥要做這個(gè)項(xiàng)目
首先玩過RL的人應(yīng)該都知道，如果把1個(gè)env和1個(gè)agent交互的過程變成n個(gè)env和1個(gè)agent交互，那么訓(xùn)練速度會(huì)大大提高（如果調(diào)參調(diào)的好的話），計(jì)算資源利用率也會(huì)提高，這樣就可以幾分鐘train出來一個(gè)不錯(cuò)的AI而不是原來的幾小時(shí)甚至幾天。（其實(shí)這也是 tianshou （https://github.com/thu-ml/tianshou）立項(xiàng)原因之一）
其次是目前主流的環(huán)境并行解決方案是python for-loop和python subprocess+share memory。雖然很方便寫，但是……太慢了。for-loop慢是很正常的，subprocess慢是因?yàn)閹讉€(gè)原因，一個(gè)是process的切換需要context switching，導(dǎo)致overhead比較大，第二個(gè)是python有GIL，三是數(shù)據(jù)傳輸也有額外開銷，四是python有點(diǎn)慢（所以正常的環(huán)境為了確保性能都是C++寫的）
那么我們?yōu)樯恫荒苤苯釉贑++這一層把并行給做了呢？對(duì)外只暴露batch的接口
然后有人要說了，為啥不分布式呢？ray其實(shí)也做了這個(gè)而且看起來挺靈活。但是根據(jù)測(cè)試，ray的計(jì)算資源利用率不太行，通常加速比是負(fù)數(shù)（，這樣會(huì)花很多冤枉錢在訓(xùn)練上面
然后有人要說了，為啥不GPU用CUDA做呢？其實(shí)有項(xiàng)目已經(jīng)做這個(gè)方面的優(yōu)化了比如 Brax / Isaac-gym，但是GPU的話其實(shí)不夠通用，要重寫整個(gè)env代碼，意味著不能很好兼容生態(tài)以及不能用一些受商業(yè)保護(hù)的代碼。所以現(xiàn)在他們只重寫了幾個(gè)類似mujoco的env，雖然能達(dá)到幾千萬FPS但畢竟不能適用于所有RL環(huán)境。
Brax：https://github.com/google/brax
?Isaac-gym：https://developer.nvidia.com/isaac-gym
然后有人要說了，為啥之前沒人做這個(gè)C++層面并行的idea？其實(shí)是有的，只不過沒做好。舉幾個(gè)例子是openai的 gym3 和 procgen 也是這么做的，但是我當(dāng)時(shí)順便測(cè)了下速發(fā)現(xiàn)他們 實(shí) 在 不 行。看這里有個(gè)討論：https://github.com/thu-ml/tianshou/issues/409
?gym3?：https://github.com/openai/gym3
procgen：https://github.com/openai/procgen
（一個(gè)小插曲：我還在寫EnvPool的時(shí)候，678月連著放了Brax/Isaac-gym/WarpDrive，雖然都是通過GPU進(jìn)行加速，但是感覺非?；?，害怕連這種代碼量巨大的項(xiàng)目也要卷……不過看起來是我多慮了）
（還有一個(gè)小插曲：某天看到票圈里面有個(gè)人說，花了三個(gè)月時(shí)間做Atari實(shí)驗(yàn)一共用了60萬RMB，說應(yīng)該把Atari剔除出RL benchmark。那我負(fù)責(zé)喊666然后順便可以推銷一波EnvPool，能省40萬的話我要的也不多，就40萬吧（霧
那么EnvPool性能究竟如何？
首先先解釋一下EnvPool是如何區(qū)分同步step和異步step的。實(shí)際上在內(nèi)部實(shí)現(xiàn)的時(shí)候這兩種直接被看做是同一種方法，簡(jiǎn)單來講舉個(gè)例子，有8個(gè)env（num_envs=8）但是我每次設(shè)置只返回其中4個(gè)的結(jié)果（batch_size=4），每次step的時(shí)候把剩下的env放在background繼續(xù)跑，這就是異步step；同步step就是num_envs == batch_size，每次等全部step好再返回。理論上這種異步step能夠有效提高CPU利用率（事實(shí)確實(shí)如此）
然后直接上圖
昨天和今天剛跑出來的。
首先EnvPool可以很方便的定義是同步版（每次step所有env，和vector_env一樣）還是異步版（每次異步step指定env），這里面可以認(rèn)為Sample Factory和EnvPool (async/numa) 是一組都是異步step，剩下的都是同步step。公平起見這里來算下每一組的倍數(shù)：
| Highest FPS | Laptop (12) | TPU-VM (96) | Apollo (96) | DGX-A100 (256) |
|---|---|---|---|---|
| For-loop | 8.37x | 46.09x | 39.28x | 108.43x |
| Subprocess | 2.24x | 4.10x | 8.14x | 5.91x |
| EnvPool (sync) | 1x | 1x | 1x | 1x |
| Sample Factory | 1.87x | 1.83x | 1.56x | 1.32x |
| EnvPool (async) | 1x | 1x | 1x | 1x |
為了不傷害Subprocess的自尊心，我這里就不拿async EnvPool和它比了。NUMA太優(yōu)秀了就先一邊站著。
值得一提的是Sample Factory是異步實(shí)現(xiàn)的env.step，但是和EnvPool同步版的性能卻差不多，更加證明了python-level parallalization是多么不行）
其實(shí)還有vizdoom的benchmark結(jié)果的，但是還沒從內(nèi)部版checkout到開源版本上面，到時(shí)候會(huì)更的！
那么……生態(tài)呢？
開源項(xiàng)目必須得在生態(tài)營(yíng)造的地方下功夫才不至于會(huì)死。
首先是安裝，使用 pip install envpool 就能在linux機(jī)子上安裝成功（需要python版本3.7以上）（時(shí)間太趕還要上課做作業(yè)暫時(shí)先沒考慮別的操作系統(tǒng)了……）
其次是用法，使用
env_gym = envpool.make_gym("Pong-v5", ...)
env_dm = envpool.make_dm("Pong-v5", ...)
# obs/act shape
observation_space = env_gym.observation_space
action_space = env_gym.action_space
就能輕松得到一個(gè)gym.Env或者dm_env，在同步模式下API是和已有的標(biāo)準(zhǔn)API完全一致，當(dāng)然也可以只step/reset部分env，只需要多傳一個(gè)參數(shù)就行：
while True:
act = policy(obs)
obs, rew, done, info = env.step(act, env_id)
env_id = info["env_id"]
（這其實(shí)已經(jīng)是異步step了，如果每次env_id都只有一部分的話）
以及在設(shè)計(jì)新版EnvPool的時(shí)候，我們還考慮了結(jié)構(gòu)化數(shù)據(jù)（比如nested observation/action）、multi-player env、以及二者融合的情況，接口都預(yù)留好了；
以及對(duì)于潛在的contributor而言，只需要在C++一側(cè)按照envpool/core/env.h的API把環(huán)境接好，然后實(shí)例化注冊(cè)一下就能成功用上EnvPool的加速。
對(duì)于開源項(xiàng)目最重要的測(cè)試和文檔，現(xiàn)在已經(jīng)搭了一套比較成熟的CI pipeline進(jìn)行C++和Python兩方面的測(cè)試，并且文檔已經(jīng)就位了：
里面有各種API解釋和各種可配置參數(shù)。以及順帶提供了一些examples在envpool/examples at master · sail-sg/envpool（https://github.com/sail-sg/envpool/tree/master/examples），理論上可以直接跑的（雖然確實(shí)還比較少但是……我不想掉頭發(fā)，會(huì)做的會(huì)做的 咕咕咕咕）
最后日常不要臉求star，再放一遍鏈接
https://github.com/sail-sg/envpool
推薦閱讀
關(guān)于程序員大白
程序員大白是一群哈工大，東北大學(xué)，西湖大學(xué)和上海交通大學(xué)的碩士博士運(yùn)營(yíng)維護(hù)的號(hào)，大家樂于分享高質(zhì)量文章，喜歡總結(jié)知識(shí)，歡迎關(guān)注[程序員大白]，大家一起學(xué)習(xí)進(jìn)步！