English

JaxUED: A simple and useable UED library in Jax

Machine Learning 2024-03-21 v1 Artificial Intelligence

Abstract

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.

Keywords

Cite

@article{arxiv.2403.13091,
  title  = {JaxUED: A simple and useable UED library in Jax},
  author = {Samuel Coward and Michael Beukman and Jakob Foerster},
  journal= {arXiv preprint arXiv:2403.13091},
  year   = {2024}
}

Comments

11 pages, 5 figures

R2 v1 2026-06-28T15:26:24.539Z