JAX (Flax) RL

本仓库包含使用 JAX (Flax) 实现的强化学习算法：

• 可学习温度参数的软演员-评论家算法（SAC）[1812.05905]
• 优势加权演员-评论家算法（AWAC）[2006.09359]
• 图像增强就是全部需求 [2004.13649]（仅支持 K=1, M=1）
• 带有剪裁双 Q 学习的深度确定性策略梯度（DDPG）[1509.02971]，结合剪裁双 Q 学习 [1802.09477]
• 随机集成双 Q 学习：无需模型即可快速学习 [2101.05982]
• 行为克隆

本仓库的目标是提供简单、清晰的实现，以便在此基础上开展研究。请勿将本仓库用于基线实验，请改用原始实现（SAC：softlearning，AWAC：rlkit，DrQ：drq)。

如果您在工作中使用了 JAXRL，请在出版物中引用本仓库：

@misc{jaxrl,
  author = {Kostrikov, Ilya},
  doi = {10.5281/zenodo.5535154},
  month = {10},
  title = {{JAXRL: Implementations of Reinforcement Learning algorithms in JAX}},
  url = {https://github.com/ikostrikov/jaxrl},
  year = {2021}
}

您可以在 这里 找到本仓库的更新版本。

更改记录

2022年1月10日

• 添加了随机集成双 Q 学习：无需模型即可快速学习 [2101.05982] 的实现

2021年7月20日

• 添加了带有剪裁双 Q 学习的深度确定性策略梯度（DDPG）[1509.02971] 的实现

2021年5月19日

• 添加了软演员-评论家 v1 [1801.01290] 的实现

2021年4月29日

• 添加了来自“图像增强就是全部需求” [2004.13649] 的数据增强实现

安装

先决条件：

• Python 3.8–3.9（暂不支持 3.10）
• Poetry
• patchelf

建议的构建环境：

# 通用构建依赖
sudo apt-get update; sudo apt-get install make build-essential libssl-dev zlib1g-dev \
libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
# Mujoco 依赖
apt-get -y install wget unzip software-properties-common \
    libgl1-mesa-dev \
    libgl1-mesa-glx \
    libglew-dev \
    libosmesa6-dev patchelf
# Mujoco 安装
curl -OL https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
mkdir ~/.mujoco
tar -zxf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco
rm mujoco210-linux-x86_64.tar.gz

安装步骤如下：

poetry install
# 若需 GPU 支持，请运行
pip install "jax[cuda]==0.3.10" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

如需进一步了解如何在 GPU 上运行代码，请参考 官方仓库 中的相关说明。

对于 MuJoCo 的安装，您可能需要在 .bashrc 文件中添加以下内容：

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/costa/.mujoco/mujoco210/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia

开发

如果您希望修改代码，请按照上述说明进行安装。

示例

故障排除

如果遇到内存不足错误，尤其是在启用视频保存功能时，请参阅 JAX 关于 GPU 内存分配的文档：docs。此外，您可以尝试使用以下环境变量运行程序：

XLA_PYTHON_CLIENT_MEM_FRACTION=0.80 python ...

如果您在远程机器上运行代码，并希望为 DeepMind Control Suite 保存视频，请使用 EGL 进行渲染：

MUJOCO_GL=egl python train.py --env_name=cheetah-run --save_dir=./tmp/ --save_video

TensorBoard

启动 TensorBoard 查看训练和评估日志：

tensorboard --logdir=./tmp/

结果

基于状态的连续控制

基于像素的连续控制

Docker

构建

将您的 MuJoCo 密钥复制到 ./vendor 目录下：

cd remote
docker build -t ikostrikov/jaxrl . -f Dockerfile 

测试

sudo docker run -v <examples-dir>:/jaxrl/ ikostrikov/jaxrl:latest python /jaxrl/train.py --env_name=HalfCheetah-v2 --save_dir=/jaxrl/tmp/

# 使用 GPU
sudo docker run --rm --gpus all -v <examples-dir>:/jaxrl/ --gpus=all ikostrikov/jaxrl:latest python /jaxrl/train.py --env_name=HalfCheetah-v2 --save_dir=/jaxrl/tmp/

贡献

在向本仓库贡献代码之前，请先通过 issue 讨论您计划进行的更改。如果您不熟悉 pull request，请阅读 这篇文档。

致谢

感谢 @evgenii-nikishin 在 JAX 方面的帮助。同时也感谢 @dibyaghosh 在 vmapped ensembles 方面的支持。

JAXRL 快速上手指南

JAXRL 是一个基于 JAX (Flax) 的强化学习算法库，提供了简洁、干净的实现，适合用于科研和二次开发。支持 SAC、AWAC、DrQ、DDPG、REDQ 及行为克隆等主流算法。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: 3.8 或 3.9 (暂不支持 3.10+)
• 硬件: 支持 CUDA 的 NVIDIA GPU (可选，用于加速训练)

前置依赖

在安装前，请确保系统已安装以下基础工具和 MuJoCo 依赖：

# 更新源并安装通用构建依赖
sudo apt-get update
sudo apt-get install make build-essential libssl-dev zlib1g-dev \
libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev

# 安装 MuJoCo 图形渲染依赖
apt-get -y install wget unzip software-properties-common \
    libgl1-mesa-dev \
    libgl1-mesa-glx \
    libglew-dev \
    libosmesa6-dev patchelf

# 安装 Poetry (Python 包管理工具)
curl -sSL https://install.python-poetry.org | python3 -

MuJoCo 安装

下载并配置 MuJoCo 210：

# 下载 MuJoCo
curl -OL https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
mkdir -p ~/.mujoco
tar -zxf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco
rm mujoco210-linux-x86_64.tar.gz

# 配置环境变量 (添加到 ~/.bashrc 末尾)
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/$USER/.mujoco/mujoco210/bin' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia' >> ~/.bashrc
source ~/.bashrc

注意：请将上述命令中的 /home/$USER/ 替换为你实际的用户目录路径，并将你的 MuJoCo License Key (mjkey.txt) 放入 ~/.mujoco/ 目录下。

安装步骤

1. 克隆仓库

   git clone https://github.com/ikostrikov/jaxrl.git
   cd jaxrl
   

2. 使用 Poetry 安装项目依赖

   poetry install
   

3. 安装 GPU 支持 (可选) 如果需要利用 GPU 加速，请在激活虚拟环境后运行以下命令安装特定版本的 JAX：

   poetry run pip install "jax[cuda]==0.3.10" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   

   注：若无 GPU，可跳过此步，JAX 将自动使用 CPU 运行。

基本使用

以下是在 HalfCheetah-v2 环境中训练 SAC 算法的最简示例。

启动训练

在项目根目录下运行：

poetry run python train.py --env_name=HalfCheetah-v2 --save_dir=./tmp/

常用参数说明

• --env_name: 指定 Gym/MuJoCo 环境名称 (例如 Hopper-v2, Walker2d-v2)。
• --save_dir: 日志和模型保存路径。
• --save_video: 启用视频录制 (需配合 MUJOCO_GL=egl 在无头服务器上使用)。

查看训练日志

使用 TensorBoard 可视化训练曲线：

tensorboard --logdir=./tmp/

然后在浏览器访问显示的地址 (通常为 http://localhost:6006)。

无头服务器运行提示

如果在远程服务器运行且需要保存视频，请使用 EGL 渲染后端：

MUJOCO_GL=egl poetry run python train.py --env_name=cheetah-run --save_dir=./tmp/ --save_video