JaxUED: A simple and useable UED library in Jax
This provides a tool for researchers in reinforcement learning to accelerate UED experiments, but it is incremental as it focuses on implementation efficiency rather than new algorithmic insights.
The authors introduced JaxUED, an open-source library implementing Unsupervised Environment Design algorithms in Jax, achieving about 100x speedups over prior CPU-based methods.
We present JaxUED, an open-source library providing minimal dependency implementations of modern Unsupervised Environment Design (UED) algorithms in Jax. JaxUED leverages hardware acceleration to obtain on the order of 100x speedups compared to prior, CPU-based implementations. Inspired by CleanRL, we provide fast, clear, understandable, and easily modifiable implementations, with the aim of accelerating research into UED. This paper describes our library and contains baseline results. Code can be found at https://github.com/DramaCow/jaxued.