minimax: Efficient Baselines for Autocurricula in JAX
This work addresses a major obstacle to rapid innovation in the RL community by providing efficient baselines for autocurricula, though it is incremental as it focuses on implementation improvements rather than new algorithmic breakthroughs.
The paper tackles the high computational cost of unsupervised environment design (UED) in reinforcement learning by introducing the minimax library, which uses JAX for hardware acceleration and achieves over 120× speedups in training time compared to previous implementations.
Unsupervised environment design (UED) is a form of automatic curriculum learning for training robust decision-making agents to zero-shot transfer into unseen environments. Such autocurricula have received much interest from the RL community. However, UED experiments, based on CPU rollouts and GPU model updates, have often required several weeks of training. This compute requirement is a major obstacle to rapid innovation for the field. This work introduces the minimax library for UED training on accelerated hardware. Using JAX to implement fully-tensorized environments and autocurriculum algorithms, minimax allows the entire training loop to be compiled for hardware acceleration. To provide a petri dish for rapid experimentation, minimax includes a tensorized grid-world based on MiniGrid, in addition to reusable abstractions for conducting autocurricula in procedurally-generated environments. With these components, minimax provides strong UED baselines, including new parallelized variants, which achieve over 120$\times$ speedups in wall time compared to previous implementations when training with equal batch sizes. The minimax library is available under the Apache 2.0 license at https://github.com/facebookresearch/minimax.