Package: jumanji
Environments|Installation|Quickstart|Training|Citation|DocsWelcome to the Jungle! 🌴Jumanji is a diverse suite of scalable reinforcement learning environments written in JAX.Jumanji is helping pioneer a new wave of hardware-accelerated research and development in the
field of RL. Jumanji's high-speed environments enable faster iteration and large-scale
experimentation while simultaneously reducing complexity. Originating in the Research Team atInstaDeep, Jumanji is now developed jointly with the open-source
community. To join us in these efforts, reach out, raise issues and read ourcontribution guidelinesor juststar🌟 to stay up to date with the latest developments!Goals 🚀Provide a simple, well-tested API for JAX-based environments.Make research in RL more accessible.Facilitate the research on RL for problems in the industry and help close the gap between
research and industrial applications.Provide environments whose difficulty can be scaled to be arbitrarily hard.Overview 🦜🥑Environment API: core abstractions for JAX-based environments.🕹️Environment Suite: a collection of RL environments ranging from simple games to NP-hard
combinatorial problems.🍬Wrappers: easily connect to your favourite RL frameworks and libraries such asAcme,Stable Baselines3,RLlib,OpenAI GymandDeepMind-Envthrough ourdm_envandgymwrappers.🎓Examples: guides to facilitate Jumanji's adoption and highlight the added value of
JAX-based environments.🏎️Training:example agents that can be used as inspiration for the agents one may implement
in their research.Environments 🌍Jumanji provides a diverse range of environments ranging from simple games to NP-hard combinatorial
problems.EnvironmentCategoryRegistered Version(s)SourceDescription🔢 Game2048LogicGame2048-v1codedoc🎨 GraphColoringLogicGraphColoring-v0codedoc💣 MinesweeperLogicMinesweeper-v0codedoc🎲 RubiksCubeLogicRubiksCube-v0RubiksCube-partly-scrambled-v0codedoc✏️ SudokuLogicSudoku-v0Sudoku-very-easy-v0codedoc📦 BinPack (3D BinPacking Problem)PackingBinPack-v2codedoc🏭 JobShop (Job Shop Scheduling Problem)PackingJobShop-v0codedoc🎒 KnapsackPackingKnapsack-v1codedoc▒ TetrisPackingTetris-v0codedoc🧹 CleanerRoutingCleaner-v0codedoc:link: ConnectorRoutingConnector-v2codedoc🚚 CVRP (Capacitated Vehicle Routing Problem)RoutingCVRP-v1codedoc🚚 MultiCVRP (Multi-Agent Capacitated Vehicle Routing Problem)RoutingMultiCVRP-v0codedoc:mag: MazeRoutingMaze-v0codedoc:robot: RobotWarehouseRoutingRobotWarehouse-v0codedoc🐍 SnakeRoutingSnake-v1codedoc📬 TSP (Travelling Salesman Problem)RoutingTSP-v1codedocMulti Minimum Spanning Tree ProblemRoutingMMST-v0codedocInstallation 🎬You can install the latest release of Jumanji from PyPI:pipinstalljumanjiAlternatively, you can install the latest development version directly from GitHub:pipinstallgit+https://github.com/instadeepai/jumanji.gitJumanji has been tested on Python 3.8 and 3.9.
Note that because the installation of JAX differs depending on your hardware accelerator,
we advise users to explicitly install the correct JAX version (see theofficial installation guide).Rendering:Matplotlib is used for rendering all the environments. To visualize the environments
you will need a GUI backend. For example, on Linux, you can install Tk via:apt-get install python3-tk, or using conda:conda install tk. Check outMatplotlib backendsfor a list of
backends you can use.Quickstart ⚡RL practitioners will find Jumanji's interface familiar as it combines the widely adoptedOpenAI GymandDeepMind Environmentinterfaces. From OpenAI Gym, we adopted
the idea of aregistryand therendermethod, while ourTimeStepstructure is inspired by
DeepMind Environment.Basic Usage 🧑‍💻importjaximportjumanji# Instantiate a Jumanji environment using the registryenv=jumanji.make('Snake-v1')# Reset your (jit-able) environmentkey=jax.random.PRNGKey(0)state,timestep=jax.jit(env.reset)(key)# (Optional) Render the env stateenv.render(state)# Interact with the (jit-able) environmentaction=env.action_spec().generate_value()# Action selection (dummy value here)state,timestep=jax.jit(env.step)(state,action)# Take a step and observe the next state and time stepstaterepresents the internal state of the environment: it contains all the information required
to take a step when executing an action. This shouldnotbe confused with theobservationcontained in thetimestep, which is the information perceived by the agent.timestepis a dataclass containingstep_type,reward,discount,observationandextras. This structure is similar todm_env.TimeStepexcept for theextrasfield that was added to allow users to log environments metrics that are neither part of
the agent's observation nor part of the environment's internal state.Advanced Usage 🧑‍🔬Being written in JAX, Jumanji's environments benefit from many of its features including
automatic vectorization/parallelization (jax.vmap,jax.pmap) and JIT-compilation (jax.jit),
which can be composed arbitrarily.
We provide an example of a more advanced usage in theadvanced usage guide.Registry and Versioning 📖Like OpenAI Gym, Jumanji keeps a strict versioning of its environments for reproducibility reasons.
We maintain a registry of standard environments with their configuration.
For each environment, a version suffix is appended, e.g.Snake-v1.
When changes are made to environments that might impact learning results,
the version number is incremented by one to prevent potential confusion.
For a full list of registered versions of each environment, check outthe documentation.Training 🏎️To showcase how to train RL agents on Jumanji environments, we provide a random agent and a vanilla
actor-critic (A2C) agent. These agents can be found injumanji/training/.Because the environment framework in Jumanji is so flexible, it allows pretty much any problem to
be implemented as a Jumanji environment, giving rise to very diverse observations. For this reason,
environment-specific networks are required to capture the symmetries of each environment.
Alongside the A2C agent implementation, we provide examples of such environment-specific
actor-critic networks injumanji/training/networks.⚠️ The example agents injumanji/trainingareonlymeant to serve as inspiration for how one
can implement an agent. Jumanji is first and foremost a library of environments - as such, the
agents and networks willnotbe maintained to a production standard.For more information on how to use the example agents, see thetraining guide.Contributing 🤝Contributions are welcome! See our issue tracker forgood first issues. Please read
ourcontributing guidelinesfor
details on how to submit pull requests, our Contributor License Agreement, and community guidelines.Citing Jumanji ✏️If you use Jumanji in your work, please cite the library using:@misc{bonnet2023jumanji,
    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 Vincent Coyette and
        Paul Duckworth and Laurence I. Midgley and Tristan Kalloniatis and Sasha Abramowitz and
        Cemlyn N. Waters and Andries P. Smit and Nathan Grinsztajn and Ulrich A. Mbou Sob and
        Omayma Mahjoub and Elshadai Tegegn and Mohamed A. Mimouni and Raphael Boige and
        Ruan de Kock and Daniel Furelos-Blanco and Victor Le and Arnu Pretorius and
        Alexandre Laterre
    },
    year={2023},
    eprint={2306.09884},
    url={https://arxiv.org/abs/2306.09884},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}See Also 🔎Other works have embraced the approach of writing RL environments in JAX.
In particular, we suggest users check out the following sister repositories:🤖Qdaxis a library to accelerate
Quality-Diversity and neuro-evolution algorithms through hardware accelerators and parallelization.🌳Evojaxprovides tools to enable neuroevolution algorithms
to work with neural networks running across multiple TPU/GPUs.🦾Braxis a differentiable physics engine that simulates
environments made up of rigid bodies, joints, and actuators.🏋️‍Gymnaximplements classic environments including
classic control, bsuite, MinAtar and a collection of meta RL tasks.🎲Pgxprovides classic board game environments like
Backgammon, Shogi, and Go.Acknowledgements 🙏The development of this library was supported with Cloud TPUs
from Google'sTPU Research Cloud(TRC) 🌤.