【Google最新成果】使用新的物理模拟引擎加速强化学习

深度强化学习实验室

官网：http://www.neurondance.com/

论坛：http://deeprl.neurondance.com/

来源：GoogleAI Blog

上一篇文章我们介绍了Google AI 开放的【最新】如何降低深度强化学习研究的计算成本(Reducing the Computational Cost of DeepRL)， 其通过对强化学习的计算成本进行考量，综合来提升强化学习的计算成本。本篇文章Google AI 再次带来《使用新的物理模拟引擎加速强化学习》，两者从降低成本，加速学习的模拟环境，其实早在2019年的时候伯克利Pieter Abbeel就曾经提出过《Accelerated Methods For Deep Reinforcement Learning》https://arxiv.org/pdf/1803.02811.pdf ，本文再一次将加速作为重点进行了研究，话不多少放内容：

强化学习(RL) 是一种流行的教学机器人导航和操纵物理世界的方法，其本身可以简化并表示为刚体之间的交互¹（即，当对它们施加力时不会变形的固体物理对象）。为了便于在实际时间内收集训练数据，RL 通常利用模拟，其中任意数量的复杂对象的近似值由许多由关节连接并由执行器提供动力的刚体组成。但这带来了一个挑战：RL 智能体通常需要数百万到数十亿的模拟帧才能精通简单的任务，例如步行、使用工具或组装玩具积木。

虽然通过回收模拟帧在提高训练效率方面取得了进展，但一些 RL 工具通过将模拟帧的生成分布到许多模拟器来回避这个问题。这些分布式模拟平台产生了令人印象深刻的结果，训练速度非常快，但它们必须在具有数千个 CPU 或 GPU 的计算集群上运行，而大多数研究人员无法访问这些集群。

本文作者展示了一种新的物理仿真引擎“Brax”，它与仅具有单个 TPU 或 GPU 的大型计算集群的性能相匹配。该引擎旨在在单个加速器上高效运行数千个并行物理模拟以及机器学习 (ML) 算法，并在互连加速器的 pod 中无缝扩展数百万个模拟。我们已经开源了引擎以及参考 RL 算法和模拟环境，这些都可以通过Colab访问。使用这个新平台，我们展示了比传统工作站设置快 100-1000 倍的训练。

三个典型的 RL 工作流程。该留下显示了一个典型的工作站流程：一台机器上，对CPU的环境，训练需要数小时或数天。在中间显示了一个典型的分布式仿真流程：通过训练模拟耕出上千台机器只需要几分钟。的右示出了BRAX流动：学习和大批量模拟发生并排单个CPU / GPU芯片上。

物理仿真引擎设计机会
刚体物理用于视频游戏、机器人、分子动力学、生物力学、图形和动画以及其他领域。为了对此类系统进行准确建模，模拟器集成了来自重力、电机驱动、关节约束、物体碰撞等的力，以模拟物理系统随时间的运动。

仔细研究当今大多数物理模拟引擎的设计方式，有一些提高效率的巨大机会。正如我们上面提到的，典型的机器人学习管道将单个学习器置于紧密反馈中，同时进行许多模拟，但在分析此架构后，人们发现：

1. 这种布局带来了巨大的延迟瓶颈。由于数据必须通过数据中心内的网络传输，因此学习者必须等待 10,000 多纳秒才能从模拟器中获取经验。如果这种体验已经与学习者的神经网络在同一设备上，延迟将降至 <1 纳秒。

2. 训练智能体所需的计算（一个模拟步骤，然后更新智能体的神经网络）被打包数据（即，在引擎内编组数据，然后转换为诸如protobuf 之类的有线格式，然后进入TCP缓冲区，然后在学习者端撤消所有这些步骤）。

3. 每个模拟器中发生的计算非常相似，但并不完全相同。

Brax 设计
针对这些观察结果，Brax 的设计使其物理计算在其数千个并行环境中的每一个环境中都完全相同，方法是确保模拟没有分支（即，模拟“ if”逻辑因此而发散环境状态）。物理引擎中分支的一个示例是在球和墙壁之间施加接触力：将根据球是否接触墙壁执行不同的代码路径。也就是说，如果球接触墙壁，则将执行用于模拟球从墙壁反弹的单独代码。Brax 混合使用以下三种策略来避免分支：

• 用连续函数替换离散分支逻辑，例如使用带符号距离函数近似球壁接触力。这种方法会带来最大的效率收益。

• 在 JAX 的实时编译期间评估分支。许多基于环境静态属性的分支，例如两个对象是否有可能发生碰撞，可以在模拟时间之前进行评估。

• 在仿真过程中运行分支的两侧，然后只选择所需的结果。因为这会执行一些最终没有使用的代码，所以与上面相比，它浪费了操作。

一旦保证计算完全一致，就可以降低整个训练架构的复杂性，以便在单个 TPU 或 GPU 上执行。这样做可以消除跨机器通信的计算开销和延迟。在实践中，这些变化将可比工作负载的训练成本降低了 100 到 1000 倍。

Brax 环境
环境是微小的封装世界，定义了 RL 智能体要学习的任务。环境不仅包含模拟世界的手段，还包含功能，例如如何观察世界以及在该世界中定义目标。

近年来出现了一些标准的基准环境，用于测试新的 RL 算法并使用研究科学家普遍理解的指标评估这些算法的影响。Brax 包括来自流行的OpenAI 健身房的 四个此类即用型环境：Ant、HalfCheetah、Humanoid和Reacher。

从左到右：Ant、HalfCheetah、Humanoid 和 Reacher 是 RL 研究的流行基线环境。

Brax 还包括三种新颖的环境：对象的灵巧操纵（机器人技术中的一个流行挑战）、广义运动（一种智能体前往放置在其周围任何位置的目标）以及工业机器人手臂的模拟。

左图： Grap，一只学习灵巧操作的爪手。中间： Fetch，一只玩具，盒状狗，学习一般的基于目标的运动策略。右图：工业机器人手臂UR5e 的仿真。

性能基准
分析 Brax 性能的第一步是测量它模拟大批量环境的速度，因为这是要克服的关键瓶颈，以便学习者消耗足够的经验来快速学习。

下面的这两个图显示了 Brax 可以产生多少物理步骤（环境状态的更新），因为它的任务是并行模拟越来越多的环境。左图显示，Brax 将每秒步数与并行环境的数量成线性比例，仅在 10,000 个环境时遇到内存带宽瓶颈，这不仅足以训练单个智能体，还适合训练整个群体智能体。右图显示了两件事：第一，Brax 不仅在 TPU 上表现良好，而且在高端 GPU 上也表现良好（参见V100和P100曲线），第二，通过利用 JAX 的设备并行原语, Brax 可在多个设备之间无缝扩展，达到每秒数亿个物理步骤（参见TPUv3 8x8曲线，即 64 个TPUv3芯片通过高速互连直接相互连接）

左图： 4x2 TPU v3 上每个 Brax 环境每秒模拟步数的缩放比例。右图： Ant 环境中多个加速器每秒模拟步数的缩放比例。

分析 Brax 性能的另一种方法是衡量其对在单个工作站上运行强化学习实验所需时间的影响。在这里，我们将 Brax 训练流行的Ant基准测试环境与其OpenAI 对应物进行比较，后者由MuJoCo 物理引擎提供支持。

在下图中，蓝线代表标准工作站设置，其中学习器在 GPU 上运行，模拟器在 CPU 上运行。我们看到训练蚂蚁以合理的熟练度（y 轴上的分数为 4000）运行所需的时间从蓝线的大约 3 小时下降到在加速器硬件上使用 Brax 的大约 10 秒。有趣的是，即使仅在 CPU 上（灰线），Brax 的执行速度也快了一个数量级，这得益于处于同一进程中的学习器和模拟器。

Brax 的优化PPO与标准 GPU 支持的 PPO 学习 MuJoCo-Ant-v2 环境，评估了 1000 万步。请注意，x 轴是以秒为单位的 log-wallclock-time。阴影区域表示超过 5 个副本的最低和最高性能种子，实线表示平均值。

Physics Fidelity
设计一个与现实世界行为相匹配的模拟器是一个已知的难题，这项工作没有解决。尽管如此，将 Brax 与参考模拟器进行比较以确保其产生至少同样有效的输出是有用的。在这种情况下，我们再次将 Brax 与MuJoCo进行比较，后者因其模拟质量而广受好评。我们希望看到，在其他条件相同的情况下，无论是在 MuJoCo 还是 Brax 中训练，策略都具有相似的奖励轨迹。

MuJoCo-Ant-v2 与 Brax Ant，显示根据环境获得的平均情节分数绘制的环境步骤数。两种环境都使用相同的SAC标准实现进行了训练。阴影区域表示五次运行中表现最低和最高的种子，实线表示平均值。

这些曲线表明，随着两个模拟器的奖励以大致相同的速度增长，两个引擎计算物理的复杂性或解决难度相当。由于两条曲线都以大致相同的奖励达到顶峰，因此我们相信相同的一般物理限制适用于在任一模拟中尽其所能操作的智能体。

我们还可以测量 Brax 保持线性动量、角动量和能量守恒的能力。

Brax 和其他几个物理引擎的线性动量（左）、角动量（中）和能量（右）非守恒缩放。y 轴表示预期计算的漂移（漂移越小越好），x 轴表示模拟的时间量。

这种物理模拟质量的衡量标准是由 MuJoCo 的作者首先提出的，目的是了解模拟如何在计算越来越大的时间步长时偏离轨道。在这里，Brax 的表现与其邻居相似。

结论
我们邀请研究人员通过在Brax Training Colab 中训练他们自己的策略来对 Brax 的物理保真度进行更定性的测量。学习到的轨迹与 OpenAI Gym 中看到的轨迹非常相似。

我们的工作使快速、可扩展的 RL 和机器人研究变得更加容易——以前只能通过大型计算集群才能实现的研究现在可以在工作站上运行，或者通过托管的 Google Colaboratory免费运行。我们的 Github 存储库不仅包括 Brax 模拟引擎，还包括大量用于快速训练的参考 RL 算法。我们迫不及待地想看看 Brax 能够实现什么样的新研究。

致谢
我们要感谢我们的论文合著者：Anton Raichuk、Sertan Girgin、Igor Mordatch 和 Olivier Bachem。我们还要感谢 Erwin Coumans 关于构建物理引擎的建议、Blake Hechtman 和 James Bradbury 为 JAX 和 XLA 提供优化帮助，以及 Luke Metz 和 Shane Gu 的建议。我们还要感谢 Vijay Sundaram、Wright Bagwell、Matt Leffler、Gavin Dodd、Brad Mckee 和 Logan Olson，感谢他们帮助孵化这个项目。

原文链接：

https://ai.googleblog.com/2021/07/speeding-up-reinforcement-learning-with.html

完

总结1：周志华 || AI领域如何做研究-写高水平论文

总结2：全网首发最全深度强化学习资料(永更)

总结3:  《强化学习导论》代码/习题答案大全

总结4：30+个必知的《人工智能》会议清单

总结5：2019年-57篇深度强化学习文章汇总

总结6:   万字总结 || 强化学习之路

总结7：万字总结 || 多智能体强化学习(MARL)大总结

总结8：深度强化学习理论、模型及编码调参技巧

完

第113篇：如何降低深度强化学习研究的计算成本

第112篇：基于Novelty-Pursuit的高效探索方法

第111篇：Reward is enough奖励机制。

第110篇：163篇ICML2021强化学习领域论文汇总

第109篇：【Easy-RL】200页强化学习总结笔记

第108篇：清华大学李升波老师《强化学习与控制》

第107篇：阿里巴巴2022届强化学习实习生招聘

第106篇：奖励机制不合理：内卷，如何解决？

第105篇：FinRL: 一个量化金融自动交易RL库

第104篇：RPG: 通过奖励发现多智能体多样性策略

第103篇：解决MAPPO（Multi-Agent PPO）技巧

第102篇：82篇AAAI2021强化学习论文接收列表

第101篇：OpenAI科学家提出全新强化学习算法

第100篇：Alchemy: 元强化学习(meta-RL)基准环境

第99篇：NeoRL:接近真实世界的离线强化学习基准

第98篇：全面总结(值函数与优势函数)的估计方法

第97篇：MuZero算法过程详细解读

第96篇:  值分布强化学习（Distributional RL）总结

第95篇：如何提高"强化学习算法模型"的泛化能力?

第94篇：多智能体强化学习《星际争霸II》研究

第93篇：MuZero在Atari基准上取得了新SOTA效果

第92篇：谷歌AI掌门人Jeff Dean获冯诺依曼奖

第91篇：详解用TD3算法通关BipedalWalker环境

第90篇：Top-K Off-Policy  RL论文复现

第89篇：腾讯开源分布式多智能TLeague框架

第88篇：分层强化学习(HRL)全面总结

第87篇：165篇CoRL2020 accept论文汇总

第86篇：287篇ICLR2021深度强化学习论文汇总

第85篇：279页总结"基于模型的强化学习方法"

第84篇：阿里强化学习领域研究助理/实习生招聘

第83篇：180篇NIPS2020顶会强化学习论文

第82篇：强化学习需要批归一化(Batch Norm)吗？

第81篇：《综述》多智能体强化学习算法理论研究

第80篇：强化学习《奖励函数设计》详细解读

第79篇: 诺亚方舟开源高性能强化学习库“刑天”

第78篇：强化学习如何tradeoff"探索"和"利用"？

第77篇：深度强化学习工程师/研究员面试指南

第76篇：DAI2020 自动驾驶挑战赛(强化学习)

第75篇：Distributional Soft Actor-Critic算法

第74篇：【中文公益公开课】RLChina2020

第73篇：Tensorflow2.0实现29种深度强化学习算法

第72篇：【万字长文】解决强化学习"稀疏奖励"

第71篇：【公开课】高级强化学习专题

第70篇：DeepMind发布"离线强化学习基准“

第69篇：深度强化学习【Seaborn】绘图方法

第68篇：【DeepMind】多智能体学习231页PPT

第67篇：126篇ICML2020会议"强化学习"论文汇总

第66篇：分布式强化学习框架Acme，并行性加强

第65篇：DQN系列(3): 优先级经验回放(PER)

第64篇：UC Berkeley开源RAD来改进强化学习算法

第63篇：华为诺亚方舟招聘 || 强化学习研究实习生

第62篇：ICLR2020- 106篇深度强化学习顶会论文

第61篇：David Sliver 亲自讲解AlphaGo、Zero

第60篇：滴滴主办强化学习挑战赛:KDD Cup-2020

第59篇：Agent57在所有经典Atari 游戏中吊打人类

第58篇：清华开源「天授」强化学习平台

第57篇：Google发布"强化学习"框架"SEED RL"

第56篇：RL教父Sutton实现强人工智能算法的难易

第55篇：内推 ||  阿里2020年强化学习实习生招聘

第54篇：顶会 || 65篇"IJCAI"深度强化学习论文

第53篇：TRPO/PPO提出者John Schulman谈科研

第52篇：《强化学习》可复现性和稳健性，如何解决？

第51篇：强化学习和最优控制的《十个关键点》

第50篇：微软全球深度强化学习开源项目开放申请

第49篇：DeepMind发布强化学习库 RLax

第48篇：AlphaStar过程详解笔记

第47篇：Exploration-Exploitation难题解决方法

第46篇：DQN系列(2): Double DQN 算法

第45篇：DQN系列(1): Double Q-learning

第44篇：科研界最全工具汇总

第43篇：起死回生|| 如何rebuttal顶会学术论文?

第42篇：深度强化学习入门到精通资料综述

第41篇：顶会征稿 ||  ICAPS2020: DeepRL

第40篇：实习生招聘 || 华为诺亚方舟实验室

第39篇：滴滴实习生|| 深度强化学习方向

第38篇：AAAI-2020 || 52篇深度强化学习论文

第37篇：Call For Papers# IJCNN2020-DeepRL

第36篇：复现"深度强化学习"论文的经验之谈

第35篇：α-Rank算法之DeepMind及Huawei改进

第34篇：从Paper到Coding, DRL挑战34类游戏

第33篇：DeepMind-102页深度强化学习PPT

第32篇：腾讯AI Lab强化学习招聘(正式/实习)

第31篇：强化学习，路在何方？

第30篇：强化学习的三种范例

第29篇：框架ES-MAML：进化策略的元学习方法

第28篇：138页“策略优化”PPT--Pieter Abbeel

第27篇：迁移学习在强化学习中的应用及最新进展

第26篇：深入理解Hindsight Experience Replay

第25篇：10项【深度强化学习】赛事汇总

第24篇：DRL实验中到底需要多少个随机种子？

第23篇：142页"ICML会议"强化学习笔记

第22篇：通过深度强化学习实现通用量子控制

第21篇：《深度强化学习》面试题汇总

第20篇：《深度强化学习》招聘汇总(13家企业）

第19篇：解决反馈稀疏问题之HER原理与代码实现

第18篇："DeepRacer" —顶级深度强化学习挑战赛

第17篇：AI Paper | 几个实用工具推荐

第16篇：AI领域：如何做优秀研究并写高水平论文？

第15篇：DeepMind开源三大新框架！

第14篇：61篇NIPS2019DeepRL论文及部分解读

第13篇：OpenSpiel(28种DRL环境+24种DRL算法)

第12篇：模块化和快速原型设计Huskarl DRL框架

第11篇：DRL在Unity自行车环境中配置与实践

第10篇：解读72篇DeepMind深度强化学习论文

第9篇：《AutoML》：一份自动化调参的指导

第8篇：ReinforceJS库（动态展示DP、TD、DQN）

第7篇：10年NIPS顶会DRL论文(100多篇)汇总

第6篇：ICML2019-深度强化学习文章汇总

第5篇：深度强化学习在阿里巴巴的技术演进

第4篇：深度强化学习十大原则

第3篇：“超参数”自动化设置方法---DeepHyper

第2篇：深度强化学习的加速方法

第1篇：深入浅出解读"多巴胺（Dopamine）论文"、环境配置和实例分析