大规模分布式强化学习架构

今天，我们要介绍的是Menger——一种具有本地化推理能力的大规模分布式RL架构，可通过多个处理集群（如Borg单元）扩展数千个Actors，从而减少了芯片放置任务的训练时间。在接下来的章节，我们介绍了如何使用Google TPU配置Menger，从而提高训练速度，进一步我们通过芯片放置任务验证框架的性能和可扩展性。可以发现，与基准模型相比，Menger将培训时间减少了8.6倍。

Menger设计思路

当前有各种各样的分布式RL系统，如Acme和SEED RL，然而，这些系统往往只从一个特定角度对分布式强化学习系统进行优化。例如，Acme从频繁的Learner获取模型，使每个Actor都进行本地推理，而SEED RL则通过分配一部分TPU内核执行批量调用，进行集中推理。对通信成本和推理成本的衡量是不同优化系统的区别，具体包括：（1）向/从集中式推理服务器发送/接收观察和动作的通信成本，或从Larner获取模型的通信成本；（2）相比加速器（TPU/GPU）成本，Actor的推理成本大小。考虑到观察值、动作和模型大小等目标程序要求，Menger使用类似Acme的局部推理，但同时尽可能的增加Actor的可扩展性。要实现良好扩展性和训练速度，主要挑战包括以下两点：

Actor向Learner进行大量读取请求以进行模型检索，这就造成Learner的负担，随着Actor数量的增加模型表现明显受限（如收敛时间的显著增长）。

在将训练数据输送给TPU计算核心时，TPU性能通常受到输入管道效率的限制。随着TPU计算核心数量的增加（如TPU Pod），输入管道的性能对于训练时间的影响更加明显。

高效的模型检索

为应对第一个挑战，在TensorFlow代码中，我们在Learner和Actor之间引入了透明的分布式缓存组件，并通过Reverb进行优化（类似于Dota中使用的方法）。缓存组件的主要职责是对Actor的大量请求和Learner的处理能力进行平衡。通过添加这些缓存组件，不仅显着减轻了过多请求对Learner的压力，而且以少量的通信成本将Actor分配给多个Borg单元。我们的研究表明，对有512个Actors、大小为16MB的模型，引入缓存组件可以将平均读取延迟降低约4.0倍，从而实现更快的训练迭代，在PPO等策略算法中效果更加明显。

（编辑：常州站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!