pith. sign in

arxiv: 2605.19926 · v1 · pith:3GB7FEJKnew · submitted 2026-05-19 · 💻 cs.LG

JAXenstein: Accelerated Benchmarking for First-Person Environments

Ruo Yu Tao , George Konidaris This is my paper

Pith reviewed 2026-05-20 08:00 UTC · model grok-4.3

classification 💻 cs.LG
keywords JAXreinforcement learningbenchmarkfirst-person environmentsWolfenstein 3Dvisual RLpartial observabilityexploration
0
0 comments X

The pith

JAXenstein ports the 1992 Wolfenstein 3D engine into JAX to create a fast, scalable benchmark for visual first-person reinforcement learning tasks.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces JAXenstein to address the lack of visual first-person benchmarks in the JAX RL ecosystem. It implements the classic Wolfenstein 3D rendering engine in JAX so that researchers can run large-scale experiments on tasks that require handling partial observability and exploration. The resulting environment runs several times faster than comparable vision-based benchmarks while remaining easy to extend toward more complex domains. If the port works as claimed, it removes a key bottleneck that currently slows algorithm iteration in visual RL.

Core claim

JAXenstein is an open-source JAX-based benchmark that implements the Wolfenstein 3D rendering engine for fast and scalable experimentation in visual first-person tasks, delivering several times the speed of existing vision-based RL environments and supporting straightforward extension to richer first-person domains.

What carries the argument

The JAX-native implementation of the Wolfenstein 3D rendering engine, which vectorizes and accelerates first-person visual simulation to support high-throughput RL training.

If this is right

  • RL researchers can iterate on algorithms for partial-observability tasks at several times the previous speed.
  • Visual first-person domains become practical for large-scale JAX-based training runs without specialized hardware.
  • New exploration methods can be prototyped and scaled quickly before moving to more complex simulators.
  • The benchmark can be extended to additional first-person settings while retaining the same performance advantage.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Adoption of JAXenstein could shift early-stage RL development toward simpler, faster visual testbeds before scaling to photorealistic environments.
  • The approach suggests that reimplementing classic game engines in JAX may be a general pattern for creating efficient RL benchmarks.
  • Faster iteration on first-person tasks may accelerate progress on embodied agents that must act under limited visual information.

Load-bearing premise

A JAX port of the 1992 Wolfenstein 3D engine supplies a sufficiently rich and representative testbed for modern RL challenges around exploration and partial observability.

What would settle it

Measure wall-clock time to reach a fixed exploration or navigation performance level on JAXenstein versus a modern first-person benchmark such as DeepMind Lab, using the same JAX-based RL algorithm and hardware.

Figures

Figures reproduced from arXiv: 2605.19926 by George Konidaris, Ruo Yu Tao.

Figure 1
Figure 1. Figure 1: Example environments in the JAXenstein benchmark, entirely implemented in end-to-end [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Speed comparisons between JAXenstein and similar benchmarks. [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Baseline results across JAXenstein environments. Runs were conducted on the recurrent [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Steps per second without image resizing. [PITH_FULL_IMAGE:figures/full_fig_p008_4.png] view at source ↗
read the original 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The manuscript introduces JAXenstein, a JAX-based implementation of the 1992 Wolfenstein 3D raycasting engine, as an open-source benchmark for visual first-person reinforcement learning tasks. It positions the tool as filling a gap in the JAX RL ecosystem by enabling fast, scalable experimentation on domains that test exploration and partial observability, with the central claims being that it is several times faster than comparable vision-based benchmarks and easily extensible to more complex first-person settings.

Significance. If the performance and extensibility claims are substantiated with quantitative evidence, JAXenstein could provide a valuable, high-throughput testbed that accelerates iteration on RL algorithms for partially observable visual environments within the JAX ecosystem. This would be particularly useful if the environment is shown to surface non-trivial exploration and belief-maintenance challenges beyond those in simpler discrete domains.

major comments (2)
  1. Abstract: the claim that 'JAXenstein is several times faster than comparable vision-based benchmarks' is presented without any quantitative runtime measurements, error bars, baseline comparisons, or implementation details, leaving the central performance advantage unverified and load-bearing for the paper's contribution.
  2. Introduction (or equivalent motivation section): the positioning of the benchmark as 'crucial for testing both exploration and an agent's ability to overcome partial observability' is undermined by reliance on fixed 64×64 ray-cast maps, binary wall/door geometry, and a small discrete action space; this setup risks allowing agents to solve tasks via short memorized sequences rather than long-horizon belief or large-scale exploration, directly engaging the stress-test concern.
minor comments (2)
  1. The manuscript would benefit from a dedicated experiments or results section that includes concrete runtime tables, comparisons to other JAX RL environments or vision benchmarks (e.g., Atari ports), and discussion of extensibility mechanisms with example code or modifications.
  2. Notation and environment description could be clarified by specifying the exact state representation, observation space dimensionality, and reward structure to aid reproducibility.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive feedback on our manuscript. We address each major comment below and indicate where revisions will be made to strengthen the presentation of results and clarify the benchmark's capabilities.

read point-by-point responses

  1. Referee: [—] Abstract: the claim that 'JAXenstein is several times faster than comparable vision-based benchmarks' is presented without any quantitative runtime measurements, error bars, baseline comparisons, or implementation details, leaving the central performance advantage unverified and load-bearing for the paper's contribution.

    Authors: We agree that the abstract's performance claim would benefit from explicit quantitative support to stand alone. The full manuscript reports runtime benchmarks in the Experiments section, including wall-clock times and throughput comparisons against vision-based baselines such as Atari environments in JAX and other raycasting implementations, with means and standard deviations over multiple runs. To address this directly, we will revise the abstract to incorporate specific speedup figures (e.g., 4-8x faster depending on resolution and batch size) along with a brief reference to the detailed evaluation. revision: yes

  2. Referee: [—] Introduction (or equivalent motivation section): the positioning of the benchmark as 'crucial for testing both exploration and an agent's ability to overcome partial observability' is undermined by reliance on fixed 64×64 ray-cast maps, binary wall/door geometry, and a small discrete action space; this setup risks allowing agents to solve tasks via short memorized sequences rather than long-horizon belief or large-scale exploration, directly engaging the stress-test concern.

    Authors: This is a fair point about the current setup's potential limitations. While the fixed 64×64 binary maps and discrete actions were chosen to enable high-throughput JAX vectorization and focus on rendering speed, the first-person raycast view inherently creates partial observability, as agents receive only local observations and must integrate information across timesteps. Our preliminary experiments indicate that memoryless agents perform poorly, suggesting some requirement for belief maintenance. Nevertheless, we recognize the risk of short-horizon memorization and will add an explicit limitations discussion in the revised manuscript, along with new results on randomized starting positions and procedurally varied map elements to better stress long-horizon exploration. revision: partial

Circularity Check

0 steps flagged

No circularity: empirical benchmark introduction with no derivations or self-referential claims

full rationale

The paper introduces JAXenstein as a new JAX-based software artifact implementing a 1992 raycasting engine for RL benchmarking. Its central claims concern measured speed improvements and extensibility, which are presented as empirical outcomes of the implementation rather than any first-principles derivation, fitted parameter, or prediction that reduces to the inputs by construction. The provided text contains no equations, no self-citations invoked as load-bearing uniqueness theorems, and no renaming of known results. The contribution is self-contained as a tool release whose validity rests on external benchmarking rather than internal logical loops.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The contribution is primarily an engineering port rather than new theory; no free parameters, invented physical entities, or non-standard axioms are introduced in the abstract.

axioms (1)
  • domain assumption JAX's vectorized execution model can deliver real-time rendering performance for first-person 3D environments.
    The speed claim rests on this unstated premise about JAX's suitability for graphics workloads.

pith-pipeline@v0.9.0 · 5675 in / 1139 out tokens · 57500 ms · 2026-05-20T08:00:54.383789+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Reference graph

Works this paper leans on

23 extracted references · 23 canonical work pages

  1. [1]

    and Barto, Andrew G

    Sutton, Richard S. and Barto, Andrew G. , publisher=. Reinforcement Learning:. 1998 , address=

    work page 1998

  2. [2]

    R. S. Sutton and D. McAllester and S. Singh and Y. Mansour. Policy Gradient Methods for Reinforcement Learning with Function Approximation. Advances in Neural Information Processing Systems 12. 2000

    work page 2000

  3. [3]

    R. J. Williams. Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. Machine Learning. 1992

    work page 1992

  4. [4]

    and Sutton, Richard S

    Barto, Andrew G. and Sutton, Richard S. and Anderson, Charles W. , journal=. Neuronlike adaptive elements that can solve difficult learning control problems , year=

    work page

  5. [5]

    Andrew William Moore , title =

    work page

  6. [6]

    The Arcade Learning Environment: An Evaluation Platform for General Agents , journal =. 2013

    work page 2013

  7. [7]

    and Veness, Joel and Bellemare, Marc G

    Mnih, Volodymyr and Kavukcuoglu, Koray and Silver, David and Rusu, Andrei A. and Veness, Joel and Bellemare, Marc G. and Graves, Alex and Riedmiller, Martin and Fidjeland, Andreas K. and Ostrovski, Georg and Petersen, Stig and Beattie, Charles and Sadik, Amir and Antonoglou, Ioannis and King, Helen and Kumaran, Dharshan and Wierstra, Daan and Legg, Shane ...

    work page

  8. [8]

    MuJoCo: A physics engine for model-based control , year=

    Todorov, Emanuel and Erez, Tom and Tassa, Yuval , booktitle=. MuJoCo: A physics engine for model-based control , year=

    work page

  9. [9]

    Proceedings of the 31st International Conference on International Conference on Machine Learning - Volume 32 , pages =

    Silver, David and Lever, Guy and Heess, Nicolas and Degris, Thomas and Wierstra, Daan and Riedmiller, Martin , title =. Proceedings of the 31st International Conference on International Conference on Machine Learning - Volume 32 , pages =. 2014 , publisher =

    work page 2014

  10. [10]

    and Hunt, Jonathan J

    Lillicrap, Timothy P. and Hunt, Jonathan J. and Pritzel, Alexander and Heess, Nicolas and Erez, Tom and Tassa, Yuval and Silver, David and Wierstra, Daan , booktitle =

    work page

  11. [11]

    2017 , eprint=

    Proximal Policy Optimization Algorithms , author=. 2017 , eprint=

    work page 2017

  12. [12]

    Proceedings of The 33rd International Conference on Machine Learning , pages =

    Asynchronous Methods for Deep Reinforcement Learning , author =. Proceedings of The 33rd International Conference on Machine Learning , pages =. 2016 , editor =

    work page 2016

  13. [13]

    Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem , title =

    C. Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem , title =

    work page

  14. [14]

    IEEE Transactions on Games , year =

    Marek Wydmuch and Micha. IEEE Transactions on Games , year =

    work page

  15. [15]

    Leibo and Denis Teplyashin and Tom Ward and Marcus Wainwright and Heinrich K

    Charles Beattie and Joel Z. Leibo and Denis Teplyashin and Tom Ward and Marcus Wainwright and Heinrich K. DeepMind Lab , journal =. 2016 , eprinttype =

    work page 2016

  16. [16]

    James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander

    work page

  17. [17]

    Alexander Mordvintsev , title=

    work page

  18. [18]

    2000 , publisher=

    3D Computer Graphics , author=. 2000 , publisher=

    work page 2000

  19. [19]

    CoRR , volume =

    Maxime Chevalier-Boisvert and Bolun Dai and Mark Towers and Rodrigo de Lazcano and Lucas Willems and Salem Lahlou and Suman Pal and Pablo Samuel Castro and Jordan Terry , title =. CoRR , volume =

    work page

  20. [20]

    Advances in Neural Information Processing Systems , volume=

    Discovered policy optimisation , author=. Advances in Neural Information Processing Systems , volume=

    work page

  21. [21]

    Journal of Machine Learning Research , year =

    Antonin Raffin and Ashley Hill and Adam Gleave and Anssi Kanervisto and Maximilian Ernestus and Noah Dormann , title =. Journal of Machine Learning Research , year =

    work page

  22. [22]

    International Conference on Learning Representations , year=

    Exploration by random network distillation , author=. International Conference on Learning Representations , year=

    work page

  23. [23]

    and Darrell, Trevor , title =

    Pathak, Deepak and Agrawal, Pulkit and Efros, Alexei A. and Darrell, Trevor , title =. International Conference on Machine Learning (ICML) , year =

    work page