Google 开源可大规模扩展的深度强化学习新架构 SEED RL

2020 年 4 月 3 日 CSDN

图源：Google AI 官博

作者 | Jesus Rodriguez

译者 | 弯月，责编 | 夕颜

出品 | CSDN（ID：CSDNnews）

如今，深度强化学习（Deep reinforcement learning，DRL）是深度学习领域发展最快的一个方向。AI领域近年来的一些重大里程碑都源自DRL，例如AlphaGo、Dota2 Five或Alpha Star，DRL似乎是最接近人类智能的学科。然而，虽然我们取得了如此大的进展，但DRL方法在现实世界中的实现仍局限于大型人工智能实验室。部分原因是DRL体系结构依赖巨量的训练，对大多数组织而言，这笔费用超出了可承受的范围且不切实际。

最近，Google Research发表的一篇论文中提出了SEED RL，这是一种可大规模扩展的DRL模型新架构。

在现实世界中实现DRL模型的难度与它们的架构有着直接的关系。本质上，DRL包含各种任务，例如运行环境、模型推断、模型训练或重放缓冲等等。大多数现代DRL架构都无法有效地分配这类任务的计算资源，从而导致实现成本不合理。在AI硬件加速器等组件的帮助下，我们已经克服了其中的一些限制，但仍然很有限。近年来出现的新架构已被市场上许多成功的DRL实现所采用。

IMPALA带来的启示

在当前的DRL架构中，IMPALA为该领域树立了新标准。IMPALA最初是由DeepMind在2018年的一篇研究论文中提出的，它引入了一种模型，利用专门的数值计算加速器，充分发挥了监督学习早已享有多年的速度和效率。IMPALA的核心是角色模型（Actor model），常常用于最大化并发和并行。

基于IMPALA的DRL代理架构主要包含两大组件：Actor和Learner。在这个模型中，通常Actor都在CPU上运行，并在环境中采取的步骤与对该模型的推断之间进行迭代，以预测下一个动作。Actor会不断更新推理模型的参数，并在收集到足够数量的观测后，将观测和动作的轨迹发送给Learner，以完成对Learner的优化。在这种架构中，Learner使用来自数百台机器的分布式推理输入在GPU上训练模型。从计算的角度来看，IMPALA架构可使用GPU加速Learner的学习，而Actor可在许多机器上扩展。

图源：Google AI 官博

IMPALA在DRL架构中建立了新的标准。但是，该模型具有一些固有的局限性。

使用CPU进行神经网络推断：Actor机器通常基于CPU。当模型的计算需求增加时，推理所花费的时间会超过环境步骤的计算。解决方案是增加Actor的数量，而这会增加成本并影响融合。

资源利用效率低下：Actor在两个任务之间交替进行：环境步骤和推断步骤。通常这两个任务的计算要求并不一样，因此会导致资源利用率低下或Actor迟缓。

带宽要求：模型参数、循环状态和观测在Actor和Learner之间传递。此外，基于内存的模型需要发送很大的状态数据，因此增加了带宽需求。

Google 受 IMPALA actor模型的启发，开发了一种新架构，用以解决之前的架构在DRL模型缩放方面的一些局限性。

SEED RL

总体而言，Google 的 SEED RL 架构看起来与IMPALA极为相似，但也引入了一些变化，解决了DeepMind模型的一些主要限制。在SEED RL中，神经网络推理由Learner在专用硬件（GPU或TPU）上集中完成，通过确保将模型参数和状态保持在本地来加快推理速度，并避免数据传输瓶颈。每个环境步骤和推断步骤都会发送到Learner，并在推理完成后将动作发送回Actor。这个聪明的解决方案解决了IMPALA等模型的推理局限性，但可能会带来延迟的问题。

为了最大程度地减少延迟影响，SEED RL依靠gPRC进行消息传递和流传输。具体来说，SEED RL利用流式RPC，建立从Actor到Learner的连接，而且元数据仅发送一次。此外，该框架还包括一个批处理模块，能够有效地将多个Actor的推理调用放到批处理中。

图源：Google AI 官博

深入研究IMPALA架构，我们就会发现其中包含三种基本类型的线程：

推理
数据预取
训练

推理线程会收到一批观测、奖励和episode终止标志。它们负责加载循环状态，并将数据发送到推理TPU内核。接收采样的动作和新的重复状态，并在存储最新的重复状态的同时，将动作发送回Actor。轨迹完全展开后，会被添加到FIFO队列或重播缓冲区中，然后由数据预取线程进行采样。最后，将轨迹推入设备缓冲区，以供每个参加培训的TPU内核使用。训练线程（Python主线程）采用预取的轨迹，使用训练的TPU核心计算梯度，并将梯度同步应用于所有TPU核心的模型（推理和训练）。我们可以通过调整推理和训练核心的比率，实现最大的吞吐量和利用率。

图源：Google AI 官博

在SEED RL架构中，Actor可以扩展到成千上万个内核，而Learner的数量也可以扩展到成千上万台机器，因此可以实现每秒数百万帧的训练速度。SEED RL基于TensorFlow 2 API，而且性能也经过了TPU加速。

为了评估SEED RL，Google使用了常见的DRL基准测试环境，例如cade学习环境、DeepMind Lab环境以及最近发布的Google Research Football环境。所有环境下的结果都很惊人。例如，在DeepMind Lab环境中，SEED RL使用64个Cloud TPU内核实现了每秒240万帧的数据传输，比以前的最新分布式代理IMPALA提升了80倍。此外，我们也看到了速度和CPU利用率的提高。

图源：Google AI 官博

SEED RL代表了可大规模扩展的DRL模型的进步。Google Research在GitHub上开源了最初的SEED RL体系架构。不难想象，在不久的将来，这将成为许多DRL实现的基础模型。

原文链接：

https://towardsdatascience.com/google-open-sources-a-new-architecture-for-massively-scalable-deep-reinforcement-learning-3af21adfcfb1

本文为CSDN翻译文章，转载请注明出处。

作为“百万人学AI”的重要组成部分，2020 AIProCon 开发者万人大会将于6月26日通过线上直播形式，让开发者们一站式学习了解当下 AI 的前沿技术研究、核心技术与应用以及企业案例的实践经验，同时还可以在线参加精彩多样的开发者沙龙与编程项目。参与前瞻系列活动、在线直播互动，不仅可以与上万名开发者们一起交流，还有机会赢取直播专属好礼，与技术大咖连麦。

评论区留言入选，可获得价值299元的「2020 AI开发者万人大会」在线直播门票一张。快来动动手指，写下你想说的话吧！

推荐阅读

☞ 华为集齐 AI 龙珠，“召唤神龙”为期不远

☞ 中国无人机“老炮儿”回忆录

☞ 京东商城背后的AI技术能力揭秘 - 基于关键词自动生成摘要

☞互联网之父确诊新冠，一代传奇：任谷歌副总裁、NASA 访问科学家

☞从未如此简单：10分钟带你逆袭Kafka！

☞无代码时代来临，程序员如何保住饭碗？