JAX（Just After eXecution）是一个由DeepMind开发的机器/深度学习库。所有JAX操作都是基于XLA或加速线性代数。XLA由谷歌开发，是一个特定领域的线性代数编译器，使用整个程序优化来加速计算。它使BERT的训练速度几乎快了7.3倍。
它是为高性能数值计算而设计的。JAX于2018年推出，目前由Alphabet子公司DeepMind使用。它类似于数字计算库NumPy，是Python编程的另一个库。其数值函数的API是基于NumPy的。
在这篇文章中，我们将探讨使用JAX进行强化学习的不同库和框架。
JAX强化学习代理
1.RLax
RLax（发音为'放松'）是一个关于JAX的简单库。它为实现强化学习代理提供了有用的构建模块。
它可以直接从GitHub或PyPI安装。
所有RLax代码都可以使用jax.jit为不同的硬件（如CPU、GPU、TPU）进行编译。
欲了解更多信息，请点击这里。
2.Haiku
Haiku是一个用于JAX的库。该神经网络允许用户使用熟悉的编程模型，并提供JAX的纯函数转换。
提供的两个核心工具包括一个模块抽象，hk.Module，和一个简单的函数转换，hk.transform。
它是用纯Python写的，但依赖于C++代码。
存储库可以在这里找到。
- Gymnax
Gymnax是Open AI的gym环境的JAX兼容版本。Gym是一个开源的Python库，用于开发和比较RL算法。它提供了一个标准的API来在学习算法和环境之间进行交流。在发布之后，该API已经成为这样做的领域标准。
欲了解更多信息，请点击这里。
4.Dopamine
Dopamine是一个用于原型设计RL算法的研究框架。它的目的是满足对一个代码库的需求，在这个代码库中，用户可以自由地试验一些疯狂的想法（理论研究）。
使用Dopamine最简单的方法是从源头安装并直接修改源代码。2020年发布的版本支持JAX代理，其中包括一个量化回归代理（QR-DQN）的实现。
它也可以用pip安装。此外，Dopamine支持Atari环境和Mujoco环境。
欲了解更多信息，请点击这里。
5.JAX FLAX (RL)
2020年推出的Flax是一个为JAX设计的高性能神经网络库，它的设计非常灵活。通过分叉一个例子和修改训练循环来尝试新的训练形式，而不是通过向框架添加功能。
它是与JAX团队紧密合作开发的。FLAX背后的基本理念是库代码应该易于阅读和理解。
在其核心部分，Flax是围绕着被称为模块的参数化函数建立的，它可以作为普通函数使用。
谷歌研究:Flax资源库在GitHub上。
6.coax
coax是一个RL的python包，用于用基于JAX的函数近似器解决GymAI环境。
它的设计与RL的核心概念一致，而不是与代理的高级概念一致。这使得coax更加模块化，对RL用户更加友好。
欲了解更多信息，请点击这里。
7.7.Acme
Acme是一个强化学习（RL）构件库，致力于暴露简单的代理。首先，这些代理作为参考实现，为算法性能提供强有力的基线。Acme的构建模块是以这样一种方式设计的，即代理可以在多种规模上编写。
它同时支持TensorFlow v2和JAX。
作者：今天你吸猫了么
链接：https://juejin.cn/post/7117552893523132430
来源：稀土掘金
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
