NAVIX: Scaling MiniGrid Environments with JAX

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 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.