JAXenstein: Accelerated Benchmarking for First-Person Environments
For reinforcement learning researchers, it addresses the bottleneck of lacking fast, scalable first-person benchmarks in the JAX ecosystem.
JAXenstein introduces a JAX-based benchmark implementing the Wolfenstein 3D rendering engine for visual first-person tasks, achieving several times faster speed than comparable benchmarks.
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.