Package: xminigrid
XLand-MiniGridMeta-Reinforcement Learning in JAX🥳XLand-MiniGrid wasacceptedtoIntrinsically Motivated Open-ended Learningworkshop at NeurIPS 2023.We look forward to seeing everyone at the poster session!XLand-MiniGridis a suite of tools, grid-world environments and benchmarks for meta-reinforcement learning research inspired by
the diversity and depth ofXLandand the simplicity and minimalism ofMiniGrid. Despite the similarities,
XLand-MiniGrid is written in JAX from scratch and designed to be highly scalable, democratizing large-scale experimentation
with limited resources. Ever wanted to reproduce aDeepMind AdAagent? Now you can and not in months, but days!Features🔮 System of rules and goals that can be combined in arbitrary ways to produce
diverse task distributions🔧 Simple to extend and modify, comes with example environments ported from the originalMiniGrid🪄 Fully compatible with all JAX transformations, can run on CPU, GPU and TPU📈 Easily scales to $2^{16}$ parallel environments and millions of steps per second on a single GPU🔥 Multi-GPU PPO baselines in thePureJaxRLstyle, which can achieve1 trillionenvironment steps under two daysHow cool is that? For more details, take a look at thetechnical paperorexamples, which will walk you through the basics and training your own adaptive agents in minutes!Installation 🎁⚠️ XLand-MiniGrid is currently in alpha stage, so expect breaking changes! ⚠️The latest release of XLand-MiniGrid can be installed directly from PyPI:pip install xminigrid
# or, from github directly
pip install "xminigrid @ git+https://github.com/corl-team/xland-minigrid.git"Alternatively, if you want to install the latest development version from the GitHub and run provided algorithms or scripts,
install the source as follows:git clone git@github.com:corl-team/xland-minigrid.git
cd xland-minigrid
# additional dependencies for baselines
pip install -e ".[dev,benchmark]"Note that the installation of JAX may differ depending on your hardware accelerator!
We advise users to explicitly install the correct JAX version (see theofficial installation guide).Basic Usage 🕹️Most users who are familiar with other popular JAX-based environments
(such asgymnaxorjumnaji),
will find that the interface is very similar.
On the high level, current API combinesdm_envand gymnax interfaces.importjaximportxminigridfromxminigrid.wrappersimportGymAutoResetWrapperkey=jax.random.PRNGKey(0)reset_key,ruleset_key=jax.random.split(key)# to list available benchmarks: xminigrid.registered_benchmarks()benchmark=xminigrid.load_benchmark(name="trivial-1m")# choosing ruleset, see section on rules and goalsruleset=benchmark.sample_ruleset(ruleset_key)# to list available environments: xminigrid.registered_environments()env,env_params=xminigrid.make("XLand-MiniGrid-R9-25x25")env_params=env_params.replace(ruleset=ruleset)# auto-reset wrapperenv=GymAutoResetWrapper(env)# fully jit-compatible step and reset methodstimestep=jax.jit(env.reset)(env_params,reset_key)timestep=jax.jit(env.step)(env_params,timestep,action=0)# optionally render the stateenv.render(env_params,timestep)Similar to gymnasium or jumanji, users can register new environment
variations withregisterfor convenient further usage withmake.timestepis a dataclass containingstep_type,reward,discount,observation, as well as the internal environmentstate.For a bit more advanced introduction see providedwalkthrough notebook.On environment interfaceCurrently, there are a lot of new JAX-based environments appearing, each offering its own variant of API. Initially, we tried to reuse Jumaji, but it turned out
that its designis not suitable for meta learning. The Gymnax design appeared to be more appropriate, but unfortunately it is not actively supported and
often departs from the idea that parameters should only be contained inenv_params. Furthermore, splittingtimestepinto multiple entities seems suboptimal to us, as it complicates many things, such as envpool or dm_env
style auto reset, where the reset occurs on the next step (we need access to done of previous step).Therefore, we decided that we would make a minimal interface that would cover just our needs without the
goal of making it generic. The core of our library is interface independent, and we plan
to switch to the new one when/if a better design becomes available
(e.g. when stable GymnasiumFuncEnvis released).Rules and Goals 🔮In XLand-MiniGrid, the system of rules and goals is the cornerstone of the
emergent complexity and diversity. In the original MiniGrid
some environments have dynamic goals, but the dynamics are never changed.
To train and evaluate highly adaptive agents, we need to be able to change
the dynamics in non-trivial ways.Rulesare the functions that can change the environment state in some deterministic
way according to the given conditions.Goalsare similar to rules, except they do
not change the state, they only test conditions. Every task should be described with a goal, rules and initial objects. We call theserulesets.
Currently, we support only one goal per task.To illustrate, we provide visualization for specific ruleset. To solve this task, agent should take blue pyramid and put it near the purple square to transform both
objects into red circle. To complete the goal, red circle should be placed near
green circle. However, placing purple square near yellow circle will make it unsolvable in this trial. Initial objects positions will be randomized on each reset.For more advanced introduction, see corresponding section in the providedwalkthrough notebook.Benchmarks 🎲While composing rules and goals by hand is flexible, it can quickly become cumbersome.
Besides, it's hard to express efficiently in a JAX-compatible way due to the high number of heterogeneous computationsTo avoid significant overhead during training and facilitate reliable comparisons between agents,
we pre-sampled several benchmarks with up tothree million unique tasks, following the procedure used to train DeepMind
AdA agent from the original XLand. These benchmarks differ in the generation configs, producing distributions with
varying levels of diversity and average difficulty of the tasks. They can be used for different purposes, for example
thetrivial-1mbenchmark can be used to debug your agents, allowing very quick iterations. However, we would caution
against treating benchmarks as a progression from simple to complex. They are just different 🤷.Pre-sampled benchmarks are hosted onHuggingFaceand will be downloaded and cached on the first use:importjax.randomimportxminigridfromxminigrid.benchmarksimportBenchmark# downloading to path specified by XLAND_MINIGRID_DATA,# ~/.xland_minigrid by defaultbenchmark:Benchmark=xminigrid.load_benchmark(name="trivial-1m")# reusing cached on the second usebenchmark:Benchmark=xminigrid.load_benchmark(name="trivial-1m")# users can sample or get specific rulesetsbenchmark.sample_ruleset(jax.random.PRNGKey(0))benchmark.get_ruleset(ruleset_id=benchmark.num_rulesets()-1)# or split them for train & testtrain,test=benchmark.shuffle(key=jax.random.PRNGKey(0)).split(prop=0.8)We also provide thescriptused to generate these benchmarks. Users can use it for their own purposes:python scripts/ruleset_generator.py --helpIn depth description of all available benchmarks is providedin the technical paper(Section 3).P.S.Be aware, that benchmarks can change, as we are currently testing and balancing them!Environments 🌍We provide environments from two domains.XLandis our main focus for meta-learning. For this domain we provide single
environment and numerous registered variants with different grid layouts and sizes. All of them can be combined
with arbitrary rulesets.To demonstrate the generality of our library we also port majority of
non-language based tasks from originalMiniGrid. Similarly, some environments come with multiple registered variants.
However, we have no current plans to actively develop and support them (but that may change).NameDomainVisualizationGoalXLand-MiniGridXLandspecified by the provided rulesetMiniGrid-EmptyMiniGridgo to the green goalMiniGrid-EmptyRandomMiniGridgo the green goal from different starting positionsMiniGrid-FourRoomsMiniGridgo the green goal, but goal and starting positions are randomizedMiniGrid-LockedRoomMiniGridfind the key to unlock the door, go to the green goalMiniGrid-MemoryMiniGridremember the initial object and choose it at the end of the corridorMiniGrid-PlaygroundMiniGridgoal is not specifiedMiniGrid-UnlockMiniGridunlock the door with the keyMiniGrid-UnlockPickUpMiniGridunlock the door and pick up the object in another roomMiniGrid-BlockedUnlockPickUpMiniGridunlock the door blocked by the object and pick up the object in another roomMiniGrid-DoorKeyMiniGridunlock the door and go to the green goalUsers can get all registered environments withxminigrid.registered_environments(). We also provide manual control to easily explore the environments:python -m xminigrid.manual_control --env-id="MiniGrid-Empty-8x8"Baselines 🚀In addition to the environments, we provide high-qualityalmostsingle-file
implementations of recurrent PPO baselines in the style ofPureJaxRL. With the help of magicaljax.pmaptransformation
they can scale to multiple accelerators, achieving impressive FPS of millions during training.Agents can be trained from the terminal and default arguments can be overwritten from the command line or from the yaml config:# for meta learning
python training/train_meta_task.py \
    --config-path='some-path/config.yaml' \
    --env_id='XLand-MiniGrid-R1-9x9'

# for minigrid envs
python training/train_singe_task.py \
    --config-path='some-path/config.yaml' \ 
    --env_id='XLand-MiniGrid-R1-9x9'For the source code and hyperparameters available see/trainingor runpython training/train_meta_task.py --help.
Furthermore, we provide standalone implementations that can be trained in Colab:xland,minigrid.P.S.Do not expect that provided baselines will solve the hardest environments or benchmarks
available. How much fun would that be 🤔? However, we hope that they will
help to get started quickly!Contributing 🔨We welcome anyone interested in helping out! Please take a look at ourcontribution guidefor further instructions and open an issue if something is not clear.See Also 🔎A lot of other work is going in a similar direction, transforming RL through JAX. Many of them have inspired us,
and we encourage users to check them out as well.Brax- fully differentiable physics engine used for research and development of robotics.Gymnax- implements classic environments including classic control, bsuite, MinAtar and simplistic meta learning tasks.Jumanji- a diverse set of environments ranging from simple games to NP-hard combinatorial problems.Pgx- JAX implementations of classic board games, such as Chess, Go and Shogi.JaxMARL- multi-agent RL in JAX with wide range of commonly used environments.Let's build together!Citation 🙏@inproceedings{nikulin2023xlandminigrid,title={{XL}and-MiniGrid: Scalable Meta-Reinforcement Learning Environments in {JAX}},author={Alexander Nikulin and Vladislav Kurenkov and Ilya Zisman and Viacheslav Sinii and Artem Agarkov and Sergey Kolesnikov},booktitle={Intrinsically-Motivated and Open-Ended Learning Workshop, NeurIPS2023},year={2023},url={https://openreview.net/forum?id=xALDC4aHGz}}