JAXenstein: Accelerated Benchmarking for First-Person Environments

TL;DR

JAXenstein is an open-source, fast, and scalable JAX-based benchmark for visual first-person reinforcement learning tasks, addressing a key gap in existing benchmarking tools.

Contribution

It introduces JAXenstein, a novel benchmark using Wolfenstein 3D, enabling rapid experimentation in visual first-person environments with improved speed and extensibility.

Findings

JAXenstein is several times faster than comparable benchmarks.

It is easily extensible to more complex first-person domains.

Provides a scalable platform for visual RL research.

Abstract

The progression of reinforcement learning algorithms have been driven by challenging benchmarks. The rate in which a researcher can iterate on a problem setting directly impacts the speed of algorithm development. Modern machine learning has produced tools that allow for fast and scalable algorithm development like the JAX library. With the availability of these tools, a serious bottleneck in algorithm development is the availability of large and complex domains for experimentation. Most notably, the JAX reinforcement learning ecosystem does not have any benchmarks that test visual first-person tasks; these domains are crucial for testing both exploration and an agent's ability to overcome partial observability. We introduce JAXenstein: an open-source JAX-based benchmark that implements the Wolfenstein 3D rendering engine for fast and scalable experimentation in visual first-person tasks. JAXenstein is several times faster than comparable vision-based benchmarks, and is easily extensible to more complex first-person domains.