EvoJAX：硬件加速神经进化

EvoJAX是一个可扩展、通用、硬件加速的神经进化工具包。它基于JAX库构建，使神经进化算法能够在多个TPU/GPU上并行运行神经网络。EvoJAX通过将进化算法、神经网络和任务全部用NumPy实现，并即时编译以在加速器上运行，从而实现极高的性能。

本仓库还包含了EvoJAX在广泛任务中的几个可扩展示例，包括监督学习、强化学习和生成艺术，展示了EvoJAX如何在单个加速器上几分钟内运行您的进化实验，相比之下使用CPU可能需要数小时或数天。

EvoJAX论文：https://arxiv.org/abs/2202.05008（介绍[视频](https://youtu.be/TMkft3wWpb8)）

如果您希望在您的出版物中引用本项目，请使用以下BibTeX：

@article{evojax2022,
  title={EvoJAX: Hardware-Accelerated Neuroevolution},
  author={Tang, Yujin and Tian, Yingtao and Ha, David},
  journal={arXiv preprint arXiv:2202.05008},
  year={2022}
}

使用EvoJAX的出版物列表（请通过PR添加缺失的条目）：

现代创造力进化策略：拟合具体图像和抽象概念（NeurIPS创造力研讨会2021，EvoMUSART 2022）

安装

EvoJAX是在JAX上实现的，需要先安装JAX。

安装JAX：请首先按照JAX的安装说明进行安装，可选择GPU/TPU后端支持。如果未设置JAX，EvoJAX安装仍会尝试拉取仅CPU版本的JAX。请注意，Colab运行时预装了JAX。

安装EvoJAX：

# 从PyPI安装
pip install evojax

# 或者，从我们的GitHub仓库安装
pip install git+https://github.com/google/evojax.git@main

如果您还想安装某些可选功能所需的额外依赖项，请使用

pip install evojax[extra]
# 或
pip install git+https://github.com/google/evojax.git@main#egg=evojax[extra]

代码概览

EvoJAX是一个由三个主要组件组成的框架，我们期望用户对其进行扩展。

神经进化算法 所有神经进化算法都应实现evojax.algo.base.NEAlgorithm接口，并位于evojax/algo/中。有关EvoJAX中可用算法的信息，请参见此处。
策略网络 所有神经网络都应实现evojax.policy.base.PolicyNetwork接口，并保存在evojax/policy/中。在本仓库中，我们给出了MLP、ConvNet、Seq2Seq和PermutationInvariant模型的示例实现。
任务所有任务都应实现evojax.task.base.VectorizedTask，并位于evojax/task/中。

这些组件可以独立使用，也可以由管理训练流程的evojax.trainer和evojax.sim_mgr协调使用。虽然它们应该足以满足当前提供的策略和任务的需求，但我们计划在未来根据需要扩展其功能。

示例

作为快速入门，我们提供了一些非平凡的示例（examples/中的脚本和examples/notebooks中的笔记本）来说明EvoJAX的用法。我们在每个脚本的顶部提供了启动训练过程的示例命令。这些脚本和笔记本在TPU和/或NVIDIA V100 GPU上运行：

监督学习任务

虽然在实践中显然会使用梯度下降来解决此类任务，但重点是要表明神经进化也能在短时间内以一定程度的准确性解决这些问题，这在将这些模型适应到更复杂的任务中时会很有用，因为在那些情况下基于梯度的方法可能不起作用。

MNIST分类 - 我们展示了EvoJAX在单个GPU上5分钟内训练ConvNet策略，达到>98%的测试准确率。
Seq2Seq学习 - 我们演示了EvoJAX能够学习一个具有数十万参数的大型网络来完成seq2seq任务。

经典控制任务

包含控制任务有两个目的：1) 与监督学习任务不同，EvoJAX中的控制任务步数不确定，因此我们使用这些示例来展示我们任务展开循环的效率。2) 我们希望展示在JAX中实现任务的加速优势，并说明如何从头开始实现一个任务。

运动 - Brax是一个用JAX实现的可微分物理引擎。我们将其包装为一个任务，并在GPU/TPU上使用EvoJAX进行训练。EvoJAX只需几十分钟就能解决Brax中的运动任务。
倒立摆摆动 - 我们说明了如何在JAX中实现经典控制任务，并将其集成到EvoJAX的流程中，以显著加快训练速度。

新颖任务

在最后这一类中，我们超越简单的说明，展示了更实用和吸引遗传和进化计算领域研究人员的新颖任务示例，目的是帮助他们在EvoJAX中尝试想法。

<table width="100%"> <tr> <td width="30%"> <img width="100%" src="https://yellow-cdn.veclightyear.com/835a84d5/50a8dcdd-e122-455a-acba-97749f42d71a.gif"></img> </td> <td width="30%"> <img width="100%" src="https://yellow-cdn.veclightyear.com/835a84d5/c4e2cee2-baa7-49de-98ce-ee8e39ea4293.gif"></img> </td> <td width="40%"> <img width="100%" src="https://yellow-cdn.veclightyear.com/835a84d5/b19ed966-cd39-4541-b9ce-65626241e07f.gif"></img> </td> </tr> <tr> <td> 多智能体水世界 </td> <td> ES-CLIP：<i>"一幅猫的画"</i> </td> <td> 史莱姆排球 </td> </tr> </table> * [水世界](https://github.com/google/evojax/blob/main/examples/train_waterworld.py) - 在这个[任务](https://cs.stanford.edu/people/karpathy/reinforcejs/waterworld.html)中,智能体试图获取尽可能多的食物,同时避开毒药。EvoJAX能够在单个GPU上用数十分钟训练出智能体。此外,我们证明在EvoJAX中可以进行[多智能体训练](https://github.com/google/evojax/blob/main/examples/train_waterworld_ma.py),这有利于学习能够应对环境复杂性和不确定性的策略。

抽象画(笔记本1和笔记本2) - 我们复现了这项计算创造力工作的结果,并展示了如何使用EvoJAX在单个GPU上高效地加速原本需要多个CPU和GPU的实现,这在之前是不可能的。此外,使用多个GPU/TPU,EvoJAX可以进一步将上述工作的速度提升近乎线性。我们还展示了EvoJAX的模块化设计允许其组件被独立使用 - 在这种情况下,可以只使用EvoJAX的ES算法,同时利用自己的训练循环和环境实现。
神经网络史莱姆排球 - 在这个任务中,智能体的目标是让球落在对手一侧的地面上,使对手失去一条生命。当任一智能体失去全部五条生命或达到时间限制时,回合结束。当对手失去生命时,智能体获得+1的奖励,当自己失去生命时获得-1的奖励。 EvoJAX能够在单个GPU上用不到5分钟训练出智能体,相比之下在多个CPU上需要数小时。这个实现基于Slime Volleyball Gym环境,它是原始JavaScript版游戏的Python移植版。在所有这些版本中,内置的AI对手和不太理想的物理引擎都是相同的。