PureJaxRL（纯Jax实现的端到端强化学习训练）

PureJaxRL是一个高性能的、端到端的Jax强化学习（RL）实现。在GPU上并行运行多个智能体时，我们的实现比标准的PyTorch RL实现快1000多倍。与其他Jax RL实现不同，我们将整个训练流程（包括环境）都在JAX中实现。这使我们能够通过JIT编译和避免CPU-GPU数据传输获得显著的速度提升。它还使得调试更加容易，因为系统是完全同步的。更重要的是，这段代码允许你使用jax对整个RL训练流程进行jit、vmap、pmap和scan操作。有了这个，我们可以：

• 🏃 在一个GPU上高效地并行运行大量种子
• 💻 进行快速超参数调优
• 🦎 通过元进化发现新的RL算法

更多详情，请访问附带的博客文章：https://chrislu.page/blog/meta-disco/

这个笔记本演示了基本用法：

性能

在不进行向量化的情况下，我们的实现比CleanRL的PyTorch基准快10倍，如单线程性能图所示。

Cartpole	Minatar-Breakout

通过向量化训练，我们可以在单个GPU上训练2048个PPO智能体，用时仅为训练单个PyTorch PPO智能体的一半。向量化智能体训练允许同时跨多个种子进行训练，快速调整超参数，甚至进行进化元强化学习。

向量化Cartpole	向量化Minatar-Breakout

代码理念

PureJaxRL受CleanRL的启发，提供高质量的单文件实现，具有研究友好的特性。与CleanRL一样，这不是一个模块化的库，也不是为了被导入而设计的。该仓库专注于实现的简洁性和清晰度，是研究人员和实践者的优秀资源。

安装

使用requirements.txt文件安装依赖项：

pip install -r requirements.txt

要在您的加速器上使用JAX，可以在JAX文档中找到更多详细信息。

使用示例

examples/walkthrough.ipynb介绍了基本用法。

examples/brax_minatar.ipynb介绍了如何使用PureJaxRL进行Brax和MinAtar。

待办事项

PureJaxRL仓库计划进行以下改进：

1. 更高效的内存日志记录
2. 与Weights & Biases (WandB)集成以进行实验跟踪
3. 连接到非Jax环境，如envpool

相关工作

PureJaxRL基于Jax和RL生态系统中的其他工具。请查看以下项目：

• Gymnax
• Evosax
• CleanRL
• Brax
• Jumanji
• Pgx

以下仓库与PureJaxRL相关，并扩展了其用途！

• JaxMARL
• JAX-CORL
• JaxIRL
• POPJym

以下仓库和项目是purejaxrl的前身：

• 无模型对手塑造 (ICML 2022) (https://github.com/luchris429/Model-Free-Opponent-Shaping)

• 发现式策略优化 (NeurIPS 2022) (https://github.com/luchris429/discovered-policy-optimisation)

• 对抗性廉价对话 (ICML 2023) (https://github.com/luchris429/adversarial-cheap-talk)

引用

如果您在研究中使用了PureJaxRL，请引用以下论文：

@article{lu2022discovered,
    title={发现式策略优化},
    author={Lu, Chris and Kuba, Jakub and Letcher, Alistair and Metz, Luke and Schroeder de Witt, Christian and Foerster, Jakob},
    journal={神经信息处理系统进展},
    volume={35},
    pages={16455--16468},
    year={2022}
}