WARNING: Rljax is currently in a beta version and being actively improved. Any contributions are welcome :)

RL Algorithms in JAX

Rljax is a collection of RL algorithms written in JAX.

Setup

You can install dependencies simply by executing the following. To use GPUs, nvidia-driver and CUDA must be installed.

pip install --upgrade https://storage.googleapis.com/jax-releases/`nvcc -V | sed -En "s/.* release ([0-9]*)\.([0-9]*),.*/cuda\1\2/p"`/jaxlib-0.1.55-`python3 -V | sed -En "s/Python ([0-9]*)\.([0-9]*).*/cp\1\2/p"`-none-manylinux2010_x86_64.whl jax
pip install -e .

If you don't have a GPU, please executing the following instead.

pip install --upgrade jaxlib jax
pip install -e .

If you want to use a MuJoCo physics engine, please install mujoco-py.

Algorithms

Currently, following algorithms have been implemented.

• Proximal Policy Optimization(PPO)
• Deep Deterministic Policy Gradient(DDPG)
• Twin Delayed DDPG(TD3)
• Soft Actor-Critic(SAC)
• Deep Q Network(DQN)
• N-step return
• Dueling Network
• Double Q-Learning
• Prioritized Experience Replay(PER)
• Soft Actor-Critic for Discrete Settings(SAC-Discrete)

We plan to implement the following algorithms in the future.

• Quantile Regression DQN(QR-DQN)
• Implicit Quantile Network(IQN)

Below shows that our algorithms successfully learning the discrete action environment CartPole-v0 and the continuous action environment InvertedPendulum-v2.