Research map / Reinforcement learning / Visual and representation RL
XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in 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.
Presents XLand-MiniGrid, a JAX-based grid-world meta-RL environment suite with procedurally generated rule/goal benchmarks (up to 3M rulesets) and RL² recurrent-PPO baselines that reach millions of steps per second.
Problem
Meta-RL research lacks environments that are both non-trivial in task diversity and computationally affordable, since XLand is closed-source and simpler benchmarks dominate.
Approach
A JAX-native, MiniGrid-style environment with an array-encoded system of production rules and goals (inspired by XLand) that can be combined into procedurally generated task distributions and vectorized/jitted across GPUs/TPUs.
Key results
- Simulation throughput reaches tens of millions of steps per second on multiple A100 devices with random policy; single-device training saturates near 1M steps per second and reaches almost 6M with multiple devices.
- RL² baseline performs far from optimal on high-1m even after one trillion transitions, especially at the 20th score percentile.
- On held-out goal-type generalization, there remains a large gap on all benchmarks, including the minimal-diversity one.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use bf16 precision during training (ablation) | Recurrent PPO training on H100 · PPO training throughput | steps per second: 1.2M -> 1.5M | Helped |
| Use bf16 precision during training (ablation) | Recurrent PPO training on RTX 4090 · PPO training throughput | steps per second: around 800k -> 2.1M | Helped |
| Use rendered image observations instead of symbolic observations (ablation) | RL² recurrent PPO · XLand-MiniGrid meta-RL training | steps per second: about 20k | Hurt |
| Use rendered image observations instead of symbolic observations (ablation) | RL² recurrent PPO · XLand-MiniGrid meta-RL benchmarks | generalization | Hurt |
Surprising findings
- Performance on high-1m remains sub-optimal even after one trillion training transitions.
- Image observations make benchmarks harder and reduce training throughput to about 20k steps per second on a single GPU.
Limitations
- No multi-agent support, procedural grid generation, multi-output rules, or goal composition compared with full XLand.
- Grid layouts are fixed in advance and cannot change under jit; new rules and goals require modifying source code.
- Only the simple recurrent PPO/RL² baseline is provided, with GRU rather than modern architectures.
Experimental setup
Datasets: trivial-1m small-1m medium-1m medium-3m high-1m high-3m
Benchmarks: XLand-MiniGrid meta-RL benchmarks (trivial, small, medium, high)
Metrics: return 20th score percentile simulation throughput (steps per second) training throughput (steps per second)
Baselines: RL² with recurrent PPO Ada Jumanji EnvPool
Builds on: XLand MiniGrid PureJaxRL Anakin RL²
Closest papers by meaning
- NAVIX: Scaling MiniGrid Environments with JAX (NeurIPS 2025 · similarity 0.89)
- XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.88)
- JaxMARL: Multi-Agent RL Environments and Algorithms in JAX (NeurIPS 2024 · similarity 0.88)
- MetaBox-v2: A Unified Benchmark Platform for Meta-Black-Box Optimization (NeurIPS 2025 · similarity 0.87)
- Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX (ICLR 2024 · similarity 0.86)
- Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning (ICML 2024 · similarity 0.86)
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.