instadeepai/jumanji

What it solves

Jumanji 提供了一套可扩展的强化学习 (RL) 环境套件，旨在加速 RL 研究与开发。它解决了对高速、硬件加速环境的需求，这些环境允许更快的迭代与大规模实验，同时降低复杂度。

How it works

该库完全使用 JAX 编写，利用了 JAX 的自动向量化、并行化 (jax.vmap, jax.pmap) 以及 JIT 编译 (jax.jit) 能力。它结合了受 OpenAI Gym 启发的接口（使用注册机制与渲染方法）和 DeepMind Environment 的接口（使用 TimeStep 结构）。

Who it’s for

它专为需要高性能环境来测试智能体（agent）的 RL 从业者与研究人员设计，特别是那些专注于组合问题、逻辑游戏与路径规划问题的研究者。

Highlights

• Diverse Environment Suite: 包括跨越逻辑 (例如 Sudoku, Rubik's Cube)、装箱 (例如 BinPack, Tetris)、路径规划 (例如 TSP, CVRP) 以及集群 (例如 Search and Rescue) 等类别的 22 个环境。
• Hardware Acceleration: 完全基于 JAX，能够在 GPU/TPU 上实现大规模并行化与高速执行。
• Framework Agnostic: 为流行的 RL 库如 Acme, Stable Baselines3, RLlib, 和 Gymnasium 提供封装器。
• Scalable Difficulty: 环境设计旨在使其难度可以被扩展至任意困难的程度。