# JaxUED: A simple and useable UED library in Jax

arXiv: https://arxiv.org/abs/2403.13091 · HTML: https://mlsanity.com/paper/2403.13091/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 20 references that are in the map.
- First posted on arXiv 19 Mar 2024.

## Builds on (most cited first, 13 of 13)

- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [Emergent Complexity and Zero-shot Transfer via Unsupervised Environment Design](https://mlsanity.com/paper/2012.02096/) (2020)
- [Prioritized Level Replay](https://mlsanity.com/paper/2010.03934/) (2020)
- [Evolving Curricula with Regret-Based Environment Design](https://mlsanity.com/paper/2203.01302/) (2022)
- [Human-Timescale Adaptation in an Open-Ended Task Space](https://mlsanity.com/paper/2301.07608/) (2023)
- [Replay-Guided Adversarial Environment Design](https://mlsanity.com/paper/2110.02439/) (2021)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [Stabilizing Unsupervised Environment Design with a Learned Adversary](https://mlsanity.com/paper/2308.10797/) (2023)
- [minimax: Efficient Baselines for Autocurricula in JAX](https://mlsanity.com/paper/2311.12716/) (2023)

## Built on by (most cited first, 11 of 11)

- [Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks](https://mlsanity.com/paper/2410.23208/) (2024)
- [No Regrets: Investigating and Improving Regret Approximations for Curriculum Discovery](https://mlsanity.com/paper/2408.15099/) (2024)
- [The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Design](https://mlsanity.com/paper/2406.17949/) (2024)
- [gfnx: Fast and Scalable Library for Generative Flow Networks in JAX](https://mlsanity.com/paper/2511.16592/) (2025)
- [Syllabus: Portable Curricula for Reinforcement Learning Agents](https://mlsanity.com/paper/2411.11318/) (2024)
- [Video Game Level Design as a Multi-Agent Reinforcement Learning Problem](https://mlsanity.com/paper/2510.04862/) (2025)
- [Improving Regret Approximation for Unsupervised Dynamic Environment Generation](https://mlsanity.com/paper/2601.14957/) (2026)
- [TABX: A High-Throughput Sandbox Battle Simulator for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2602.01665/) (2026)
- [NASimJax: A GPU-Accelerated Policy Learning Framework for Penetration Testing](https://mlsanity.com/paper/2603.19864/) (2026)
- [PACE: Parameter Change for Unsupervised Environment Design](https://mlsanity.com/paper/2605.01358/) (2026)
- [Jaxolotl: A Unified High-Performance Benchmark Suite for LTL-Based Multi-Task RL](https://mlsanity.com/paper/2609.38065/) (2026)

## Abstract (arXiv)

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.