[论文解读] EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine

Jiayi Weng, Min Lin+10 more|arXiv (Cornell University)|Jun 21, 2022

Reinforcement Learning in Robotics被引用 14

一句话总结

EnvPool 是一个高度并行的强化学习环境执行引擎，通过基于 C++ 的线程池执行器优化环境模拟，显著加速强化学习训练，在 MuJoCo 上实现高达每秒 300 万个物理步，在 Atari 上实现每秒 100 万个帧，比基线 Python 实现快 19.2 倍。它可与现有强化学习库无缝集成，支持在笔记本电脑上五分钟内完成 Atari Pong 或 MuJoCo Ant 的快速训练。

ABSTRACT

There has been significant progress in developing reinforcement learning (RL) training systems. Past works such as IMPALA, Apex, Seed RL, Sample Factory, and others, aim to improve the system's overall throughput. In this paper, we aim to address a common bottleneck in the RL training system, i.e., parallel environment execution, which is often the slowest part of the whole system but receives little attention. With a curated design for paralleling RL environments, we have improved the RL environment simulation speed across different hardware setups, ranging from a laptop and a modest workstation, to a high-end machine such as NVIDIA DGX-A100. On a high-end machine, EnvPool achieves one million frames per second for the environment execution on Atari environments and three million frames per second on MuJoCo environments. When running EnvPool on a laptop, the speed is 2.8x that of the Python subprocess. Moreover, great compatibility with existing RL training libraries has been demonstrated in the open-sourced community, including CleanRL, rl_games, DeepMind Acme, etc. Finally, EnvPool allows researchers to iterate their ideas at a much faster pace and has great potential to become the de facto RL environment execution engine. Example runs show that it only takes five minutes to train agents to play Atari Pong and MuJoCo Ant on a laptop. EnvPool is open-sourced at https://github.com/sail-sg/envpool.

研究动机与目标

• 解决强化学习训练系统中并行环境执行的低效瓶颈问题。
• 在从笔记本电脑到高端 DGX-A100 系统的各类硬件上提升模拟吞吐量。
• 通过提供 OpenAI Gym 和 dm_env API 的即插即用替代方案，使研究人员在极少代码修改下即可更快训练智能体。
• 支持同步和异步执行模式，以最大化硬件利用率。
• 确保与主流强化学习训练库（如 CleanRL、rl_games 和 DeepMind Acme）的广泛兼容性。

提出的方法

• 实现基于 C++ 的线程池执行器，以并行运行多个强化学习环境，最大限度减少 Python GIL 和 I/O 开销。
• 在 C++ 端优化 Python 包装器，降低环境交互和观测处理的延迟。
• 通过标准 OpenAI Gym 和与 dm_env 兼容的 Python API 暴露环境交互，实现无缝集成。
• 支持同步和异步执行模式，后者可实现更高吞吐量并减少数据陈旧性。
• 设计模块化架构，使开发人员可轻松将基于 C++ 的自定义强化学习环境集成到 EnvPool 中。
• 采用高效的内存管理与 C++ 和 Python 之间的零拷贝数据传输，减少序列化开销。

实验结果

研究问题

• RQ1 一个高度并行化的环境执行引擎是否能显著降低强化学习训练中的模拟延迟？
• RQ2 EnvPool 的性能在从笔记本电脑到高端 DGX 系统的不同硬件配置下如何扩展？
• RQ3 在保持标准强化学习基准中样本效率的前提下，EnvPool 能在多大程度上提升训练吞吐量？
• RQ4 EnvPool 是否能无需修改智能体代码，高效集成到现有强化学习训练流水线中？
• RQ5 在环境批处理中，与同步执行相比，异步执行可实现多大的性能提升？

主要发现

• 在配备 256 个 CPU 核心的 NVIDIA DGX-A100 上，EnvPool 在 Atari 上实现每秒 100 万个帧，在 MuJoCo 上实现每秒 300 万个物理步，分别比基线 Python 实现快 14.9 倍和 19.2 倍。
• 在配备 12 个 CPU 核心的笔记本电脑上，EnvPool 的运行速度比标准 Python 子进程式环境执行快 2.8 倍。
• 在笔记本电脑上，Atari Pong 和 MuJoCo Ant 的端到端训练均在五分钟内完成，证明其在快速研究迭代中的实际可用性。
• 在 MuJoCo HalfCheetah-v3 的相同 Acme PPO 训练设置下，EnvPool 的训练壁时延小于 DummyVecEnv 的一半，同时保持了相近的最终回合回报。
• 训练曲线显示，只要保持 num_envs 和 batch_size 的乘积恒定，增加并行环境数量（num_envs）可显著缩短训练时间，且不降低样本效率。
• 该系统与主流强化学习框架（包括 CleanRL、rl_games、Ray 和 DeepMind Acme）表现出强大的兼容性，支持即插即用的集成。