mansicer/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 -r requirements.txt
pip install git+https://github.com/deepmind/dm-haiku

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