Google发布"强化学习"框架"SEED RL"

深度强化学习实验室报道

来源：Google Research

作者：DeepRL

本框架是Google发布于ICLR2020顶会上，这两天发布于Google Blog上

论文Arxiv：https://arxiv.org/abs/1910.06591

|| 关于Seed RL

现阶段在云端训练复杂的机器学习模型十分昂贵。报告显示，华盛顿大学（University of Washington）的Grover专为生成和检测假新闻而设计，该模型在两周的在训练中共花费了2.5万美元；OpenAI训练其GPT-2语言模型，每小时需花费256美元；谷歌训练其BERT双向转换器模型，预估也花费了6912美元。而这些模型前期的训练成本，对一些大型AI实验室或科技巨头来说也许不足为提，但对于许多的AI创企和小型AI实验室而言，却带来了压力.

当前的框架

上一代的分布式强化学习代理（例如IMPALA）利用专门用于数值计算的加速器，充分利用了（无）监督学习多年来受益的速度和效率。RL代理的体系结构通常分为actor和learner。actor通常在CPU上运行，并且在环境中采取的步骤与对模型进行推断之间进行迭代，以预测下一个动作。通常，actor会更新推理模型的参数，并且在收集到足够数量的观察结果之后，会将观察结果和动作的轨迹发送给learner，从而对learner进行优化。在这种架构中，learner使用来自数百台机器上的分布式推理的输入在GPU上训练模型。

RL Agent（例如IMPALA）的体系结构具有许多缺点：

• 与使用加速器相比，使用CPU进行神经网络推理的效率和速度要慢得多，并且随着模型变得越来越大且计算量越来越大，问题变得越来越严重。
• 在actor和learner之间发送参数和中间模型状态所需的带宽可能是瓶颈。
• 在一台机器上处理两个完全不同的任务（即环境渲染和推理）不可能最佳地利用机器资源。

SEED RL体系结构

SEED RL体系结构旨在解决这些缺点。通过这种方法，learner可以在专用硬件（GPU或TPU）上集中进行神经网络推理，从而通过确保模型参数和状态保持局部状态来加快推理速度并避免数据传输瓶颈。尽管观察结果在每个环境步骤都发送给learner，但由于基于gRPC框架和异步流RPC的非常高效的网络库，延迟保持在较低水平。这使得在一台机器上每秒最多可以实现一百万个查询。learner可以扩展到数千个核心（例如，在Cloud TPU上最多2048个），actor的数量可以扩展到数千台机器，以充分利用learner，从而可以以每秒数百万帧的速度进行训练。SEED RL基于TensorFlow 2 API，在我们的实验中，是通过TPU加速的。

SEED RL的特点与性能

• 基于谷歌的TensorFlow 2.0框架，SEED RL的特点是能通过集中模型推理，来利用图形卡和TPU（张量处理单元）。
• 为了避免数据传输瓶颈，SEED RL还使用学习器组件来集中执行AI推理，而该组件也使用来自分布式推理的输入来训练模型。
• 目标模型的变量和状态信息将保持在本地，并将每个环境步骤的观察结果发送给学习器组件。同时，由于该模型使用了基于开放源代码通用RPC框架的网络库，因此它的延迟也将保持在最低水平。
• SEED RL的学习器组件能够扩展到成千上万个核心，例如在Cloud TPU上最多可扩展到2048个，而actor的数量可扩展多达数千台机器。

实验结果如下

SEED RL在常用的Arcade学习环境，DeepMind Lab环境以及最近发布的Google Research Football环境中进行了基准测试。

在DeepMind Lab上，作者使用64个Cloud TPU内核实现了每秒240万帧的数据传输速度，与以前的最新分布式代理IMPALA相比，提高了80倍。这样可以显着提高挂钟时间和计算效率。对于相同的速度，IMPALA需要的CPU是SEED RL的3-4倍。

通过针对现代加速器进行优化的架构，自然会增加模型大小，以提高数据效率。我们表明，通过增加模型的大小和输入分辨率，我们可以解决以前未解决的Google Research Football任务“困难”。

更多参数请查看Google blog:

https://ai.googleblog.com/2020/03/massively-scaling-reinforcement.html

|| 代码使用

1. 克隆仓库

git clone https://github.com/google-research/seed_rl.git
cd seed_rl

2. 单机本地机器训练

./run_local.sh [Game] [Agent] [Num. actors]
./run_local.sh atari r2d2 4
./run_local.sh football vtrace 4
./run_local.sh dmlab vtrace 4

3. 使用AI平台进行分布式训练

第一步是配置GCP和一个将用于培训的Cloud项目：

• 按照https://cloud.google.com/sdk/install上的说明安装Cloud SDK，并设置您的GCP项目。
• 确保为您的项目启用了计费。
• 启用AI平台（“云机器学习引擎”）和Compute Engine API。
• 如https://cloud.google.com/ml-engine/docs/ working-with-cloud-storage所述，授予对AI Platform服务帐户的访问权限。
• 在您的shell脚本中进行云身份验证，以便SEED脚本可以使用您的项目：

gcloud auth login
gcloud config set project [YOUR_PROJECT]
gcp/train_[scenario_name].sh