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目前已发布在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倍。

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

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仓库了解更多信息,并尝试在你的项目中使用它。