Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

JaxUED: A simple and useable UED library in Jax

Samuel Coward, Michael Beukman, Jakob Foerster

TL;DR

JaxUED introduces a fast, minimal-dependency JAX-based library for Unsupervised Environment Design (UED), offering nearly single-file reference implementations of core UED algorithms to accelerate research. It formalizes the UnderspecifiedEnv interface, provides a LevelSampler, and includes reference implementations for Domain Randomisation (DR), Replay-based methods (PLR, ACCEL), and PAIRED, all paired with a fully in-JAX Maze environment for benchmarking. Empirical results show that Domain Randomisation is unexpectedly competitive with state-of-the-art UED methods on maze tasks, and JaxUED delivers substantial speedups over CPU baselines. The work aims to enable rapid, reproducible UED experimentation and to spur the development of more general and robust RL agents by lowering the barrier to implementing and comparing UED algorithms.

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.

JaxUED: A simple and useable UED library in Jax

TL;DR

JaxUED introduces a fast, minimal-dependency JAX-based library for Unsupervised Environment Design (UED), offering nearly single-file reference implementations of core UED algorithms to accelerate research. It formalizes the UnderspecifiedEnv interface, provides a LevelSampler, and includes reference implementations for Domain Randomisation (DR), Replay-based methods (PLR, ACCEL), and PAIRED, all paired with a fully in-JAX Maze environment for benchmarking. Empirical results show that Domain Randomisation is unexpectedly competitive with state-of-the-art UED methods on maze tasks, and JaxUED delivers substantial speedups over CPU baselines. The work aims to enable rapid, reproducible UED experimentation and to spur the development of more general and robust RL agents by lowering the barrier to implementing and comparing UED algorithms.

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.
Paper Structure (19 sections, 3 figures, 3 tables)

This paper contains 19 sections, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Who is this for?
  3. Unsupervised Environment Design
  4. The JaxUED Library
  5. Environment Interface
  6. Wrappers
  7. Level Sampler
  8. Maze Environment
  9. Reference Implementations
  10. Replay-Based Methods
  11. Domain Randomisation
  12. PAIRED
  13. Results
  14. Performance
  15. Runtime Speed
  16. ...and 4 more sections

Figures (3)

  • Figure 1: Illustration of the training process of replay-based methods. Training can be viewed as a deterministic Markov Decision Process (MDP) controlled by a fixed, stochastic meta-policy defined by hyperparameters. Transitions determine what kind of PPO update is performed.
  • Figure 2: Visualization of an example batch of holdout levels used for evaluation, generated using minimaxjiang2023Minimax. Such levels were used to evaluate the performance of minimax and DCD in jiang2023Minimax.
  • Figure 3: IQM of mean solve rate across over 100 trials of minimax evaluation levels, measured over 10 random seeds. Error bars correspond to min-max performance over the seeds. The number after the method name indicates the maximum number of walls in the base DR distribution, either 25 or 60, and in the case of PAIRED indicates the number of editor environment steps taken by the adversary.