cs.LGOct 26, 2026

Jaxnasium: Embracing PyTrees For Any-Agent, Cross-Suite Reinforcement Learning

Authors: Koen Ponse, Aske Plaat, Thomas M. Moerland

Organizations: Leiden Institute of Advanced Computer Science Leiden University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Hardware-accelerated reinforcement learning in JAX has gained rapid popularity in recent years and has proven to be a successful method of scaling online reinforcement learning. The resulting ecosystem, however, is fragmented: environment suites define their own interfaces, and algorithm implementations are typically written for a single suite, a single type of space, and either JAX-native or host-based environments. Here we present Jaxnasium, a library for JAX-based reinforcement learning, aiming to make the field more accessible. Jaxnasium represents observation and action spaces as PyTrees, and uses this structure to build networks and support multi-agent processing. Through this, its near-single-file implementations run on discrete, continuous, image, and composite spaces, as well as in multi-agent environments, without any changes to the algorithm code. The same files serve as readable, research-friendly algorithms, as well as importable algorithms for more applied workloads. Through a minimal Gymnasium-style interface and wrappers for 11 suites, over 300 environments are available on which a small PPO implementation with default hyperparameters is able to run. Additionally, through a simple wrapper, we can run Gymnasium and EnvPool environments with the same algorithm and training loop, with no performance loss compared to architectures specifically built for training JAX algorithms on host-based environments. Jaxnasium further provides tools for multi-seed evaluation, hyperparameter sweeps, and project scaffolding. Jaxnasium is open-sourced at https://github.com/ponseko/jaxnasium

Figures & tables

Figure 1: Jaxnasium registers 333 environments and includes algorithm implementations that can train on over 300 of them without any changes. The functionality is familiar to those who have used well-known libraries such as Gymnasium and Stable-Baselines.
Figure 3: The default agent (policy) network. An observation network creates an encoder for each observation leaf, a shared body processes the concatenated features, and an output network creates a head for each action leaf, with a distribution matching the leaf’s space. While the structure is useful for mapping spaces onto networks, it remains optional. Critic networks are created similarly, but with different output heads.
Figure 5: Average throughput of Jaxnasium and other existing established libraries on various environment backends with a small PPO network for each. Not all libraries have an implementation available to run on each environment backend, hence these bars are omitted from the graph. Jaxnasium, despite using a single small implementation to run on all types of environments, matches or exceeds the throughput of other libraries. Because not many environments have been implemented in all these backends, this graph only averages over 6 simple environments, with a cheap step function. Appendix C lists extra environments, which paint roughly the same picture.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

SuiteDomainEnvs
Brax ( Freeman et al., 2021 )Continuous control10
Craftax ( Matthews et al., 2024 )Open-ended survival4
Gymnax ( Lange, 2022 )Classic control, bsuite, MinAtar24
JaxMARL ( Rutherford et al., 2024 )Multi-agent coordination29
Jumanji ( Bonnet et al., 2024 )Combinatorial problems26
Navix ( Pignatelli et al., 2025 )Grid worlds41
Appendix
Table 1: Registered environments per suite currently available in Jaxnasium. For convenience, Jaxnasium bundles the 5 classic-control environments; other libraries have to be installed and can subsequently be initialized with jym.make(...) . More environments will likely be added as they become available.
Figure 7: Throughput tested for various backends and different libraries on the Jumanji suite. Averaged over 5 runs. Standard deviation is plotted, but is negligible.
Figure 8: Throughput tested for various backends and different libraries on the MinAtar suite. Averaged over 5 runs. Standard deviation is plotted, but is negligible.
Figure 9: Throughput tested for various backends and different libraries on the MuJoCo and Brax suites. Averaged over 5 runs. Standard deviation is plotted, but is negligible.
More in this topic

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,742 papers

Explore similar work

CardsList
Nov 16, 2023

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. The code is available at https://github.com/flairox/jaxmarl.

Alexander Rutherford, Benjamin Ellis, Matteo Gallici +18

University of Oxford · Universitat Politècnica de Catalunya · University College London +2

RL BenchmarksReinforcement LearningGame-Playing AgentsGPU-Accelerated RLMulti-Agent Reinforcement Learning
AskPDF
Oct 2, 2025

Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX

Reinforcement learning (RL) research requires diverse, challenging environments that are both tractable and scalable. While modern video games may offer rich dynamics, they are computationally expensive and poorly suited for large-scale experimentation due to their CPU-bound execution. We introduce Octax, a high-performance suite of classic arcade game environments implemented in JAX, based on CHIP-8 emulation, a predecessor to Atari, which is widely adopted as a benchmark in RL research. Octax provides the JAX community with a long-awaited end-to-end GPU alternative to Atari games, offering image-based environments, spanning puzzle, action, and strategy genres, all executable at massive scale on modern GPUs. Our JAX-based implementation achieves orders-of-magnitude speedups over traditional CPU emulators. We demonstrate Octax's capabilities by training RL agents across multiple games, showing significant improvements in training speed and scalability compared to existing solutions. The environment's modular design enables researchers to easily extend the suite with new games or generate novel environments using large language models, making it an ideal platform for large-scale RL experimentation. Our open-source framework is available at https://github.com/riiswa/octax/.

Waris Radji, Thomas Michel, Hector Piteau

Univ. Lille, Inria, CNRS, Centrale Lille, UMR 9189-CRIStAL, France · Independent Researcher

RL BenchmarksReinforcement LearningGPU-Accelerated RLRL for Game AI
AskPDF
May 19, 2026

JAXenstein: Accelerated Benchmarking for First-Person Environments

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.

Ruo Yu Tao, George Konidaris

Department of Computer Science, Brown University, USA

RL BenchmarksPartially Observable RLRL ExplorationVisual RL
AskPDF