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

7 月 22 日消息，机器人要如何完成这样一个动作？我们一般会基于强化学习，在仿真环境中进行模拟训练。

这时，如果在一台机器的 CPU 环境下进行模拟训练，那么需要几个小时到几天。

但现在，只需一个 TPU/GPU，就能和数千个 CPU 或 GPU 的计算集群的速度一样快，直接将所需时间缩短到几分钟，相当于将强化学习的速度提升了 1000 倍！

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

三种策略避免逻辑分支

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

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

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

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

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

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

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

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

主要分为以下三个方法：

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

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

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

否则，就执行其他代码；

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

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

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

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

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

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

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

速度最高提升 1000 倍

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

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

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

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

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

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

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

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

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