# XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX

arXiv: https://arxiv.org/abs/2312.12044 · HTML: https://mlsanity.com/paper/2312.12044/

## Key facts

- Cited by 55 papers in the map (data through 14 Oct 2026).
- Builds on 54 of its 71 references that are in the map.
- First posted on arXiv 19 Dec 2023.

## Builds on (most cited first, 30 of 54)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Deep Reinforcement Learning that Matters](https://mlsanity.com/paper/1709.06560/) (2017)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [Distributed Prioritized Experience Replay](https://mlsanity.com/paper/1803.00933/) (2018)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Open-Ended Learning Leads to Generally Capable Agents](https://mlsanity.com/paper/2107.12808/) (2021)
- [What Matters In On-Policy Reinforcement Learning? A Large-Scale Empirical Study](https://mlsanity.com/paper/2006.05990/) (2020)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Human-Timescale Adaptation in an Open-Ended Task Space](https://mlsanity.com/paper/2301.07608/) (2023)
- [Revisiting the Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2305.09836/) (2023)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Sample Factory: Egocentric 3D Control from Pixels at 100000 FPS with Asynchronous Reinforcement Learning](https://mlsanity.com/paper/2006.11751/) (2020)
- [A Dissection of Overfitting and Generalization in Continuous Reinforcement Learning](https://mlsanity.com/paper/1806.07937/) (2018)

## Built on by (most cited first, 30 of 55)

- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks](https://mlsanity.com/paper/2410.23208/) (2024)
- [Meta-World+: An Improved, Standardized, RL Benchmark](https://mlsanity.com/paper/2505.11289/) (2025)
- [No Regrets: Investigating and Improving Regret Approximations for Curriculum Discovery](https://mlsanity.com/paper/2408.15099/) (2024)
- [In-Context Reinforcement Learning for Variable Action Spaces](https://mlsanity.com/paper/2312.13327/) (2023)
- [Learning to Reason at the Frontier of Learnability](https://mlsanity.com/paper/2502.12272/) (2025)
- [Retrieval-Augmented Decision Transformer: External Memory for In-context RL](https://mlsanity.com/paper/2410.07071/) (2024)
- [NAVIX: Scaling MiniGrid Environments with JAX](https://mlsanity.com/paper/2407.19396/) (2024)
- [AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers](https://mlsanity.com/paper/2411.11188/) (2024)
- [XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning](https://mlsanity.com/paper/2406.08973/) (2024)
- [A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks](https://mlsanity.com/paper/2410.22391/) (2024)
- [Cross-environment Cooperation Enables Zero-shot Multi-agent Coordination](https://mlsanity.com/paper/2504.12714/) (2025)
- [Sample-Efficient Online Learning in LM Agents via Hindsight Trajectory Rewriting](https://mlsanity.com/paper/2510.10304/) (2025)
- [The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Design](https://mlsanity.com/paper/2406.17949/) (2024)
- [Yes, Q-learning Helps Offline In-Context RL](https://mlsanity.com/paper/2502.17666/) (2025)
- [An Optimisation Framework for Unsupervised Environment Design](https://mlsanity.com/paper/2505.20659/) (2025)
- [Partially Observable Reinforcement Learning with Memory Traces](https://mlsanity.com/paper/2503.15200/) (2025)
- [SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmas](https://mlsanity.com/paper/2503.14576/) (2025)
- [Zero-Shot Adaptation of Behavioral Foundation Models to Unseen Dynamics](https://mlsanity.com/paper/2505.13150/) (2025)
- [Mixture-of-Experts Meets In-Context Reinforcement Learning](https://mlsanity.com/paper/2506.05426/) (2025)
- [N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs](https://mlsanity.com/paper/2411.01958/) (2024)
- [Benchmarking General-Purpose In-Context Learning](https://mlsanity.com/paper/2405.17234/) (2024)
- [Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX](https://mlsanity.com/paper/2510.01764/) (2025)
- [ARM-FM: Automated Reward Machines via Foundation Models for Compositional Reinforcement Learning](https://mlsanity.com/paper/2510.14176/) (2025)
- [Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds](https://mlsanity.com/paper/2502.02869/) (2025)
- [Deep RL Needs Deep Behavior Analysis: Exploring Implicit Planning by Model-Free Agents in Open-Ended Environments](https://mlsanity.com/paper/2506.06981/) (2025)
- [PuzzleJAX: A Benchmark for Reasoning and Learning](https://mlsanity.com/paper/2508.16821/) (2025)
- [ARLBench: Flexible and Efficient Benchmarking for Hyperparameter Optimization in Reinforcement Learning](https://mlsanity.com/paper/2409.18827/) (2024)
- [The Yokai Learning Environment: Tracking Beliefs Over Space and Time](https://mlsanity.com/paper/2508.12480/) (2025)
- [gfnx: Fast and Scalable Library for Generative Flow Networks in JAX](https://mlsanity.com/paper/2511.16592/) (2025)

## Abstract (arXiv)

Inspired by the diversity and depth of XLand and the simplicity and minimalism of MiniGrid, we present XLand-MiniGrid, a suite of tools and grid-world environments for meta-reinforcement learning research. Written in JAX, XLand-MiniGrid is designed to be highly scalable and can potentially run on GPU or TPU accelerators, democratizing large-scale experimentation with limited resources. Along with the environments, XLand-MiniGrid provides pre-sampled benchmarks with millions of unique tasks of varying difficulty and easy-to-use baselines that allow users to quickly start training adaptive agents. In addition, we have conducted a preliminary analysis of scaling and generalization, showing that our baselines are capable of reaching millions of steps per second during training and validating that the proposed benchmarks are challenging. XLand-MiniGrid is open-source and available at https://github.com/dunnolab/xland-minigrid.