JaxUED: A simple and useable UED library in Jax

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,...

Full description

Saved in:
Bibliographic Details
Main Authors Coward, Samuel, Beukman, Michael, Foerster, Jakob
Format Journal Article
LanguageEnglish
Published 19.03.2024
Subjects
Computer Science - Artificial Intelligence
Computer Science - Learning
Online AccessGet full text

Cover

More Information
Summary: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.
DOI:10.48550/arxiv.2403.13091