RLlib: Abstractions for Distributed Reinforcement Learning

A1 主要贡献

本文旨在解决强化学习（RL）领域在系统和抽象设计方面进展缓慢的问题，与深度学习（DL）领域形成鲜明对比。尽管RL社区受益于DL的系统进步，但RL算法固有的计算模式高度不规则且嵌套复杂，给并行化和组件复用带来了巨大挑战。现有分布式框架（如MPI、分布式TensorFlow）依赖于长时运行的程序副本间的通信，这种模式不自然地封装并行性和资源需求，导致组件复用困难，每次开发新算法时都需要重新实现复杂的分布式通信和执行逻辑。

核心问题：如何为具有异构、分布式和复杂嵌套并行特性的RL算法设计一个可组合的、可扩展的并行计算模型，以促进代码复用和研究效率。

研究目标与创新点：
本文主张围绕逻辑上集中的程序控制（logically centralized program control）和并行性封装（parallelism encapsulation）原则来构建分布式RL组件。作者提出了一种分层并行任务模型（hierarchical parallel task model），其中一个主控程序（driver）将算法子任务（如rollouts、梯度计算）委托给工作进程并行执行，而工作进程本身也可以进一步将任务委托给自己的子工作进程，从而支持嵌套计算。这种模型具有以下优势：
1. 简化开发：将分布式控制逻辑完全封装在单个进程中，比多进程并发执行更容易实现。
2. 促进代码复用：将算法组件（如策略评估、策略优化）分离为子程序，使其可以在不同的执行模式中复用。
3. 高效资源利用：允许将具有不同资源需求（如CPU vs. GPU）的子任务放置在不同机器上。
4. 无缝嵌套：以该模型编写的分布式算法可以无缝地相互嵌套。

基于上述模型，本文设计并实现了RLlib，一个可扩展的RL库。其主要贡献如下：
1. 提出了一种通用且可组合的分层编程模型，用于RL训练，解决了现有框架在组件组合和嵌套并行方面的局限性。
2. 详细介绍了RLlib库，它建立在该模型之上，为广泛的RL算法提供了可扩展的抽象（如策略评估器和策略优化器），从而实现了快速开发。
3. 展示了该模型的性能优势，并通过实验证明RLlib在多种RL工作负载上达到或超过了现有最先进实现（state-of-the-art）的性能。

下图展示了一个RL算法如何组合使用无导数优化、策略评估、基于梯度的优化和基于模型的规划，体现了RL算法在多个层次和物理设备上利用并行性的复杂性。

图1. 与深度学习不同，RL算法在多个层次和物理设备上利用并行性。这里我们展示了一个RL算法，它组合了无导数优化、策略评估、基于梯度的优化和基于模型的规划（表2）。
图1. 与深度学习不同，RL算法在多个层次和物理设备上利用并行性。这里我们展示了一个RL算法，它组合了无导数优化、策略评估、基于梯度的优化和基于模型的规划（表2）。

A3 背景知识/关键Observation/设计原则

1.1 RL训练工作负载的不规则性

现代RL算法的计算模式高度不规则。与深度学习中主导的张量代数等计算模式不同，RL算法的计算模式不规则性体现在多个层面，这对流行的分布式框架所支持的计算模型构成了挑战。

不规则性的具体表现。
1. 任务持续时间和资源需求差异巨大：根据算法不同，任务的持续时间和资源需求可能相差几个数量级。例如，A3C【【索引30，Asynchronous methods for deep reinforcement learning，Mnih, V., et al.，2016，International Conference on Machine Learning】】的更新可能只需几毫秒，而像PPO【【索引40，Proximal policy optimization algorithms，Schulman, J., et al.，2017，arXiv】】这样的算法则将rollouts分批处理，粒度要大得多。
2. 通信模式多样：通信模式各不相同，从同步到异步的基于梯度的优化，再到高吞吐量离策略算法（如Ape-X【【索引21，Distributed prioritized experience replay，Horgan, D., et al.，2018，International Conference on Learning Representations】】和IMPALA【【索引14，Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures，Espeholt, L., et al.，2018，arXiv】】）中存在的多种类型的异步任务。
3. 嵌套计算：嵌套计算由基于模型的混合算法（见表2）、与RL或DL训练结合的超参数调优，或在单个算法内结合无导数和基于梯度的优化（如AlphaGo Zero【【索引41，Mastering the game of Go without human knowledge，Silver, D., et al.，2017】】）产生。
4. 大量的状态维护：RL算法通常需要维护和更新大量状态，包括策略参数、重放缓冲区甚至外部模拟器。

现有框架的局限性。由于上述不规则性，开发者不得不混合使用多种框架，如参数服务器、类MPI框架中的集体通信原语、任务队列等。对于更复杂的算法，通常需要构建定制的分布式系统，其中进程独立计算并通过无中央控制的方式进行协调（如图2(a)所示）。虽然这种方法可以实现高性能，但开发和评估成本高昂，不仅因为需要实现和调试分布式程序，还因为组合这些算法会进一步使其实现复杂化（如图3所示）。此外，现有的计算框架（如Spark【【索引43，Spark: Cluster computing with working sets，Zaharia, M., et al.，2010】】、MPI）通常假设计算模式是规则的，当子任务具有不同的持续时间、资源需求或嵌套结构时，它们难以应对。

表1. RL涵盖了广泛的计算需求。

表1. RL涵盖了广泛的计算需求。
表1. RL涵盖了广泛的计算需求。

1.2 用于分布式RL的逻辑集中控制

提出逻辑集中控制模型。一个理想的编程模型应能捕捉RL训练的所有需求，而这并不需要放弃构建计算的高级框架。本文的关键洞见是，对于每个分布式RL算法，都可以编写一个等效的、表现出逻辑上集中的程序控制（logically centralized program control）的算法（如图2(b)所示）。也就是说，替代让独立执行的进程（图2(a)中的A, B, C, D）相互协调（例如，通过RPC、共享内存、参数服务器或集体通信），可以由一个单一的驱动程序（图2(b)和2(c)中的D）将算法子任务委托给其他进程并行执行。在这种模式下，工作进程A、B和C被动地持有状态（例如，策略或模拟器状态），但在被D调用之前不执行任何计算。

扩展至分层控制模型。为了支持嵌套计算，本文提出将集中控制模型扩展为分层委托控制（hierarchical delegation of control，如图2(c)所示）。这允许工作进程（如B, C）在执行任务时，能够进一步将工作（如模拟、梯度计算）委托给它们自己的子工作进程。

图2. 目前大多数RL算法采用完全分布式的方式编写(a)，其中复制的进程根据其角色（如果有）独立计算并相互协调。我们提出了一种分层控制模型(c)，它扩展了(b)以支持RL和超参数调优工作负载中的嵌套，从而简化和统一了用于实现的编程模型。
图2. 目前大多数RL算法采用完全分布式的方式编写(a)，其中复制的进程根据其角色（如果有）独立计算并相互协调。我们提出了一种分层控制模型(c)，它扩展了(b)以支持RL和超参数调优工作负载中的嵌套，从而简化和统一了用于实现的编程模型。

该模型的优势。构建于这样一个逻辑集中和分层控制模型之上，具有几个重要优势。首先，等效算法在实践中通常更容易实现，因为分布式控制逻辑完全封装在单个进程中，而不是多个并发执行的进程中。其次，将算法组件分离为子程序（例如，执行rollouts、计算关于某个策略损失的梯度），使得代码可以在不同的执行模式中复用。具有不同资源需求（例如，CPU vs GPU）的子任务可以被放置在不同的机器上，从而降低计算成本，这一点在第5节中得到了验证。最后，用这种模型编写的分布式算法可以无缝地相互嵌套，满足了并行性封装原则。

该模型的性能。逻辑集中控制模型可以实现高性能，而本文提出的分层变体性能更佳。这是因为进程之间的大部分数据传输（图2中的蓝色箭头）发生在驱动程序的带外（out of band），不会通过任何中央瓶颈。实际上，许多高度可扩展的分布式系统（例如Spark【【索引43，Spark: Cluster computing with working sets，Zaharia, M., et al.，2010】】、Bigtable【【索引8，Bigtable: A distributed storage system for structured data，Chang, F., et al.，2008，ACM Transactions on Computer Systems (TOCS)】】、MapReduce【【索引11，MapReduce: simplified data processing on large clusters，Dean, J. and Ghemawat, S.，2008，Communications of the ACM】】）在设计中都利用了集中控制。在单个可微的张量图内，像TensorFlow这样的框架也实现了将张量计算逻辑上集中调度到可用物理设备上。本文的提议将这一原则扩展到更广泛的机器学习系统设计领域。

2. 分层并行任务模型

现有框架在组合性上的缺陷。如图3所示，使用像MPI【【索引18，A highperformance, portable implementation of the MPI message passing interface standard，Gropp, W., et al.，1996，Parallel computing】】和分布式Tensorflow【【索引1，TensorFlow: Large-scale machine learning on heterogeneous distributed systems，Abadi, M., et al.，2016，arXiv】】这样的框架对整个程序进行并行化，在尝试组合两个程序或组件时，通常需要对算法进行显式修改以插入协调点。这限制了快速原型化新颖分布式RL应用的能力。尽管图3中的例子很简单，但新的针对长时训练任务的超参数调优算法，例如HyperBand【【索引25，Hyperband: Bandit-based configuration evaluation for hyperparameter optimization，Li, L., et al.，2016】】和基于种群的训练（PBT）【【索引22，Population based training of neural networks，Jaderberg, M., et al.，2017，arXiv】】，越来越要求对训练过程进行细粒度控制。

基于任务的编程模型。我们建议在像Ray【【索引31，Ray: A distributed framework for emerging AI applications，Moritz, P., et al.，2017，arXiv】】这样灵活的基于任务的编程模型之上，构建具有分层和逻辑集中控制的RL库。基于任务的系统允许子程序以细粒度的方式被异步地调度和执行在工作进程上，并且结果可以在进程之间被检索或传递。

图3. 将分布式超参数搜索与一个同样需要分布式计算的函数组合在一起，涉及到复杂的嵌套并行计算模式。使用MPI(a)，必须从头编写一个混合了两者元素的新程序。而使用分层控制(b)，组件可以保持不变，只需作为远程任务被调用即可。
图3. 将分布式超参数搜索与一个同样需要分布式计算的函数组合在一起，涉及到复杂的嵌套并行计算模式。使用MPI(a)，必须从头编写一个混合了两者元素的新程序。而使用分层控制(b)，组件可以保持不变，只需作为远程任务被调用即可。

2.1 与现有分布式ML抽象的关系

在逻辑集中控制模型中使用现有抽象。尽管通常是为分布式控制而设计的，像参数服务器和集体通信操作这样的抽象也可以在逻辑集中控制模型中使用。例如，RLlib在其某些策略优化器中使用了allreduce和参数服务器（如图4所示），我们将在第5节评估它们的性能。

2.2 Ray对分层控制的实现

实现选择。我们注意到，在单台机器内，所提出的编程模型可以简单地用线程池和共享内存来实现，但理想情况下，底层框架应能按需扩展到更大的集群。

选择Ray作为底层框架。我们选择在Ray框架之上构建RLlib，因为它允许Python任务分布在大型集群上。Ray的分布式调度器天然适合分层控制模型，因为在Ray中实现嵌套计算没有中央任务调度瓶颈。

Ray Actor实现逻辑集中控制。要实现逻辑集中控制模型，首先需要一种机制来启动新进程并在其上调度任务。Ray通过Ray actors满足了这一要求，Ray actors是可以在集群中创建并接受远程方法调用（即任务）的Python类。Ray允许这些actor在方法调用中继而启动更多的actor并在这些actor上调度任务，从而满足了我们对分层委托的需求。

Ray的性能特性。为了性能，Ray提供了标准的通信原语，如聚合和广播，并且关键地通过共享内存对象存储实现了大型数据对象的零拷贝共享。如第5节所示，这保证了RLlib算法的性能。我们将在第4节进一步讨论框架的性能。

A2 方法细节

3. 强化学习的抽象

RLlib的核心抽象。为了利用RLlib进行分布式执行，算法必须声明其策略（policy）$\pi$、经验后处理器（experience postprocessor）$\rho$和损失（loss）$L$。这些可以在任何深度学习框架中指定，包括TensorFlow和PyTorch。RLlib提供了策略评估器（policy evaluators）和策略优化器（policy optimizers），它们实现了分布式策略评估和训练的策略。

3.1 定义策略图

策略模型$\pi$。RLlib的抽象如下。开发者指定一个策略模型$\pi$，它将当前观测值$o_t$和（可选的）RNN隐藏状态$h_t$映射到一个动作$a_t$和下一个RNN状态$h_{t+1}$。也可以返回任意数量的用户定义值$y_{it}$（例如，价值预测、TD误差）：

公式1
公式1

轨迹后处理器$\rho$。大多数算法还会指定一个轨迹后处理器$\rho$，它转换一个从$t$开始的包含$K$个元组 $\{ (o_t, h_t, a_t, h_{t+1}, y_{1t} \dots y_{Nt}, r_t, o_{t+1}) \}$ 的批次$X_{t,K}$。这里$r_t$和$o_{t+1}$是在执行动作后得到的奖励和新观测值。典型的应用包括优势估计（advantage estimation）【【索引39，High-dimensional continuous control using generalized advantage estimation，Schulman, J., et al.，2015，arXiv】】和目标重标记（goal relabeling）【【索引4，Hindsight experience replay，Andrychowicz, M., et al.，2017，arXiv】】。为了同时支持多智能体环境，环境中其他$P$个智能体的经验批次$X^p_{t,K}$也可以被访问：

公式2
公式2

损失函数$L$。基于梯度的算法定义一个组合损失$L$，通过下降该损失可以改进策略和辅助网络：

公式3
公式3

效用函数$u_i$。最后，开发者还可以指定任意数量的效用函数$u_i$，在训练期间根据需要调用，例如，返回训练统计数据$s$，更新目标网络，或调整退火调度：

公式4
公式4

PolicyGraph接口。为了与RLlib对接，这些算法函数应在一个PolicyGraph类中定义，该类具有以下方法：

PolicyGraph接口定义
PolicyGraph接口定义

3.2 策略评估

PolicyEvaluator类。为了收集经验，RLlib提供了一个PolicyEvaluator类，该类包装了策略图和环境，并增加了一个用于sample()经验批次的方法。策略评估器实例可以作为Ray远程actor创建，并在集群中复制以实现并行化。为了使其用法具体化，我们考虑一个最小的TensorFlow策略梯度实现，该实现扩展了rllib.TFPolicyGraph帮助模板：

一个最小的TensorFlow策略梯度实现
一个最小的TensorFlow策略梯度实现

并行采样。基于这个策略图定义，开发者可以创建多个策略评估器副本ev，并在每个副本上调用ev.sample.remote()来并行地从环境中收集经验。RLlib支持OpenAI Gym【【索引6，OpenAI gym，Brockman, G., et al.，2016】】、用户定义的环境，以及像ELF【【索引42，ELF: an extensive, lightweight and flexible research platform for real-time strategy games，Tian, Y., et al.，2017，CoRR】】这样的批处理模拟器：

evaluators = [rllib.PolicyEvaluator.remote(
    env=SomeEnv, graph=PolicyGradient)
    for _ in range(10)]
print(ray.get([ev.sample.remote() for ev in evaluators]))

并行采样代码示例
并行采样代码示例

图4. 四种RLlib策略优化器step方法的伪代码。每个step()操作本地策略图和远程评估器副本数组。Ray远程调用以橙色突出显示；其他Ray原语以蓝色突出显示（第4节）。Apply是更新权重的简写。小批量代码和辅助函数已省略。RLlib中的参数服务器优化器还实现了此处未显示的流水线操作。
图4. 四种RLlib策略优化器step方法的伪代码。每个step()操作本地策略图和远程评估器副本数组。Ray远程调用以橙色突出显示；其他Ray原语以蓝色突出显示（第4节）。Apply是更新权重的简写。小批量代码和辅助函数已省略。RLlib中的参数服务器优化器还实现了此处未显示的流水线操作。

3.3 策略优化

PolicyOptimizer抽象。RLlib将算法的实现分为特定于算法的策略图的声明和与算法无关的策略优化器的选择。策略优化器负责性能关键的任务，如分布式采样、参数更新和管理重放缓冲区。为了分配计算，优化器操作一组策略评估器副本。

使用示例。为了完成示例，开发者选择一个策略优化器，并用现有评估器的引用来创建它。异步优化器使用评估器actor在多个CPU上并行计算梯度（图4(c)）。每个optimizer.step()运行一轮远程任务来改进模型。在步骤之间，可以直接查询策略图副本，例如，打印训练统计信息：

optimizer = rllib.AsyncPolicyOptimizer(
    graph=PolicyGradient, workers=evaluators)
while True:
    optimizer.step()
    print(optimizer.foreach_policy(
        lambda p: p.get_train_stats()))

与梯度下降优化器的类比。策略优化器将众所周知的梯度下降优化器抽象扩展到RL领域。一个典型的梯度下降优化器实现step(L(θ), X, θ) ⇒ θ_opt。RLlib的策略优化器则操作本地策略图G和一组远程评估器副本，即step(G, ev1...evn, θ) ⇒ θ_opt，将RL的采样阶段作为优化的一部分（即在策略评估器上调用sample()以产生新的模拟数据）。

PolicyOptimizer抽象的优势。将执行策略与策略和损失定义分开，具有以下优势。首先，可以换入专门的优化器以利用可用的硬件或算法特性，而无需更改算法的其余部分。其次，PolicyGraph类封装了与深度学习框架的交互，使算法作者可以避免将分布式系统代码与数值计算混合在一起，并使得优化器实现可以在不同的深度学习框架中得到改进和重用。

性能与多样性。如图4所示，通过利用集中控制，策略优化器简洁地捕捉了RL优化中的广泛选择：同步与异步、allreduce与参数服务器、以及使用GPU与CPU。RLlib的策略优化器提供了与优化的参数服务器（图5(a)）和基于MPI的实现（第5节）相当的性能。在逻辑集中控制模型中，提取出这种优化器抽象是容易的，因为每个策略优化器都完全控制它所实现的分布式计算。

图5. RLlib的集中控制策略优化器性能与专业系统中的实现相当或更优。使用8个内部shard的RLlib参数服务器优化器与在相似条件下测试的分布式TensorFlow实现具有竞争力。RLlib的Ape-X策略优化器在256个worker和4的帧跳过率下，可扩展至每秒16万帧，远超256个worker下约4.5万fps的参考吞吐量，表明单线程Python控制器可以高效地扩展到高吞吐量。
图5. RLlib的集中控制策略优化器性能与专业系统中的实现相当或更优。使用8个内部shard的RLlib参数服务器优化器与在相似条件下测试的分布式TensorFlow实现具有竞争力。RLlib的Ape-X策略优化器在256个worker和4的帧跳过率下，可扩展至每秒16万帧，远超256个worker下约4.5万fps的参考吞吐量，表明单线程Python控制器可以高效地扩展到高吞吐量。

3.4 抽象的完备性和通用性

完备性论证。我们通过在API中制定表2中列出的算法族来展示RLlib抽象的完备性。在适用时，我们还描述了在RLlib中的具体实现：

表2. RLlib的策略优化器和评估器在一个逻辑集中控制模型中捕获了通用组件（评估、回放、基于梯度的优化器），并利用Ray的分层任务模型来支持其他分布式组件。

表2. RLlib的策略优化器和评估器在一个逻辑集中控制模型中捕获了通用组件
表2. RLlib的策略优化器和评估器在一个逻辑集中控制模型中捕获了通用组件

图6. 复杂的RL架构可以很容易地在RLlib的分层控制模型中捕获。这里蓝线表示数据传输，橙线表示开销较轻的方法调用。每个train()调用都包含了一批组件之间的远程调用。
图6. 复杂的RL架构可以很容易地在RLlib的分层控制模型中捕获。这里蓝线表示数据传输，橙线表示开销较轻的方法调用。每个train()调用都包含了一批组件之间的远程调用。

4. 框架性能

本节我们讨论对RLlib至关重要的Ray【【索引31，Ray: A distributed framework for emerging AI applications，Moritz, P., et al.，2017，arXiv】】的属性和其他优化。

4.1 单节点性能

4.2 分布式性能

A4 实验环境

A4 实验结果

A5 结论

RLlib是一个用于强化学习的开源库，它利用细粒度的嵌套并行性，在广泛的RL工作负载上实现了最先进的性能。它既提供了一系列参考算法的实现，也提供了可扩展的抽象，使得用户可以轻松地组合和构建新的算法。RLlib的成功证明了以逻辑集中控制和分层任务模型为基础构建分布式RL系统的有效性，为该领域的研究和应用开发提供了强大的工具。