RobertTLange/gymnax

What it solves

gymnax 解决了基于 CPU 的强化学习 (RL) 环境速度缓慢的瓶颈。通过在 JAX 中重新实现经典 RL 环境，它允许研究人员在相同的加速器 (GPU/TPU) 上运行环境和策略，从而消除 CPU 与 GPU 之间的通信开销，并实现大规模的向量化。

How it works

它将经典的 OpenAI Gym API 引入 JAX 生态系统。由于它是用 JAX 编写的，环境是完全功能性的，并且与 JAX 原语如 jit (即时编译)、vmap (向量化映射) 和 pmap (并行映射) 兼容。这使得能够同时在数千个环境实例上进行高吞吐量的 rollouts，并使用 lax.scan 将整个回合 (episode) 循环编译成单个加速器操作。

Who it’s for

它专为需要高吞吐量实验的 RL 研究人员和开发者、从事 meta-RL (需要将不同的环境配置并行化) 的研究人员，以及进化优化策略的使用者所设计。

Highlights

• High-Throughput Vectorization: 利用 vmap 和 pmap 实现大规模并行 rollouts。
• Accelerator-Native: 在 GPU/TPU 上同时运行环境和策略，以促进 Anakin sub-architecture 等架构的实现。
• Broad Environment Support: 包含经典控制、bsuite、MinAtar 和各种 meta-RL 任务的实现。
• Functional Control: 允许对随机种子和超参数进行显式控制，以加速并行化。
• Visualization Tools: 包含一个 Visualizer 用于生成智能体 (agent) 行为的 GIF 动画。