Skip to content

kwk2696/sb3-jax-haiku

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

70 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Stable Baslines with JAX & Haiku

Implementation of Stable Baselines based on JAX & Haiku.

This library is based on Stable Baselines 3 (https://github.com/DLR-RM/stable-baselines3).

Implemented Algorithms

Name Online_learning Box Discrete MultiDiscrete MultiBinary
BC ✔️ ✔️
OnlineBC ✔️ ✔️ ✔️
DT ✔️
DU ✔️ ✔️
SAC ✔️ ✔️ ✔️
PPO ✔️ ✔️ ✔️

Install

git clone https://github.com/kwk2696/sb3-jax-haiku.git
pip install -e .

Benchmark

We use Intel i9-10940X, RTX 3090 to benchmark Decision Transformer (DT) on MuJoCo Ant environment.

Name sb3-torch sb3-jax-haiku
SAC 163 step / sec 236 step / sec
DT 0.03 step / sec 3 step / sec

Example

Example codes are available in tests directory.

Currently Working On ...

TD3, Generative modeling for RL (e.g. diffuser)