EnvPool: 高性能并行强化学习环境执行引擎

envpool

EnvPool:高性能并行强化学习环境执行引擎

EnvPool是一个基于C++的高性能批处理环境池,它利用pybind11和线程池技术,可以大幅提升强化学习环境的并行执行效率。EnvPool具有出色的性能表现,在高端设备上可以实现每秒约100万帧的Atari游戏模拟速度,以及每秒约300万步的Mujoco物理引擎模拟速度。同时,EnvPool还提供了兼容性很好的API,支持gym和dm_env两种接口规范,支持同步和异步执行模式,支持单玩家和多玩家环境。

主要特性

EnvPool的主要特性包括:

兼容OpenAI gym、DeepMind dm_env以及gymnasium的API
默认以批处理方式管理和交互多个环境
支持同步和异步两种执行模式
支持单玩家和多玩家环境
提供简单的C++ API用于添加新环境
单环境情况下可实现约2倍的加速
在256个CPU核心上可实现每秒100万Atari游戏帧或300万Mujoco模拟步数,比Python子进程方案快约20倍
在12核PC上比Python子进程方案快约3倍
相比GPU方案如Brax/Isaac-gym,EnvPool是一个更通用的RL环境并行化解决方案
支持JAX jit函数的XLA接口
兼容多个主流RL库,如Stable-Baselines3、Tianshou、ACME、CleanRL、rl_games等

支持的环境

目前EnvPool支持以下环境:

Atari游戏
Mujoco(gym版本)
经典控制问题:CartPole、MountainCar、Pendulum、Acrobot
文本类RL环境:Catch、FrozenLake、Taxi、NChain、CliffWalking、Blackjack
ViZDoom单人模式
DeepMind控制套件
Box2D物理引擎
Procgen程序生成环境
Minigrid网格世界

EnvPool支持的部分环境示例

安装

EnvPool目前已发布在PyPI上,支持Python 3.7及以上版本。可以通过以下命令安装:

pip install envpool

安装完成后,可以在Python中导入并检查版本:

import envpool
print(envpool.__version__)

如果没有报错,说明安装成功。

性能基准测试

EnvPool在不同硬件配置下进行了性能测试,包括笔记本电脑、工作站、TPU虚拟机和NVIDIA DGX-A100等。测试环境包括Atari游戏中的Pong和Mujoco环境中的Ant。

与基准方案相比,EnvPool在各种配置下都实现了显著的性能提升:

在高端256核DGX-A100上,EnvPool可以达到每秒100万Atari帧或300万Mujoco步数,比gym.vector_env基准快14.9/19.6倍。
在普通12核PC上,EnvPool的吞吐量是gym.vector_env的3.1/2.9倍。

下面是详细的性能数据对比:

Atari最高FPS	笔记本(12核)	工作站(32核)	TPU-VM(96核)	DGX-A100(256核)
For循环	4,893	7,914	3,993	4,640
子进程	15,863	47,699	46,910	71,943
Sample-Factory	28,216	138,847	222,327	707,494
EnvPool(同步)	37,396	133,824	170,380	427,851
EnvPool(异步)	49,439	200,428	359,559	891,286
EnvPool(NUMA+异步)	/	/	373,169	1,069,922

Mujoco最高FPS	笔记本(12核)	工作站(32核)	TPU-VM(96核)	DGX-A100(256核)
For循环	12,861	20,298	10,474	11,569
子进程	36,586	105,432	87,403	163,656
Sample-Factory	62,510	309,264	461,515	1,573,262
EnvPool(同步)	66,622	380,950	296,681	949,787
EnvPool(异步)	105,126	582,446	887,540	2,363,864
EnvPool(NUMA+异步)	/	/	896,830	3,134,287

EnvPool性能对比图

API使用示例

EnvPool提供了同步和异步两种API使用方式。以下是一些简单的使用示例:

同步API

import envpool
import numpy as np

# 创建gym环境
env = envpool.make("Pong-v5", env_type="gym", num_envs=100)
obs = env.reset()  # 形状应为(100, 4, 84, 84)
act = np.zeros(100, dtype=int)
obs, rew, term, trunc, info = env.step(act)

在同步模式下,EnvPool的使用方式与OpenAI Gym非常相似,主要区别在于默认采用批处理交互方式。

异步API

import envpool
import numpy as np

# 创建异步环境
num_envs = 64
batch_size = 16
env = envpool.make("Pong-v5", env_type="gym", num_envs=num_envs, batch_size=batch_size)
action_num = env.action_space.n
env.async_reset()  # 向所有环境发送初始重置信号
while True:
    obs, rew, term, trunc, info = env.recv()
    env_id = info["env_id"]
    action = np.random.randint(action_num, size=batch_size)
    env.send(action, env_id)

在异步模式下,step函数被拆分为send和recv两个函数。send函数不会等待环境执行完毕,而是立即返回。recv函数用于获取"下一个状态",但不保证获取的是刚刚send的那些环境的状态。

贡献

EnvPool仍在持续开发中,欢迎社区贡献更多环境和功能。如果你想参与贡献,可以查看项目的贡献指南。

许可证

EnvPool采用Apache 2.0许可证。其他第三方源代码和数据遵循其各自的许可证。

引用

如果你在研究中使用了EnvPool,请引用以下论文:

@inproceedings{weng2022envpool,
 author = {Weng, Jiayi and Lin, Min and Huang, Shengyi and Liu, Bo and Makoviichuk, Denys and Makoviychuk, Viktor and Liu, Zichen and Song, Yufan and Luo, Ting and Jiang, Yukun and Xu, Zhongwen and Yan, Shuicheng},
 booktitle = {Advances in Neural Information Processing Systems},
 editor = {S. Koyejo and S. Mohamed and A. Agarwal and D. Belgrave and K. Cho and A. Oh},
 pages = {22409--22421},
 publisher = {Curran Associates, Inc.},
 title = {Env{P}ool: A Highly Parallel Reinforcement Learning Environment Execution Engine},
 url = {https://proceedings.neurips.cc/paper_files/paper/2022/file/8caaf08e49ddbad6694fae067442ee21-Paper-Datasets_and_Benchmarks.pdf},
 volume = {35},
 year = {2022}
}

EnvPool为强化学习研究提供了一个高效的环境执行引擎,可以大幅提升训练速度,让研究人员能够更快地迭代想法。它有望成为强化学习环境执行的事实标准。如果你对EnvPool感兴趣,可以访问GitHub仓库了解更多信息,并尝试在你的项目中使用它。