Metadata-Version: 2.3
Name: jaxnasium
Version: 0.1.1
Summary: A lightweight utility library for reinforcement learning projects in JAX and Equinox.
Author: Koen Ponse
Author-email: Koen Ponse <k.ponse@liacs.leidenuniv.nl>
Requires-Dist: create-rl-app>=0.1.0
Requires-Dist: equinox>=0.12.2,!=0.13.3
Requires-Dist: jax>=0.5.2,!=0.7.*,!=0.9.*
Requires-Dist: jax>=0.8.2 ; extra == 'algs'
Requires-Dist: distrax>=0.1.7 ; extra == 'algs'
Requires-Dist: optax>=0.2.5 ; extra == 'algs'
Requires-Dist: brax>=0.14.2 ; extra == 'brax'
Requires-Dist: envpool>=1.2.5,!=1.2.6 ; extra == 'envpool'
Requires-Dist: jaxnasium[brax,envpool,factoriax,gymnasium,gymnax,jaxmarl,jumanji,navix,octax,pgx,playground,xminigrid] ; extra == 'envs'
Requires-Dist: factoriax>=0.1.0 ; extra == 'factoriax'
Requires-Dist: jax[cuda] ; extra == 'gpu'
Requires-Dist: gymnasium>=1.1.0 ; extra == 'gymnasium'
Requires-Dist: gymnax>=0.0.9 ; extra == 'gymnax'
Requires-Dist: jaxmarl[mabrax]>=0.2.0 ; extra == 'jaxmarl'
Requires-Dist: jumanji>=1.1.2 ; extra == 'jumanji'
Requires-Dist: kinetix-env>=3.0.2 ; extra == 'kinetix'
Requires-Dist: navix>=0.7.4 ; extra == 'navix'
Requires-Dist: octax>=0.1.1 ; extra == 'octax'
Requires-Dist: opencv-python-headless>=4.10.0.84 ; extra == 'octax'
Requires-Dist: pillow>=10.0.0 ; extra == 'octax'
Requires-Dist: pgx>=2.6.0 ; extra == 'pgx'
Requires-Dist: playground>=0.2.0 ; extra == 'playground'
Requires-Dist: xminigrid>=0.9.3 ; extra == 'xminigrid'
Requires-Python: >=3.11
Provides-Extra: algs
Provides-Extra: brax
Provides-Extra: envpool
Provides-Extra: envs
Provides-Extra: factoriax
Provides-Extra: gpu
Provides-Extra: gymnasium
Provides-Extra: gymnax
Provides-Extra: jaxmarl
Provides-Extra: jumanji
Provides-Extra: kinetix
Provides-Extra: navix
Provides-Extra: octax
Provides-Extra: pgx
Provides-Extra: playground
Provides-Extra: xminigrid
Description-Content-Type: text/markdown

# Jaxnasium: A Lightweight Utility Library for JAX-based RL Projects

[![PyPI version](https://badge.fury.io/py/jaxnasium.svg)](https://badge.fury.io/py/jaxnasium)
[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
[![Documentation Status](https://img.shields.io/badge/docs-latest-brightgreen.svg)](https://ponseko.github.io/jaxnasium/)

Jaxnasium is your one-stop shop for your JAX RL code: a single environment API that existing suites are wrapped into, a set of general algorithms you can either import or copy into your project, and the tooling to train, evaluate and sweep them.

1. 🕹️ **One environment API.** Import environments from Gymnax, Jumanji, Brax, Crax, Pgx, JaxMARL, xMinigrid, Navix, Playground, PopJym or Craftax through `jym.make(...)`, automatically wrapped to a common standard.
2. 🤖 **Readable cross-suite algorithms.** Various algorithms that can operate on any of those environments in near-single-file philosophy, built in [Equinox](https://github.com/patrick-kidger/equinox) with a familiar [Stable-Baselines](https://github.com/DLR-RM/stable-baselines3)-like API and end-to-end JIT training in the spirit of [PureJaxRL](https://github.com/luchris429/purejaxrl).
3. 👥 **Multi-agent for free.** Single-agent algorithm code transparently upgrades to multi-agent environments through PyTrees and function transformations.
4. 📊 **Sweeps and evaluation.** Grid, random, Sobol and one-at-a-time searches over many seeds, locally or across a Slurm job array.
5. 🚀 **Project scaffolding.** Bootstrap a complete codebase with a single CLI command.

For more details, see the [Documentation](https://ponseko.github.io/jaxnasium/).

## Installation

```bash
pip install "jaxnasium[algs]"   # [algs] pulls in optax + distrax, needed for jaxnasium.algorithms
```

Third-party environment suites are _not_ dependencies; install the ones you want to use
(e.g. `pip install gymnax`). For a brand-new project, let the [CLI](https://ponseko.github.io/jaxnasium/cli/) set everything up for you:

```bash
uvx jaxnasium init <projectname>   # or: pipx run jaxnasium init <projectname>
cd <projectname> && uv run train.py
```

## Quickstart

```python
import jax
import jaxnasium as jym
from jaxnasium.algorithms import PPO

env = jym.make("CartPole-v1")
key = jax.random.PRNGKey(0)

algorithm = PPO(total_timesteps=500_000, learning_rate_start=2.5e-3)
agent, metrics = algorithm.train(key, env)
returns = agent.evaluate(key, env, num_eval_episodes=10)
```

`train` returns the trained agent alongside its training metrics: by default the mean
return of the episodes that finished in each training iteration, one scalar per iteration
(iterations without a finished episode yield `NaN`).

## 🏠 Environments

Jaxnasium does not aim to deliver a full environment suite. Instead, `jym.make(...)` imports
environments from existing suites (provided they are installed) and wraps them to the
Jaxnasium API standard:

```python
env = jym.make("Breakout-MinAtar")  # from Gymnax
env = jym.FlattenObservationWrapper(env)

algorithm = PPO(**some_good_hyperparameters)
agent, metrics = algorithm.train(jax.random.PRNGKey(0), env)

# > Using an environment from Gymnax via gymnax.make(Breakout-MinAtar).
# > Wrapping Gymnax environment with GymnaxWrapper
# >  Control this behavior by passing wrappers=[...] to jym.make
# > Wrapping environment in VecEnvWrapper
# > ... training results
```

By default `make` applies the adapter wrapper for the environment's library and a
[`LogWrapper`](https://ponseko.github.io/jaxnasium/api/Wrappers/) for episode statistics; the algorithms add a
`VecEnvWrapper` themselves when training.

> For convenience, Jaxnasium bundles the 5 [classic-control environments](https://gymnasium.farama.org/environments/classic_control/), which need no external dependencies.

See [Available Environments](https://ponseko.github.io/jaxnasium/api/Available-Environments/) for a complete list of available environments.

### Environment API

The Jaxnasium API stays close to the _somewhat_ established [Gymnax](https://github.com/RobertTLange/gymnax) API for the `reset()` and `step()` functions, but allows for truncated episodes in a manner closer to [Gymnasium](https://gymnasium.farama.org/).

```python
obs, env_state = env.reset(key)  # <-- Mirroring Gymnax

# env.step(): Gymnasium Timestep tuple with state information
(obs, reward, terminated, truncated, info), env_state = env.step(key, env_state, action)
```

Writing your own environment means subclassing [`jym.Environment`](https://ponseko.github.io/jaxnasium/api/Environment/) and
implementing `step_env`, `reset_env`, `observation_space` and `action_space`.

## 🤖 Algorithms

Algorithms in `jaxnasium.algorithms` follow a near-single-file implementation philosophy. In contrast to implementations in [CleanRL](https://github.com/vwxyzjn/cleanrl) or [PureJaxRL](https://github.com/luchris429/purejaxrl), Jaxnasium algorithms are built in Equinox and follow a class-based design with a familiar [Stable-Baselines](https://github.com/DLR-RM/stable-baselines3) API.

Each algorithm is split into an **algorithm**, which holds the hyperparameters and the training
logic, and the **agent**, holding the trainable state along with the trainer that produced it.

```python
agent, metrics = PPO(**some_good_hyperparameters).train(key, env)

agent.get_action(key, obs)                          # act
agent, more_metrics = agent.train(key, env)         # keep training
agent.save("ppo_agent.eqx")                         # checkpoint, trainer included
```

Networks are configurable without touching the algorithm, observation and action spaces may
be arbitrary PyTrees of spaces, and observation/reward normalization is built in and
checkpointed with the agent. See [Algorithms](https://ponseko.github.io/jaxnasium/algorithms/Algorithms/) for the details.

| Algorithm | Multi-Agent<sup>1</sup> | Observation Spaces                 | Action Spaces                           | Composite (nested) Spaces<sup>2</sup> |
| --------- | ----------------------- | ---------------------------------- | --------------------------------------- | ------------------------------------- |
| **PPO**   | ✅                      | `Box`, `Discrete`, `MultiDiscrete` | `Box`, `Discrete`, `MultiDiscrete`      | ✅                                    |
| **DQN**   | ✅                      | `Box`, `Discrete`, `MultiDiscrete` | `Discrete`, `MultiDiscrete`<sup>3</sup> | ✅                                    |
| **PQN**   | ✅                      | `Box`, `Discrete`, `MultiDiscrete` | `Discrete`, `MultiDiscrete`<sup>3</sup> | ✅                                    |
| **SAC**   | ✅                      | `Box`, `Discrete`, `MultiDiscrete` | `Box`, `Discrete`, `MultiDiscrete`      | ✅                                    |

<sup>1</sup> All algorithms support automatic multi-agent transformation through the `auto_upgrade_multi_agent` parameter. See [Multi-Agent](https://ponseko.github.io/jaxnasium/algorithms/Multi-Agent/) for more information.

<sup>2</sup> Algorithms support composite (nested) spaces. See [Spaces](https://ponseko.github.io/jaxnasium/api/Spaces/) for more information.

<sup>3</sup> MultiDiscrete action spaces in PQN and DQN are only supported when flattening to a Discrete action space. E.g. via the `FlattenActionSpaceWrapper`.

## 👥 Multi-agent

Multi-agent support is a **core pillar** of Jaxnasium's design rather than a separate API.
Multi-agent Environments simply expose a PyTree of spaces whose first level is the agent dimension, and
algorithms build one agent per entry automatically:

```python
env = jym.make("MPE_simple_spread_v3")  # a JaxMARL environment
agent, metrics = PPO().train(key, env)  # one PPO agent per agent in the environment

agent.get_action(key, obs)  # {"agent_0": ..., "agent_1": ..., ...}
```

Hyperparameters may even differ per agent, e.g. `PPO(gamma={"agent_0": 0.9, "agent_1": 0.99})`.
See [Multi-Agent](https://ponseko.github.io/jaxnasium/algorithms/Multi-Agent/).

## 📊 Sweeps and evaluation

`jaxnasium.eval` turns any function into a parameter sweep, and provides a standard
train-and-evaluate trial for Jaxnasium algorithms. Jobs are addressed by index, so the same
script runs locally in a loop or as a Slurm job array:

```python
from jaxnasium.eval import AlgorithmEvaluation, GridSearch, SobolSearch, Sweep

sweep = Sweep(
    AlgorithmEvaluation(jax.random.split(key, 10), batch_size=5, save_path="results/"),
    GridSearch({"env": ["CartPole-v1", "Acrobot-v1"], "algorithm": ["PPO", "DQN"]}),
    SobolSearch({"learning_rate_start": (1e-4, 1e-2, "log")}, num_samples=32, seed=key),
)

len(sweep)          # 2 envs * 2 algorithms * 32 learning rates
result = sweep[0]()  # run one job
```

See [Sweeps](https://ponseko.github.io/jaxnasium/eval/Sweeps/) and [Searches](https://ponseko.github.io/jaxnasium/eval/Searches/).

## Where to go next

|                                                                                                                                                                                                          |                                       |
| -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------- |
| [Environment](https://ponseko.github.io/jaxnasium/api/Environment/) · [Spaces](https://ponseko.github.io/jaxnasium/api/Spaces/) · [Wrappers](https://ponseko.github.io/jaxnasium/api/Wrappers/)          | Build or adapt an environment         |
| [Algorithms](https://ponseko.github.io/jaxnasium/algorithms/Algorithms/) · [Multi-Agent](https://ponseko.github.io/jaxnasium/algorithms/Multi-Agent/)                                                    | Train agents                          |
| [Networks](https://ponseko.github.io/jaxnasium/algorithms/networks/Networks/) · [Architectures](https://ponseko.github.io/jaxnasium/algorithms/networks/Architectures/)                                  | Swap in your own model                |
| [Sweeps](https://ponseko.github.io/jaxnasium/eval/Sweeps/) · [Searches](https://ponseko.github.io/jaxnasium/eval/Searches/)                                                                              | Run experiments at scale              |
| [Compilation](https://ponseko.github.io/jaxnasium/api/Compilation/)                                                                                                                                      | Ahead-of-time compilation and caching |
| [Checkpointing](https://ponseko.github.io/jaxnasium/algorithms/core/Checkpointing/) · [Tree utilities](https://ponseko.github.io/jaxnasium/tree/Tree/) · [CLI](https://ponseko.github.io/jaxnasium/cli/) | Everything else                       |
