sbx-rl 0.21.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)
• Simplicity Bias for Scaling Up Parameters in Deep Reinforcement Learning (SimBa)

Example

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

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