单机超越分布式?!强化学习新姿势,并行环境模拟器EnvPool实现速度成本双赢

简介: 单机超越分布式?!强化学习新姿势,并行环境模拟器EnvPool实现速度成本双赢

在训练强化学习智能体的时候,你是否为训练速度过慢而发愁?又是否对昂贵的大规模分布式系统加速望而却步?来自 Sea AI Lab 团队的最新研究结果表明,其实鱼和熊掌可以兼得:对于强化学习标准环境 Atari 与 Mujoco,如果希望在短时间内完成训练,需要采用数百个 CPU 核心的大规模分布式解决方案;而使用 EnvPool,只需要一台游戏本就能完成相同体量的训练任务,并且用时不到 5 分钟,极大地降低了训练成本。



目前,EnvPool 项目已在 GitHub 开源,收获超过 500 Stars,并且受到众多强化学习研究者的关注。


项目地址:https://github.com/sail-sg/envpool

在线文档:https://envpool.readthedocs.io/en/latest/

arXiv 链接:https://arxiv.org/abs/2206.10558


EnvPool 是一个基于 C++ 、高效、通用的强化学习并行环境(vectorized environment)模拟器,不仅能够兼容已有的 gym/dm_env API,还支持了多智能体环境。除了 OpenAI Gym 本身拥有的环境外,EnvPool 还支持一些额外的复杂环境。

目前支持的环境有:

Atari games

Mujoco(gym)

Classic control RL envs: CartPole, MountainCar, Pendulum, Acrobot

Toy text RL envs: Catch, FrozenLake, Taxi, NChain, CliffWalking, Blackjack

ViZDoom single player

DeepMind Control Suite


追求极致速度

EnvPool 采取了 C++ 层面的并行解决方案。根据现有测试结果,使用 EnvPool 并行运行多个强化学习环境,能在正常笔记本上比主流的 Python Subprocess 解决方案快近 3 倍;使用多核 CPU 服务器能够达到更好的性能。

例如在 NVIDIA DGX-A100(256 核 CPU 服务器)上的测试结果表明,Atari 游戏能够跑出每秒一百多万帧的惊人速度,Mujoco 物理引擎的任务更是能跑出每秒三百多万模拟步数的好成绩,比 Python Subprocess 快近二十倍,比之前最快的 CPU 异步模拟器 Sample Factory 还快两倍。

与此同时,EnvPool + CleanRL 的整系统测试表明,使用原始的 PPO 算法,直接把原来基于 Python Subprocess 的主流解决方案替换成 EnvPool,整体系统在标准的 Atari 基准测试中能快近三倍!

标准测试结果表明,对于数量稍大(比如超过 32)的并行环境,Subprocess 的运行效率十分堪忧。

为此有研究者提出分布式解决方案(比如 Ray)和基于 GPU 的解决方案(比如 Brax 和 Isaac-gym)进行加速。分布式方案经过测试,计算资源利用率其实并不高;基于 GPU 的解决方案虽然可以达到千万 FPS,但并不是所有环境都能使用 CUDA 重写,不能很好兼容生态以及不能复用一些受商业保护的代码。

EnvPool 由于采用了 C++ 层面的并行解决方案,并且大部分强化学习环境都使用 C++ 实现来保证运行效率,因此只需要在 C++ 层面实现接口即可完成高效的并行。

打造开放生态

EnvPool 对上下游的生态都有着良好的支持:

对于上游的强化学习环境而言,目前最多使用的是 OpenAI Gym,其次是 DeepMind 的 dm_env 接口。EnvPool 对两种环境 API 都完全支持,并且每次 env.step 出来的数据都是经过 numpy 封装好的,用户不必每次手动合并数据,同时也提高了吞吐量;

对于下游的强化学习算法库而言,EnvPool 支持了目前 PyTorch 最为流行的两个算法库 Stable-baselines3 和 Tianshou,同时还支持了 ACME、CleanRL 和 rl_games 等强化学习算法库,并且达到了令人惊艳的效果(在笔记本电脑上,2 分钟训练完 Atari Pong、5 分钟训练完 Mujoco Ant/HalfCheetah,并且通过了 Gym 原本环境的验证)。


为了更好地营造生态,EnvPool 采用了 Bazel 进行构建,拥有完善的软件工程标准,也提供了高质量的代码、单元测试和在线文档

使用示例

以下是一些简单的 EnvPool 使用示例。首先导入必要的包:


import numpy as np, envpool


以 gym.Env 接口为例,初始化 100 个 Atari Pong 的并行环境,只需要一行代码:


env = envpool.make_gym("Pong-v5", num_envs=100)


访问 observation_space 和 action_space 和 Gym 如出一辙:



observation_space = env.observation_spaceaction_space = env.action_space


在同步模式下,API 与已有的 Gym API 无缝衔接,只不过第一维大小是并行环境个数:




obs = env.reset()  # should be (100, 4, 84, 84)act = np.zeros(100, dtype=int)obs, rew, done, info = env.step(act)


当然也可以只 step/reset 部分环境,只需多传一个参数 env_id 即可:





while True:  act = policy(obs)  obs, rew, done, info = env.step(act, env_id)  env_id = info["env_id"]


为了追求极致性能,EnvPool 还支持异步模式的 step/reset,于在线文档和论文中对此进行了详细阐述与实验。

实际体验效果

EnvPool 只需要一句命令 pip install envpool 就能安装,目前仅支持 Linux 操作系统并且 Python 版本为 3.7/3.8/3.9。EnvPool 在短期内会对其他操作系统(Windows 和 macOS)进行支持。

在 EnvPool 的论文中,作者们给出了如下 rl_games 的惊艳结果:只用一台游戏本,在其他条件都完全相同的情况下,把基于 ray 的 vectorized env 实现直接换成 EnvPool,能够直接获得免费的加速。

Atari Pong 使用了 EnvPool,可以在大约 5 分钟的时候训练完成,相比而言 ray 的方案需要在大约 15 分钟的时候训练完成(达到接近 20 的 reward);Mujoco Ant 更为明显,使用原始 PPO 算法在不到 5 分钟的时间内达到了超过 5000 的 reward,而基于 ray 的解决方案运行了半小时还没达到 5000。



为了验证上述结果的真实性,我们尝试着运行了一下在 EnvPool README 提供的 colab 示例:

Pong:https://colab.research.google.com/drive/1iWFv0g67mWqJONoFKNWUmu3hdxn_qUf8?usp=sharing

Ant:https://colab.research.google.com/drive/1C9yULxU_ahQ_i6NUHCvOLoeSwJovQjdz?usp=sharing


我们找了台和论文 appendix 中与 rl_games 的实验类似配置的机子:

OS: Debian GNU/Linux 11 (bullseye) x86_64

Kernel: 5.16.0-0.bpo.4-amd64

CPU: AMD Ryzen 9 5950X (32) @ 3.400GHz

GPU: NVIDIA GeForce RTX 3080 Ti

Memory: 128800MiB


从 colab 的结果来看,Pong 跑到 200 个 epoch 结果已经收敛了,于是设置 200 个 epoch 跑测试,发现在这台测试机上运行代码,居然只需要要 2 分 17 秒就能跑完,并且 reward 能超过 19:


Ant 跑到 1000 个 epoch 结果差不多收敛了。设置 max_epochs 为 1000 之后运行,在 2 分 36 秒之后运行完毕,并且 reward 超过了 5300:



相比 RLLib 和 SeedRL 等工作,使用几百个 CPU 核心的分布式计算,EnvPool + rl_games 只用了单机做到了同样的效果(甚至更好),确实不错。

相关文章
|
6月前
|
机器学习/深度学习 分布式计算 数据处理
分布式计算框架:并行力量的交响乐章
分布式计算框架如Apache Spark解决单机计算挑战,通过拆分任务到多机并行处理提升效率。Spark以其内存计算加速处理,支持批处理、查询、流处理和机器学习。以下是一个PySpark统计日志中每日UV的示例,展示如何利用SparkContext、map和reduceByKey进行数据聚合分析。这些框架的运用,正改变大数据处理领域,推动数据分析和机器学习的边界。【6月更文挑战第18天】
250 2
|
1月前
|
机器学习/深度学习 边缘计算 人工智能
第二届边缘计算与并行、分布式计算国际学术会议(ECPDC 2025) 2025 2nd international Conference on Edge Computing, Parallel and Distributed Computing
第二届边缘计算与并行、分布式计算国际学术会议(ECPDC 2025) 2025 2nd international Conference on Edge Computing, Parallel and Distributed Computing 机器学习 计算学习理论 数据挖掘 科学计算 计算应用 数字图像处理 人工智能
59 6
|
2月前
|
SQL NoSQL 安全
分布式环境的分布式锁 - Redlock方案
【10月更文挑战第2天】Redlock方案是一种分布式锁实现,通过在多个独立的Redis实例上加锁来提高容错性和可靠性。客户端需从大多数节点成功加锁且总耗时小于锁的过期时间,才能视为加锁成功。然而,该方案受到分布式专家Martin的质疑,指出其在特定异常情况下(如网络延迟、进程暂停、时钟偏移)可能导致锁失效,影响系统的正确性。Martin建议采用fencing token方案,以确保分布式锁的正确性和安全性。
54 0
|
2月前
|
存储 SQL 消息中间件
Hadoop-26 ZooKeeper集群 3台云服务器 基础概念简介与环境的配置使用 架构组成 分布式协调框架 Leader Follower Observer
Hadoop-26 ZooKeeper集群 3台云服务器 基础概念简介与环境的配置使用 架构组成 分布式协调框架 Leader Follower Observer
52 0
|
4月前
|
存储 分布式计算 算法
探索Hadoop的三种运行模式:单机模式、伪分布式模式和完全分布式模式
在配置Hadoop集群之前,了解这三种模式的特点、适用场景和配置差异是非常重要的。这有助于用户根据个人需求和资源情况,选择最适合自己的Hadoop运行模式。在最初的学习和开发阶段,单机模式和伪分布式模式能为用户提供便利和成本效益。进而,当用户要处理大规模数据集时,完全分布式模式将是理想的选择。
261 2
|
4月前
|
分布式计算 并行计算 大数据
NumPy 并行计算与分布式部署
【8月更文第30天】随着数据量的不断增长,传统的单机计算模型已经难以满足对大规模数据集处理的需求。并行和分布式计算成为了处理这些大数据集的关键技术。虽然 NumPy 本身并不直接支持并行计算,但可以通过结合其他库如 Numba 和 Dask 来实现高效的并行和分布式计算。
42 1
|
4月前
|
机器学习/深度学习 分布式计算 PyTorch
大规模数据集管理:DataLoader在分布式环境中的应用
【8月更文第29天】随着大数据时代的到来,如何高效地处理和利用大规模数据集成为了许多领域面临的关键挑战之一。本文将探讨如何在分布式环境中使用`DataLoader`来优化大规模数据集的管理与加载过程,并通过具体的代码示例展示其实现方法。
253 1
|
4月前
|
消息中间件 Java Kafka
如何在Kafka分布式环境中保证消息的顺序消费?深入剖析Kafka机制,带你一探究竟!
【8月更文挑战第24天】Apache Kafka是一款专为实时数据管道和流处理设计的分布式平台,以其高效的消息发布与订阅功能著称。在分布式环境中确保消息按序消费颇具挑战。本文首先介绍了Kafka通过Topic分区实现消息排序的基本机制,随后详细阐述了几种保证消息顺序性的策略,包括使用单分区Topic、消费者组搭配单分区消费、幂等性生产者以及事务支持等技术手段。最后,通过一个Java示例演示了如何利用Kafka消费者确保消息按序消费的具体实现过程。
167 3
|
4月前
|
资源调度 Java 调度
项目环境测试问题之Schedulerx2.0通过分布式分片任务解决单机计算瓶颈如何解决
项目环境测试问题之Schedulerx2.0通过分布式分片任务解决单机计算瓶颈如何解决
项目环境测试问题之Schedulerx2.0通过分布式分片任务解决单机计算瓶颈如何解决
|
4月前
|
监控 负载均衡 Java
(九)漫谈分布式之微服务组件篇:探索分布式环境下各核心组件的必要性!
本文将深入探讨微服务中各个组件的必要性,以此帮助各位更好地加深对分布式系统的掌握度。
156 1
下一篇
DataWorks