在过去的十年中，强化学习（RL）已经成为机器学习中最受关注的研究领域之一，RL的应用可以很好地解决芯片布局和资源等问题管理。复杂的问题，以及围棋/Dota2/捉迷藏等具有挑战性的游戏。简单来说，RL基础设施就是数据收集和训练的循环。Actor根据环境收集样本数据，然后传输给Learner来训练和更新模型。当前大多数RL实现都需要环境中数千个样本的多次迭代来学习目标任务，例如Dota2每2秒学习数千帧样本。这样一来，强化学习架构不仅要具备强大的数据处理能力，比如添加演员来收集大量样本，还??要能够在训练过程中快速迭代这些样本。Actor和Learner交互的RL架构。Learner使用采样数据训练模型并将更新后的模型传递给Actor（例如TF-Agents、IMPALA）。今天，我们将推出Menger——一种具有本地化推理的大规模分布式RL架构，可扩展到跨多个处理集群（如Borg单元）的数千个Actor，从而减少芯片放置任务的训练时间。在接下来的章节中，我们将介绍如何使用GoogleTPU配置Menger以提高训练速度，并进一步通过芯片放置任务验证框架的性能和可扩展性。可以发现，与基线模型相比，Menger将训练时间减少了8.6倍。Menger设计思路目前有多种分布式强化学习系统，如Acme和SEEDRL，但这些系统往往只是从特定角度优化分布式强化学习系统。例如，Acme从频繁的Learner中获取模型，使每个Actor可以进行本地推理，而SEEDRL则通过分配一部分TPU核执行批量调用来进行集中推理。通信成本和推理成本的度量是不同优化系统之间的差异，包括：（1）向/从集中式推理服务器发送/接收观察和动作的通信成本，或从Larner获取模型的通信成本；（2）与加速器（TPU/GPU）成本相比，Actor的推理成本较大。考虑到目标程序的观察、动作、模型大小等需求，Menger采用了类Acme的局部推理，但同时尽可能增加了Actors的可扩展性。要实现良好的可扩展性和训练速度，主要挑战包括以下两点：Actors向Learner发出大量读取请求以进行模型检索，这给Learner造成了负担。随着Actor数量的增加，模型的性能明显受到限制（比如收敛时间显着增加）。将训练数据馈送到TPU计算核心时，TPU性能通常受到输入管道效率的限制。随着TPU计算核心数量的增加（如TPUPod），输入管道的性能对训练时间的影响更加明显。高效的模型检索为了解决第一个挑战，在TensorFlow代码中，我们在Learner和Actor之间引入了一个透明的分布式缓存组件，并通过Reverb对其进行优化（类似于Dota中使用的方法）。缓存组件的主要职责是平衡来自Actors的大量请求和Learners的处理能力。通过添加这些缓存组件，不仅可以显着缓解Learner过多请求的压力，还可以通过少量的通信成本将Actors分发到多个Borg单元。我们的研究表明，对于一个512个Actors，大小为16MB的模型，缓存组件的引入可以将平均读取延迟降低约4.0倍，从而实现更快的训练迭代，在策略算法中效果更加明显，例如PPO。一个分布式RL系统，其中多个Actor被放置在不同的Borg单元中。不同Borg单元中大量Actors的频繁模型更新请求限制了Learner的性能以及Learner和Actors之间的通信网络，导致整体收敛时间显着增加。虚线表示不同机器之间的gRPC通信。引入透明分布式缓存服务的分布式RL系统。多个Actor放在不同的Borg单元中，Learner只将更新后的模型发送给分布式缓存组件。每个缓存组件都应该为相邻的参与者和缓存更新模型请求。缓存区不仅减轻了Learner服务于模型更新请求的负担，还降低了Actors的平均读取延迟。高吞吐量输入管道为了提高输入数据管道的吞吐量，Menger使用了Reverb，这是一种专为机器学习应用程序设计的新型开源数据存储系统。在Reverb中，可以使用在线或离线算法进行体验回放。然而，目前单个Reverb无法扩展到具有数万个actor的分布式RL系统，并且actor的写入吞吐量效率低下。具有单个播放缓冲区的分布式RL系统。在Actor发出大量经过深思熟虑的写入请求后，回放缓冲区可能会受到限制并降低整体吞吐量。此外，一旦我们将Learner扩展为具有多个计算引擎(TPUPod)，单个回放缓冲区就无法高效地为这些引擎提供数据，从而严重影响整体收敛时间。为了更好地了解分布式RL系统中重放缓冲区的效率，我们评估了不同负载大小(16MB-512MB)和不同参与者(16-2048)下的平均写入延迟。我们将播放缓冲区和Actor放在同一个Borg单元中，可以发现随着actor数量的增加，平均写入延迟显着增加。将Actor的数量从16扩展到2048会使大小为16MB和512MB的有效负载的平均写入延迟分别增加约6.2倍和约18.9倍。这种写入延迟的增加会影响数据收集时间，从而导致训练效率低下。当不同大小(16MB-512MB)和Actor数量（16到2048）的有效载荷放置在同一个Borg单元上时，单个Reverb重放缓冲区的平均写入延迟。为了缓解这种情况，我们使用Reverb的分片功能来增加Actor、Learner和播放缓冲区之间的吞吐量。Sharding可以在多个播放缓冲服务器之间平衡大量Actor的写入负载，而不是只作用于单个播放缓冲服务器。同时，由于少数Actor共享同一个服务器，可以最小化每个播放缓冲服务器的平均写入延迟。这允许Menger跨多个Borg单元扩展数千个Actor。具有切片播放缓冲区的分布式RL系统。每个播放缓冲区用于存储位于同一Borg单元上的特定Actor。此外，分块播放缓冲区为加速器内核提供了具有更高吞吐量的输入管道。示例验证：芯片布局我们在大型网表的芯片布局任务中测试了Menger。与baseline相比，Menger使用了512个TPU核心，在训练时间上有显着提升（高达约8.6倍，即在最优配置下，训练时间可以从约8.6小时减少到1小时）。虽然Menger针对TPU进行了优化，但框架才是性能提升的关键因素，我们希望在GPU上进行实验时能看到类似的提升。与芯片放置的基线相比，使用Menger对不同数量的TPU核心进行训练的时间得到改善。我们相信Menger架构及其在贴片任务中的优异表现，为进一步缩短芯片设计周期提供了方向。同时，我们也可以应用这个架构来实现其他具有挑战性的现实世界问题。本文转载自雷锋网。如需转载，请在雷锋网官网申请授权。
