Rejax is a library of RL algorithms which are implemented in pure Jax.
It allows you to accelerate your RL pipelines by using jax.jit
, jax.vmap
, jax.pmap
or any other transformation on whole training algorithms.
Use it to quickly search for hyperparameters, evaluate agents for multiple seeds in parallel, or run meta-evolution experiments on your GPUs and TPUs.
If you're new to rejax and want to learn more about it,
- Install via pip:
pip install rejax
- Install from source:
pip install git+https://github.com/keraJLi/rejax
- Use
jax.jit
on the whole train function to run training exclusively on your GPU! - Use
jax.vmap
andjax.pmap
on the initial seed or hyperparameters to train a whole batch of agents in parallel!
from rejax.algos import get_agent
# Get train function and initialize config for training
train_fn, config_cls = get_agent("sac")
train_config = config_cls.create(env="CartPole-v1", learning_rate=0.001)
# Jit the training function
jitted_train_fn = jax.jit(train_fn)
# Vmap training function over 300 initial seeds
vmapped_train_fn = jax.vmap(jitted_train_fn, in_axes=(None, 0))
# Train 300 agents!
keys = jax.random.split(jax.random.PRNGKey(0), 300)
train_state, evaluation = vmapped_train_fn(train_config, keys)
Benchmark on an A100 80G and a Intel Xeon 4215R CPU. Note that the hyperparameters were set to the default values of cleanRL, including buffer sizes. Shrinking the buffers can yield additional speedups due to better caching, and enables training of even more agents in parallel.
Algorithm | Link | Discrete | Continuous | Notes |
---|---|---|---|---|
PPO | here | โ | โ | |
SAC | here | โ | โ | discrete version as in Christodoulou, 2019 |
DQN | here | โ | incl. DDQN, Dueling DQN | |
DDPG | here | โ | ||
TD3 | here | โ |
The implementations focus on clarity! Easily modify the implemented algorithms by overwriting isolated parts, such as the loss function, trajectory generation or parameter updates. For example, easily turn DQN into DDQN by writing
class DoubleDQN(DQN):
@classmethod
def update(cls, config, state, minibatch):
# Calculate DDQN-specific targets
targets = ddqn_targets(config, state, minibatch)
# The loss function predicts Q-values and returns MSBE
def loss_fn(params):
...
return jnp.mean((targets - q_values) ** 2)
# Calculate gradients
grads = jax.grad(loss_fn)(state.q_ts.params)
# Update train state
q_ts = state.q_ts.apply_gradients(grads=grads)
state = state.replace(q_ts=q_ts)
return state
Using callbacks, you can run logging to the console, disk, wandb, and much more. Even when the whole train function is jitted! For example, run a jax.experimental.io_callback regular intervals during training, or print the current policies mean return:
def print_callback(config, state, rng):
policy = make_act(config, state) # Get current policy
episode_returns = evaluate(policy, ...) # Evaluate it
jax.debug.print( # Print results
"Step: {}. Mean return: {}",
state.global_step,
episode_returns.mean(),
)
return () # Must return PyTree (None is not a PyTree)
config = config.replace(eval_callback=print_callback)
Callbacks have the signature callback(config, train_state, rng) -> PyTree
, which is called every eval_freq
training steps with the config and current train state. The output of the callback will be aggregated over training and returned by the train function. The default callback runs a number of episodes in the training environment and returns their length and episodic return, such that the train function returns a training curve.
Importantly, this function is jit-compiled along with the rest of the algorithm. However, you can use one of Jax's callbacks such as jax.experimental.io_callback
to implement model checkpoining, logging to wandb, and more, all while maintaining the advantages of a completely jittable training function.
Libraries:
- Brax along with several environments, brax implements PPO and SAC within their environment interface
Single file implementations:
- PureJaxRL implements PPO, recurrent PPO and DQN
- Stoix features DQN, DDPG, TD3, SAC, PPO, as well as popular extensions and more
@misc{rejax,
title={rejax},
url={https://github.com/keraJLi/rejax},
journal={keraJLi/rejax},
author={Liesen, Jarek and Lu, Chris and Lange, Robert},
year={2024}
}