Jax-Baseline

Jax-Baseline is a Reinforcement Learning implementation using JAX and Flax/Haiku libraries, mirroring the functionality of standard baselines.

Features

• 2-3 times faster than previous Torch and Tensorflow implementations
• Optimized using JAX's Just-In-Time (JIT) compilation
• Flexible solution for Gym and Unity ML environments

Installation

pip install -r requirement.txt
pip install .

Implementation Status

• ✔️ : Optional implemented
• ✅ : Defualt implemented at papers
• ❌ : Not implemeted yet or can not implemented
• 💤 : Implemented but didn't update a while (can not guarantee working well now)

Supported Environments

Name	Q-Net based	Actor-Critic based	DPG based
Gymnasium	✔️	✔️	✔️
MultiworkerGym with Ray	✔️	✔️	✔️
Unity-ML Environments	💤	💤	💤