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/

本笔记本将带你了解基本用法：

请查看 RESOURCES.MD，了解属于 Jax RL 生态系统的其他 GitHub 仓库！

性能

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

Cartpole	Minatar-Breakout

启用向量化训练后，我们可以在一半时间内训练 2048 个 PPO 智能体，而同样的时间仅够在单个 GPU 上训练一个 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。

相关工作

请查看 RESOURCES 列表，了解与 PureJaxRL 密切相关的库！

以下仓库和项目是 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={Discovered policy optimisation},
    author={Lu, Chris and Kuba, Jakub and Letcher, Alistair and Metz, Luke and Schroeder de Witt, Christian and Foerster, Jakob},
    journal={Advances in Neural Information Processing Systems},
    volume={35},
    pages={16455--16468},
    year={2022}
}

PureJaxRL 快速上手指南

PureJaxRL 是一个高性能、端到端的 JAX 强化学习（RL）实现。其核心特点是将**整个训练流程（包括环境）**完全用 JAX 编写，从而利用 JIT 编译和避免 CPU-GPU 数据传输，实现比传统 PyTorch 实现快 1000 倍以上的训练速度。它特别适合在单块 GPU 上并行运行大量智能体、进行超参数快速调优以及元进化研究。

环境准备

• 系统要求：Linux 或 macOS（Windows 支持有限，建议使用 WSL2），配备 NVIDIA GPU 以获得最佳性能。
• 前置依赖：
  • Python 3.8+
  • CUDA/cuDNN（如需使用 GPU 加速）
  • 已配置好的 JAX 环境（支持 GPU/TPU）

注意：请确保您的 JAX 版本已正确安装并能在加速器上运行。具体 JAX 安装细节请参考 JAX 官方文档。国内用户若遇到下载问题，可尝试配置清华源或阿里源安装基础 Python 包，但 JAX 本身的二进制包建议参照官方指引或使用国内镜像站提供的 wheel 包。

安装步骤

克隆仓库并安装依赖：

git clone https://github.com/luchris429/purejaxrl.git
cd purejaxrl
pip install -r requirements.txt

基本使用

PureJaxRL 的设计理念是“单文件实现”，代码清晰简洁，不适合像常规库那样通过 import 调用模块化组件，而是直接运行脚本或参考 Notebook 示例。

1. 运行官方示例 (推荐)

最快速的上手方式是运行官方的 Colab Notebook 或本地 Jupyter Notebook，其中包含了从环境初始化到训练完成的完整流程。

在线体验 (无需本地配置): 点击以下链接直接在 Google Colab 中运行：

本地运行: 如果您已在本地安装好环境，可以下载并运行示例 Notebook：

# 启动 Jupyter Lab 或 Notebook
jupyter lab examples/walkthrough.ipynb

2. 核心代码逻辑简述

在纯 JAX 模式下，训练循环通常包含以下关键变换操作，这也是 PureJaxRL 的核心优势：

• jit: 编译整个训练步骤以加速执行。
• vmap: 向量化处理，同时并行训练多个随机种子（seeds）或智能体。
• pmap: 跨多设备（多卡）并行。
• scan: 高效地展开时间步长的循环。

以下是一个概念性的代码结构示例（具体变量名请参考 examples/walkthrough.ipynb）：

import jax
import jax.numpy as jnp
from purejaxrl.ppo import make_train # 假设导入对应的算法文件

# 1. 配置训练参数
config = {
    "SEED": 0,
    "NUM_ENVS": 16, # 并行环境数量
    "TOTAL_TIMESTEPS": 1e5,
    # ... 其他超参数
}

# 2. 创建训练函数
train_fn = make_train(config)

# 3. 使用 vmap 并行运行多个种子 (例如同时跑 100 个种子)
rngs = jax.random.split(jax.random.PRNGKey(0), 100)
vmapped_train = jax.vmap(train_fn)

# 4. 执行训练 (JIT 会自动应用于内部逻辑)
out = vmapped_train(rngs)

# 5. 获取结果
metrics = out["metrics"]

3. 更多示例

• Brax 与 MinAtar 环境: 查看 examples/brax_minatar.ipynb 了解如何在更复杂的环境中应用 PureJaxRL。

提示：PureJaxRL 的代码旨在作为研究参考和实验基石，建议直接阅读其单文件源码（如 purejaxrl/ppo.py）以理解算法实现的每一个细节。