sbx-rl 0.15.0

Jax version of Stable Baselines, implementations of reinforcement learning algorithms.

Stable Baselines Jax (SB3 + JAX = SBX)

See https://github.com/araffin/sbx

Proof of concept version of Stable-Baselines3 in Jax.

Implemented algorithms:

• Soft Actor-Critic (SAC) and SAC-N
• Truncated Quantile Critics (TQC)
• Dropout Q-Functions for Doubly Efficient Reinforcement Learning (DroQ)
• Proximal Policy Optimization (PPO)
• Deep Q Network (DQN)
• Twin Delayed DDPG (TD3)
• Deep Deterministic Policy Gradient (DDPG)
• Batch Normalization in Deep Reinforcement Learning (CrossQ)

Example

from sbx import DDPG, DQN, PPO, SAC, TD3, TQC, DroQ, CrossQ

model = TQC("MlpPolicy", "Pendulum-v1", verbose=1)
model.learn(total_timesteps=10_000, progress_bar=True)