# NAVIX: Scaling MiniGrid Environments with JAX

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

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Placed in RL training of language and agent models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 35 of its 48 references that are in the map.
- First posted on arXiv 28 Jul 2024.

## Placed in

- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

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

- [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)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [The NetHack Learning Environment](https://mlsanity.com/paper/2006.13760/) (2020)
- [Prioritized Level Replay](https://mlsanity.com/paper/2010.03934/) (2020)
- [Behaviour Suite for Reinforcement Learning](https://mlsanity.com/paper/1908.03568/) (2019)
- [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)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Simplifying Deep Temporal Difference Learning](https://mlsanity.com/paper/2407.04811/) (2024)
- [EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine](https://mlsanity.com/paper/2206.10558/) (2022)
- [POPGym: Benchmarking Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2303.01859/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [Improving Intrinsic Exploration with Language Abstractions](https://mlsanity.com/paper/2202.08938/) (2022)
- [Adversarially Guided Actor-Critic](https://mlsanity.com/paper/2102.04376/) (2021)
- [History Compression via Language Models in Reinforcement Learning](https://mlsanity.com/paper/2205.12258/) (2022)
- [Podracer architectures for scalable Reinforcement Learning](https://mlsanity.com/paper/2104.06272/) (2021)
- [Interesting Object, Curious Agent: Learning Task-Agnostic Exploration](https://mlsanity.com/paper/2111.13119/) (2021)
- [Alchemy: A benchmark and analysis toolkit for meta-reinforcement learning agents](https://mlsanity.com/paper/2102.02926/) (2021)
- [BeBold: Exploration Beyond the Boundary of Explored Regions](https://mlsanity.com/paper/2012.08621/) (2020)
- [Rank the Episodes: A Simple Approach for Exploration in Procedurally-Generated Environments](https://mlsanity.com/paper/2101.08152/) (2021)

## Built on by (most cited first, 18 of 18)

- [Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks](https://mlsanity.com/paper/2410.23208/) (2024)
- [Benchmarking Partial Observability in Reinforcement Learning with a Suite of Memory-Improvable Domains](https://mlsanity.com/paper/2508.00046/) (2025)
- [In-Context Reinforcement Learning via Communicative World Models](https://mlsanity.com/paper/2508.06659/) (2025)
- [Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX](https://mlsanity.com/paper/2510.01764/) (2025)
- [Investigating Memory in Model-Free RL with POPGym Arcade](https://mlsanity.com/paper/2503.01450/) (2025)
- [From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learning](https://mlsanity.com/paper/2510.02484/) (2025)
- [Multi-Agent Craftax: Benchmarking Open-Ended Multi-Agent Reinforcement Learning at the Hyperscale](https://mlsanity.com/paper/2511.04904/) (2025)
- [EconoJax: A Fast & Scalable Economic Simulation in Jax](https://mlsanity.com/paper/2410.22165/) (2024)
- [Inducing, Detecting and Characterising Neural Modules: A Pipeline for Functional Interpretability in Reinforcement Learning](https://mlsanity.com/paper/2501.17077/) (2025)
- [Stackelberg Coupling of Online Representation Learning and Reinforcement Learning](https://mlsanity.com/paper/2508.07452/) (2025)
- [Chargax: A JAX Accelerated EV Charging Simulator](https://mlsanity.com/paper/2507.01522/) (2025)
- [Assistax: A Multi-Agent Hardware-Accelerated Reinforcement Learning Benchmark for Assistive Robotics](https://mlsanity.com/paper/2507.21638/) (2025)
- [CoGrid & the Multi-User Gymnasium: A Framework for Multi-Agent Experimentation](https://mlsanity.com/paper/2604.15044/) (2026)
- [Self-Supervised On-Policy Reinforcement Learning via Contrastive Proximal Policy Optimisation](https://mlsanity.com/paper/2605.13554/) (2026)
- [Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX](https://mlsanity.com/paper/2605.20577/) (2026)
- [Jaxolotl: A Unified High-Performance Benchmark Suite for LTL-Based Multi-Task RL](https://mlsanity.com/paper/2609.38065/) (2026)
- [Scalable Multi-Task Inverse Reinforcement Learning](https://mlsanity.com/paper/2610.00758/) (2026)
- [Adaptive Expert Guidance for Efficient On-Policy Reinforcement Learning](https://mlsanity.com/paper/2610.06019/) (2026)

## Abstract (arXiv)

As Deep Reinforcement Learning (Deep RL) research moves towards solving large-scale worlds, efficient environment simulations become crucial for rapid experimentation. However, most existing environments struggle to scale to high throughput, setting back meaningful progress. Interactions are typically computed on the CPU, limiting training speed and throughput, due to slower computation and communication overhead when distributing the task across multiple machines. Ultimately, Deep RL training is CPU-bound, and developing batched, fast, and scalable environments has become a frontier for progress. Among the most used Reinforcement Learning (RL) environments, MiniGrid is at the foundation of several studies on exploration, curriculum learning, representation learning, diversity, meta-learning, credit assignment, and language-conditioned RL, and still suffers from the limitations described above. In this work, we introduce NAVIX, a re-implementation of MiniGrid in JAX. NAVIX achieves over 200 000x speed improvements in batch mode, supporting up to 2048 agents in parallel on a single Nvidia A100 80 GB. This reduces experiment times from one week to 15 minutes, promoting faster design iterations and more scalable RL model development.