[论文解读] TorchBeast: A PyTorch Platform for Distributed RL

Heinrich Küttler, Nantas Nardelli+5 more|arXiv (Cornell University)|Oct 8, 2019

Model-Driven Software Engineering Techniques参考文献 13被引用 29

一句话总结

TorchBeast 是一个基于 PyTorch 的分布式强化学习平台，实现了 IMPALA 算法，重点关注简洁性、可扩展性和性能。它提供了纯 Python 版本（MonoBeast）和高性能、多机版本（PolyBeast），后者使用 C++ 实现系统级组件，在 Atari 环境中实现了与基于 TensorFlow 的 IMPALA 相当的性能。

ABSTRACT

TorchBeast is a platform for reinforcement learning (RL) research in PyTorch. It implements a version of the popular IMPALA algorithm for fast, asynchronous, parallel training of RL agents. Additionally, TorchBeast has simplicity as an explicit design goal: We provide both a pure-Python implementation ("MonoBeast") as well as a multi-machine high-performance version ("PolyBeast"). In the latter, parts of the implementation are written in C++, but all parts pertaining to machine learning are kept in simple Python using PyTorch, with the environments provided using the OpenAI Gym interface. This enables researchers to conduct scalable RL research using TorchBeast without any programming knowledge beyond Python and PyTorch. In this paper, we describe the TorchBeast design principles and implementation and demonstrate that it performs on-par with IMPALA on Atari. TorchBeast is released as an open-source package under the Apache 2.0 license and is available at \url{https://github.com/facebookresearch/torchbeast}.

研究动机与目标

• 解决当前缺乏可访问、高性能且可扩展的基于 PyTorch 的分布式强化学习实现的问题。
• 通过将所有机器学习代码保留在 Python 中使用 PyTorch，使研究人员能够在无需深入了解系统知识的情况下快速原型设计并扩展强化学习算法。
• 提供一个最小化、可轻松 fork 的代码库，支持单机原型设计和大规模多机训练。
• 确保与 OpenAI Gym 环境的兼容性，并保持与成熟的基于 TensorFlow 的 IMPALA 实现相当的性能。

提出的方法

• 使用基于 V-trace 的离策略校正方法，实现基于 PyTorch 的 IMPALA 算法，采用演员-评论家学习框架。
• 将系统划分为演员（数据收集）、学习者（模型更新）和批处理机制，将轨迹存储为结构化张量。
• 采用多进程架构，使用共享内存和套接字（MonoBeast）或 RPC 与 C++ 线程（PolyBeast），绕过 GIL 限制，实现高吞吐量训练。
• 以每个组件一个主文件（如 polybeast.py）的方式组织代码库，以最大化可读性和可修改性。
• 将系统级组件（如批处理、队列）与学习逻辑解耦，使研究人员仅需修改学习和环境代码即可。
• 使用基于 Python 的推理和训练循环，支持动态批处理和梯度更新，隐藏底层并行性实现细节。

实验结果

研究问题

• RQ1 基于 PyTorch 的 IMPALA 实现是否能在标准强化学习基准测试中达到与原始基于 TensorFlow 的 IMPALA 相当的性能？
• RQ2 纯 Python 实现在不牺牲性能的前提下，能在多大程度上支持可扩展的、高吞吐量的分布式强化学习训练？
• RQ3 如何设计一个面向研究、最小化代码量的框架，以支持快速原型设计，同时支持大规模多机训练？
• RQ4 在保持高性能的分布式强化学习系统中，为兼顾简洁性和可读性，需要做出哪些工程权衡？
• RQ5 与纯 Python 或其他系统相比，使用 C++ 实现系统组件（如批处理）在性能和可维护性方面有何影响？

主要发现

• 在 Atari 套件上，TorchBeast 在吞吐量、数据效率、稳定性以及最终智能体性能方面均达到与原始基于 TensorFlow 的 IMPALA 实现相当的水平。
• 纯 Python 的 MonoBeast 变体无需掌握底层系统编程知识，即可实现快速原型设计和实验。
• PolyBeast 变体通过使用 C++ 实现系统级组件，在多台机器上实现了有效的可扩展性，同时保持所有机器学习组件在 PyTorch 中，实现了高吞吐量。
• 该设计成功隔离了系统级复杂性，使研究人员仅需修改智能体和环境代码，而无需接触底层基础设施。
• 该平台通过多进程共享内存和 RPC 有效缓解了 Python 全局解释器锁（GIL）的影响，实现了真正的并行性。
• 以 Apache 2.0 许可证开源发布，确保了研究社区的广泛可访问性和可扩展性。