PureJaxRL

The experimental/s5 directory contains a re-implementation of the architecture described in the paper Structured State Space Models for In-Context Reinforcement Learning.

Note: This implementation is currently a work-in-progress and the code is undergoing cleaning. Use with caution for production or stable research environments.

PureJaxRL is designed for high-performance, end-to-end Reinforcement Learning (RL) by implementing the entire training pipeline—including the environment—in JAX.

Key Characteristics:

• End-to-End JAX: By keeping the environment and training logic in JAX, the system avoids CPU-GPU data transfers and leverages JIT compilation for massive speedups.
• Vectorization: The implementation is built to support jit, vmap, pmap, and scan across the entire pipeline, enabling the simultaneous training of thousands of agents (e.g., 2048 PPO agents) on a single GPU.
• Non-Modular Design: Inspired by CleanRL, PureJaxRL provides high-quality, single-file implementations. It is intended to be read and used as a research resource rather than being imported as a modular library.

The following resources provide deep dives and educational content for mastering JAX-based RL:

• Achieving 4000x Speedups with PureJaxRL: Explains how JAX vectorization speeds up training.
• Breaking down State-of-the-Art PPO Implementations in JAX: An in-depth look at PureJaxRL's PPO implementation.
• A Gentle Introduction to Deep Reinforcement Learning in JAX: A general tutorial on Deep RL in JAX.
• Writing an RL Environment in JAX: A guide on creating your own environments.
• Getting started with JAX (MLPs, CNNs & RNNs): A basic tutorial on JAX neural networks.
• awesome-jax: A curated list of useful libraries in JAX.

PureJaxRL provides interactive notebooks to demonstrate usage. For a general walkthrough of the basic usage, use the walkthrough.ipynb notebook. For specific implementations involving Brax and MinAtar environments, use the brax_minatar.ipynb notebook.

- Basic usage: examples/walkthrough.ipynb
- Brax and MinAtar: examples/brax_minatar.ipynb

Install the necessary dependencies using the provided requirements.txt file. Additionally, ensure that JAX is correctly configured for your specific hardware accelerators (GPUs/TPUs) by following the official JAX installation documentation.

pip install -r requirements.txt

Use these specialized libraries to build or enhance your RL workflows:

• evosax / evojax: Evolution strategies in JAX.
• flashbax / dejax: Accelerated replay buffers in JAX.
• rlax: RL components and building blocks in JAX.
• mctx: Monte Carlo tree search in JAX.
• distrax: Distributions in JAX.
• optax: Gradient-based optimizers in JAX.
• flax: Neural Networks in JAX.

A variety of environments are available for training and benchmarking in JAX:

• gymnax: Classic RL environments in JAX.
• brax: Continuous control environments in JAX.
• JaxMARL: Multi-agent algorithms and environments in pure JAX.
• jumanji: Suite of unique RL environments in JAX.
• pgx: Suite of popular board games in JAX.
• popjaxrl: Partially-observed RL environments (POPGym) in JAX.
• waymax: Self-driving car simulator in JAX.
• Craftax: A challenging crafter-like and nethack-inspired benchmark in JAX.
• xland-minigrid: A large-scale meta-RL environment in JAX.
• navix: Classic minigrid environments in JAX.
• autoverse: A fast, evolvable description language for RL environments.
• qdx: Quantum Error Corection with JAX.
• matrax: Matrix games in JAX.
• AlphaTrade: Limit Order Book (LOB) in JAX.

If you are looking for RL algorithms implemented entirely end-to-end in JAX for maximum performance (up to 4000x speedups), consider these repositories:

• purejaxrl: Classic and simple end-to-end RL training in pure JAX.
• rejax: Modular and importable end-to-end JAX RL training.
• Stoix: End-to-end JAX RL training with advanced logging, configs, and more.
• purejaxql: Simple single-file end-to-end JAX baselines for Q-Learning.
• jym: Educational and beginner-friendly end-to-end JAX RL training.

To use Brax environments within the PureJaxRL training loop, you must wrap them to match the Gymnax interface. Key wrappers include:

• BraxGymnaxWrapper: Converts a Brax environment into a Gymnax-compatible interface (providing reset, step, observation_space, and action_space).
• VecEnv: Uses jax.vmap to vectorize the reset and step methods of the environment.
• ClipAction: Clips actions to a specified low and high range.
• NormalizeVecObservation: Performs running normalization of vectorized observations.
• NormalizeVecReward: Performs running normalization of vectorized rewards using a specified gamma.

The PPO implementation in this repository uses a specialized actor-critic loss function.

Value Loss

The value loss uses a clipped objective to prevent large updates to the value function: value_losses = jnp.square(value - targets) value_losses_clipped = jnp.square(value_pred_clipped - targets) value_loss = 0.5 * jnp.maximum(value_losses, value_losses_clipped).mean()

Actor Loss (Drift-based)

Instead of a standard PPO ratio loss, this implementation uses a drift-based approach with DPO_ALPHA and DPO_BETA parameters to regularize the policy update based on the advantage (gae). It calculates a drift term using nn.relu and nn.tanh to stabilize training.

Total Loss

The total loss is a weighted sum of the actor loss, value loss, and entropy: total_loss = loss_actor + config["VF_COEF"] * value_loss - config["ENT_COEF"] * entropy

Install the required dependencies for PureJaxRL using the provided requirements file via pip:

!pip install -r https://raw.githubusercontent.com/luchris429/purejaxrl/main/requirements.txt