ipsec / Stoix

🏛️A research-friendly codebase for fast experimentation of single-agent reinforcement learning in JAX • End-to-End JAX RL

Distributed Single-Agent Reinforcement Learning End-to-End in JAX

stoic - a person who can endure pain or hardship without showing their feelings or complaining.

Welcome to Stoix! 🏛️

Stoix provides simplified code for quickly iterating on ideas in single-agent reinforcement learning with useful implementations of popular single-agent RL algorithms in JAX allowing for easy parallelisation across devices with JAX's pmap. All implementations are fully compilable with JAX's jit thus making training and environment execution very fast. However, this requires environments written in JAX. Algorithms and their default hyperparameters have not been hyper-optimised for any specific environment and are useful as a starting point for research and/or for initial baselines.

To join us in these efforts, please feel free to reach out, raise issues or read our contribution guidelines (or just star 🌟 to stay up to date with the latest developments)!

Stoix is fully in JAX with substantial speed improvement compared to other popular libraries. We currently provide native support for the Jumanji environment API and wrappers for popular JAX-based RL environments.

Code Philosophy 🧘

The current code in Stoix was initially largely taken and subsequently adapted from Mava. As Mava develops, Stoix will hopefully adopt their optimisations that are relevant for single-agent RL. Like Mava, Stoix is not designed to be a highly modular library and is not meant to be imported. Our repository focuses on simplicity and clarity in its implementations while utilising the advantages offered by JAX such as pmap and vmap, making it an excellent resource for researchers and practitioners to build upon. Stoix follows a similar design philosophy to CleanRL and PureJaxRL, where we allow for some code duplication to enable readability, easy reuse, and fast adaptation. A notable difference between Stoix and other single-file libraries is that Stoix makes use of abstraction where relevant. It is not intended to be purely educational with research utility as the primary focus. In particular, abstraction is currently used for network architectures, environments, logging, and evaluation.

Overview 🦜

Stoix currently offers the following building blocks for Single-Agent RL research:

Implementations of Algorithms 🥑

• Deep Q-Network (DQN) - Paper
• Double DQN (DDQN) - Paper
• Dueling DQN - Paper
• Categorical DQN (C51) - Paper
• Munchausen DQN (M-DQN) Paper
• Quantile Regression DQN (QR-DQN) - Paper
• DQN with Regularized Q-learning (DQN-Reg) Paper
• REINFORCE With Baseline - Paper
• Deep Deterministic Policy Gradient (DDPG) - Paper
• Twin Delayed DDPG (TD3) - Paper
• Distributed Distributional DDPG (D4PG) - Paper
• Soft Actor-Critic (SAC) - Paper
• Proximal Policy Optimization (PPO) - Paper
• Discovered Policy Optimization (DPO) Paper
• Maximum a Posteriori Policy Optimisation (MPO) - Paper
• On-Policy Maximum a Posteriori Policy Optimisation (V-MPO) - Paper
• Advantage-Weighted Regression (AWR) - Paper
• AlphaZero - Paper
• MuZero - Paper
• Sampled Alpha/Mu-Zero - Paper

Environment Wrappers 🍬

Stoix offers wrappers for Gymnax, Jumanji, Brax, XMinigrid, Craftax and even JAXMarl (although using Centralised Controllers).

Statistically Robust Evaluation 🧪

Stoix natively supports logging to json files which adhere to the standard suggested by Gorsane et al. (2022). This enables easy downstream experiment plotting and aggregation using the tools found in the MARL-eval library.

Performance and Speed 🚀

As the code in Stoix (at the time of creation) was in essence a port of Mava, for further speed comparisons we point to their repo. Additionally, we refer to the PureJaxRL blog post here where the speed benefits of end-to-end JAX systems are discussed.

Below we provide some plots illustrating that Stoix performs equally to that of PureJaxRL but with the added benefit of the code being already set up for pmap distribution over devices as well as the other features provided (algorithm implementations, logging, config system, etc).

I've also included a plot of the training time for 5e5 steps of PPO as one scales the number of environments. PureJaxRL does not pmap and thus runs on a single a device.

Lastly, please keep in mind for practical use that current networks and hyperparameters for algorithms have not been tuned.

Installation 🎬

At the moment Stoix is not meant to be installed as a library, but rather to be used as a research tool.

You can use Stoix by cloning the repo and pip installing as follows:

git clone https://github.com/EdanToledo/Stoix.git
cd Stoix
pip install -e .

We have tested Stoix on Python 3.10. Note that because the installation of JAX differs depending on your hardware accelerator, we advise users to explicitly install the correct JAX version (see the official installation guide).

Quickstart ⚡

To get started with training your first Stoix system, simply run one of the system files. e.g.,

python stoix/systems/ppo/ff_ppo.py

Stoix makes use of Hydra for config management. In order to see our default system configs please see the stoix/configs/ directory. A benefit of Hydra is that configs can either be set in config yaml files or overwritten from the terminal on the fly. For an example of running a system on the CartPole environment, the above code can simply be adapted as follows:

python stoix/systems/ppo/ff_ppo.py env=gymnax/cartpole

Contributing 🤝

Please read our contributing docs for details on how to submit pull requests, our Contributor License Agreement and community guidelines.

Roadmap 🛤️

We plan to iteratively expand Stoix in the following increments:

• 🌴 Support for more environments as they become available.
• 🔁 More robust recurrent systems.
  • Add recurrent variants of all systems
  • Allow easy interchangability of recurrent cells/architecture via config
• 📊 Benchmarks on more environments.
  • Create leaderboard of algorithms
• 🦾 More algorithm implementations:
  • Muesli - Paper
  • DreamerV3 - Paper
  • R2D2 - Paper
  • Rainbow - Paper
• 🎮 Self-play 2-player Systems for board games.

Please do follow along as we develop this next phase!

Citing Stoix 📚

If you use Stoix in your work, please cite us:

@software{toledo2024stoix,
author = {Toledo, Edan},
doi = {10.5281/zenodo.10916258},
month = apr,
title = {{Stoix: Distributed Single-Agent Reinforcement Learning End-to-End in JAX}},
url = {https://github.com/EdanToledo/Stoix},
version = {v0.0.1},
year = {2024}
}

Acknowledgements 🙏

We would like to thank the authors and developers of Mava as this was essentially a port of their repo at the time of creation. This helped set up a lot of the infracstructure of logging, evaluation and other utilities.

See Also 🔎

Related JAX Libraries In particular, we suggest users check out the following repositories:

• 🦁 Mava: Distributed Multi-Agent Reinforcement Learning in JAX.
• 🔌 OG-MARL: datasets with baselines for offline MARL in JAX.
• 🌴 Jumanji: a diverse suite of scalable reinforcement learning environments in JAX.
• 😎 Matrax: a collection of matrix games in JAX.
• 🔦 Flashbax: accelerated replay buffers in JAX.
• 📈 MARL-eval: standardised experiment data aggregation and visualisation for MARL.
• 🦊 JaxMARL: accelerated MARL environments with baselines in JAX.
• 🌀 DeepMind Anakin for the Anakin podracer architecture to train RL agents at scale.
• ♟️ Pgx: JAX implementations of classic board games, such as Chess, Go and Shogi.
• 🔼 Minimax: JAX implementations of autocurricula baselines for RL.

Disclaimer: This is not an official InstaDeep product nor is any of the work putforward associated with InstaDeep in any official capacity.