jumanji

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

Jumanji @ ICLR 2024

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

欢迎来到丛林！🌴

Jumanji是一个用JAX编写的多样化、可扩展的强化学习环境套件。它现在拥有22个环境！

Jumanji正在帮助开创强化学习领域硬件加速研究和开发的新浪潮。Jumanji的高速环境使得更快的迭代和大规模实验成为可能，同时降低了复杂性。Jumanji起源于InstaDeep的研究团队，现在与开源社区共同开发。要加入我们的努力，请联系我们、提出问题并阅读我们的贡献指南，或者只需加星 🌟 以随时了解最新进展！

目标 🚀

1. 为基于JAX的环境提供简单、经过充分测试的API。
2. 使强化学习研究更容易获得。
3. 促进工业问题的强化学习研究，帮助缩小研究和工业应用之间的差距。
4. 提供难度可以任意扩展的环境。

概览 🦜

• 🥑 环境API：JAX环境的核心抽象。
• 🕹️ 环境套件：包含从简单游戏到NP难组合问题的各种强化学习环境。
• 🍬 包装器：轻松连接您喜欢的强化学习框架和库，如Acme、Stable Baselines3、RLlib、OpenAI Gym和DeepMind-Env，通过我们的dm_env和gym包装器。
• 🎓 示例：指南以促进Jumanji的采用并突出基于JAX环境的附加价值。
• 🏎️ 训练：示例智能体，可作为研究中实现智能体的灵感来源。

环境 🌍

Jumanji提供了多种环境，从简单游戏到NP难组合问题。

安装 🎬

您可以从PyPI安装Jumanji的最新版本：

pip install -U jumanji

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

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

Jumanji已在Python 3.8和3.9上进行了测试。 请注意，由于JAX的安装因硬件加速器而异，我们建议用户明确安装正确的JAX版本（请参阅官方安装指南）。

渲染： 所有环境的渲染都使用Matplotlib。要可视化环境，您需要一个GUI后端。例如，在Linux上，您可以通过以下方式安装Tk： apt-get install python3-tk，或使用conda：conda install tk。查看 Matplotlib后端以获取可用后端列表。

快速入门 ⚡

基本用法 🧑‍💻

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环境中训练强化学习代理，我们提供了一个随机代理和一个普通的演员-评论家（A2C）代理。这些代理可以在jumanji/training/中找到。

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

⚠️ jumanji/training中的示例代理仅用作如何实现代理的灵感。Jumanji首先是一个环境库 - 因此，代理和网络不会维护到生产标准。

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

贡献 🤝

欢迎贡献！请查看我们的问题跟踪器以找到适合新手的问题。请阅读我们的贡献指南，了解如何提交拉取请求、我们的贡献者许可协议和社区指南的详细信息。

引用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提供了像西洋双陆棋、将棋和围棋等经典棋盘游戏环境。

致谢 🙏

这个库的开发得到了来自Google的TPU Research Cloud（TRC）🌤的Cloud TPU支持。