NAVIX: Scaling MiniGrid Environments with JAX
Eduardo Pignatelli, Jarek Liesen, Robert Tjarko Lange, Chris Lu, Pablo, Samuel Castro, Laura Toni

TL;DR
NAVIX is a JAX-based re-implementation of MiniGrid that dramatically accelerates environment simulation, enabling large-scale parallel RL experiments and reducing training times from weeks to minutes.
Contribution
We developed NAVIX, a scalable, high-performance MiniGrid environment in JAX, achieving over 200,000x speedup and supporting thousands of agents in parallel.
Findings
Achieves 200,000x speedup in batch mode
Supports up to 2048 agents on a single GPU
Reduces experiment time from one week to 15 minutes
Abstract
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…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Code & Models
Videos
Taxonomy
TopicsDistributed and Parallel Computing Systems · Peer-to-Peer Network Technologies · Multimedia Communication and Technology
MethodsSPEED: Separable Pyramidal Pooling EncodEr-Decoder for Real-Time Monocular Depth Estimation on Low-Resource Settings
