araffin/sbx

SBX: Stable Baselines Jax (SB3 + Jax) RL algorithms

What it solves

SBX is a proof-of-concept implementation of Stable-Baselines3 (SB3) using the Jax framework. It provides a set of reliable reinforcement learning (RL) algorithms that are optimized for performance and compatible with the SB3 API, allowing users to switch between PyTorch-based SB3 and Jax-based SBX without changing their high-level code structure.

How it works

By rewriting the SB3 algorithms in Jax, the project leverages Jax's hardware acceleration and functional programming model. It implements several key RL algorithms including PPO, DQN, SAC, TD3, and TQC, as well as specialized configurations like DroQ and SimBa. It maintains API compatibility with SB3, which means it can be integrated into existing tools like the RL Zoo for training and evaluation.

Who it’s for

Researchers and developers working in reinforcement learning who want the performance benefits of Jax while maintaining the same user-friendly API as Stable-Baselines3.

Highlights

  • SB3 API Compatibility: Seamlessly integrates with the RL Zoo and other SB3-compatible tools.
  • Detailed Algorithm Suite: Includes a wide range of algorithms such as Soft Actor-Critic (SAC), Proximal Policy Optimization (PPO), and Deep Q Network (DQN).
  • Jax-powered: Built on Jax for improved computational efficiency.
  • Advanced RL Techniques: Supports parameter resets for off-policy algorithms and specialized architectures like SimBa.

Related

  • Project
    DLR-RM/stable-baselines3

    A set of reliable PyTorch implementations of reinforcement learning algorithms designed to provide stable baselines for research and industry.

  • Project
    Stable-Baselines-Team/stable-baselines3-contrib

    A contribution package for Stable-Baselines3 containing experimental reinforcement learning algorithms and niche tools for researchers and developers.

  • Project
    DLR-RM/rl-baselines3-zoo

    A training framework for Stable Baselines3 that provides scripts for training, evaluating, and tuning RL agents, along with a collection of pre-tuned hyperparameters for common environments.

  • Project
    RobertTLange/gymnax

    A collection of reinforcement learning environments reimplemented in JAX to enable high-throughput, vectorized experiments on GPUs and TPUs.

  • Project
    danijar/dreamerv3

    A scalable reinforcement learning algorithm that uses world models to master diverse control tasks with a fixed set of hyperparameters.