过去几年,强化学习领域取得了令人印象深刻的进展,不仅可以下围棋,还可以玩 StarCraft、Dota 2 等各种游戏,并且还获得了超越人类玩家的傲人成绩。
但强化学习本质上是由异构任务组成的,即便是目前最先进的分布式算法也无法有效地使用已有的计算资源来完成任务。大量的数据,和对资源的低效利用,使得实验的成本高得令人望而却步。由此一来,现在强化学习技术越来越多地被用来训练系统玩简单游戏,这无疑与强化学习领域的上升发展是“背道而驰”的。
一种解决之道,便是采用分布式方法,将计算任务分散到不同的机器上。目前最为优秀的分布式强化学习是 DeepMind 于 2018 年推出的 IMPALA 架构,然而 IMPALA 存在着一系列的缺点,例如资源利用率低、无法大规模扩展等。
针对这些问题,谷歌在最新工作《SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference》(为 ICLR 2020 Oral 论文)中,提出了一个能够扩展到数千台机器的强化学习架构——SEED RL(Scalable and Efficient Deep-RL,即可扩展且高效的深度强化学习),同时该架构还能够以每秒数百万帧的速度进行训练,计算效率显著提高。
谷歌研究者在 Google Research Football、DeepMind Lab 等常用的强化学习基准上测试了 SEED RL 的性能,结果非常令人惊喜。其中,在 DeepMind Lab 基准上比较了 SEED RL 和 IMPALA 的性能。
结果显示,SEED RL 使用 64 个 Cloud TPU 的计算速度达到了每秒 240 万帧,比当前最好的分布式强化学习架构 IMPALA 提高了 80 倍。在相同速度下,IMPALA 所需要的 CPU 数量是 SEED RL 的 3 到 4 倍。
SEED RL 本质上是在 IMPALA 基础上的一个改进。因此我们先简单来看下 IMPALA 是如何做分布式强化学习的。
IMPALA 的体系架构分为 Actor 和 Learner。其中 actor 一般运行在 CPU 上,它在环境中采取步骤,然后在模型上运行推断,进行迭代,然后预测下一个动作。Actor 会更新推理模型的参数,在收集到足够多的观察数据后,actor 会将观察和动作的结果发送给 learner,learner 根据这些反馈结果进行优化。
SEED RL 框架同样是由 Learner 和 Actor 两部分组成。但不同的是,SEED RL 依托于一种通过中心化模型推理,并引入快速通信层来大规模利用加速器(GPU 或 TPU)的新型架构。SEED RL 架构一览。与 IMPALA 架构不同,Actor 仅在环境中采取动作,而推理集中由 Learner 使用来自多个 Actor 的批量数据在加速器上执行。
实验的结果,我们在文章前面已经展示。这里总结一下谷歌这项工作的意义:1、相比 IMPALA,SEED RL 框架将推理部分转移到 Learner 中。这样调整后,Actor 中就不再有任何关于神经网络相关的计算了。这样做的好处是,模型的大小不会影响 Actor,于是便可扩大 Actor 的数量,做到可扩展。
2、在 SEED RL 的架构下,Learner(以 GPU 或 TPU 为主)可以专注于批量推理,而 Actor(以 CPU 为主)可以更加适应多环境。整体来说,这种结构会降低实验成本。3、涉及模型的所有内容都留在 Learner 本地,只有观测结果和动作会在 Learner 和 Actor 之间进行传输,这可以将带宽需求降低多达 99%。
4、使用了具有最小延迟和最小开销的 streaming gRPC,并将批处理集成到服务器模块当中。
据了解,目前 SEED RL 框架已经开源,打包后可以非常容易地在 Google Cloud 上运行。开源地址是:https://github.com/google-research/seed_rl