跳到主要内容

论文解读：RLlib — 分布式强化学习的抽象与实现

· 阅读需 14 分钟
Zhiyuan Pan
Blog Author

解读 ICML 2018 论文，RLlib 基于 Ray 框架提供分布式强化学习的通用抽象，支持多种 RL 算法的高效并行训练。

论文来源：ICML 2018, Eric Liang, Richard Liaw, Philipp Moritz, Robert Nishihara, Roy Fox, Ken Goldberg, Joseph E. Gonzalez, Michael I. Jordan, Ion Stoica (UC Berkeley) 项目地址：http://rllib.io | 基于 Ray 开源框架

1. 论文要解决什么问题？​

强化学习（RL）算法的计算模式和深度学习有着根本性的不同。深度学习的并行化相对"规整"——本质上就是把同一个模型复制到多张卡上做数据并行或模型并行。但 RL 算法的计算模式极其不规则，体现在以下四个维度：

  1. 任务粒度差异巨大：A3C 的一次更新可能只需要毫秒级，但 PPO 的一次 rollout 批次可能是分钟级的大颗粒任务。
  2. 通信模式多样：从同步梯度聚合，到异步梯度更新，再到 Ape-X 那种多种异步任务并存的高吞吐 off-policy 架构，模式五花八门。
  3. 嵌套计算随处可见：model-based 方法需要在内部嵌套规划；超参搜索（HyperBand、PBT）需要嵌套整个训练过程；AlphaGo 更是在一个算法里同时组合了 MCTS、自我对弈、模型训练等多个分布式组件。
  4. 大量有状态对象：策略参数、replay buffer、模拟器状态等，必须跨多个物理设备和进程管理。

当时的现状是：大家不得不用 MPI、Distributed TensorFlow、Parameter Server 等底层工具"手搓"分布式逻辑，每实现一个新算法就要重写大量分布式通信代码。不同算法之间几乎没有可复用的组件，组合两个算法更是噩梦（Figure 3 展示了用 MPI 做嵌套超参搜索有多痛苦）。

核心问题：能不能设计一套通用的抽象，让各种 RL 算法都能以可组合、可扩展的方式实现分布式训练？

2. 核心思路：逻辑集中控制 + 分层委托​

论文提出的编程模型核心理念是 Logically Centralized Control with Hierarchical Delegation（逻辑集中控制 + 分层委托）。

传统做法的问题​

传统的分布式 RL 通常采用"全分布式"架构（Figure 2(a)）：多个进程 A、B、C、D 各自独立运行，通过 MPI、参数服务器或共享内存相互协调。这种做法的问题是：

  • 分布式控制逻辑散落在所有进程中，调试困难
  • 每加一个新组件就要修改已有进程的协调代码
  • 组合两个算法时，协调逻辑爆炸性增长

RLlib 的方案​

RLlib 的做法是：用一个单线程的 driver 程序 D来做逻辑上的集中控制（Figure 2(b)）。Worker 进程 A、B、C 被动持有状态（策略参数、模拟器状态等），但不主动执行计算——只有被 driver 调用时才工作。

关键在于，这并不意味着所有数据都要经过中心节点。实际的数据传输（蓝色箭头）大部分在 worker 之间直接发生，不经过 driver，所以性能不受中心瓶颈限制。

为了支持嵌套计算（比如 AlphaGo 里套 MCTS），RLlib 进一步引入了分层委托（Figure 2(c)）：worker 进程 B、C 可以继续向自己的 sub-worker 委派任务，形成树状的控制层级。

这种设计的三大优势：

  • 易于实现：等价的算法逻辑集中在一个进程中，比分散在多个进程中简单得多
  • 代码复用：rollout、梯度计算、策略更新等子任务可以独立实现，跨算法共享
  • 天然支持嵌套：分层委托模型下，算法可以无缝嵌套组合

3. Ray：底层执行引擎​

RLlib 选择构建在 Ray 框架之上。Ray 提供了几个关键能力，恰好满足分层集中控制模型的需求：

  • Ray Actor：Python 类可以被创建为集群中的远程 actor，接受远程方法调用。Actor 可以进一步创建子 actor 并调度子任务——完美匹配分层委托的需求。
  • 共享内存对象存储（Object Store）：同一台机器上的 worker 可以通过共享内存零拷贝读取大数据对象（如 rollout 数据、神经网络权重），不需要序列化/反序列化。
  • 通信原语：aggregate 和 broadcast 等操作支持高效的集体通信。
  • 灵活的任务调度：任何 Python 函数/方法都可以通过 func.remote() 异步远程执行，返回一个 future 占位符，支持细粒度的嵌套并行。

4. 三大核心抽象​

RLlib 的 API 围绕三个层次分明的抽象构建：

4.1 Policy Graph（策略图）​

Policy Graph 定义了算法的核心数学逻辑，是最底层的抽象。一个 Policy Graph 需要指定：

  • 策略模型 $\pi_\theta(o_t, h_t) \Rightarrow (a_t, h_{t+1}, y_t^1 \ldots y_t^N)$：输入观测和 RNN 隐状态，输出动作、下一隐状态、以及任意辅助值（如 value prediction、TD error）
  • 后处理器 $\rho_\theta$：对一个 batch 的轨迹数据做变换，例如计算 advantage estimation、goal relabeling
  • 损失函数 $L(\theta; X)$：基于梯度的算法用来更新策略和辅助网络
  • 工具函数 $u^i$：如更新 target network、设置探索率 $\epsilon$ 等
abstract class rllib.PolicyGraph:
    def act(self, obs, h): action, h, y*
    def postprocess(self, batch, b*): batch
    def gradients(self, batch): grads
    def get_weights; def set_weights;
    def u*(self, args*)

Policy Graph 可以用任何深度学习框架实现（TensorFlow、PyTorch 等），它封装了与深度学习框架的交互，让上层不需要关心具体的框架细节。

4.2 Policy Evaluator（策略评估器）​

Policy Evaluator 将 Policy Graph 和环境（Environment）包装在一起，提供 sample() 方法来收集经验数据。

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

关键特性：

  • 可以创建为 Ray remote actor，分布到集群中并行采样
  • 支持 OpenAI Gym、用户自定义环境、以及 ELF 这样的 batched simulator
  • 支持向量化（vectorization），一个 evaluator 内部可以同时跑多个环境实例
  • 支持多智能体场景，允许同一环境中运行多个策略

4.3 Policy Optimizer（策略优化器）​

Policy Optimizer 是最上层的抽象，负责分布式训练策略的编排。它操作一个本地 Policy Graph 和一组远程 Policy Evaluator，实现 step() 方法来推进一步训练。

核心公式：$step(G, ev_1 \ldots ev_n, \theta) \Rightarrow \theta_{opt}$

这个抽象将执行策略和算法逻辑解耦：同一个算法（如 PPO）可以搭配不同的 optimizer 来适应不同的硬件条件，而不需要修改算法代码本身。

论文给出了四种 Policy Optimizer 实现（Figure 4）：

Optimizer 类型工作方式适用场景
Allreduce所有 evaluator 并行采样+计算梯度，聚合后更新，广播新权重同步训练，多 CPU/GPU
Local Multi-GPUevaluator 采样后 pin 到本地 GPU 显存，做多轮 SGDPPO 等需要多 epoch 训练的算法，GPU 显存充足时
Asynchronousevaluator 异步提交梯度，driver 即时应用A3C 等异步算法
Sharded Parameter Server梯度和参数分片到多个 PS shard，evaluator 异步更新大规模异步训练

5. 算法覆盖：一套抽象统一所有主流 RL 算法​

论文用 Table 2 系统性地展示了 RLlib 的三层抽象如何覆盖几乎所有主流 RL 算法家族：

DQN 系列：

  • Policy Graph 中 $y^1$ 存储 TD error，$\rho_\theta$ 计算 n-step return，$L$ 定义 Q loss
  • $u^1$ 实现 target network 更新，$u^2$ 设置探索率 $\epsilon$
  • DQN 实现使用内置 replay buffer 的 policy optimizer
  • Ape-X 变体：创建带有不同 $\epsilon$ 分布的 evaluator，编写高吞吐 policy optimizer（约 200 行代码），流水线化 replay buffer actor 之间的数据采样和传输

Policy Gradient / Off-policy PG：

  • $y^1$ 存储 value prediction，$\rho_\theta$ 实现 advantage estimation，$L$ 组合 actor 和 critic loss

PPO：

  • 利用 Local Multi-GPU optimizer，将样本 pin 到 GPU 显存中做多轮 SGD
  • 每轮：从 evaluator 收集样本 -> 本地多 GPU 优化 -> 广播新权重

A3C：

  • 使用 Asynchronous optimizer 或 Sharded Parameter Server optimizer
  • 从 evaluator 收集梯度更新策略参数 $\theta$ 的标准副本

DDPG：

  • 复用 DQN 的 replay policy optimizer，$L$ 包含 actor 和 critic 双损失
  • 也可搭配 Ape-X optimizer 使用

Model-based / Hybrid：

  • $\pi_\theta(o_t, h_t)$ 扩展为基于模型 rollout 做决策，模型 rollout 可通过 Ray 并行化
  • 模型损失可捆绑到 $L$ 中联合训练，或通过 Ray 原语单独并行训练并通过 $u^i$ 定期更新权重

Multi-Agent：

  • Policy Evaluator 支持在同一环境中运行多个策略，为每个 agent 产出经验 $X_{t,K}^P$
  • $\rho_\theta$ 可以汇总来自多个 agent 的经验，支持 centralized critic / value function

Evolutionary Methods：

  • 通过 non-gradient-based 的 policy optimizer 支持
  • ES 实现：移植单线程 ES 只需少量改动，再用 aggregation tree of actors 进一步扩展

AlphaGo Zero：

  • 逻辑集中控制多个分布式组件：model optimizer、self-play evaluator、candidate evaluator、shared replay buffer
  • 所有组件作为 Ray actor 由顶层 AlphaGo policy optimizer 统一调度
  • 约 70 行伪代码即可 sketch 主算法循环

6. 框架性能优化​

RLlib 在单节点和分布式两个层面都做了细致的性能优化：

6.1 单节点性能​

  • 有状态计算（Stateful Computation）：通过 Ray actor 让任务共享可变状态，对操作第三方模拟器或神经网络权重至关重要
  • 共享内存对象存储：rollout 数据和网络权重等大对象通过共享内存在 worker 之间传递，无中心瓶颈；同机 worker 可以零拷贝读取
  • 向量化（Vectorization）：batch 化 policy evaluation，支持 batched environment，使用列式数组格式传递经验数据
  • 数据压缩：使用 LZ4 算法压缩经验 batch，对图像观测可将网络流量和内存占用降低一个数量级，压缩速率约 1 GB/s/core

6.2 分布式性能​

  • 轻量级任务：Ray 的远程调用开销在同机约 200 微秒，跨节点约 1 毫秒，支持细粒度并行
  • 嵌套并行（Nested Parallelism）：组件可以动态创建子任务，调用图天然是动态的，Ray 支持任意 Python 函数/方法作为轻量级远程任务
  • 资源感知调度：远程调用可以指定资源需求（CPU/GPU），Ray 使用资源感知调度器保证组件性能
  • 容错与掉队者缓解（Fault Tolerance & Straggler Mitigation）：利用 Ray 内置容错机制，支持使用 preemptible/spot 实例降低成本；通过 ray.wait() 原语丢弃最慢的 evaluator 任务（如 PPO 中），以牺牲少量 bias 换取吞吐量

7. 实验评估​

7.1 采样效率​

Policy evaluation 吞吐量近线性扩展（Figure 7）：

  • Pendulum-CPU：1 到 128 核，从 15k 扩展到 1.5M actions/s（64x64 全连接网络）
  • Pong-GPU：1 到 128 核，从 2.4k 扩展到约 200k actions/s（DQN 卷积架构）

7.2 大规模测试​

ES（Evolution Strategies）：

  • 在 AWS m4.16xl 实例上扩展到 8192 核
  • Humanoid-v1 任务达到 reward 6000 的中位时间为 3.7 分钟
  • 比当时最佳公开结果（Salimans et al., 2017）快两倍以上

PPO：

  • 从 1 台 p2.16xl GPU 实例开始，逐步加入 m4.16xl CPU 实例
  • 在 512 CPU x 64 GPU 的配置下，超越了高度优化的 MPI 参考实现（Figure 8(b)）
  • Local Multi-GPU optimizer 在数据能完全放入 GPU 显存时，SGD 吞吐量显著优于 Allreduce（Table 3: Humanoid-v1 上 4 GPU 达 2.1M samples/s vs 330k samples/s）

A3C：

  • 在 x1.16xl 机器上，12 分钟解决 PongDeterministic-v4（async optimizer），9 分钟（sharded param-server optimizer），匹配 well-tuned baseline

Ape-X：

  • 256 workers 时达到 160k frames/sec，frameskip 为 4
  • 远超参考实现在 256 workers 时约 45k fps 的吞吐量（Figure 5(b)）

7.3 Parameter Server 对比​

RLlib 使用 8 个内部 shard 的 parameter server optimizer，在 A3C 上与 Distributed TensorFlow PS 实现性能相当（Figure 5(a)），验证了集中控制模型不会引入性能瓶颈。

8. 可组合性：50 行 PPO-ES，70 行 AlphaGo Zero​

RLlib 抽象的真正威力在于算法的可组合性。

PPO-ES 混合算法​

在 ES 优化的内循环中跑 PPO 更新，随机扰动 PPO 模型。实现只需约 50 行代码，且不需要修改 PPO 的任何代码。实验表明：

  • 在 Walker2d-v1 上比纯 PPO 收敛更快，最终 reward 更高
  • 类似的 A3C-ES 实现在 PongDeterministic-v4 上比纯 ES 快 30%

这展示了"封装并行性"（encapsulating parallelism）的价值——组件边界清晰，组合时互不干扰。

AlphaGo Zero Sketch​

论文 sketch 了用 Ray + RLlib 可扩展地实现 AlphaGo Zero 的方案（约 70 行主循环伪代码），涉及四大分布式组件：

  1. 逻辑集中控制：model optimizer、self-play evaluator、candidate evaluator、shared replay buffer 都作为 Ray actor，由顶层 policy optimizer 统一管理
  2. 共享 replay buffer：self-play 数据通过 Ray 对象引用路由到共享 buffer
  3. Best player 跟踪：候选模型必须以 >= 55% 胜率替换当前最佳模型，只需在主循环加一个 if 判断
  4. MCTS：作为 policy graph 的子程序处理，可选择性地通过 Ray 并行化

9. 与已有工作的对比​

维度传统 RL 库RLlib
并行模型长驻进程间通信（MPI, PS, etc.）短任务 + 分层集中控制
代码复用每个算法自带分布式逻辑Policy Graph / Evaluator / Optimizer 三层解耦，跨算法共享
可组合性组合两个算法需要修改两者的协调代码组件边界封装并行性，无缝嵌套
容错通常需要手动处理Ray 原生容错 + straggler mitigation
框架绑定通常绑定 TensorFlow 或特定框架Policy Graph 可用任意 DL 框架实现

相比当时的其他 RL 库（Intel Coach, OpenAI Baselines, TensorForce 等），RLlib 的独特贡献不在于实现了更多算法，而在于提出了一套可扩展、可组合的抽象层次，使得实现新算法和组合已有算法都变得简单。

10. 个人总结与思考​

这篇论文的核心贡献可以用一句话概括：为分布式 RL 找到了正确的抽象层次。

为什么这套抽象好用？

关键在于它抓住了 RL 算法的结构性共性：几乎所有 RL 算法都可以分解为"定义策略和损失"（Policy Graph）、"收集经验"（Policy Evaluator）、"组织分布式训练"（Policy Optimizer）三个层次。这三层的耦合度很低——换 optimizer 不需要改算法，换算法不需要改 optimizer——这使得每一层都可以独立优化和复用。

逻辑集中控制的洞察非常深刻。表面上看，让一个单线程 driver 控制整个分布式系统似乎会成为瓶颈。但实际上：(1) driver 只负责"控制流"，真正的"数据流"在 worker 之间直接传输；(2) 短任务模型意味着调度本身是轻量级的；(3) 分层委托让控制可以层层下发，不会都堆在顶层。这个设计思路后来在 Ray 生态的很多系统中都被沿用。

对今天的启示：虽然 RLlib 的具体 API 经过多年演化已经有很大变化（比如后来引入了 Algorithm、RolloutWorker 等新抽象），但论文中提出的核心原则——封装并行性、逻辑集中控制、可组合的层次化抽象——至今仍然是设计分布式 ML 系统的黄金法则。特别是在当下 LLM + RL（如 RLHF/RLAIF）的场景中，如何优雅地组合 LLM 推理、reward model 评估、PPO 训练等异构计算，RLlib 当年的思路依然极具参考价值。