1个GPU几分钟完成强化学习训练，谷歌新引擎将深度学习提速1000倍机器人是如何完成这样的动作的？我们一般是在基于强化学习的模拟环境中进行模拟训练。这时候如果在一台机器的CPU环境下进行模拟训练，需要几个小时到几天的时间。而现在，只需要一个TPU/GPU就可以和拥有数千个CPU或GPU的计算集群一样快，直接将所需时间缩短到几分钟！相当于将强化学习的速度提高了1000倍！这是Brax，一款由谷歌科学家开发的物理模拟引擎。避免逻辑分支的三种策略如今大多数物理模拟引擎是如何设计的？将重力、电机驱动、关节约束、物体碰撞等任务集成在一个模拟器中，并行进行多个模拟，逼近真实的运动系统。△对于每个模拟时间步长，力和力矩都被整合。在这种情况下，每个模拟器的计算都不同，数据必须通过数据中心内的网络传输。这种并行布局也会导致高延迟——也就是说，学习者可能需要等待超过10,000纳秒才能从模拟器中获得经验。那么如何才能缩短这个延迟时间呢？Brax选择通过避免模拟中的分支来降低整个训练架构的复杂性，以确保在数千个并行环境中计算的完全一致性。直到复杂度降低到可以在单个TPU或GPU上执行，跨机器通信的计算开销降低，延迟才能有效消除。主要分为以下三种方法：连续函数代替离散分支逻辑。例如，在计算小球与墙壁的接触力时，会产生一个分支：如果小球接触到墙壁，则执行模拟小球从墙壁弹起的独立函数。代码;否则，执行其他代码；这里可以利用符号距离函数来避免这个if/else产生离散的分支逻辑。使用JAX即时编译中的评估分支，根据模拟时间之前环境的静态属性评估分支，例如两个对象是否可能发生碰撞。在仿真中只选择需要的分支结果。使用这三种策略后，我们得到了一个模拟由刚体、关节和执行器组成的环境的物理引擎。同时，它也是一种在此环境下实现各种操作（如进化策略、直接轨迹优化等）的学习算法。那么Brax的性能如何呢？速度提升高达1000倍。Brax测试中使用的基准是OpenAIGym中的Ant、HalfCheetah、Humanoid和Reacher。同时，增加了三个新环境：包括灵巧的物理操作、通用运动（例如去放置物体周围的任何地方）和模拟工业机器人手臂：研究人员首先在Brax进行了测试并行模拟。当有很多环境时，可以发生多少物理步骤（即对环境状态的更新）。测试结果中的TPUv38x8曲线表明，Brax可以在多设备间无缝扩展，达到每秒数亿物理步数：不仅在TPU上，从V100和P100曲线来看，Brax在高结束GPU。然后是Brax在单个工作站上运行强化学习实验所需的时间。在这里，研究人员将在Ant基准环境中训练的Brax引擎与MuJoCo物理引擎进行了比较：可以看出，与MuJoCo（蓝线）所需的近3小时相比，仅使用Brax的加速器硬件仅需10秒.使用Brax不仅可以提高单核训练的效率，还可以扩展到大规模并行模拟训练。论文地址：https://arxiv.org/abs/2106.13281下载地址：https://github.com/google/brax
