1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍
本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。
机器人要如何完成这样一个动作?
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s2.51cto.com/oss/202107/27/5e09b4a6c22d6872308b81595c10b726.jpg)
我们一般会基于强化学习,在仿真环境中进行模拟训练。
这时,如果在一台机器的CPU环境下进行模拟训练,那么需要几个小时到几天。
但现在,只需一个TPU/GPU,就能和数千个CPU或GPU的计算集群的速度一样快,直接将所需时间缩短到几分钟!
相当于将强化学习的速度提升了1000倍!
这就是来自谷歌的科学家们开发的物理模拟引擎Brax。
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s6.51cto.com/oss/202107/27/1bb04481ee365e08481374ce4cbf49f6.jpg)
三种策略避免逻辑分支
现在大多数的物理模拟引擎都是怎么设计的呢?
将重力、电机驱动、关节约束、物体碰撞等任务都整合在一个模拟器中,并行地进行多个模拟,以此来逼近现实中的运动系统。
△对于每个模拟时间步长,力和力矩被整合在一起
这种情况下,每个模拟器中的计算都不相同,且数据必须在数据中心内通过网络传输。
这种并行布局也就导致了较高的延迟时间——即学习者可能需要超过10000纳秒的等待时间,才能从模拟器中获得经验。
那么怎样才能缩短这种延迟时间呢?
Brax选择通过避免模拟中的分支来保证数千个并行环境中的计算完全统一,进而降低整个训练架构的复杂度。
直到复杂度降低到可以在单一的TPU或GPU上执行,跨机器通信的计算开销就随之降低,延迟也就能被有效消除。
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s3.51cto.com/oss/202107/27/2366a0d6f80d567aac3a5bb99dad9a6d.jpg)
主要分为以下三个方法:
- 连续函数替换离散分支逻辑
比如,在计算一个小球与墙壁之间的接触力时,就产生了一个分支:
如果球接触墙壁,就执行模拟球从墙壁反弹的独立代码;
否则,就执行其他代码;
这里就可以通过符号距离函数来避免这种if/else的离散分支逻辑的产生。
- 使用JAX即时编译中评估分支
在仿真时间之前评估基于环境静态属性的分支,例如两个物体是否有可能发生碰撞。
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s3.51cto.com/oss/202107/27/364c58bb52d6e03e0f97de16b157e6fd.jpg)
- 在模拟中只选择需要的分支结果
在使用了这三种策略之后,我们就得到了一个模拟由刚体、关节、执行器组成环境的物理引擎。
同时也是一种实现在这种环境中各类操作(如进化策略,直接轨迹优化等)的学习算法。
那么Brax的性能究竟如何呢?
速度最高提升1000倍
Brax测试所用的基准是OpenAI Gym中Ant、HalfCheetah、Humanoid、Reacher四种。
同时也增加了三个新环境:包括对物理的灵巧操作、通用运动(例如前往周围任何一个放置了物体的地点)、以及工业机器人手臂的模拟:
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s6.51cto.com/oss/202107/27/e8adca5cfb5725cdf46fca346bca389c.jpg)
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s2.51cto.com/oss/202107/27/9fb35da4803f4bff8dd701df55c34927.jpg)
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s3.51cto.com/oss/202107/27/b5301b2edf6afb51938431fe30265c88.jpg)
研究人员首先测试了Brax在并行模拟越来越多的环境时,可以产生多少次物理步骤(也即对环境状态的更新)。
测试结果中的TPUv3 8x8曲线显示,Brax可以在多个设备之间进行无缝扩展,每秒可达到数亿个物理步骤:
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s5.51cto.com/oss/202107/27/1760129e9937cc23c87d75bccfa87cf2.jpg)
而不仅是在TPU上,从V100和P100曲线也能看出,Brax在高端GPU上同样表现出色。
然后就是Brax在单个工作站(workstation)上运行一个强化学习实验所需要的时间。
在这里,研究人员将基于Ant基准环境训练的Brax引擎与MuJoCo物理引擎做了对比:
![1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍](https://s4.51cto.com/oss/202107/27/ad91dbd618db9ce7079d3b917088e6c5.jpg)
可以看到,相对于MuJoCo(蓝线)所需的将近3小时时间,使用了Brax的加速器硬件最快只需要10秒。
使用Brax,不仅能够提高单核训练的效率,还可以扩展到大规模的并行模拟训练。
论文地址:
https://arxiv.org/abs/2106.13281
下载:
https://github.com/google/brax
相关文章
- 微软为 Windows 10 21H2 适配了新的用户离开自动锁屏检测功能
- Windows 10 21H1功能更新即将登场:新功能前瞻
- 微软 Windows 10 周二更新:移除旧版 Edge 浏览器
- 坐看大数据之道:51CTO专访大数据专家郑玮
- 微软 Windows 10任务栏菜单新变化:隐藏触摸键盘等选项将移至 “设置”应用
- Windows 10 21H2功能更新将进一步改善任务栏交互
- 面试官问我zookeeper选举过程,我当场给他讲了源码
- 微软计划对Windows 10驱动程序发起另一项重要改变
- 微软 Windows1021H2 “太阳谷”将获得新个性化设置:聚焦壁纸、触摸键盘、设备特定用途
- Windows10打印蓝屏后续:微软推新补丁已修复错误
- 即将迎来大改版!Windows10全新文件管理器上手体验
- Windows 10进一步精简:右键菜单去掉画图3D编辑功能
- Google Dremel让大数据不再“恐怖”
- 微软 Windows10 记事本未来将变得更加出色
- Windows 10“设备使用”扩大测试:根据用途进行优化
- 微软为在WindowsUI 3中剔除了Pivot控件而向开发者致歉
- 微软为 Windows101909 推送 KB5000850 更新,修复资源管理器搜索等问题
- 曝微软正将 Windows10「任务栏」从 Explorer.exe 中移出:响应将更流畅稳定
- Windows 10批量换装新图标:五彩缤纷 舒服多了
- Windows10再爆新UI!这样的文件管理器你喜欢吗