谷歌推出分布式强化学习框架SEED RL，计算速度比IMPALA提高了80倍！
近年来，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
- 发表于:
- 本文为 InfoQ 中文站特供稿件
- 首发地址：https://www.infoq.cn/article/p9fSANJZKfRuOjdMQoAd
- 如有侵权，请联系 cloudcommunity@tencent.com 删除。
相关快讯
扫码
添加站长 进交流群
领取专属 10元无门槛券
私享最新 技术干货
腾讯云开发者
