PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Koen Ponse, Aske Plaat, Thomas M. Moerland

arXiv:2610.36075v1cs.LGcs.DL

TL;DR

JAX reinforcement learning spans a fragmented set of environment interfaces, space types, and execution settings, while algorithms are often written for only a subset. Jaxnasium uses PyTrees and a common interface to support shared algorithms across these settings; its host-environment wrapper matches or exceeds the fastest baseline in the tested backends.

  • Problem

    JAX reinforcement learning algorithms are often tied to particular suites, space types, or JAX-native versus host-based environments, limiting the range of tasks a single implementation supports.

    Source: §1 · pp. 1–2

  • Method

    Jaxnasium represents spaces as PyTrees and provides a shared environment interface and wrappers so near-single-file algorithms can operate across space types, agents, and JAX-native or host environments.

    Source: §1 · p. 2, §5 · p. 7

  • Results

    Jaxnasium matches or exceeds the fastest baseline in every tested environment backend, including host-based settings.

    Source: §5 · p. 8

  • Takeaways & Limitations

    The same algorithm files serve as readable research references and importable practitioner algorithms across over 300 registered environments, including Gymnasium and EnvPool environments.

    Source: §8 · p. 10

    Jaxnasium currently implements four algorithms and is aimed at research-sized labs; large distributed production experiments are likely better served by other libraries.

    Source: §7 · p. 9

Abstract

from arXiv · show

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

1 INTRODUCTION

JAX environment suites and algorithm implementations have proliferated, but algorithms often remain tied to particular suites, space types, or execution settings. Jaxnasium uses PyTrees and a shared environment interface to make near-single-file algorithms work across these settings, with research and applied use in view.

1 INTRODUCTION

  • Algorithms are commonly reimplemented across suites, observation and action spaces, and single- versus multi-agent settings, while production implementations can be harder for researchers to modify.Research-oriented libraries favor readable implementations; production-oriented libraries favor broad importable support through more distributed components.Source: §1 · p. 1
  • JAX-native and host-based environments require different training architectures, leaving researchers with implementations that support only subsets of tasks or are difficult to modify.Source: §1 · p. 2
  • Jaxnasium represents observation and action spaces as PyTrees, letting near-single-file algorithms build networks that support discrete, continuous, image, composite, and multi-agent settings without code changes.Source: §1 · p. 2
  • A minimal Gymnasium-style interface and wrappers let one small PPO implementation run on over 300 environments, including host-based Gymnasium and EnvPool environments in the same compiled loop.The paper reports matching or exceeding the throughput of implementations built specifically for host-based environments.Source: §1 · p. 2
  • Jaxnasium’s contributions also include project scaffolding, multi-seed evaluation, hyperparameter sweeps, and PyTree utilities.Source: §1 · p. 2

2 BACKGROUND & LANDSCAPE

Accelerator-based JAX training can compile and vectorize rollouts and updates, but the ecosystem spans both JAX-native and host-based environments. This section describes the execution and interface landscape, along with JAX, PyTrees, and Equinox concepts relevant to Jaxnasium.

2 BACKGROUND & LANDSCAPE

  • Host-based environments still generate trajectories on the CPU while policy actions and network updates run on the accelerator, requiring repeated data transfers.Source: §2 · p. 2
  • Anakin describes end-to-end accelerator training for JIT-compatible environments, while Sebulba retains host-side data generation for traditional environments.Source: §2 · p. 3
  • PureJAXRL popularized simple single-device JAX implementations and reported that training could be multiple orders of magnitude faster, potentially reducing experiment times from days to minutes.Source: §2 · p. 3
  • Traditional stateful environment APIs conflict with JAX’s pure functional requirements, and Gymnasium’s experimental functional interface arrived after several JAX-native suites and has seen little adoption.Source: §2 · p. 3
  • JAX transformations such as jax.jit and jax.vmap support compiled training loops and batched environment steps, but require pure functions and compile-time-known array shapes.Source: §2 · p. 3
  • PyTrees represent nested structures of arrays, which JAX transformations can flatten and reconstruct while operating on their leaves.Source: §2 · p. 4
  • Equinox modules are PyTrees that combine array attributes with static or filtered metadata, allowing module methods and instances to cross JAX transformation boundaries.Jaxnasium builds its environments, wrappers, networks, and algorithms as Equinox modules.Source: §2 · p. 4

3 ENVIRONMENTS AS PYTREES

Jaxnasium’s environment interface gives algorithms one functional, JIT-compatible target while preserving familiar Gymnasium conventions. Its central abstraction treats composite observation and action spaces as PyTrees, which also represents multi-agent spaces without a separate interface.

3 ENVIRONMENTS AS PYTREES

  • Existing JAX environment suites use varied interfaces, so combining them requires wrappers that bridge differences in conventions.Source: §3 · p. 4
  • The interface targets algorithms with functional JIT-compatible calls, separate termination and truncation signals, and arbitrary PyTree observation and action spaces.It stays close to Gymnasium, and environments themselves are PyTrees so their parameters can pass through JAX transformations.Source: §3 · p. 4

3.1 ENVIRONMENT INTERFACE

  • Reset returns an initial observation and state; step returns separate termination and truncation signals, and automatic resets preserve the finished episode’s final observation in info for bootstrapping after truncation.Source: §3.1 · p. 4

3.2 PYTREES OF SPACES

  • Primitive spaces include Box, Discrete, and MultiDiscrete, while any PyTree of primitive spaces can represent a composite space using standard Python containers.Algorithms can map over the structure to construct matching networks; top-level entries can also represent agents with different, potentially composite spaces.Source: §3.2 · p. 5

4 ALGORITHMS OVER PYTREES

Jaxnasium’s PyTree-based algorithms combine readable, near-single-file implementations with importable classes and networks that adapt to varied observation and action spaces. The same single-agent logic can be dispatched across agents at the PyTree boundary, with independent learners as the default and centralized-critic support also available.

4 ALGORITHMS OVER PYTREES

  • Jaxnasium offers small implementations of PPO, SAC, DQN, and PQN that aim to combine readable research code with importable algorithms for applied use.Networks, replay buffers, and normalization classes are separated from the algorithm implementations.Source: §4 · p. 6
  • Importable Equinox classes retain a fully JIT-compatible training loop and can operate over varied action spaces and nearly any environment type.The library is available through pip and offers an API similar to Stable-Baselines.Source: §4 · p. 6

4.1 PYTREES OF NETWORKS

  • PyTree-mapped encoders process observation leaves, a shared body combines their features, and action-leaf heads produce distributions matching each action space.The default shared body is a SimBa network; the default architecture is optional.Source: §4.1 · p. 6, Figure 2 · p. 6

4.2 PYTREES OF AGENTS (MULTI-AGENT)

  • Multi-agent environments expose agents at the first PyTree level, allowing the algorithm to create a single-agent view for each entry.The transformation dispatches over agents at the PyTree boundary, leaving algorithm and agent logic in single-agent form.Source: §4.2 · pp. 6–7
  • By default, each agent is an independent learner with its own objective and networks; centralized critic observations and per-agent hyperparameters are also supported.Identically shaped agent operations are vectorized with jax.vmap; otherwise, they run per agent.Source: §4.2 · p. 7

5 ENVIRONMENT INTEROPERABILITY

Jaxnasium uses wrappers to present JAX-native and host-based environments through one research-friendly training implementation. In the tested throughput comparison, it matched or exceeded the fastest baseline on every backend.

5 ENVIRONMENT INTEROPERABILITY

  • Wrappers let JAX-native suites, Gymnasium, and EnvPool environments use the expected interface within the same compiled training loop.For host-based environments, only the environment transition runs on the host; rollout, update, and network code remain unchanged.Source: §5 · p. 7
  • Jaxnasium registers 333 environments from eleven JAX suites, and included algorithms run on over 300 of them in an end-to-end training loop.The wrappers translate suite signatures, time steps, and spaces into the Jaxnasium interface.Source: §5 · p. 7
  • For Gymnasium and EnvPool, wrappers convert host spaces and use JAX callbacks to step environments while handling finished-episode observations for resets and bootstrapping.EnvPool can alternatively be stepped through its own XLA interface.Source: §5 · p. 7
  • Jaxnasium matched or exceeded the fastest baseline on every tested backend, so a separate host-environment implementation was not necessary for those settings.Figure 3 averages throughput over six simple environments with cheap step functions; multi-device setups were not tested.Source: Figure 3 · p. 8, §5 · p. 8

6 RESEARCH TOOLING

Jaxnasium includes lightweight tools for project setup, evaluation, hyperparameter sweeps, and common reinforcement-learning utilities. Its evaluation tooling supports multiple search methods and metrics, while wrappers and agent utilities cover recurring training needs.

6 RESEARCH TOOLING

  • Jaxnasium provides these recurring tools partly to support more controlled and deterministic experiments, since seeding and aggregation can affect conclusions.The paper notes that project setup and evaluation setup are increasingly generated by large language models.Source: §6 · p. 8
  • The create-rl-app CLI scaffolds projects and adds templates for PPO and multi-environment sweeps.It can also be used as a standalone package to create a CLI for sharing configured templates.Source: §6 · pp. 8–9
  • The evaluation module sweeps arguments of functions, including training runs across algorithms, hyperparameters, or environments, with grid, random, and Sobol search.It also provides standard reinforcement-learning evaluation metrics.Source: §6 · p. 9
  • The jaxnasium.tree utilities support tree stacking, unstacking, agent-level mapping, and matching PRNG-key splits; wrappers provide vectorization, logging, and normalization.Agents can be saved and loaded with normalization built in, and training functions can be compiled ahead of runtime.Source: §6 · p. 9

7 LIMITATIONS & FUTURE WORK

Jaxnasium prioritizes maintainable algorithms and research-sized setups over a large algorithm catalog or large-scale distributed production experiments. Its automatic multi-agent transformation is limited to individual learning; more complex multi-agent structures require algorithm changes.

7 LIMITATIONS & FUTURE WORK

  • The library currently implements four algorithms and is unlikely to match other libraries’ algorithm counts, prioritizing maintainability and updating algorithms with improvements.Its main focus is simple setups for research-sized labs rather than large distributed, production-ready experiments.Source: §7 · p. 9
  • Automatic multi-agent transformation is limited to individual multi-agent learning; complex structures require algorithm changes, although CTDE examples are documented.The paper cites de Witt et al. (2020) for individual learning and Rashid et al. (2020) for complex structures.Source: §7 · p. 9
  • Future work aims to incorporate engineering improvements and updated algorithm defaults, with the goal of providing baselines and starting points across varied environments.The authors frame Jaxnasium as a research library.Source: §7 · p. 9

8 CONCLUSION

Jaxnasium organizes observation and action spaces as PyTrees so near-single-file algorithms can handle diverse spaces and multi-agent environments without code changes. The same implementations serve researchers and practitioners across wrapped environments, including host-based Gymnasium and EnvPool environments.

8 CONCLUSION

  • PyTree-structured spaces let near-single-file algorithms train on discrete, continuous, image, and composite spaces and extend to multi-agent settings without code changes.The components around an algorithm, rather than the algorithm itself, follow the PyTree structure.Source: §8 · pp. 9–10
  • The implementation files serve as readable research references and importable algorithms for practitioners.Jaxnasium also provides project scaffolding, multi-seed evaluation, and hyperparameter-sweep tools.Source: §8 · p. 10
  • Wrappers make the algorithms usable across over 300 registered environments and on host-based Gymnasium and EnvPool environments without sacrificing comparative performance.The conclusion also presents Jaxnasium as a starting point for researchers working with different environment types.Source: §8 · p. 10

A INCLUDED ENVIRONMENTS

The table lists the environments currently registered per suite and notes which environments are bundled or require installation.

A INCLUDED ENVIRONMENTS

  • Table 1 lists the registered environments available in Jaxnasium for each suite.The caption says additional environments may be added as they become available.Source: Table 1 · p. 14
  • Jaxnasium bundles the five classic-control environments; other libraries must be installed before initialization with jym.make(...).Source: Table 1 · p. 14
  • The listed environment coverage is expected to grow as more environments become available.Source: Table 1 · p. 14

B EXAMPLE OF A SWEEP CONFIGURATION

The example configures sweeps from reusable search stages, which can be nested or branched by algorithm. It illustrates independent jobs that can run locally or through a cluster job array.

B EXAMPLE OF A SWEEP CONFIGURATION

11 SobolSearch({"tau": (1e-3, 5e-2)}, 32, seed=key)],

7 GridSearch({"algorithm": {

C EXTRA ENVIRONMENTS FOR BACKEND PERFORMANCE COMPARISON

The extra backend benchmarks show a recurring pattern: JAX-native environments are typically fastest, while the wrapper avoids diminished performance against agents not built for end-to-end JAX training. The figures report throughput for Jumanji, MinAtar, and MuJoCo and Brax.

C EXTRA ENVIRONMENTS FOR BACKEND PERFORMANCE COMPARISON

  • Across the extra benchmarks, JAX-native environments are typically fastest, while the wrapper does not diminish performance relative to JAX agents not built for end-to-end training.Source: C EXTRA ENVIRONMENTS FOR BACKEND… · p. 15, Figure 4 · p. 15, Figure 5 · p. 16
  • Figure 4 reports throughput for different backends and libraries on Jumanji, averaged over five runs.The plotted standard deviation is described as negligible.Source: Figure 4 · p. 15
  • Figures 5 and 6 report throughput for different backends and libraries on MinAtar and on MuJoCo and Brax, respectively.Each figure averages five runs and describes the plotted standard deviation as negligible.Source: Figure 5 · p. 16, Figure 6 · p. 16

Related papers

Loading 2610.36075v1…