參考鏈接：100萬幀數據僅1秒！AI大牛顏水成團隊強化學習新作，代碼已開源 (qq.com)
GTC 2021演講：https://events.rainfocus.com/widget/nvidia/nvidiagtc/sessioncatalog/session/1630239583490001Z5dE
論文：
開源代碼：
EnvPool是一種高度並行的強化學習環境執行引擎，其性能明顯優於現有環境執行器。通過專用於RL用例的策划設計，我們利用通用異步執行模型的技術，在環境執行中使用C++線程池實現。
以下是EnvPool的幾個亮點：
- 兼容OpenAI gym API和DeepMind dm_env API；
- 管理一個envs池，默認與批處理API中的envs交互；
- 同步執行API和異步執行API；
- 用於添加新環境的簡單C++開發人員API；
- 每秒模擬100萬Atari幀，具有256個CPU內核，吞吐量是基於Python子進程的向量環境的約13倍；
- 在低資源設置(如12個CPU內核)上，吞吐量是基於Python子進程的向量環境的吞吐量的約3倍；
- 與現有的基於GPU的解決方案(Brax/Isaac-gym)相比，EnvPool是各種加速RL環境並行化的通用解決方案；
- 兼容一些現有的強化學習庫，例如Tianshou。
Installation
PyPI
EnvPool目前托管在PyPI上。它需要Python≥3.7。您可以使用以下命令簡單地安裝EnvPool：
$ pip install envpool
安裝后，打開Python控制台並輸入：
import envpool print(envpool.__version__)
From Source
請參考guideline.
Documentation
教程和API文檔托管在envpool.readthedocs.io.
示例代碼位於文件夾examples/.
Supported Environments
我們正在從我們的內部版本中開源所有可用的環境，敬請期待。
Benchmark Results
我們使用ALE Atari環境(帶有環境包裝器)在不同硬件設置上執行基准測試，包括具有96個CPU內核和2個NUMA節點1的TPUv3-8虛擬機(VM)，以及具有256個CPU內核和8個NUMA節點的NVIDIA DGX-A100節點。基准包括1) 簡單的Python for循環；2) 最流行的RL環境並行化執行Python子進程，例如，gym.vector_env；3) 據我們所知，EnvPool之前最快的RL環境執行器Sample Factory。
我們報告了同步模式、異步模式和NUMA + 異步模式下的EnvPool性能，與不同數量的工作線程(即CPU內核數量)的基准進行比較。從結果中可以看出，EnvPool在所有設置的基准上都取得了顯著的改進。在高端設置上，EnvPool在256個CPU內核上實現了每秒100萬幀，這是gym.vector_env基准的13.3倍。在具有12個CPU內核的典型PC設置上，EnvPool的吞吐量是gym.vector_env的2.8倍。
我們的基准測試腳本在examples/benchmark.py中。4類系統的詳細配置為：
- Personal laptop: 12 core
Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz - TPU-VM: 96 core
Intel(R) Xeon(R) CPU @ 2.00GHz - Apollo: 96 core
AMD EPYC 7352 24-Core Processor - DGX-A100: 256 core
AMD EPYC 7742 64-Core Processor
1 NUMA (Non Uniform Memory Access)技術可以使眾多服務器像單一系統那樣運轉，同時保留小系統便於編程和管理的優點。非統一內存訪問(NUMA)是一種用於多處理器的電腦內存體設計，內存訪問時間取決於處理器的內存位置。在NUMA下，處理器訪問它自己的本地存儲器的速度比非本地存儲器(存儲器的地方到另一個處理器之間共享的處理器或存儲器)快一些。
基准測試腳本(examples/benchmark.py)：
import argparse import time import numpy as np import envpool if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--task", type=str, default="Pong-v5") parser.add_argument("--num-envs", type=int, default=645) parser.add_argument("--batch-size", type=int, default=248) # num_threads == 0 means to let envpool itself determine parser.add_argument("--num-threads", type=int, default=0) # thread_affinity_offset == -1 means no thread affinity parser.add_argument("--thread-affinity-offset", type=int, default=0) parser.add_argument("--total-iter", type=int, default=50000) args = parser.parse_args() env = envpool.make_gym( "Pong-v5", num_envs=args.num_envs, batch_size=args.batch_size, num_threads=args.num_threads, thread_affinity_offset=args.thread_affinity_offset, ) env.async_reset() action = np.ones(args.batch_size, dtype=np.int32) t = time.time() for _ in range(args.total_iter): info = env.recv()[-1] env.send(action, info["env_id"]) duration = time.time() - t fps = args.total_iter * args.batch_size / duration * 4 print(f"Duration = {duration:.2f}s") print(f"EnvPool FPS = {fps:.2f}")
envpool.list_all_envs():
>>> envpool.list_all_envs() ['Solaris-v5', 'BeamRider-v5', 'Earthworld-v5', 'YarsRevenge-v5', 'VideoCube-v5', 'MsPacman-v5', 'MazeCraze-v5', 'Tennis-v5', 'Klax-v5', 'Berzerk-v5', 'TicTacToe3d-v5', 'Asteroids-v5', 'Superman-v5', 'Seaquest-v5', 'VideoPinball-v5', 'MrDo-v5', 'Surround-v5', 'Kaboom-v5', 'AirRaid-v5', 'IceHockey-v5', 'Adventure-v5', 'MiniatureGolf-v5', 'ChopperCommand-v5', 'Robotank-v5', 'UpNDown-v5', 'JourneyEscape-v5', 'Atlantis2-v5', 'Blackjack-v5', 'WordZapper-v5', 'VideoCheckers-v5', 'Hero-v5', 'Pitfall-v5', 'SirLancelot-v5', 'Jamesbond-v5', 'SpaceWar-v5', 'Alien-v5', 'Frostbite-v5', 'Hangman-v5', 'DemonAttack-v5', 'Qbert-v5', 'LostLuggage-v5', 'StarGunner-v5', 'HumanCannonball-v5', 'Othello-v5', 'Galaxian-v5', 'BasicMath-v5', 'Pacman-v5', 'Combat-v5', 'KungFuMaster-v5', 'Darkchambers-v5', 'Et-v5', 'Gravitar-v5', 'Koolaid-v5', 'MarioBros-v5', 'Skiing-v5', 'Atlantis-v5', 'FishingDerby-v5', 'Freeway-v5', 'BankHeist-v5', 'RoadRunner-v5', 'Boxing-v5', 'Crossbow-v5', 'Zaxxon-v5', 'Warlords-v5', 'HauntedHouse-v5', 'Assault-v5', 'Breakout-v5', 'WizardOfWor-v5', 'FlagCapture-v5', 'Casino-v5', 'DonkeyKong-v5', 'Gopher-v5', 'ElevatorAction-v5', 'MontezumaRevenge-v5', 'Tutankham-v5', 'Krull-v5', 'Phoenix-v5', 'Enduro-v5', 'TimePilot-v5', 'Trondead-v5', 'KeystoneKapers-v5', 'Turmoil-v5', 'Venture-v5', 'LaserGates-v5', 'Pooyan-v5', 'Amidar-v5', 'Entombed-v5', 'Tetris-v5', 'Pong-v5', 'Pitfall2-v5', 'Frogger-v5', 'Joust-v5', 'Defender-v5', 'Riverraid-v5', 'NameThisGame-v5', 'Kangaroo-v5', 'VideoChess-v5', 'BattleZone-v5', 'Asterix-v5', 'CrazyClimber-v5', 'Centipede-v5', 'SpaceInvaders-v5', 'PrivateEye-v5', 'KingKong-v5', 'Bowling-v5', 'DoubleDunk-v5', 'Backgammon-v5', 'Carnival-v5']
envpool.make_gym(args):
- task_id：task id, use
envpool.list_all_envs()to see all support tasks; - num_envs：how many envs are in the envpool;how many envs are in the envpool, default to
1; - batch_size：async configuration, see the last section, default to
num_envs; - num_threads：the maximum thread number for executing the actual
env.step, default tobatch_size; - thread_affinity_offset：the start id of binding thread.
-1means not to use thread affinity in thread pool, and this is the default behavior;
>>> envpool.make_gym( ... "Pong-v5", ... num_envs=645, ... batch_size=248, ... num_threads=0, ... thread_affinity_offset=0, ... ) AtariGymEnvPool(num_envs=645, batch_size=248, num_threads=0, max_num_players=1, thread_affinity_offset=0, base_path='/opt/conda/lib/python3.8/site-packages/envpool', seed=42, max_episode_steps=25000, stack_num=4, frame_skip=4, noop_max=30, zero_discount_on_life_loss=False, episodic_life=False, reward_clip=False, img_height=84, img_width=84, task='pong')
API Usage
下面的內容展示了EnvPool的同步和異步API用法。您還可以在以下位置運行完整腳本examples/env_step.py
Synchronous API
import envpool import numpy as np # make gym env env = envpool.make("Pong-v5", env_type="gym", num_envs=100) # or use envpool.make_gym(...) obs = env.reset() # should be (100, 4, 84, 84) act = np.zeros(100, dtype=int) obs, rew, done, info = env.step(act)
在同步模式下，envpool與openai-gym/dm-env非常相似。具有同義的reset和step函數。但是有一個例外，在envpool中批處理交互是默認設置。因此，在創建envpool期間，有一個num_envs參數表示您希望並行運行多少個envs。
env = envpool.make("Pong-v5", env_type="gym", num_envs=100)
傳遞給step函數的action的第一個維度應該等於num_envs。
act = np.zeros(100, dtype=int)
當任何一個done為true時，您不需要手動重置一個環境，而是默認情況下envpool中的所有envs都啟用了自動重置。
Asynchronous API
import envpool import numpy as np # make asynchronous env = envpool.make("Pong-v5", env_type="gym", num_envs=64, batch_size=16) env.async_reset() # send the initial reset signal to all envs while True: obs, rew, done, info = env.recv() action = np.random.randint(batch_size, size=len(info.env_id)) env.send(action, env_id)
在異步模式下，step函數分為兩部分，即send/recv函數。send接受兩個參數，一批動作，以及每個動作應發送到的相應env_id。與step不同的是，send不等待envs執行並返回下一個狀態，它在將操作饋送到envs后立即返回。(之所以稱為異步模式的原因)。
env.send(action, env_id)
要獲得"下一個狀態"，我們需要調用recv函數。但是，recv並不能保證您會取回您剛剛調用send的環境的"下一個狀態"。取而代之的是，無論envs首先完成執行什么，都會首先被recv。
state = env.recv()
除了num_envs，還有一個參數batch_size。num_envs定義了envpool總共管理了多少個envs，而batch_size定義了每次與envpool交互時涉及的envs數量。例如envpool中有64個envs在執行，每次與一批16個envs 交互時，send和recv。
envpool.make("Pong-v5", env_type="gym", num_envs=64, batch_size=16)
envpool.make還有其他可配置的參數，請查看envpool interface introduction.
Contributing
EnvPool仍在開發中。將添加更多環境，我們始終歡迎貢獻以更好地幫助EnvPool。如果你想貢獻，請查看我們的contribution guideline.
API使用總結：
目前來說，文檔工作相對不足，很多細節並沒有在樣例中進行展示，例如如何在實際強化學習算法中調用等。所以在實際強化學習實驗中，我總結了以下幾點：
- env.reset()：只重置那些未初始化的或者已經done的環境，返回的狀態是0-255的像素值；
- env.recv()：返回的狀態是狀態是0-255的像素值，獎勵未經過裁剪；只有env剛經歷過reset或者send之類的操作且還未調用過recv()函數，recv()才會返回值，否則是一個空數組；