BepfCp/baselax

Baselax (Baselines + JAX) provides stable-baselines-style implementations of reinforcement learning (RL) algorithms with Google JAX framework.

Status: in development, no practical use.

Baselax

Baselax (Baselines + jax) provides stable-baselines-style implementations of reinforcement learning (RL) algorithms with Google JAX framework, supported by the following frameworks:

• Google JAX
• DeepMind Optax
• DeepMind Haiku
• DeepMind RLax

Installation

Install via pip:

pip install -e .

# optional
pip install -e .["sb3"]
pip install -e .["all"]

Additnionally, you can download and install envpool wheels from GitHub releases for better sample efficiency.