# SBX runs Stable Baselines3 algorithms on Jax, and calls itself a proof of concept

> SBX is a Jax rewrite of the Stable Baselines3 reinforcement learning algorithms by the same maintainers. It keeps the SB3 API, ships SAC, TQC, DQN, PPO, TD3, DDPG, CrossQ, DroQ and SimBa, and still depends on the PyTorch based SB3 package.

**araffin/sbx** — SBX: Stable Baselines Jax (SB3 + Jax) RL algorithms

- Repository: https://github.com/araffin/sbx
- Stars: 613 · Forks: 61
- Language: Python
- License: MIT
- Published: 2026-10-22 · Updated: 2026-10-22 · Language: en
- Canonical page: https://hysenlabs.com/projects/araffin-sbx

## The project calls itself a proof of concept, and that shapes everything else

SBX stands for Stable Baselines Jax, and the project's own first line describes it as a proof of concept version of Stable Baselines3 in Jax. That framing is worth taking seriously rather than reading past, because it tells you what the code is for: exploring a Jax based reimplementation of algorithms that already exist in a mature PyTorch library, not shipping a drop-in replacement that the maintainers promise to keep in step with upstream.

The implemented list is long. It covers Soft Actor-Critic and SAC-N, Truncated Quantile Critics, DroQ, Proximal Policy Optimization, Deep Q Network, Twin Delayed DDPG, Deep Deterministic Policy Gradient, CrossQ and SimBa. Each of those corresponds to a named paper, and the repository links the arXiv or OpenReview entry for each one.

What makes the port attractive is the API shape. Because SBX shares the SB3 API, scripts written against Stable Baselines3 read the same way here: a policy string, an environment, a learn call with total_timesteps, a predict call. The code below in the RL Zoo section shows how far that compatibility goes, since the zoo's algorithm map is simply overwritten.

## Jax and Flax arrive, but so does the PyTorch based SB3

The dependency list in setup.py is worth reading before you install anything, because it is not a pure Jax stack. The declared requirements include stable_baselines3 in the range >=2.9.0,<3.0, jax in the range >=0.4.24,<0.12.0, and then flax, optax, tqdm and rich.

The stable_baselines3 pin is the interesting part. Stable Baselines3 is the PyTorch implementation, so a machine running SBX ends up with both frameworks installed even though the algorithms themselves are executed through Jax, Flax and Optax. A team whose reason for looking at SBX is to avoid PyTorch in a build will not get that from this package.

The Jax range is generous but bounded, and a mismatch surfaces at install time rather than at runtime. The package is published as sbx-rl, with the version read at build time from a version.txt file inside the package directory rather than hard coded in setup.py.

## Two ways to install, one for the master branch and one for the published package

The project gives two installation routes and nothing more elaborate. For the latest master version it points at the git repository, and for everyone else it points at the package index.

```
pip install git+https://github.com/araffin/sbx
```

```
pip install sbx-rl
```

The difference in practice is which code you get. The git route tracks the master branch, which is where the most recent work lands, while the package route gives you a tagged release. The recent tags include v0.28.1 and v0.28.0 on 2026-07-24, and v0.27.0 on 2026-06-15, the last of which widened the accepted Jax version range.

For development, the Makefile exposes the same linting and typing rules the project uses on itself. Formatting is Black based with a line length of 127, import sorting is handled by Ruff, and typing is checked with mypy over the sbx/, tests/ and setup.py paths.

## A first training run on Pendulum-v1, then a rollout

The shortest complete example creates a Gymnasium environment with rendering enabled, trains a TQC agent with an MLP policy for ten thousand timesteps, and then replays the learned policy in a loop.

```python
import gymnasium as gym

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

env = gym.make("Pendulum-v1", render_mode="human")

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

vec_env = model.get_env()
obs = vec_env.reset()
for _ in range(1000):
    vec_env.render()
    action, _states = model.predict(obs, deterministic=True)
    obs, reward, done, info = vec_env.step(action)
```

Two details in that snippet matter. The agent is constructed with the same three part signature used by SB3, a policy name, an environment and keyword arguments, and the environment is passed as an object rather than as a string id, which is the older style still used in the setup.py description. Second, the rollout goes through model.get_env() and vec_env.step, so the action returned by predict is applied to the vectorised environment that the model wraps, and rendering is a call on that wrapper rather than on the raw environment.

## DroQ and SimBa are configurations layered onto SAC and TQC

Two entries on the implemented list are not standalone algorithms in the way DQN or PPO are, and the project explains how to reach them. DroQ is a special configuration of SAC, obtained by passing the paper's hyperparameters through the RL Zoo config format.

```yaml
HalfCheetah-v4:
  n_timesteps: !!float 1e6
  policy: 'MlpPolicy'
  learning_starts: 10000
  gradient_steps: 20
  policy_delay: 20
  policy_kwargs: "dict(dropout_rate=0.01, layer_norm=True)"
```

The dropout_rate and layer_norm entries inside policy_kwargs are what turn ordinary SAC into the Dropout Q-Functions variant. The project adds two practical notes: raising the q-value function learning rate with qf_learning_rate: !!float 1e-3 helps, and when the same configuration is used with CrossQ you must set layer_norm=False because CrossQ already applies batch normalization.

SimBa works the same way but through the network architecture rather than the hyperparameters. The recommended setup shown in the project uses a policy named SimbaPolicy with a residual block of 128 units for the actor and 256 by 256 for the critics, two critics, AdamW through optax, and observation normalisation through VecNormalize with norm_obs set to True and norm_reward set to False.

## The RL Zoo path works by overwriting its algorithm map

Since SBX shares the SB3 API, the RL Zoo can drive it, and the mechanism is a dictionary override rather than a plugin system. The script imports the zoo modules, replaces the entries in rl_zoo3.ALGOS with the SBX classes, and then reassigns the same mapping onto rl_zoo3.train.ALGOS and rl_zoo3.exp_manager.ALGOS so that the zoo's experiment manager resolves the same classes.

```python
import rl_zoo3
import rl_zoo3.train
from rl_zoo3.train import train
from sbx import DDPG, DQN, PPO, SAC, TD3, TQC, CrossQ

rl_zoo3.ALGOS["ddpg"] = DDPG
rl_zoo3.ALGOS["dqn"] = DQN
# See note below to use DroQ configuration
# rl_zoo3.ALGOS["droq"] = DroQ
rl_zoo3.ALGOS["sac"] = SAC
rl_zoo3.ALGOS["ppo"] = PPO
rl_zoo3.ALGOS["td3"] = TD3
rl_zoo3.ALGOS["tqc"] = TQC
rl_zoo3.ALGOS["crossq"] = CrossQ
```

Once that script exists, it is invoked the way a zoo script normally is, with the algorithm and environment given on the command line.

```
python train.py --algo sac --env HalfCheetah-v4 -params train_freq:4 gradient_steps:4 -P
```

An enjoy script follows the same shape with rl_zoo3.enjoy.ALGOS and the enjoy entry point. Partial benchmark results for SBX are published on the OpenRL Benchmark wandb project, with several reports linked from there.

## Parameter resets exist for long off-policy runs, and the project keeps moving

One mechanism is not obvious from the API and deserves attention if you plan long runs. For off-policy algorithms, parameter resets can be activated by passing a list of timesteps to the model constructor, for example param_resets=[int(1e5), int(5e5)], which resets parameters and optimizers after 100,000 and 500,000 timesteps. This addresses the well known failure mode where an off-policy replay buffer fills with stale data early in training, and it is opt in rather than default.

On maintenance, the repository is not archived and the last push was on 2026-10-19, three days before this article's date, so the code is moving quickly. The tag history shows a rapid patch cadence on 2026-07-24 with both v0.28.0 and v0.28.1, where the latter was a fix for the version in the file.

The licence is MIT, which is permissive enough for commercial and research use without the obligations that copyleft licences impose. Attribution in publications should cite the Stable-Baselines3 paper, and the repository carries a CITATION.bib file with the bibtex entry, a NOTICE file and a CODE_OF_CONDUCT.md.

## Conclusion

SBX suits researchers who already know the SB3 API and want the same algorithms running through Jax and Flax without rewriting their training scripts. It is the wrong dependency if you are trying to drop PyTorch from a stack, since the package still requires stable_baselines3, and it is the wrong choice for production work because the project describes itself as a proof of concept. Verify first that your pinned Jax version falls inside the declared range, that the algorithm you need is on the implemented list, and that the DroQ and SimBa configurations you plan to copy actually match your environment setup.

## FAQ

### How do I install SBX?

Run pip install sbx-rl for the published package, or pip install git+https://github.com/araffin/sbx for the latest master version of the repository.

### Which reinforcement learning algorithms does SBX implement?

The implemented list covers SAC and SAC-N, TQC, DroQ, PPO, DQN, TD3, DDPG, CrossQ and SimBa. DroQ is a configuration of SAC and SimBa is a network architecture for off-policy algorithms.

### How do I run the DroQ configuration with the RL Zoo?

DroQ is a special configuration of SAC, so you use the algo sac with a config file that sets learning_starts, gradient_steps, policy_delay and a policy_kwargs string containing dropout_rate and layer_norm, and you invoke python train.py --algo sac --env HalfCheetah-v4 -c droq.yml -P.

### How do I run the SBX test suite?

The Makefile pytest target runs ./scripts/run_tests.sh. Tests are marked with an expensive marker that can be deselected with -m "not expensive" when you want a faster local loop.

## Sources

- [araffin/sbx on GitHub](https://github.com/araffin/sbx)
- [Issues](https://github.com/araffin/sbx/issues)
- [License: MIT](https://github.com/araffin/sbx/blob/master/LICENSE)
- [README](https://github.com/araffin/sbx/blob/master/README.md)
- [Releases](https://github.com/araffin/sbx/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/araffin-sbx