RL-X

A framework for Reinforcement Learning research.

│ Overview │ Getting Started │ Documentation │ Citation │

Overview

Highlights

• 💡 Perfect to understand and prototype algorithms:
  • One algorithm = One directory -> No backtracking through parent classes
  • Algorithms can be easily copied out of RL-X
• ⚒️ Known DL libraries: Implementations in PyTorch and mainly JAX (Flax)
• ⚡ Maximum speed: Just-In-Time (JIT) compilation and parallel environments
• 🧪 Mix and match and extend: Generic interfaces between algorithms and environments
• ⛰️​ Custom environments: Examples for custom environments with MuJoCo or pure socket communication
• 🚀​ GPU environments: MuJoCo XLA (MJX) supports running 4000 parallel environments on a GPU
• 📈 Experiments: Checkpoints, Evaluation, Console log, Tensorboard, Weights & Biases, SLURM, Docker

Implemented Algorithms

• Proximal Policy Optimization (PPO) in PyTorch, Flax
• Early Stopping Policy Optimization (ESPO) in PyTorch, Flax
• Deep Deterministic Policy Gradient (DDPG) in Flax
• Twin Delayed Deep Deterministic Gradient (TD3) in Flax
• Soft Actor Critic (SAC) in PyTorch, Flax
• Randomized Ensembled Double Q-Learning (REDQ) in Flax
• Dropout Q-Functions (DroQ) in Flax
• CrossQ in Flax
• Truncated Quantile Critics (TQC) in Flax
• Aggressive Q-Learning with Ensembles (AQE) in Flax
• Maximum a Posteriori Policy Optimization (MPO) in Flax
• Deep Q-Network (DQN) in Flax
• Deep Q-Network with Histogram Loss using Gaussians (DQN HL-Gauss) in Flax
• Double Deep Q-Network (DDQN) in Flax

Usable Environments

• Gymnasium
  • MuJoCo
  • Atari
  • Classic control
• EnvPool
  • MuJoCo
  • Atari
  • Classic control
  • DeepMind Control Suite
• Custom MuJoCo
  • Example of a custom MuJoCo environment
  • Example of a custom MuJoCo XLA (MJX) environment
• Custom Interface
  • Prototype of a custom environment interface with socket communication

All listed environments are directly embedded in RL-X and can be used out-of-the-box.

For further information on the environments (README) and algorithms (README) and how to add your own, read the respective README files.

Getting Started

Install

Default installation for a Linux system with a NVIDIA GPU:

conda create -n rlx python=3.11.4
conda activate rlx
git clone git@github.com:nico-bohlinger/RL-X.git
cd RL-X
pip install -e .[all] --config-settings editable_mode=compat
pip uninstall $(pip freeze | grep -i '\-cu12' | cut -d '=' -f 1) -y
pip install "torch>=2.4.1" --index-url https://download.pytorch.org/whl/cu118 --upgrade
pip install -U "jax[cuda12]"

For other configurations, see the detailed installation guide in the documentation.

Example

cd experiments
python experiment.py

Detailed instructions for running experiments can be found in the README file in the experiments directory or in the documentation.

Example for Google Colab:

Citation

If you use RL-X in your research, please cite the following paper:

@incollection{bohlinger2023rlx,
      title={RL-X: A Deep Reinforcement Learning Library (not only) for RoboCup}, 
      author={Nico Bohlinger and Klaus Dorer},
      booktitle={Robot World Cup},
      pages={228--239},
      year={2023},
      publisher={Springer}
}