Metadata-Version: 2.5
Name: crax-bench
Version: 0.1.0
Summary: A GPU-accelerated benchmark and algorithm suite for Safe RL.
Project-URL: Homepage, https://ttomilin.github.io/CRAX/
Project-URL: Repository, https://github.com/TTomilin/CRAX
Project-URL: Paper, https://arxiv.org/abs/2606.20376
Project-URL: Issues, https://github.com/TTomilin/CRAX/issues
Author: Mourad Boustani, Mickey Beurskens, Thiago D. Simão
Author-email: Tristan Tomilin <t.tomilin@tue.nl>
Maintainer-email: Tristan Tomilin <t.tomilin@tue.nl>
License-Expression: Apache-2.0
License-File: LICENSE
License-File: NOTICE
Keywords: JAX,constrained reinforcement learning,physics,reinforcement learning,rigidbody,safe reinforcement learning
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: absl-py
Requires-Dist: etils
Requires-Dist: flax
Requires-Dist: jax>=0.6.0
Requires-Dist: jaxopt
Requires-Dist: ml-collections
Requires-Dist: mujoco-mjx<3.12,>=3.11.0
Requires-Dist: mujoco<3.12,>=3.11.0
Requires-Dist: numpy
Requires-Dist: pillow
Requires-Dist: typing-extensions
Provides-Extra: all
Requires-Dist: dm-env; extra == 'all'
Requires-Dist: gym>=0.26.0; extra == 'all'
Requires-Dist: imageio[ffmpeg]; extra == 'all'
Requires-Dist: matplotlib; extra == 'all'
Requires-Dist: mujoco-warp<3.12,>=3.11.0; extra == 'all'
Requires-Dist: optax; extra == 'all'
Requires-Dist: orbax-checkpoint; extra == 'all'
Requires-Dist: pandas; extra == 'all'
Requires-Dist: pygame; extra == 'all'
Requires-Dist: wandb; extra == 'all'
Requires-Dist: warp-lang>=1.16.0; extra == 'all'
Provides-Extra: analysis
Requires-Dist: matplotlib; extra == 'analysis'
Requires-Dist: pandas; extra == 'analysis'
Requires-Dist: wandb; extra == 'analysis'
Provides-Extra: cuda
Requires-Dist: jax[cuda12]>=0.6.0; extra == 'cuda'
Provides-Extra: develop
Requires-Dist: dm-env; extra == 'develop'
Requires-Dist: gym>=0.26.0; extra == 'develop'
Requires-Dist: imageio[ffmpeg]; extra == 'develop'
Requires-Dist: matplotlib; extra == 'develop'
Requires-Dist: optax; extra == 'develop'
Requires-Dist: orbax-checkpoint; extra == 'develop'
Requires-Dist: pandas; extra == 'develop'
Requires-Dist: pytest; extra == 'develop'
Requires-Dist: wandb; extra == 'develop'
Provides-Extra: gym
Requires-Dist: dm-env; extra == 'gym'
Requires-Dist: gym>=0.26.0; extra == 'gym'
Provides-Extra: play
Requires-Dist: pygame; extra == 'play'
Provides-Extra: torch
Requires-Dist: gym>=0.26.0; extra == 'torch'
Requires-Dist: torch; extra == 'torch'
Provides-Extra: train
Requires-Dist: imageio[ffmpeg]; extra == 'train'
Requires-Dist: matplotlib; extra == 'train'
Requires-Dist: optax; extra == 'train'
Requires-Dist: orbax-checkpoint; extra == 'train'
Requires-Dist: wandb; extra == 'train'
Provides-Extra: vision
Requires-Dist: mujoco-warp<3.12,>=3.11.0; extra == 'vision'
Requires-Dist: warp-lang>=1.16.0; extra == 'vision'
Description-Content-Type: text/markdown

# CRAX: Constrained Reinforcement Learning Accelerated with JAX

> 🌐 **Check out the [project page](https://ttomilin.github.io/CRAX/)** to try out CRAX yourself and claim a spot on the leaderboard!

CRAX is a high-performance benchmark for **Constrained Reinforcement Learning (Safe RL)** built on top of [Brax](https://github.com/google/brax) and [MuJoCo XLA (MJX)](https://mujoco.readthedocs.io/en/stable/mjx.html). It provides GPU/TPU-accelerated environments with safety constraints and a suite of state-of-the-art safe RL algorithms.

<p align="center">
  <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/envs.gif" alt="CRAX Environments: Goal, Push, Reacher, Lift, Pathway and Height tasks" width="100%" />
</p>

## Features

- **Massively parallel simulation**: Train agents across thousands of environments simultaneously on GPU/TPU
- **Configurable difficulty levels**: Environments with varying numbers and types of hazards
- **Multiple constraint types**: Cylindrical hazards, cubical hazards, boundary constraints, velocity limits, and more
- **State-of-the-art algorithms**: Implementations of leading constrained RL methods

## Safe RL Algorithms

CRAX includes efficient JAX implementations of the following, selected with the
`--alg` flag:

| Algorithm | Flag | Description | Reference |
|-----------|------|-------------|-----------|
| **PPO** | `ppo` | Vanilla PPO (unconstrained on-policy baseline) | [Schulman et al., 2017](https://arxiv.org/abs/1707.06347) |
| **PPO-Cost** | `ppo_cost` | PPO with a cost penalty folded into the reward | |
| **PPO-Lagrange** | `ppo_lag` | PPO with Lagrangian relaxation for constraints | [Ray et al., 2019](https://cdn.openai.com/safexp-short.pdf) |
| **PPO-PID** | `ppo_pid` | PPO with PID-controlled Lagrange multiplier | [Stooke et al., 2020](https://arxiv.org/abs/2007.03964) |
| **PPO-Saute** | `ppo_saute` | State Augmentation for Safe RL | [Sootla et al., 2022](https://arxiv.org/abs/2202.06558) |
| **FOCOPS** | `focops` | First Order Constrained Optimization in Policy Space | [Zhang et al., 2020](https://arxiv.org/abs/2002.06506) |
| **P3O** | `p3o` | Penalized Proximal Policy Optimization | [Zhang et al., 2022](https://arxiv.org/abs/2205.11814) |
| **CRPO** | `crpo` | Constrained Rectified Policy Optimization | [Xu et al., 2021](https://arxiv.org/abs/2011.05869) |
| **SAC** | `sac` | Soft Actor-Critic (unconstrained off-policy baseline) | [Haarnoja et al., 2018](https://arxiv.org/abs/1801.01290) |
| **SAC-Lagrange** | `sac_lag` | SAC with Lagrangian relaxation for constraints | |
| **SAC-PID** | `sac_pid` | SAC with PID-controlled Lagrange multiplier | |

All algorithms share a common training infrastructure with hooks for custom loss functions and constraint handling, making it easy to implement new methods. The registry lives in `get_algorithm_train_fn` in `crax/training/run_utils.py`; each algorithm adds its own CLI section in `crax/training/config.py` (`--safety_bound`, `--pid_kp`, `--nu_lr`, `--tau`, ...), and irrelevant arguments are filtered out per algorithm by `filter_kwargs_for_fn`.

## Environments

CRAX environments follow a **task-agent** naming convention: `safe_<task>_<agent>`.

| Task | Agents Supported                                      | Description                                              |
|------|-------------------------------------------------------|----------------------------------------------------------|
| **Goal** | Point, Ant, Spider, Humanoid, Swimmer                 | Navigate to goal positions while avoiding hazards        |
| **Circle** | Point, Ant, Spider, Humanoid, Swimmer                 | Orbit a circular path within boundaries                  |
| **Button** | Point, Ant, Spider, Humanoid, Swimmer                 | Press target buttons while avoiding gremlins and hazards |
| **Push** | Point, Ant, Spider, Humanoid, Swimmer                 | Push a block to a goal while avoiding hazards            |
| **Velocity** | Ant, Humanoid, HalfCheetah, Hopper, Swimmer, Walker2D | Locomotion under maximum velocity constraints            |
| **Lift** | Ant, Humanoid, Spider                                 | Locomotion with restricted legs touching the floor        |
| **Height** | Humanoid, HalfCheetah, Hopper, Walker2D               | Locomotion under a height ceiling                        |
| **Pathway** | HalfCheetah, Hopper, Walker2D                         | Traverse a corridor with hazard gaps                     |
| **Reacher** | Reacher                                               | Robotic arm reaching targets while avoiding obstacles    |

## Difficulty Levels

Every environment offers three difficulty levels that progressively increase constraint difficulty.

| Environment | Level 1 | Level 2 | Level 3 |
|:-----------:|:-------:|:-------:|:-------:|
| **Goal** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_point_goal_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_point_goal_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_point_goal_level_3.png" width="100%" /> |
| **Circle** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_point_circle_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_point_circle_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_point_circle_level_3.png" width="100%" /> |
| **Push** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_block_push_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_block_push_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_block_push_level_3.png" width="100%" /> |
| **Reacher** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_reacher_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_reacher_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_reacher_level_3.png" width="100%" /> |
| **Pathway** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_walker_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_walker_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_walker_level_3.png" width="100%" /> |
| **Height** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_height_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_height_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_height_level_3.png" width="100%" /> |
| **Lift** | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_spider_level_1.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_spider_level_2.png" width="100%" /> | <img src="https://raw.githubusercontent.com/TTomilin/CRAX/main/assets/envs/safe_spider_level_3.png" width="100%" /> |

## Installation

CRAX is published on PyPI as `crax-bench` (the import name is `crax`):

```bash
pip install crax-bench             # environments only
pip install "crax-bench[train]"    # + the RL algorithms in crax.training
```

To work on CRAX itself, or to use the scripts in `results/` and `scripts/`,
install from source:

```bash
git clone https://github.com/TTomilin/CRAX.git
cd CRAX
```

Create and activate a virtual environment using your preferred tool:

```bash
# Option A: venv
python3 -m venv .venv && source .venv/bin/activate

# Option B: conda
conda create -n crax python=3.11 && conda activate crax
```

Then install the package. The base install contains only what is needed to
build, step and render environments — no RL dependencies:

```bash
pip install -e .
```

To train agents, install the `train` extra, which adds the `crax.training`
dependencies (optax, orbax, wandb):

```bash
pip install -e ".[train]"
```

Other extras:

| Extra | Adds | For |
|---|---|---|
| `train` | optax, orbax-checkpoint, wandb, imageio, matplotlib | Running the RL algorithms in `crax.training` |
| `analysis` | pandas, matplotlib, wandb | Downloading and plotting results (`results/`) |
| `gym` | gym, dm_env | The Gym/dm_env wrappers in `crax.envs.wrappers` |
| `torch` | torch | The PyTorch wrapper |
| `vision` | mujoco-warp, warp-lang | GPU pixel observations (`--vision`) |
| `play` | pygame | Playing environments with the keyboard (`crax.play`) |
| `cuda` | `jax[cuda12]` | GPU support on CUDA 12 |
| `all` | train + analysis + gym + vision + play | Everything except CUDA and torch |

Extras combine, e.g. a GPU training install with pixel observations:

```bash
pip install -e ".[train,vision,cuda]"
```

See [SETUP.md](https://github.com/TTomilin/CRAX/blob/main/SETUP.md) for pinned installs, GPU/headless configuration, wandb
and troubleshooting, and [docs/VISION_TRAINING.md](https://github.com/TTomilin/CRAX/blob/main/docs/VISION_TRAINING.md) for
training from pixel observations.

## Start Here

The quickest way to see CRAX working is
[`examples/01_run_episode.py`](https://github.com/TTomilin/CRAX/blob/main/examples/01_run_episode.py).

```bash
python examples/01_run_episode.py --env_name safe_goal_point --level 1
```

The core loop is just a few lines:

```python
import jax
from crax import envs

env = envs.get_environment('safe_goal_point', level=1)
reset_fn, step_fn = jax.jit(env.reset), jax.jit(env.step)

rng = jax.random.PRNGKey(0)
state = reset_fn(rng)
for _ in range(300):
    rng, action_rng = jax.random.split(rng)
    action = jax.random.uniform(action_rng, (env.action_size,), minval=-1.0, maxval=1.0)
    state = step_fn(state, action)
    reward, cost = float(state.reward), float(state.info['cost'])  # safety cost per step
```

See [`examples/README.md`](https://github.com/TTomilin/CRAX/blob/main/examples/README.md) for the full list of examples.

### Playing CRAX

Point-agent environments can be played with a keyboard (requires the `play` extra):

```bash
python -m crax.play --env_name safe_goal_point --level 2
```

W/S thrust, A/D turn, C cycles the camera, R resets, Space pauses, Esc quits.

### Training from the command line

The training entry points live in `crax.training` and are run as modules:

```bash
# Single environment
python -m crax.training.train_env --env_name safe_goal_point --alg ppo_lag --difficulty 1

# Curriculum across difficulty levels
python -m crax.training.train_curriculum --env_name safe_goal_point --alg ppo_lag

# Safety transfer: unsafe PPO pre-training, then safe fine-tuning
python -m crax.training.train_transfer --env_name safe_velocity_ant --alg ppo_lag
```

Installing the package (with the `train` extra) also puts them on your PATH as
`crax-train`, `crax-train-curriculum` and `crax-train-transfer`. All three share
the argument set defined in `crax/training/config.py`; run with `--help` for the full
list.

## Project Structure

`crax/` holds the simulation and the environments; its `crax/training/`
subpackage holds the RL algorithms. Nothing outside `crax.training` imports it,
so environments can be used on their own with a much smaller dependency set
(see [Installation](#installation)).

```
CRAX/
├── crax/                        # Environments, simulation and (in training/) RL algorithms
│   ├── base.py, math.py, ...    # Physics core (Brax/MJX pipelines)
│   ├── generalized/, positional/, spring/, mjx/
│   ├── io/                      # MJCF loading, image + HTML export
│   ├── envs/                    # Environment definitions
│   │   ├── builder.py           # Modular XML scene builder
│   │   ├── difficulty.py        # Difficulty level configurations
│   │   ├── hazards.py           # Hazard generation utilities
│   │   ├── goals.py             # Goal sampling utilities
│   │   └── wrappers/            # Env wrappers (incl. pixel observations)
│   └── training/                # RL algorithms and training infrastructure
│       ├── agents/              # Algorithm implementations (PPO, SAC, etc.)
│       ├── networks.py          # Policy/value network factories
│       ├── acting.py            # Rollout and evaluation loops
│       ├── curriculum.py        # Curriculum training loop
│       ├── transfer.py          # Transfer learning loop
│       ├── checkpoint.py        # Checkpoint save/load
│       ├── config.py            # Shared CLI argument definitions
│       ├── run_utils.py         # Shared training helpers
│       ├── train_env.py         # Entry point: single environment training
│       ├── train_curriculum.py  # Entry point: progressive difficulty training
│       └── train_transfer.py    # Entry point: safety transfer learning
├── results/                     # Result download + plotting pipeline
├── examples/                    # Small self-contained usage examples
├── scripts/                     # Utility & visualization scripts
└── tests/                       # Pytest suite
```

## Architecture

CRAX uses a modular architecture where constrained RL algorithms are wrappers around a base RL algorithm, e.g.:

```
ppo/train.py          # Base trainer with hooks (loss_fn, post_step_fn, init_aux_state_fn)
    │
    ├── ppo_lag/      # Lagrange multiplier update + lagrange loss
    ├── ppo_pid/      # PID controller + lagrange loss
    ├── focops/       # FOCOPS loss + nu update
    ├── p3o/          # P3O loss + kappa adaptation
    ├── ppo_saute/    # Environment wrapper approach
    └── crpo/         # Rectified reward/cost objective switching
```

This design minimizes code duplication and makes it easy to add new algorithms.

## Contributing

Planned work is tracked in [ROADMAP.md](https://github.com/TTomilin/CRAX/blob/main/ROADMAP.md). Contributions are welcome. 
The roadmap flags which items are self-contained enough to pick up directly. 
For anything substantial, please open an issue first so we can avoid duplicated work.

## Research Powered by CRAX

Papers that have used CRAX:

- Lindsay Spoor, Aske Plaat, Thomas Moerland. *Evaluation Metrics for Safe Reinforcement Learning*, 2026. [arXiv:2609.15315](https://arxiv.org/abs/2609.15315)

Using CRAX in your work? Feel free to open a pull request to add your paper to this list.

## Acknowledgements

CRAX builds on [Brax](https://github.com/google/brax), Google's JAX-based physics and RL library, and on [MuJoCo XLA (MJX)](https://mujoco.readthedocs.io/en/stable/mjx.html)
and [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp), which power the accelerated simulation and the GPU pixel renderer. We thank the Brax and
MuJoCo teams for these foundations. [PROVENANCE.md](https://github.com/TTomilin/CRAX/blob/main/PROVENANCE.md) outlines which sources come from Brax.
Our navigation and velocity suite tasks are close reimplementations from 
[Safety Gym](https://openai.com/index/safety-gym/) and [Safety-Gymnasium](https://github.com/PKU-Alignment/safety-gymnasium).

## Citation

If you use CRAX in your research, please cite:

```bibtex
@article{tomilin2026crax,
  title         = {CRAX: Fast Safe Reinforcement Learning Benchmarking},
  author        = {Tristan Tomilin and Mourad Boustani and Mickey Beurskens and Thiago D. Sim{\~a}o},
  journal       = {arXiv preprint arXiv:2606.20376},
  year          = {2026}
}
```
