Research map / Reinforcement learning / Visual and representation RL
NAVIX: Scaling MiniGrid Environments with JAX
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Re-implements the MiniGrid suite in JAX as NAVIX, a fully jittable and batched environment, achieving large speedups over the original CPU MiniGrid and providing PPO/DDQN/SAC baselines.
Problem
Existing RL environments such as MiniGrid are CPU-bound and do not scale to high throughput, which limits the speed of deep RL experimentation.
Approach
Reimplement MiniGrid entirely in JAX using an entity-component-system design so that every environment is stateful, jittable, vmap-batchable and runs on GPU/TPU, while matching MiniGrid's environments, observations, transitions and rewards (with a Markovian reward replacing MiniGrid's step-dependent reward).
Key results
- NAVIX is over 45x faster than MiniGrid on average in the 1K-step single-batch benchmark with 8 parallel environments.
- Training 2048 PPO agents (32,768 environments) in NAVIX takes under 50 s, versus about 240 s for a single PPO agent in MiniGrid, reaching roughly 670 million steps/s and a speedup of over 200,000x.
- The original MiniGrid cannot scale beyond 16 environments on 128 GB RAM, whereas NAVIX runs up to 2^21 environments in parallel on the same hardware.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Reimplement MiniGrid with JAX and batched environment execution | MiniGrid environments · MiniGrid suite | environment speedup: over $45\times$ faster | Helped |
| Add batched environment execution (ablation) | NAVIX · MiniGrid environments | speedup | Helped |
| Run 2048 PPO agents in parallel, each using 16 environments | NAVIX · Empty-8x8-v0 | training time: less than $50$s versus around $240$s | Helped |
| Use JAX vmap to run environments in parallel | NAVIX · Empty-8x8-v0 | maximum parallel environments: $16$ versus $2^{21}$ (over $2M$) | Helped |
Surprising findings
- The original MiniGrid uses a non-Markovian step-dependent reward, which the authors argue may call into question historical MiniGrid results; NAVIX uses a Markovian reward instead.
Limitations
- MiniGrid reward is non-Markovian and NAVIX departs from it, so results are not identical to original MiniGrid rewards
- Speed comparisons are on a single A100 and Xeon CPU setup and mostly on Empty-8x8 environment
- Baselines are limited to PPO, DDQN and SAC with random-search hyperparameter tuning
Experimental setup
Datasets: MiniGrid environments (NAVIX reimplementation)
Benchmarks: MiniGrid environments (NAVIX suite)
Metrics: wall-clock speedup throughput (steps per second) episodic return
Baselines: MiniGrid PPO Double DQN SAC AMaze XLand-MiniGrid
Builds on: MiniGrid JAX Rejax
Closest papers by meaning
- XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX (NeurIPS 2024 · similarity 0.89)
- JaxMARL: Multi-Agent RL Environments and Algorithms in JAX (NeurIPS 2024 · similarity 0.86)
- Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX (ICLR 2024 · similarity 0.86)
- TGRL: Temperature-Grouped Reinforcement Learning for Efficient Exploration in LLMs (NeurIPS 2026 · similarity 0.85)
- SRL: Scaling Distributed Reinforcement Learning to Over Ten Thousand Cores (ICLR 2024 · similarity 0.85)
- Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning (ICML 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.