Jumanji — 面向强化学习研究的环境库与基准集合

环境
| 安装
| 快速开始
| 训练
| 引用
| 文档

Jumanji @ ICLR 2024

Jumanji 已被 ICLR 2024 收录，查看我们的研究论文。

欢迎来到丛林！🌴

Jumanji 是一个基于 JAX 编写的多样化、可扩展的强化学习环境套件，目前包含 22 个环境。

Jumanji 致力于开创硬件加速的强化学习研究与开发新时代。其高速环境支持更快的迭代和大规模实验，同时降低了复杂性。该项目最初源于 InstaDeep 研究团队，现与开源社区共同开发。欢迎通过提出问题、阅读贡献指南或直接星标 🌟 关注最新进展。

目标 🚀

1. 为基于 JAX 的环境提供简单、经过充分测试的 API。
2. 使强化学习研究更易上手。
3. 促进针对工业问题的强化学习研究，弥合研究与工业应用之间的差距。
4. 提供难度可扩展至任意挑战级别的环境。

概述 🦜

• 🥑 环境 API：基于 JAX 的环境的核心抽象。
• 🕹️ 环境套件：涵盖从简单游戏到 NP-困难组合优化问题的强化学习环境集合。
• 🍬 封装器：通过 dm_env 和 gym 封装器，轻松连接您喜爱的强化学习框架和库，如 Acme、Stable Baselines3、RLlib、Gymnasium 和 DeepMind-Env。
• 🎓 示例：指南文档帮助您快速上手，并突显 JAX 环境的优势。
• 🏎️ 训练：提供示例智能体，可作为您研究实现的参考。

环境 🌍

Jumanji 提供从简单游戏到 NP-困难组合问题的多样化环境。

安装 🎬

你可以从 PyPI 安装 Jumanji 的最新发布版本：

pip install -U jumanji

或者，直接从 GitHub 安装最新的开发版本：

pip install git+https://github.com/instadeepai/jumanji.git

Jumanji 已在 Python 3.10、3.11 和 3.12 上测试通过。请注意，由于 JAX 的安装取决于你的硬件加速器，我们建议用户显式安装正确的 JAX 版本（参见官方安装指南）。

渲染： 所有环境的渲染均使用 Matplotlib。要可视化环境，你需要一个 GUI 后端。例如，在 Linux 上，你可以通过 apt-get install python3-tk 或使用 conda conda install tk 安装 Tk。请查看 Matplotlib 后端 了解你可用后端的列表。

快速开始 ⚡

强化学习从业者会对 Jumanji 的接口感到熟悉，因为它结合了广泛使用的 OpenAI Gym 和 DeepMind Environment 接口。我们借鉴了 OpenAI Gym 的 registry 和 render 方法，而我们的 TimeStep 结构则受 DeepMind Environment 启发。

基本用法 🧑‍💻

import jax
import jumanji

# 使用注册表实例化 Jumanji 环境
env = jumanji.make('Snake-v1')

# 重置你的（可 JIT 编译的）环境
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)

# （可选）渲染环境状态
env.render(state)

# 与（可 JIT 编译的）环境进行交互
action = env.action_spec.generate_value()          # 动作选择（此处为占位值）
state, timestep = jax.jit(env.step)(state, action) # 执行一步并观察下一个状态和时间步

• state 表示环境的内部状态：它包含执行动作所需的所有信息。这不应与 timestep 中包含的 observation（智能体感知到的信息）混淆。
• timestep 是一个数据类，包含 step_type、reward、discount、observation 和 extras。此结构类似于 dm_env.TimeStep，但额外添加了 extras 字段，以允许用户记录环境的指标，这些指标既不属于智能体的观察，也不属于环境的内部状态。

高级用法 🧑‍🔬

使用 JAX 编写，Jumanji 的环境受益于其许多特性，包括自动向量化/并行化（jax.vmap、jax.pmap）和 JIT 编译（jax.jit），它们可以任意组合使用。
我们在高级用法指南中提供了一个更高级用法的示例。

注册表和版本管理 📖

与 OpenAI Gym 类似，Jumanji 为了可重复性，对其环境进行严格的版本管理。我们维护了一个标准环境及其配置的注册表。每个环境都附加了版本后缀，例如 Snake-v1。当对环境进行可能影响学习结果的更改时，版本号将增加一，以防止潜在混淆。要查看每个环境所有注册版本的完整列表，请查阅文档。

训练 🏎️

为了展示如何在 Jumanji 环境中训练强化学习智能体，我们提供了一个随机智能体和一个普通的 Actor-Critic (A2C) 智能体。这些智能体可以在 jumanji/training/ 中找到。

由于 Jumanji 中的环境框架非常灵活，几乎任何问题都可以实现为 Jumanji 环境，从而产生非常多样化的观察。因此，需要特定于环境的网络来捕获每个环境的对称性。除了 A2C 智能体的实现，我们还在 jumanji/training/networks 中提供了此类特定于环境的 Actor-Critic 网络示例。

⚠️ jumanji/training 中的示例智能体仅旨在作为如何实现智能体的参考。Jumanji 首先是一个环境库——因此，这些智能体和网络将不会按照生产标准进行维护。

关于如何使用示例智能体的更多信息，请参阅训练指南。

贡献 🤝

欢迎贡献！请查看我们的 good first issues 跟踪器。在提交拉取请求之前，请阅读我们的贡献指南，了解有关如何提交拉取请求、我们的贡献者许可协议和社区指南的详细信息。

引用 Jumanji ✏️

如果您在作品中使用 Jumanji，请使用以下引用信息：

@misc{bonnet2024jumanji,
    title={Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX},
    author={Clément Bonnet and Daniel Luo and Donal Byrne and Shikha Surana and Sasha Abramowitz and Paul Duckworth and Vincent Coyette and Laurence I. Midgley and Elshadai Tegegn and Tristan Kalloniatis and Omayma Mahjoub and Matthew Macfarlane and Andries P. Smit and Nathan Grinsztajn and Raphael Boige and Cemlyn N. Waters and Mohamed A. Mimouni and Ulrich A. Mbou Sob and Ruan de Kock and Siddarth Singh and Daniel Furelos-Blanco and Victor Le and Arnu Pretorius and Alexandre Laterre},
    year={2024},
    eprint={2306.09884},
    url={https://arxiv.org/abs/2306.09884},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

其他相关项目 🔎

其他一些项目也采用了使用 JAX 编写强化学习环境的方法。特别地，我们建议用户查看以下姊妹仓库：

• 🤖 Qdax 一个通过硬件加速器和并行化来加速质量-多样性和神经进化算法的库。
• 🌳 Evojax 提供工具，使神经进化算法能够在跨多个 TPU/GPU 的神经网络上工作。
• 🦾 Brax 一个可微的物理引擎，模拟由刚体、关节和执行器组成的环境。
• 🏋️‍ Gymnax 实现了经典环境，包括经典控制、bsuite、MinAtar 和一系列的元强化学习任务。
• 🎲 Pgx 提供了经典棋盘游戏环境，如双陆棋、将棋和围棋。

致谢 🙏

本库的开发得到了来自谷歌 TPU 研究云 (TRC) 🌤 的云 TPU 支持。