JaxUED: A simple and useable UED library in Jax

Samuel Coward, Michael Beukman, Jakob Foerster
3/19/2024
cs.LGcs.AI

Official code: DramaCow/jaxued (102 stars, last commit 2026-01-21).

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.

View on arXivView PDF

Code Implementations(1)

DramaCow/jaxuedOfficial100%
10216PythonFeb 25, 20249 months ago

Similar Papers

minimax: Efficient Baselines for Autocurricula in JAX

78% match1 repo2172023

JaxDecompiler: Redefining Gradient-Informed Software Design

73% match1 repo162024

JaxARC: A High-Performance JAX-based Environment for Abstraction and Reasoning Research

71% match1 repo162026

JaxPruner: A concise library for sparsity research

71% match1 repo2362023

JAXbind: Bind any function to JAX

70% match1 repo822024

DrJAX: Scalable and Differentiable MapReduce Primitives in JAX

70% match1 repo39.0k2024

Cite this paper

@article{coward2024jaxued,
  title  = {JaxUED: A simple and useable UED library in Jax},
  author = {Samuel Coward and Michael Beukman and Jakob Foerster},
  year   = {2024},
  eprint = {2403.13091},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2403.13091v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2403.13091.svg)](https://codeofpaper.com/paper/2403.13091)

Discussion