1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍

2021年07月22日 15:19 次阅读 稿源:量子位 条评论

机器人要如何完成这样一个动作?我们一般会基于强化学习,在仿真环境中进行模拟训练。这时,如果在一台机器的CPU环境下进行模拟训练,那么需要几个小时到几天。但现在,只需 一个TPU/GPU,就能和数千个CPU或GPU的计算集群的速度一样快,直接将所需时间缩短到几分钟!

博雯 发自 凹非寺

量子位 报道 | 公众号 QbitAI

相当于将强化学习的速度提升了 1000倍

这就是来自Google的科学家们开发的物理模拟引擎 Brax

三种策略避免逻辑分支

现在大多数的物理模拟引擎都是怎么设计的呢?

将重力、电机驱动、关节约束、物体碰撞等任务都整合在一个模拟器中, 并行地进行多个模拟,以此来逼近现实中的运动系统。

对于每个模拟时间步长,力和力矩被整合在一起

这种情况下,每个模拟器中的计算都不相同,且数据必须在数据中心内通过网络传输。

这种并行布局也就导致了较高的延迟时间——即学习者可能需要超过 10000纳秒的等待时间,才能从模拟器中获得经验。

那么怎样才能缩短这种延迟时间呢?

Brax选择通过 避免模拟中的分支来保证数千个并行环境中的计算完全统一,进而降低整个训练架构的复杂度。

直到复杂度降低到可以在单一的TPU或GPU上执行,跨机器通信的计算开销就随之降低,延迟也就能被有效消除。

主要分为以下三个方法:

  • 连续函数替换离散分支逻辑

比如,在计算一个小球与墙壁之间的接触力时,就产生了一个分支:

如果球接触墙壁,就执行模拟球从墙壁反弹的独立代码;

否则,就执行其他代码;

这里就可以通过符号距离函数来避免这种if/else的离散分支逻辑的产生。

  • 使用JAX即时编译中评估分支

在仿真时间之前评估基于环境静态属性的分支,例如两个物体是否有可能发生碰撞。

  • 在模拟中只选择需要的分支结果

在使用了这三种策略之后,我们就得到了一个模拟由刚体、关节、执行器组成环境的物理引擎。

同时也是一种实现在这种环境中各类操作(如进化策略,直接轨迹优化等)的学习算法。

那么Brax的性能究竟如何呢?

速度最高提升1000倍

Brax测试所用的基准是OpenAI Gym中Ant、HalfCheetah、Humanoid、Reacher四种。

同时也增加了三个新环境:包括对物理的灵巧操作、通用运动(例如前往周围任何一个放置了物体的地点)、以及工业机器人手臂的模拟:

研究人员首先测试了Brax在并行模拟越来越多的环境时,可以产生多少次物理步骤(也即对环境状态的更新)。

测试结果中的TPUv3 8x8曲线显示,Brax可以在多个设备之间进行无缝扩展,每秒可达到 数亿个物理步骤:

而不仅是在TPU上,从V100和P100曲线也能看出,Brax在高端 GPU上同样表现出色。

然后就是Brax在单个工作站 (workstation)上运行一个强化学习实验所需要的时间。

在这里,研究人员将基于Ant基准环境训练的Brax引擎与MuJoCo物理引擎做了对比:

可以看到,相对于MuJoCo(蓝线)所需的将近3小时时间,使用了Brax的加速器硬件最快只需要 10秒

使用Brax,不仅能够提高单核训练的效率,还可以扩展到大规模的并行模拟训练。

论文地址:

https://arxiv.org/abs/2106.13281

下载:

https://github.com/google/brax

参考链接:

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

对文章打分

1个GPU几分钟搞定强化学习训练,谷歌新引擎让深度学习提速1000倍

2 (8%)
已有 条意见

    最新资讯

    加载中...

    今日最热

    加载中...

    热门评论

      Top 10

      招聘


      Advertisment ad adsense googles cpro.baidu.com
      created by ceallan