作者/机构: Lasse Espeholt∗, Raphaël Marinier∗, Piotr Stanczyk∗, Ke Wang & Marcin Michalski, Brain Team, Google Research

A1 主要贡献

本文旨在解决强化学习（RL）中的两大挑战：扩展性问题和现代计算资源（特别是加速器）的有效利用问题。当前先进的分布式算法在处理RL固有的异构任务（环境运行、模型推理、模型训练、经验回放等）时，计算资源利用效率低下，导致实验成本过高。

为应对这些挑战，本文提出了一种名为SEED（Scalable, Efficient, Deep-RL）的现代可扩展强化学习代理。其核心贡献与创新点如下：

  1. 提出一种新颖的分布式架构：该架构的核心特点是中心化推理（centralized inference）。与传统Actor-Learner架构中每个Actor在本地CPU上进行模型推理不同，SEED将所有模型的推理任务集中到中心的Learner上，并由加速器（如TPU）执行。这使得Actor仅需负责运行环境，极大地提高了加速器的利用率并降低了对Actor端计算资源的需求。

  2. 优化通信层以解决延迟问题：中心化推理会引入大量的远程调用，从而产生网络延迟。SEED通过一个基于gRPC的快速流式RPCs（streaming RPCs）通信层来最小化延迟。该框架保持Actor与Learner之间的连接，并集成了一个高效的批处理模块，从而在某些情况下实现了比本地推理更快的端到端延迟。

  3. 显著提升训练速度和降低成本：通过有效利用现代加速器，SEED能够以每秒数百万帧的速度进行训练。与常用的分布式RL算法IMPALA相比，SEED在实现同等或更优性能的同时，将实验成本降低了40%至80%。例如，在Atari-57任务上，SEED达到当前最佳水平的速度比先前研究快三倍。

  4. 通用性和SOTA性能：SEED架构具有良好的灵活性，本文成功地在其上实现了两种先进的分布式算法：IMPALA/V-trace（策略梯度）和R2D2（Q-learning）。实验证明，该架构在多个基准测试（Atari-57、DeepMind Lab和Google Research Football）上均取得了SOTA（state-of-the-art）或接近SOTA的性能，特别是在Google Research Football任务上刷新了SOTA记录。

  5. 开源实现：本文开源了SEED的实现代码以及在Google Cloud上进行大规模实验的示例，旨在帮助社区复现研究结果并在此基础上进行创新，从而推动强化学习领域的发展。

A3 背景知识与架构分析

传统Actor-Learner架构及其缺陷。在介绍SEED架构之前，我们首先分析IMPALA【9, IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures, 2018, ICML】、Ape-X、OpenAI Rapid等算法所采用的通用Actor-Learner架构。该架构的概览如图1a所示。在这种架构中，大量Actor重复地从Learner（或参数服务器）读取模型参数。然后，每个Actor使用本地模型采样动作并生成一个完整的轨迹，包括观测、动作、策略logits或Q值。最后，这个轨迹连同循环状态被传输到一个共享队列或经验回放缓冲区。Learner则异步地从队列/回放区中读取批量轨迹来优化模型。


图1: 架构概览

传统架构的四个主要缺点。该架构存在以下几个方面的不足：
1. 使用CPU进行神经网络推理：Actor通常部署在基于CPU的机器上。众所周知，CPU在执行神经网络计算时效率低下【34, Large-scale deep unsupervised learning using graphics processors, 2009, ACM】。当模型计算需求增加时，推理时间会超过环境步进的计算时间。此时的解决方案是增加Actor数量，但这会提高成本并影响收敛性【9, IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures, 2018, ICML】。
2. 资源利用效率低下：Actor在环境步进和模型推理这两个任务之间交替进行。这两项任务的计算需求通常不匹配，导致资源利用率低或Actor运行缓慢。例如，某些环境是单线程的，而神经网络则易于并行化。
3. 带宽需求高：模型参数、循环状态和观测数据需要在Actor和Learner之间传输。与模型参数相比，观测轨迹的大小通常只占几个百分点。此外，基于记忆的模型会发送大量状态，进一步增加了带宽需求。
尽管像GA3C【25, Incremental Off-policy Reinforcement Learning Algorithms, 2017, University of Alberta】和单机版IMPALA这样的单机方法避免了问题1（使用CPU进行推理）和问题3（网络带宽需求），但它们受限于问题2（资源使用）和无法满足多种环境所需的大规模扩展。

A2 方法细节

SEED架构设计。SEED中使用的架构（如图1b所示）解决了上述问题。我们将推理和轨迹累积都转移到了Learner上，这使得它在概念上变成了一个带有远程环境的单机设置（除了处理故障）。这种逻辑上的转移有效地使Actor变成了一个围绕环境的小循环。在每个环境步骤中，观测数据被发送到Learner，Learner运行推理并将动作发送回Actor。这引入了一个新问题：4. 延迟。

延迟解决方案：gRPC与批处理。为了最小化延迟，我们创建了一个使用gRPC（grpc.io）的简单框架，这是一个高性能的RPC库。具体来说，我们采用了流式RPC（streaming RPCs），其中从Actor到Learner的连接保持开放，元数据只发送一次。此外，该框架还包含一个批处理模块，可以有效地将多个Actor的推理调用批处理在一起。在Actor可以与Learner部署在同一台机器上的情况下，gRPC使用unix domain sockets，从而减少了延迟、CPU和系统调用开销。总体而言，对于我们考虑的多种模型，包括网络和推理在内的端到端延迟更快（见附录A.7）。

Off-policy行为的差异。IMPALA和SEED架构的不同之处在于，对于SEED，在任何时间点，模型只有一个副本存在，而对于分布式IMPALA，每个Actor都有自己的副本。这改变了轨迹的离策略（off-policy）方式。在IMPALA中（图2a），一个Actor在整个轨迹中使用相同的策略$π_{θ_t}$。而在SEED中（图2b），一个展开轨迹期间的策略可能会多次改变，后续步骤会使用更接近优化时所用策略的最新策略。


(a) IMPALA中的离策略。在整个轨迹中，策略保持不变。当轨迹被发送到队列进行优化时，策略已经改变了两次。


(b) SEED中的离策略。优化模型对策略有即时影响。因此，轨迹由从许多不同策略（$π_{θ_t}$, $π_{θ_{t+1}}$, ...）中采样的动作组成。

图2: 在异步优化模型参数θ时评估策略π的“近似在线策略”（near on-policy）的变体。


图3: SEED中Learner的详细架构（带有一个可选的经验回放缓冲区）。

Learner的详细架构。SEED架构中Learner的详细视图如图3所示。其中运行着三种类型的线程：1. 推理（Inference） 2. 数据预取（Data prefetching） 3. 训练（Training）。
* 推理线程接收一批观测、奖励和回合终止标志。它们加载循环状态，并将数据发送到推理TPU核心。接收到采样动作和新的循环状态后，动作被发送回Actor，同时最新的循环状态被存储起来。
* 当一个轨迹完全展开后，它被添加到一个FIFO队列或经验回放缓冲区，稍后由数据预取线程进行采样。
* 最后，这些轨迹被推送到参与训练的每个TPU核心的设备缓冲区中。训练线程（主Python线程）获取预取的数据，使用训练TPU核心计算梯度，并同步地将梯度应用于所有TPU核心（包括推理和训练核心）上的模型。
推理和训练核心的比例可以进行调整，以实现最大的吞吐量和利用率。该架构可通过轮询方式将Actor分配给TPU主机，并为每个TPU主机设置独立的推理线程，从而扩展到TPU pod（2048个核心）。当Actor等待Learner响应时，它们处于空闲状态，因此为了充分利用机器，我们在单个Actor上运行多个环境。

对传统架构问题的解决方案总结。我们通过以下方式解决了之前列出的问题：
1. 将推理移至Learner，从而消除了Actor上的任何神经网络相关计算。在这种架构中，增加模型大小不会增加对更多Actor的需求（实际上恰恰相反）。
2. 在Learner上对推理进行批处理，并在Actor上运行多个环境。这充分利用了Learner上的加速器和Actor上的CPU。TPU核心用于推理和训练的数量可以精细调整，以匹配推理和训练的工作负载。所有这些因素都有助于降低实验成本。
3. 所有与模型相关的东西都保留在Learner上，只有观测和动作在Actor和Learner之间发送。这将带宽需求降低了多达99%。
4. 使用具有最小延迟和开销的流式gRPC，并将批处理集成到服务器模块中。
我们在SEED框架中提供了以下两种算法的实现：V-trace和Q-learning。

3.1 V-TRACE

V-trace算法的适配。我们适配到该框架的算法之一是V-trace【9, IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures, 2018, ICML】。我们没有包含任何在IMPALA基础上提出的新增功能，例如【48, Representation learning with contrastive predictive coding, 2018, CoRR】或【12, Shaping belief states with generative environment models for RL, 2019, CoRR】中提出的。这些新增功能也可以应用于SEED，并且由于它们的计算成本更高，它们将从SEED架构中受益。

3.2 Q-LEARNING

Q-learning算法的实现与架构通用性。为了展示SEED架构的通用性，我们完整实现了R2D2【19, Recurrent experience replay in distributed reinforcement learning, 2018】，一种当前最先进的分布式基于价值的代理。R2D2本身建立在DQN【27, Human-level control through deep reinforcement learning, 2015, Nature】的一长串改进之上，包括：双重Q学习（double Q-learning）【49, Double Q-learning, 2010, NIPS】；【50, Deep reinforcement learning with double q-learning, 2016, AAAI】，多步自举目标（multi-step bootstrap targets）【43, Learning to predict by the methods of temporal differences, 1988, Machine learning】；【44, Reinforcement learning: an introduction mit press, 1998】，对偶网络架构（dueling network architecture）【54, Dueling network architectures for deep reinforcement learning, 2016, ICML】，优先分布式经验回放（prioritized distributed replay buffer）【37, Prioritized experience replay, 2015, ICLR】；【18, Distributed prioritized experience replay, 2018, ICLR】，价值函数重缩放（value-function rescaling）【33, Observe and look further: Achieving consistent performance on atari, 2018, arXiv】，LSTMs【17, Long short-term memory, 1997, Neural computation】和预热（burn-in）【19, Recurrent experience replay in distributed reinforcement learning, 2018】。

中心化经验回放缓冲区。我们证明了将经验回放缓冲区保留在Learner上是可行的，并且实现方式直接灵活，而非使用分布式回放缓冲区。这通过减少设置中的一种作业类型来降低复杂性。它的缺点是受限于Learner的内存，但在我们的实验中这远不是问题：一个包含$10^5$条长度为120、观测为84×84未压缩灰度图像的轨迹的回放缓冲区（遵循R2D2的超参数）占用85GB内存，而Google Cloud机器可以提供数百GB的内存。然而，在需要更大回放缓冲区的情况下，没有什么能阻止将分布式回放缓冲区与SEED的中心化推理结合使用。

A4 实验环境

  • 数据集/环境:

    • DeepMind Lab: 一个基于雷神之锤3引擎的3D环境，包含迷宫、激光枪战和记忆任务。实验在4个常用任务上进行评估。
    • Google Research Football: 一个类似FIFA视频游戏的环境。实验在“Hard”任务上进行评估。
    • Arcade Learning Environment (ALE): 包含57款Atari 2600游戏的基准测试。

  • 模型架构:

    • DeepMind Lab: 使用IMPALA中的ResNet模型。
    • Google Research Football: 使用Kurach等人【21, Google research football: A novel reinforcement learning environment, 2019, arXiv】提出的模型。输入状态SMM（Super Mini Map）有三种尺寸：96x72（默认），120x90（中），144x108（大）。
    • ALE (Atari-57): 采用与R2D2相同的神经网络结构。具体为：3个卷积层（滤波器数量[32, 64, 64]，核大小[8x8, 4x4, 3x3]，步长[4, 2, 1]，ReLU激活），接一个512单元的全连接层，再接一个512隐藏单元的LSTM层，最后是512隐藏单元的对偶网络头（dueling heads）。

  • 硬件配置:

    • SEED: 主要使用TPU v3核心（实验配置从2核到64核不等），也与Nvidia P100 GPU进行了对比。Actor运行在多核CPU机器上（例如4核或8核）。
    • IMPALA (基准): 使用Nvidia Tesla P100 GPU进行训练，Actor端的推理在CPU上完成。

  • 软件配置:

    • 代码实现: 基于TensorFlow 2。
    • 依赖库: 通信层使用gRPC。
    • 开源: 代码已在GitHub上开源 (http://github.com/google-research/seed_rl)，并提供了在Google Cloud上运行的Docker镜像和启动脚本。

A4 实验结果

DeepMind Lab 与 V-trace

  • 实验内容：评估SEED与IMPALA在4个DeepMind Lab任务上的稳定性与速度。
  • 实验结果：
    • 稳定性 (图4)：在完全相同的超参数下，SEED的稳定性略优于IMPALA，并且最终回报稍高。
    • 速度 (图5, 表1)：SEED使用2个TPU v3核心比IMPALA使用1个Nvidia P100快2.5倍。扩展到8个TPU核心后，速度进一步提升4.4倍（总计提速11倍），且样本效率保持不变。尝试通过增加批次大小来进一步提速，但发现这会损害样本效率。当SEED也使用P100时，速度比IMPALA慢1.58倍，但成本更低。
  • 分析结论：SEED的中心化推理架构显著提升了训练速度，并且在扩展时能保持良好的样本效率。


图4: 在完全相同的条件下（175个actor，相同的超参数等）比较IMPALA和SEED。图表显示了按最终性能排序的不同超参数组合。


图5: 在4个DeepMind Lab任务上的训练过程。每条曲线是24次运行中基于最终回报的最佳结果。样本效率在最多8个TPU v3核心的情况下得以保持，这使得训练速度比IMPALA基准快11倍。上排：X轴是帧数。下排：X轴是小时。

Google Research Football 与 V-trace

  • 实验内容：在“Hard”任务上评估SEED的速度和性能，并探究增大输入地图（SMM）尺寸的影响。
  • 实验结果：
    • 速度 (表1)：SEED使用2个TPU v3核心比IMPALA基准（使用2个P100）快1.6倍。扩展到8核后，再提速4.1倍。
    • 地图尺寸影响 (表1, 表2)：增大SMM尺寸会降低约60%的训练速度，但能显著提高最终得分。这表明默认的地图表示可能不够精细。SEED在两种奖励函数下都超越了Kurach等人【21, Google research football: A novel reinforcement learning environment, 2019, arXiv】的结果。将训练帧数增加到40亿帧，性能得到大幅提升（例如，在scoring奖励下，平均得分从0.46提升到4.76）。
  • 分析结论：SEED的计算效率允许使用更大的模型输入，从而在复杂任务上获得更好的性能，刷新了该任务的SOTA记录。

Arcade Learning Environment (ALE) 与 Q-learning (R2D2)

  • 实验内容：在Atari-57基准上评估SEED架构实现的R2D2。
  • 实验结果 (图6, 表1)：SEED的样本效率与原版R2D2相当，但训练速度快3.1倍。SEED仅用1.8天就达到了1880%的人类归一化得分中位数，而R2D2需要5天。这创造了Atari-57上新的最快训练时间记录。尝试通过增加actor数量和批次大小来达到更高的帧率（440K FPS），但这严重损害了样本效率，最终性能受限。
  • 分析结论：SEED架构成功地加速了先进的Q-learning算法，在保持样本效率的同时大幅缩短了训练时间。


图6: SEED及相关智能体在Atari-57上的中位数人类归一化得分。左：X轴为小时。右：X轴为环境帧数。SEED达到SOTA性能的速度（墙上时间）比R2D2快3.1倍。

成本比较

  • 实验内容：分析在DeepMind Lab和Google Research Football任务中，随着模型复杂度增加，SEED和IMPALA的训练成本变化。
  • 实验结果 (表4, 表5)：
    • DeepMind Lab (表4)：SEED的训练成本约为IMPALA的1/4。随着网络模型变大，IMPALA的成本急剧增加（因CPU推理开销大），而SEED的成本增长缓慢。
    • Google Research Football (表5)：由于环境本身计算开销较大，成本差距有所缩小，但SEED依然更便宜。随着输入地图尺寸增大，需要更多actor，SEED的成本优势更加明显。
    • ALE: 尽管没有直接的成本基准对比，但推测由于ALE环境运行速度快，推理成本占比较高，因此SEED的成本节省会更加显著。
  • 分析结论：SEED架构通过将推理任务转移到高效的加速器上，不仅提升了速度，还显著降低了训练成本，尤其是在模型或环境复杂度增加时，其成本优势更为突出。

Table 1: SEED, IMPALA 和 R2D2 的性能。

Table 2: Google Research Football “Hard”任务使用两种奖励函数的结果。

Table 3: 截至2019年9月的云资源成本。

Table 4: 在DeepMind Lab上训练10亿帧的成本。

Table 5: 在Google Research Football上训练10亿帧的成本。

A5 结论

本文介绍并分析了一种新的强化学习代理架构——SEED。通过更有效地利用现代加速器，该架构比以往的分布式架构更快、每环境帧的成本更低。实验结果表明，SEED在DeepMind Lab上实现了比强大基准IMPALA快11倍的墙上时间加速，同时保持了相同的样本效率；在Google Research Football上改进了当前最佳得分；在Atari-57上达到SOTA得分的速度比先前研究快3.1倍。

该代理已经开源，并打包以便在Google Cloud上轻松运行。我们希望这能通过让社区能够复现SOTA结果并在此基础上进行构建，从而加速强化学习的研究。

作为该新代理架构潜力的展示，我们能够在现实场景中将其扩展到每秒数百万帧（与先前研究相比提速80倍）。然而，这需要增加环境数量和使用更大的批次大小，这在测试的环境中损害了样本效率。在大批次下保持样本效率在强化学习领域已有部分研究，但我们认为这仍然是一个开放且日益重要的研究领域，以便进一步扩展强化学习。

A6 附录

A.1 DEEPMIND LAB

A.1.1 关卡缓存。我们为SEED和IMPALA都启用了DeepMind Lab的关卡缓存选项，这极大地减少了CPU使用量，并使得单个核心的actor CPU使用率稳定在接近100%。

A.1.2 超参数。
Table 6: 稳定性实验中使用的超参数范围。

A.2 GOOGLE RESEARCH FOOTBALL

A.2.1 超参数。
Table 7: 用于scoring和checkpoint奖励实验的超参数范围。

A.3 ALE

A.3.1 超参数。我们使用了与R2D2【19, Recurrent experience replay in distributed reinforcement learning, 2018】相同的超参数，除了我们使用了更多的actor以最好地利用8个TPU v3核心。为完整起见，代理超参数在表8中，环境处理参数在表9中。我们使用与R2D2相同的神经网络架构，即3个卷积层，滤波器大小为[32, 64, 64]，核大小为[8×8, 4×4, 3×3]，步长为[4, 2, 1]，ReLU激活函数和“valid”填充。它们输入到一个有512个单元的线性层，再输入到一个有512个隐藏单元的LSTM层（该层也使用one-hot编码的前一个动作和前一个环境奖励作为输入），最后输入到有512个隐藏单元的对偶头（dueling heads）。我们使用Glorot均匀初始化【10, Understanding the difficulty of training deep feedforward neural networks, 2010, AISTATS】。

Table 8: SEED在Atari-57上的代理超参数。

Table 9: Atari-57环境处理参数。

A.3.2 ATARI-57上的完整结果。
Table 10: SEED 8 TPU v3核心，610个actor（1个种子）的最终性能与R2D2（3个种子平均值）和人类的比较，每回合开始时使用最多30个随机无操作步骤。SEED的评估是在训练了400亿环境帧后，对每个游戏平均200回合的回报进行计算。


图7: SEED在57款Atari 2600游戏上的学习曲线（8个TPUv3核心，610个actor，使用1个种子进行评估）。每条曲线的每个点都是200个回合的平均回报。没有进行曲线平滑处理。曲线在约43小时的训练后结束，对应于400亿环境帧。

A.4 本地和云端运行SEED

开源与部署。SEED已开源，并附带了在本地机器上和使用AI平台（Google Cloud的一部分）大规模运行的示例。我们提供了一个公共Docker镜像，其中预编译了用C++实现的底层组件，以最小化启动SEED实验所需的时间。在云上运行的主要前提是设置一个云项目。提供的启动脚本会为您上传镜像并运行训练。更多详情请见github.com/google-research/seed_rl。

A.5 实验成本构成

Table 11: IMPALA和SEED完成10亿帧训练的成本，按组件划分。

A.6 在NVIDIA P100 GPU上对DEEPMIND LAB进行成本比较

P100上的成本对比。在4.4.1节中，我们比较了使用两个TPU v3核心的SEED和在单个Nvidia P100 GPU上的IMPALA的运行成本。在表12中，我们还比较了当两个代理都在单个Nvidia P100 GPU上运行DeepMind Lab时的成本。尽管这对SEED来说是一个次优设置，因为它在加速器上执行推理，因此从更高效的加速器（如TPU）中获益不成比例，但SEED的表现仍然不错，每帧成本比IMPALA便宜1.76倍。

Table 12: 在DeepMind Lab上，IMPALA和SEED在单个Nvidia P100 GPU上运行10亿帧的成本。

A.7 推理延迟

Table 13: IMPALA和SEED在不同环境和模型下的端到端推理延迟。