Press "Enter" to skip to content

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

     图源: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

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注