Uber开源深度神经进化加速器：单机4小时完成Atari训练
更多干货内容请关注微信公众号“AI 前线”，（ID：ai-front）
我们最近发表了一篇博文《Welcoming the Era of Deep Neuroevolution》[1]，宣布了关于深度神经进化（deep neuroevolution）的五篇论文，其中包括发现遗传算法（genetic algorithms）可以解决深度强化学习问题 [2] 以及流行的替代方法，如深度 Q 学习（deep Q-learning）和策略梯度（policy gradients）。这项工作承接了 Salimans 等人在 2017 年的研究成果 [3]，这对进化策略 (evolution strategies，ES) 也同样适用，这是另一种神经进化算法。我们进一步描述了如何通过增加对智能体的压力的形式来改进 ES，以及 ES 与梯度下降之间的关系。所有这些研究都在计算上非常昂贵：它是在分布在大型高性能计算集群上进行的，有 720~3000 个 CPU，这一点似乎让大多数研究人员、学生、公司和爱好者都无法获得深度神经进化。
今天，我们发布了开源代码 [4]，这使得我们可以更快、更便宜地开展此类研究。有了这段代码，以前训练深度神经网络来运行 Atari(在 720 个 CPU 上花费 1 小时) 的时间，现在在一台现代的台式机上只需花费 4 个小时。这一点很重要，因为它极大地影响了我们对进行此类研究所需的资源范围的认知，更多的研究人员得以能够接触到这些资源。
神经进化技术是解决具有挑战性的深度强化学习问题的竞争力的替代方案，如 Atari 和仿人运动（humanoid locomotion）。视频所示为用简单遗传算法训练的深度神经网络的行为。
事实证明，拥有数十个虚拟内核的现代高端台式机本身就像一个适度的计算集群。如果评估正确地并行执行，那么一个在 720 个 CPU 上花费 1 个小时的运行，在一台 48 核 CPU 的个人电脑上运行需要 16 个小时内，这样运行速度较慢。但并非如此，因为现代台式机也有 GPU，它在运行深度神经网络（deep neural network，DNN）时速度很快。我们的代码可以最大限度地并行使用 CPU 和 GPU。它在 GPU 上运行深度神经网络，在 CPU 上运行域（如视频游戏或物理仿真器），并且批量并行执行多个评估，从而可以有效地利用所有可用硬件。如下所述，它还包含自定义的 TensorFlow 操作，这大大提高了训练速度。
在 GPU 上进行训练需要对如何计算神经网络操作进行一些修改。在我们的设置中，使用单一 CPU 运行单个神经网络比使用 GPU 更快，但是当类似的计算（例如神经网络的前向传播）并行执行时，GPU 会大大受益。为了利用 GPU，我们将多个神经网络的前向传播聚合成批，这样的做法在神经网络研究中很常见，但通常只涉及到相同的神经网络处理一批不同的输入。然而，进化是在不同的神经网络上运行的，但即使网络不同（尽管内存需求增加），加速也会发生。我们使用基本的 TensorFlow 操作实现了这个批，它产生大约两倍的加速，将训练时间缩短到 8 个小时左右。然而，我们意识到还可以做得更好。虽然 TensorFlow 提供了所有需要的操作，但这些操作并不适用这类计算。因此，我们添加了两种类型的自定义 TensorFlow 操作，这两种操作相结合可以产生另一个 2 倍的加速，从而将单台机器上的训练时间减少到大约 4 小时，记住这个最初提到的数值。
第一次定制的 TensorFlow 操作大大加快了 GPU 的速度。它是专为不同长度的 RL 域内的异构神经网络计算而构建的，在 Atari 和许多仿真机器人学习任务中都是如此。它允许 GPU 只运行需要运行的多个网络，而不需要每次迭代运行一组固定 (大型) 的网络集合。
迄今为止所描述的改进使得 GPU 比 CPU 更为划算。事实上，即使使用多处理库来并行化计算，GPU 的速度也非常快，以至于 Atari 仿真（CPU）无法跟上。为了提高仿真性能，我们添加了第二套自定义的 TensorFlow 操作。这些将 Atari 仿真的封装从 Python 改为自定义的 TensorFlow 命令（复位、步骤、观察），这些命令充分利用了 TensorFlow 提供的快速多线程功能，而没有与 Python 和 TensorFlow 交互的典型速度拖慢现象。总体而言，这些变化导致 Atari 仿真器的速度提高了三倍。这些创新应该可以加速任何具有多个领域实例的强化学习研究，这在强化学习中，是一种日益常见的技术，如分布式深度 Q 学习 [5]（DQN）[6] 和分布式策略梯度（如 A3C[7]）。
一旦我们有能力在 GPU 上快速运行一个网络，并在 CPU 上运行速度更快的域仿真器，这一挑战就是让计算机上所有资源尽可能地运行。如果我们在每个神经网络进行前向传播，询问在当前状态下应该采取什么行为，那么当每个神经网络都计算出它的答案时，运行游戏仿真器的 CPU 在这期间什么都不做。类似的，如果我们接着采取行动，询问域仿真器“这些行为产生什么状态？”那么运行神经网络的 GPU 在这一步骤中就处于空闲状态。下面所示是多线程 CPU+GPU 选项。虽然是单线程计算的改进，但效率仍然很低。
一个更好的解决方案是，将两个或更多的神经网络子集与仿真器配对，并且保持 GPU 和 CPU 同时运行，根据不同的步骤（神经网络或仿真）准备就绪，更新来自不同集合的网络或仿真。这种方法是最右边的“流水线 CPU+GPU”选项，如下图所示。有了它，以及上面提到的其他改进，我们能够将从约 400 万参数的神经网络的训练时间降到上文提到的数值（即在单台计算上运行约 4 小时）。
优化 RL 中异构网络群体的调度问题。蓝色框是域仿真器，例如 Atari 游戏仿真器或像 MuJoCo 这样的物理引擎，它可以有不同长度的片段。使用 GPU(左) 的一种简单的方法会导致性能低下，原因有两个：1)GPU 的批大小为 1，无法利用其并行计算能力；2)GPU 等待 CPU 的空闲时间，反之亦然。一个多线程的方法 (中心) 允许多个 CPU 并行执行仿真器，从而提高 GPU 的使用率 U，但会导致 GPU 在 CPU 工作时处于空闲状态，反之亦然。我们的流水线实现 (右) 允许 GPU 和 CPU 高效运行。这种方法也适用于多个 GPU 和 CPU 同时运行，这是我们在实践中所做的。
我们的代码使研究社区中的每个人，包括学生和自学者，都能快速实验性地在迭代训练深度神经网络进行实验，以处理像 Atari 这样的具有挑战性的问题，而 Atari 迄今为止仅限于只有资金雄厚的行业和学术实验室才能进行。
更快的代码带来了研究的进展。例如，我们的新代码能够以很小的代价启动一个广泛的超参数搜索遗传算法，这导致了大多数 Atari 游戏的性能改进，而不是我们最初报道的那些情况。我们将会使用这些新的结果来更新我们最初的 arXiv 论文 [8]。更快的代码也在催化了我们目前的研究，通过缩短我们的迭代次数来改善深度神经进化，使我们能够在更多的领域尝试每一个新的想法，并延长算法的运行时间。
我们的新软件库包含了深度遗传算法的实现、Salimans 等人的进化策略算法 [9]，以及我们的（出乎意料的竞争![10]）随机搜索控制。我们希望其他人能够使用我们的代码来加速自己的研究活动，还邀请了社区建立我们的代码并改进它。例如，通过分布式 GPU 训练，以及添加为此类计算而定制的其他 TensorFlow 操作，可能会进一步提高速度。
深度神经进化过程中有很多的动力。除了我们的工作和上面提到的 OpenAI 外，最近还有提到的深度学习进展，使用了 DeepMind、Google Brain 和 Sentient 的进化算法，我们希望我们的代码，开源有助于这一势头的发展，使这个领域更容易为大众获得。
原文链接：https://eng.uber.com/accelerated-neuroevolution/
相关链接：
[1] Welcoming the Era of Deep Neuroevolution
https://eng.uber.com/deep-neuroevolution/
[2] Improving Exploration in Evolution Strategies for Deep Reinforcement Learning via a Population of Novelty-Seeking Agents
https://arxiv.org/abs/1712.06560
[3] Evolution Strategies as a Scalable Alternative to Reinforcement Learning https://arxiv.org/abs/1703.03864
[4] AI Labs - GPU Neuroevolution
https://github.com/uber-common/deep-neuroevolution/tree/master/gpu_implementation
[5] Distributed Prioritized Experience Replay
https://openreview.net/forum?id=H1Dy---0Z
[6] Human-level control through deep reinforcement learning
https://www.nature.com/articles/nature14236
[7] Asynchronous Methods for Deep Reinforcement Learning
https://arxiv.org/pdf/1602.01783.pdf
[8] Deep Neuroevolution: Genetic Algorithms Are a Competitive Alternative for Training Deep Neural Networks for Reinforcement Learning
https://arxiv.org/abs/1712.06567
[9] Evolution Strategies as a Scalable Alternative to Reinforcement Learning
https://blog.openai.com/evolution-strategies/
[10] @jeffclune的一则推文
https://twitter.com/jeffclune/status/943197998502240256
如果你喜欢这篇文章，或希望看到更多类似优质报道，记得给我留言和点赞哦！