环境 | 安装 | 快速入门 | 训练 | 引用 | 文档

Jumanji @ ICLR 2024

Jumanji已被接受至ICLR 2024，请查看我们的研究论文。

欢迎来到丛林！🌴

Jumanji是一套用JAX编写的多样化、可扩展的强化学习环境。目前已有22个环境！

Jumanji正助力开创RL领域硬件加速研发的新浪潮。Jumanji的高速环境不仅能够加快迭代速度、支持大规模实验，还能同时降低复杂度。Jumanji起源于InstaDeep的研究团队，如今已与开源社区共同开发。如您希望加入我们的行列，请随时联系我们、提交问题并阅读我们的贡献指南，或直接加星 🌟，以及时了解最新进展！

目标🚀

1. 为基于JAX的环境提供简单、经过充分测试的API。
2. 让RL研究更加普及。
3. 促进工业界RL问题的研究，帮助缩小研究与工业应用之间的差距。
4. 提供难度可任意调整的环境。

概览🦜

• 🥑 环境API：基于JAX环境的核心抽象。
• 🕹️ 环境套件：涵盖从简单游戏到NP难组合问题的各类RL环境。
• 🍬 封装器：通过我们的dm_env和gym封装器，轻松对接您喜爱的RL框架与库，如Acme、 Stable Baselines3、 RLlib、Gymnasium 以及DeepMind-Env。
• 🎓 示例：指导手册，便于Jumanji的采用，并突出基于JAX环境的附加价值。
• 🏎️ 训练：示例智能体，可作为您在研究中实现智能体的灵感来源。

环境🌍

Jumanji提供了多样化的环境，涵盖从简单游戏到NP难组合问题。

环境	类别	注册版本	源代码	描述
🔢 Game2048	逻辑	Game2048-v1	代码	文档
🎨 GraphColoring	逻辑	GraphColoring-v1	代码	文档
💣 Minesweeper	逻辑	Minesweeper-v0	代码	文档
🎲 RubiksCube	逻辑	RubiksCube-v0 RubiksCube-partly-scrambled-v0	代码	文档
🔀 SlidingTilePuzzle	逻辑	SlidingTilePuzzle-v0	代码	文档
✏️ Sudoku	逻辑	Sudoku-v0 Sudoku-very-easy-v0	代码	文档
📦 BinPack (3D BinPacking Problem)	包装	BinPack-v1	代码	文档
🧩 FlatPack (2D Grid Filling Problem)	包装	FlatPack-v0	[代码](https://github.com/instadeepai/jumanji/tree/main jumanji en environments packing flat_pack/)	[文档](https://instadeepai.github.io/jumanji en environments flat_pack/)
🏭 JobShop (Job Shop Scheduling Problem)	包装	JobShop-v0	[代码](https://github.com/instadeepai/jumanji tree main jumanji en environments packing job_shop/)	[文档](https://instadeepai.github.io jumanji en environments job_shop/)
🎒 Knapsack	包装	Knapsack-v1	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments packing knapsack/)	[文档](https://instadeepai.github.io jumanji en environments knapsack/)
▒ Tetris	包装	Tetris-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments packing tetris/)	[文档](https://instadeepai.github.io jumanji en environments tetris/)
🧹 Cleaner	路由	Cleaner-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing cleaner/)	[文档](https://instadeepai.github.io jumanji en environments cleaner/)
:link: Connector	路由	Connector-v3	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing connector/)	[文档](https://instadeepai.github.io jumanji en environments connector/)
🚚 CVRP (Capacitated Vehicle Routing Problem)	路由	CVRP-v1	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing cvrp/)	[文档](https://instadeepai.github.io jumanji en environments cvrp/)
🚚 MultiCVRP (Multi-Agent Capacitated Vehicle Routing Problem)	路由	MultiCVRP-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing multi_cvrp/)	[文档](https://instadeepai.github.io jumanji en environments multi_cvrp/)
:mag: Maze	路由	Maze-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing maze/)	[文档](https://instadeepai.github.io jumanji en environments maze/)
:robot: RobotWarehouse	路由	RobotWarehouse-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing robot_warehouse/)	[文档](https://instadeepai.github.io jumanji en environments robot_warehouse/)
🐍 Snake	路由	Snake-v1	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing snake/)	[文档](https://instadeepai.github.io jumanji en environments snake/)
📬 TSP (Travelling Salesman Problem)	路由	TSP-v1	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing tsp/)	[文档](https://instadeepai.github.io jumanji en environments tsp/)
多重最小生成树问题	路由	MMST-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing mmst)	[文档](https://instadeepai.github.io jumanji en environments mmst/)
ᗧ•••ᗣ•• PacMan	路由	PacMan-v1	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing pac_man/)	[文档](https://instadeepai.github.io jumanji en environments pac_man/)
👾 Sokoban	路由	Sokoban-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing sokoban/)	[文档](https://instadeepai.github.io jumanji en environments sokoban/)
🍎 Level-Based Foraging	路由	LevelBasedForaging-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments routing lbf/)	[文档](https://instadeepai.github.io jumanji en environments lbf/)
🚁 Search and Rescue	群体	SearchAndRescue-v0	[代码](https://github.com/instadeepai jumanji tree main jumanji en environments swarms search_and_rescue/)	[文档](https://instadeepai.github.io jumanji en environments search_and_rescue/)

安装 🎬

您可以通过 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 上，您可以使用以下命令安装 Tk： apt-get install python3-tk，或通过 conda 安装：conda install tk。有关可使用的后端列表，请参阅 Matplotlib 后端。

快速入门 ⚡

强化学习从业者会发现 Jumanji 的接口非常熟悉，因为它结合了广泛采用的 OpenAI Gym 和 DeepMind Environment 接口。我们从 OpenAI Gym 中借鉴了 “注册表”和“渲染”方法的概念，而我们的 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 字段， 以便用户可以记录既不属于智能体观测也不属于环境内部状态的环境指标。

高级用法 🧑‍🔬

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

注册表与版本管理 📖

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

训练 🏎️

为了展示如何在 Jumanji 环境上训练强化学习智能体，我们提供了随机智能体和经典的演员-评论家（A2C）智能体。这些智能体可以在 jumanji/training/ 中找到。

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

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

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

贡献 🤝

欢迎贡献！请查看我们的问题跟踪器，了解 优质首次问题。请阅读我们的 贡献指南， 以了解如何提交拉取请求、我们的贡献者许可协议以及社区准则。

引用 Jumanji ✏️

如果您在工作中使用 Jumanji，请按如下方式引用该库：

@misc{bonnet2024jumanji,
    title={Jumanji：一套基于 JAX 的多样化、可扩展强化学习环境},
    author={Clément Bonnet、Daniel Luo、Donal Byrne、Shikha Surana、Sasha Abramowitz、Paul Duckworth、Vincent Coyette、Laurence I. Midgley、Elshadai Tegegn、Tristan Kalloniatis、Omayma Mahjoub、Matthew Macfarlane、Andries P. Smit、Nathan Grinsztajn、Raphael Boige、Cemlyn N. Waters、Mohamed A. Mimouni、Ulrich A. Mbou Sob、Ruan de Kock、Siddarth Singh、Daniel Furelos-Blanco、Victor Le、Arnu Pretorius、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）提供的 Cloud TPU 支持 🌤。

Jumanji 快速上手指南

---

环境准备

系统要求

• Python 3.8 或更高版本
• 支持 JAX 的硬件环境（如 NVIDIA GPU）

前置依赖

在安装 Jumanji 之前，请确保已安装以下依赖：

• JAX
• Optax

建议使用 pip 安装这些依赖，也可以通过 Conda 或其他方式安装。

---

安装步骤

你可以通过 pip 安装 Jumanji：

pip install jumanji

如果你在中国，可以使用国内镜像源加速安装：

pip install jumanji -i https://pypi.tuna.tsinghua.edu.cn/simple

---

基本使用

以下是一个最简单的使用示例，展示如何加载并运行一个 Jumanji 环境（以 Game2048 为例）：

import jumanji
from jumanji.environments.logic.game_2048 import Game2048

# 创建环境实例
env = Game2048()

# 重置环境以获取初始状态
observation, _ = env.reset()

# 打印初始观察值
print(observation)

# 执行一个动作（例如：向右移动）
action = 1
next_observation, reward, done, _ = env.step(action)

# 打印下一步的观察值和奖励
print(next_observation)
print(reward)
print(done)

你也可以使用 gym 接口来与 Jumanji 环境交互，只需导入 gym 并注册环境即可：

import gym
env = gym.make("Game2048-v1")
observation, _ = env.reset()

---

以上就是 Jumanji 的快速上手指南。更多高级用法和环境信息，可参考 官方文档。