XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX

TL;DR

XLand-MiniGrid introduces a scalable, JAX-based suite of grid-world environments for meta-reinforcement learning, enabling large-scale experimentation with diverse tasks and efficient training on accelerators.

Contribution

It provides a new, scalable, and easy-to-use platform with extensive benchmarks for meta-RL research, combining the diversity of XLand and simplicity of MiniGrid.

Findings

Baselines reach millions of steps per second.

Benchmarks are challenging for current algorithms.

The platform facilitates rapid experimentation.

Abstract

Inspired by the diversity and depth of XLand and the simplicity and minimalism of MiniGrid, we present XLand-MiniGrid, a suite of tools and grid-world environments for meta-reinforcement learning research. Written in JAX, XLand-MiniGrid is designed to be highly scalable and can potentially run on GPU or TPU accelerators, democratizing large-scale experimentation with limited resources. Along with the environments, XLand-MiniGrid provides pre-sampled benchmarks with millions of unique tasks of varying difficulty and easy-to-use baselines that allow users to quickly start training adaptive agents. In addition, we have conducted a preliminary analysis of scaling and generalization, showing that our baselines are capable of reaching millions of steps per second during training and validating that the proposed benchmarks are challenging. XLand-MiniGrid is open-source and available at https://github.com/dunnolab/xland-minigrid.