lit-sb3

lit-sb3 is a PyTorch Lightning adaptation of Stable-Baselines3. It preserves the familiar SB3 algorithm API where practical, while moving training orchestration into Lightning so reinforcement learning runs can use Lightning callbacks, loggers, checkpointing, precision settings, and distributed strategies.

Features

• Stable-Baselines3-style reinforcement learning algorithms implemented on top of PyTorch Lightning
• Lightning-native training through RLTrainer, a subclass of lightning.Trainer
• Environment management through RLEnvDataModule, a LightningDataModule for RL stages
• Backwards-friendly model.learn(total_timesteps=..., env=...) shim for quick starts
• Support for Lightning callbacks, loggers, checkpointing, mixed precision, and accelerator strategies
• Notebook-friendly workflows, MkDocs documentation, and example notebooks in this repository

Installation

With pip:

python -m pip install lit-sb3

With uv:

uv add lit-sb3

Optional extras are available for specific workflows and environment families:

uv add lit-sb3[hub,opt,viz,atari,box2d,classic-control,mujoco,toy-text]

How to use it

Most of the library keeps the familiar sklearn-like style used by Stable-Baselines3.

Here is a quick example of how to train and run PPO on a cartpole environment:

import gymnasium as gym

from lit_sb3 import PPO

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

model = PPO("MlpPolicy", verbose=1)
model.learn(total_timesteps=10_000, env=env) # env needs to be passed here instead of algorithm

vec_env = model.get_env()
obs = vec_env.reset()
for _ in range(1000):
    action, _states = model.predict(obs, deterministic=True)
    obs, reward, done, info = vec_env.step(action)
    vec_env.render()
    # VecEnv resets automatically
    # if done:
    #   obs = env.reset()

env.close()

Lightning-native workflow

The preferred interface for anything beyond a minimal example is the Lightning-native workflow.

import gymnasium as gym

from lit_sb3 import PPO, RLEnvDataModule, RLTrainer

model = PPO("MlpPolicy")
datamodule = RLEnvDataModule(
    train_env=gym.make("CartPole-v1"),
    test_env=gym.make("CartPole-v1"),
)
trainer = RLTrainer(max_steps=10_000, run_name="ppo_cartpole")

trainer.fit(model, datamodule=datamodule)
trainer.test(model, datamodule=datamodule)

Docs

uv run mkdocs build -f ./mkdocs.yml -d ./_build/

Update template

copier update --trust -A --vcs-ref=HEAD

Credits

This project was generated with