在过去十年中，强化学习(RL) 已成为机器学习中最有前途的研究领域之一，并且在解决复杂的现实世界问题（例如芯片放置和资源管理）以及解决具有挑战性的游戏（例如，去、Dota 2和捉迷藏）。简单来说，强化学习基础设施是一个数据收集和训练循环，参与者探索环境并收集样本，然后将其发送给学习者训练和更新模型。大多数当前的 RL 技术需要对环境中数百万个样本的批次进行多次迭代才能学习目标任务（例如，Dota 2 每 2 秒从 200 万帧的批次中学习）。因此，RL 基础设施不仅应该有效地扩展（例如，增加参与者的数量）并收集大量样本，而且还应该能够在训练期间快速迭代这些大量样本。
今天我们介绍 Menger 1，这是一个具有本地化推理的大规模分布式 RL 基础设施，可在多个处理集群（例如Borg 单元）中扩展到数千个参与者，从而减少芯片放置任务的整体训练时间。在这篇文章中，我们描述了我们如何使用Google TPU 加速器实现 Menger 以进行快速训练迭代，并展示其在具有挑战性的芯片放置任务中的性能和可扩展性。门格尔减少达8.6训练时间X与基线相比实现。
Menger 系统设计
有各种分布式 RL 系统，例如Acme和SEED RL，每个都专注于优化分布式强化学习系统空间中的单个特定设计点。例如，虽然 Acme 对每个参与者使用本地推理，并从学习者那里频繁检索模型，但 SEED RL 通过分配一部分 TPU 内核来执行批量调用，从而受益于集中式推理设计。这些设计点之间的权衡是（1）支付向/从集中推理服务器发送/接收观察和动作的通信成本或支付从学习器检索模型的通信成本和（2）对参与者的推理成本（例如，CPU）与加速器（例如，TPU/GPU）的比较。由于我们的目标应用程序的要求（例如，观察的大小、动作和模型大小），Menger 以类似于 Acme 的方式使用局部推理，但将参与者的可扩展性推向了几乎无限的限制。在加速器上实现大规模可扩展性和快速培训的主要挑战包括：
1.为模型检索提供从参与者到学习器的大量读取请求很容易限制学习器，并随着参与者数量的增加迅速成为主要瓶颈（例如，显着增加收敛时间）。 2.TPU 性能通常受到将训练数据馈送到 TPU 计算核心的输入管道效率的限制。随着 TPU 计算内核数量的增加（例如TPU Pod），输入管道的性能对于整个训练运行时间变得更加重要。
高效的模型检索
为了解决第一个挑战，我们在 TensorFlow 中优化并由Reverb（在Dota 中使用的类似方法）支持的学习者和参与者之间引入了透明和分布式缓存组件。缓存组件的主要职责是在来自参与者的大量请求和学习者工作之间取得平衡。添加这些缓存组件不仅显着降低了学习器为读取请求提供服务的压力，而且还进一步将 actor 分布到多个Borg 单元中，并且通信开销很小。在我们的研究中，我们表明对于具有 512 个参与者的 16 MB 模型，引入的缓存组件将平均读取延迟降低了约 4.0 倍x导致更快的训练迭代，特别是对于PPO等基于策略的算法。
高吞吐量输入管道
为了提供高吞吐量的输入数据管道，Menger 使用Reverb，这是一个最近开源的数据存储系统，专为机器学习应用程序设计，提供了一个高效灵活的平台来实现各种 on-policy/off 的体验回放- 策略算法。然而，在具有数千个 actor 的分布式 RL 设置中，使用单个 Reverb 重播缓冲区服务目前不能很好地扩展，并且在来自 actor 的写入吞吐量方面变得效率低下。
为了更好地了解分布式设置中重播缓冲区的效率，我们评估了从 16 MB 到 512 MB 的各种有效负载大小和从 16 到 2048 个不等的参与者的平均写入延迟。我们在重播缓冲区和演员被放置在同一个博格牢房中。随着参与者数量的增加，平均写入延迟也显着增加。扩大行动者的数目为16〜2048，由〜6.2倍的平均写入延迟增加X和〜18.9 X为有效载荷大小16 MB和512 MB，分别。写入延迟的增加会对数据收集时间产生负面影响，并导致整体训练时间效率低下。
为了缓解这种情况，我们使用Reverb提供的分片功能来增加参与者、学习者和重播缓冲区服务之间的吞吐量。分片在多个重播缓冲区服务器上平衡来自大量参与者的写入负载，而不是限制单个重播缓冲区服务器，并且还最小化了每个重播缓冲区服务器的平均写入延迟（因为更少的参与者共享同一服务器）。这使 Menger 能够有效地扩展到多个 Borg 单元中的数千个参与者。
案例研究：芯片布局
我们研究了 Menger 在大型网表芯片布局的复杂任务中的优势。使用512个TPU芯，门格尔实现了在训练时间显著改进（达〜8.6 X，减少从〜训练时间8.6小时下来仅仅一小时最快配置）相比，较强的基线。虽然 Menger 针对 TPU 进行了优化，但这种性能提升的关键因素是架构，我们希望在针对 GPU 进行定制时看到类似的提升。
我们相信，Menger 基础设施及其在复杂的芯片放置任务中取得的有希望的结果展示了进一步缩短芯片设计周期的创新路径，并且不仅有可能在芯片设计过程中实现进一步创新，而且还有其他具有挑战性的现实世界任务也是如此。
更新说明：优先更新微信公众号“雨夜的博客”，后更新博客，之后才会陆续分发到各个平台，如果先提前了解更多，请关注微信公众号“雨夜的博客”。
博客来源：雨夜的博客