计算速度提高 80 倍！谷歌推出分布式强化学习框架 SEED RL

近年来，Go 和 Dota2 等博弈类游戏的火爆从侧面反映了强化学习（RL）在过去几年里取得的成就。模型或代理通过（例如在游戏中）探索环境来学习，同时针对特定的目标进行优化。然而，目前哪怕要成功学习非常简单的游戏，所需要的强化学习训练体量也非常庞大，这使得研究和产品创意的迭代计算既昂贵又耗时，因此我们亟需一个能同时解决这些问题的“超级”强化学习框架。

在论文“SEED RL：通过加速集中推理实现高效可扩展的 Deep-RL”中，我们提出了一个强化学习代理，它可以规模部署到上千台机器上，并能以每秒数百万帧的速度进行训练，显著提高了计算效率。这是通过一种新的架构实现的，它通过集中式模型推理和引入快速通信层来大规模利用加速器（GPU 和 TPU）。我们验证了 SEED RL 在流行的强化学习基准测试上的性能，如谷歌研究足球、Arcade 学习环境和 DeepMind Lab，结果表明，使用更大的模型可以提高数据效率。这些代码已经在 Github 上开源，并提供了使用 GPU 在谷歌云上运行的示例。

项目地址：https://github.com/google-research/seed_rl

当前的分布式架构

IMPALA 等上一代分布式强化学习代理，使用了专门用于数值计算的加速器，多年来，就像（非）监督学习一样，在速度和效率方面都有了提升。强化学习代理的架构通常分为行动者（actors）和学习者（learners）。行动者通常在 CPU 上运行，迭代在环境中执行动作以及在模型上运行推理来预测下一个动作的过程。行动者会经常更新推理模型的参数，在收集到足够多的观察数据后，将观察和动作的轨迹发送给学习者，学习者再对模型进行优化。在这个架构中，学习者使用来自数百台机器的分布式推理输入在 GPU 上训练模型。

早期 RL 代理 IMPALA 的架构示例。推理在行动者上进行，通常使用低效的 CPU。更新后的模型参数频繁地从学习者发送给行动者，增加了带宽需求。

强化学习代理（如 IMPALA）的架构有如下几个缺点：

• 与使用加速器相比，使用 CPU 进行神经网络推理的效率和速度要低得多，而且随着模型变大、变多，这会带来问题。

• 在行动者和学习者之间发送参数和中间模型状态所需的带宽可能成为瓶颈。

• 在一台机器上处理两个完全不同的任务（即环境渲染和推理）不太可能最优地利用机器资源。

SEED RL 架构

SEED RL 架构就是为了解决这些缺陷而设计的。使用这种方法，神经网络推理由学习者在专门的硬件（GPU 或 TPU）上集中完成，通过在本地保存模型参数和状态，实现了推理加速，避免了数据传输瓶颈。虽然在每个环境步骤中都将观察结果发送给学习者，但是得益基于 gRPC 框架和异步流 RPC 非常高效的网络库，延迟很低。这使得在一台机器上可以实现每秒一百万次查询。学习者可以扩展到数千个内核（例如，在 Cloud TPUs 上可以扩展到 2048 个），行动者的数量可以扩展到数千台机器来充分利用学习者，这使得以每秒数百万帧的速度训练成为可能。SEED RL 基于 TensorFlow 2 API，在我们的实验中，TPUs 加速了 SEED RL。

为了使该框架成功运行，两个最先进的算法被集成到 SEED RL 中。第一个是 V-trace，这是一种基于策略梯度的方法，首先由 IMPALA 引入。通常，基于策略梯度的方法可以预测一个动作的分布，从中可以对动作进行抽样。然而，由于行动者和学习者在 SEED RL 中是异步执行的，所以行动者的策略稍微落后于学习者的策略，即变成了离线策略（off-policy）。通常的基于策略梯度的方法是在线策略（on-policy），这意味着，行动者和学习者有相同的策略，在离线策略设置中会遇到收敛性和数值问题。V-trace 是一种离线策略方法，因此，在异步 SEED RL 架构中可以工作得很好。

第二种算法是 R2D2，这是一种 Q-learning 方法，它使用递归分布式重播根据预测的动作未来值来选择动作。这种方法使得 Q-learning 算法可以大规模运行，同时仍然允许使用递归神经网络，该网络可以根据一个事件中所有过去帧的信息预测未来的值。

SEED RL 架构概览。与 IMPALA 架构不同，行动者只在环境中执行动作。推理由学习者在加速器上使用来自多个行动者的批量数据集中执行。

实验

SEED RL 在常用的 Arcade 学习环境、DeepMind Lab 环境和最近发布的谷歌研究足球环境上进行了基准测试。

在 DeepMind Lab 上 IMPALA 和 SEED RL 各种配置的每秒帧数比较。SEED RL 使用 4160 个 CPU 达到每秒 2.4M 帧。假设速度相同，IMPALA 需要 14000 个 CPU。

在 DeepMind Lab 上，我们使用 64 个云 TPU 内核实现了每秒 240 万帧，这比以前最先进的分布式代理 IMPALA 提高了 80 倍。这将显著提高壁钟时间和计算效率。在同样的速度下，IMPALA 需要的 CPU 数量是 SEED RL 的 3~4 倍。

随着时间的推移，在 DeepMind Lab 游戏“explore_goal_locations_small”中使用 IMPALA 和 SEED RL 获得的奖励数比较。有了 SEED RL，训练时间大幅减少。

随着架构优化以及现代化加速器的使用，为了提高数据效率而增加模型大小是很自然的。我们证明，通过增加模型的大小和输入归结，我们能够解决之前未解决的谷歌研究足球的“困难”任务。

不同架构在谷歌研究足球“困难”任务上的得分。我们证明，通过使用输入归结和一个更大的模型，分数得到了改善，并且通过更多的训练，该模型可以显著地优于内置 AI。

我们相信 SEED RL 和它所展示的结果，证明了强化学习在利用加速器方面再一次赶上了深度学习中的其他领域。

原文链接：

https://ai.googleblog.com/2020/03/massively-scaling-reinforcement.html