RL libraries: Gymnasium, Stable-Baselines3, CleanRL, RLlib and PufferLib with speed measurements

The algorithms of the previous lessons can be written in a hundred lines, but real work uses libraries: for environments (a standard interface between agent and simulator), for algorithms (tested, comparable implementations) and for speed (simulating millions of steps per second). This lesson presents the most used Python ecosystem: Gymnasium, the standard for environments; Stable-Baselines3, to use reliable algorithms in a few lines; CleanRL, RLlib and TorchRL; and PufferLib, a recent library that bets everything on simulation speed, with a working example and measurements made for this lesson.

Gymnasium: the standard interface

In 2016 OpenAI released Gym, a set of environments with a common interface; since 2022 the project has been maintained by the Farama Foundation under the name Gymnasium. Every environment has an observation space and an action space, and two functions:

import gymnasium

env = gymnasium.make("CartPole-v1")
print(env.observation_space, env.action_space)       # Box(4,) and Discrete(2)
obs, info = env.reset(seed=0)                         # initial state
total, done = 0.0, False
while not done:
    action = env.action_space.sample()                # the policy would go here
    obs, reward, terminated, truncated, info = env.step(action)
    total += reward
    done = terminated or truncated                    # true end or time limit
print("return of a random episode:", total)

terminated indicates a true end (the pole fell), truncated a time limit: as seen in the DQN lesson, they must be treated differently. Hundreds of environments rely on this interface: classic control, Atari games, MuJoCo physics simulation (walking robots), and third-party ones (autonomous driving, power grids, finance). PettingZoo is the equivalent for multi-agent environments.

Stable-Baselines3: ready-made algorithms

Stable-Baselines3 (SB3), based on PyTorch, offers tested implementations of PPO, A2C, DQN, SAC, TD3 and DDPG, with reasonable default hyperparameters. A few lines are enough to train PPO on CartPole:

import time
import gymnasium
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy

envs = make_vec_env("CartPole-v1", n_envs=8, seed=0)      # 8 environments in parallel
model = PPO("MlpPolicy", envs, seed=0, verbose=0)          # default hyperparameters
t0 = time.time()
model.learn(total_timesteps=100_000)
mean, std = evaluate_policy(model, gymnasium.make("CartPole-v1"), n_eval_episodes=10, deterministic=True)
print(f"training {time.time() - t0:.1f} s, evaluation over 10 episodes: {mean:.1f} ± {std:.1f}")
model.save("ppo_cartpole")

With SB3 2.9 training took 48 s and the evaluation gave 453.6 ± 60.7: a good result without touching anything, a little below the hand-written PPO of the previous lesson, whose hyperparameters had been chosen for CartPole. It is the typical trade-off: the library is slower and less tuned to the single problem, but it is reliable, documented and lets you change algorithm by changing one word (PPO → SAC).

The other libraries

• CleanRL: each algorithm in a single file of a few hundred lines, without abstractions. It is made for understanding and modifying, and it is the best to read after this course. The DQN hyperparameters of lesson 4 come from there.
• RLlib (Ray project): designed to scale on clusters of many machines, with many agents and many algorithms; more complex, suited to large projects.
• TorchRL: PyTorch’s official RL library, made of “building blocks” (buffers, data collectors, losses) to combine.

PufferLib: speed first

In RL almost all the time goes not into the neural network but into the simulator: for CartPole the network does a few thousand operations, while the agent needs hundreds of thousands of steps. PufferLib, developed by Joseph Suarez (PufferAI), starts from this observation and combines three things:

• emulation: it wraps any Gymnasium or PettingZoo environment and makes it look like a “flat” environment, with observations and actions in contiguous memory arrays, the same for all environments (even those with complicated, dictionary observations);
• vectorisation: it runs many environments in parallel with different backends (Serial in the same process, Multiprocessing over several processes, Ray over several machines), writing results into shared buffers without copies;
• Ocean: a collection of environments written in C, from simple games to large multi-agent simulations, that run hundreds of times faster than Python versions (and even more on the GPU).

Versions. On PyPI version 3.0 (June 2025, pip install pufferlib) offers the Python interface described here. Version 5.0, at the time this lesson is written, is installed from the PufferAI GitHub repository and goes even further: C and CUDA environments, training on the GPU with a modified PPO, command-line use (./build.sh breakout --cu, ./puffer train, ./puffer eval latest); training requires an NVIDIA GPU. The interface changes quickly between versions: always check the documentation on puffer.ai.

The example. The PPO of the previous lesson already uses PufferLib to vectorise: the Gymnasium environment is wrapped by GymnasiumPufferEnv and eight copies are created by pufferlib.vector.make. An instructive trap met while writing it: to avoid copying data, PufferLib rewrites the same buffer of observations and rewards at every step. Storing observations with torch.as_tensor, which does not copy, made all stored steps point to the last observation, and the agent learned nothing (return stuck around 25). With torch.tensor, which copies, the same code reaches 500. In general: with high-performance libraries, whatever you want to keep must be copied.

Measurements. Steps per second with random actions, on a PC with 16 logical processors (PufferLib 3.0, Gymnasium 0.29, Python 3.11):

Repeating the measurements, the values of the first three rows change by 20-30 % (in another run: 120,000 and 118,000 steps per second), that of Ocean by 2-3 %. For an environment as simple as CartPole, parallel processes help little: the cost of communication between processes is similar to that of a step. The real jump comes from the environment written in C, which performs 1024 steps in a loop without ever returning to Python: almost 300 times faster. This is why modern libraries (PufferLib, GPU simulators such as Isaac Lab, MJX and Brax) rewrite environments in compiled languages or on the GPU: at 30 million steps per second, a training run that would take a day takes a few minutes.

import time
import numpy as np
import gymnasium
import pufferlib.emulation, pufferlib.vector, pufferlib.ocean

def make(buf=None, seed=0):
    return pufferlib.emulation.GymnasiumPufferEnv(env=gymnasium.make("CartPole-v1"), buf=buf, seed=seed)

def measure(envs, n_env, steps=2000):
    envs.reset(seed=0)
    actions = np.random.randint(0, 2, size=(steps, n_env))
    t0 = time.time()
    for k in range(steps):
        envs.step(actions[k])
    return steps * n_env / (time.time() - t0)

envs = pufferlib.vector.make(make, num_envs=8, backend=pufferlib.vector.Serial)
print(f"Serial, 8 environments: {measure(envs, 8):,.0f} steps/s"); envs.close()
envs = pufferlib.vector.make(make, num_envs=8, num_workers=4, backend=pufferlib.vector.Multiprocessing)
print(f"Multiprocessing, 8 environments on 4 processes: {measure(envs, 8):,.0f} steps/s"); envs.close()

ocean = pufferlib.ocean.env_creator("puffer_cartpole")(num_envs=1024)   # CartPole written in C
ocean.reset(seed=0)
actions = np.random.randint(0, 2, size=(500, 1024))
t0 = time.time()
for k in range(500):
    ocean.step(actions[k])
print(f"Ocean, 1024 environments in C: {500 * 1024 / (time.time() - t0):,.0f} steps/s")

Which to choose

For learning and small experiments: Gymnasium and your own code, or CleanRL. For a project that must work right away with standard algorithms: Stable-Baselines3. To train agents on complex or multi-agent environments, when simulation time dominates: PufferLib or GPU simulators. For systems distributed over several machines: RLlib. In all cases the Gymnasium interface is the meeting point: an environment written for Gymnasium can be used with all of them.