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

Abstract

We present XLand-Minigrid, a suite of tools and grid-world environments for meta-reinforcement learning research inspired by the diversity and depth of XLand and the simplicity and minimalism of MiniGrid. XLand-Minigrid is written in JAX, designed to be highly scalable, and can potentially run on GPU or TPU accelerators, democratizing large-scale experimentation with limited resources. To demonstrate the generality of our library, we have implemented some well-known single-task environments as well as new meta-learning environments capable of generating $10^8$ distinct tasks. We have empirically shown that the proposed environments can scale up to $2^{13}$ parallel instances on the GPU, reaching tens of millions of steps per second.

Cite

Text

Nikulin et al. "XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX." NeurIPS 2023 Workshops: IMOL, 2023.

Markdown

[Nikulin et al. "XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX." NeurIPS 2023 Workshops: IMOL, 2023.](https://mlanthology.org/neuripsw/2023/nikulin2023neuripsw-xlandminigrid/)

BibTeX

@inproceedings{nikulin2023neuripsw-xlandminigrid,
  title     = {{XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX}},
  author    = {Nikulin, Alexander and Kurenkov, Vladislav and Zisman, Ilya and Sinii, Viacheslav and Agarkov, Artem and Kolesnikov, Sergey},
  booktitle = {NeurIPS 2023 Workshops: IMOL},
  year      = {2023},
  url       = {https://mlanthology.org/neuripsw/2023/nikulin2023neuripsw-xlandminigrid/}
}