Technical Summary: JAXenstein

Problem Statement

The development of reinforcement learning (RL) algorithms is heavily dependent on the availability of challenging benchmarks. While modern tools like the JAX library have enabled fast, scalable, and GPU-accelerated algorithm development, a critical bottleneck remains: the lack of large, complex, visual first-person domains within the JAX ecosystem. Existing first-person benchmarks, such as those derived from ViZDoom or DeepMind Lab, are crucial for testing exploration and an agent's ability to handle partial observability, yet they have not been widely adopted in JAX due to the computational complexity of reimplementing their rendering engines. Consequently, researchers often face a trade-off between using lightweight, non-visual benchmarks (e.g., Cart Pole) or slower, computationally expensive visual environments, hindering the rapid iteration required for algorithmic advancement.

Methodology

To address this gap, the authors introduce JAXenstein, an open-source benchmark suite entirely implemented in JAX. The core of the system is a lightweight rendering engine based on the ray-casting technique originally used in Wolfenstein 3D (id Software, 1992).

The Rendering Engine

The engine utilizes Digital Differential Analyzer (DDA) ray casting to generate first-person views from tile-based maps.

• Mechanism: For each column of the image, the algorithm calculates the intersection of a ray cast from the observer with the environment (e.g., walls). Shading is applied based on the distance to the intersection to create depth.
• Optimization: By leveraging the tile-based nature of the world, the DDA algorithm "marches" along tile boundaries rather than using small, fixed intervals, significantly reducing the number of intersection checks required.
• Trade-offs: This approach prioritizes speed over visual fidelity. It requires environments to be tile-based, meaning features like stairs, jumping, or height differences are not supported. Complex objects would require more computationally expensive algorithms, potentially negating the speed benefits.

Implementation and Architecture

• JAX-Native: The entire pipeline, from environment rendering to training, is implemented in JAX. It utilizes Just-In-Time (JIT) compilation and vectorized mapping (vmap) to support massive parallelization.
• Resolution: The engine renders at lower resolutions natively, avoiding the per-timestep cost of downsampling images, which is common in deep RL.
• Extensibility: The system supports the conversion of 2D ASCII-defined mazes into first-person environments, allowing for the rapid creation of new tasks.

Key Contributions

1. First JAX-Based Visual First-Person Benchmark: JAXenstein is the first benchmark to bring visual first-person tasks to the JAX RL ecosystem, enabling end-to-end GPU-accelerated training for these domains.
2. Porting and Creation of Environments: The suite includes simplified, high-speed reimplementations of popular benchmarks:
   • ViZDoom: My Way Home and Health Gathering.
   • DeepMind Lab: Navigation mazes with static and random goal configurations (three map sizes: 01, 02, 03).
   • ASCII Mazes: Simple navigation tasks (simple, key-door) and a port of the MiniGrid-KeyCorridorS4R3-v0 environment.
3. Performance Optimization: By foregoing complex visual assets and utilizing tile-based ray casting, the system achieves significant speed improvements over existing simulators.

Results

The paper presents empirical comparisons demonstrating the efficiency of JAXenstein:

• Single-Core Speed: On a single CPU core, JAXenstein runs at approximately 6× the speed of comparable vision-based benchmarks like MiniWorld and ViZDoom.
• Parallelization: In training speed comparisons using a Recurrent Proximal Policy Optimization (PPO) algorithm, JAXenstein maintains high throughput as the number of parallel environments increases. In contrast, ViZDoom crashes when attempting to parallelize more than 64 environments.
• Baseline Performance: Experiments using recurrent PPO with various exploration strategies (including Random Network Distillation and Intrinsic Curiosity Module) on the dmlab-static-01 environment revealed that even the simplest maze navigation tasks in this benchmark are difficult for standard algorithms. This highlights the necessity of robust memory mechanisms and exploration strategies to handle the partial observability and large state spaces inherent in first-person visual domains.

Significance and Claims

The authors position JAXenstein as a tool to accelerate the iteration cycle for reinforcement learning research in visual domains. By removing the computational bottleneck of environment simulation, the benchmark allows researchers to:

• Iterate faster on algorithm development for first-person tasks.
• Conduct large-scale experiments with GPU acceleration that were previously cost-prohibitive.
• Test the limits of exploration and memory in partially observable environments without the overhead of high-fidelity rendering.

The paper modestly claims that while the current engine is limited to tile-based, flat environments, it successfully establishes a foundation for fast, scalable experimentation. Future work is envisioned to include non-playable characters and movable objects, though the current focus remains on the speed and scalability of the rendering pipeline itself.