Reinforcement Learning Handbook

Appearance

How to Choose Frameworks and Tools

On this page Stable-Baselines3, Ray RLlib, Tianshou, CleanRL, the JAX stack (Brax, PureJaxRL), and rolling your own — a comparison table of six options for different scenarios, plus the battle-tested engineering advice "copy CleanRL first, then modify."

How to Choose Frameworks and Tools ​

In a nutshell: this page helps you answer "which framework should I actually build my RL code on" — a side-by-side comparison of six options (Stable-Baselines3, Ray RLlib, Tianshou, CleanRL, Brax, and rolling your own), plus one piece of repeatedly validated engineering advice: "copy CleanRL first, then modify." By the end you'll have a concrete selection, instead of hopping back and forth between six READMEs.

First, a splash of cold water: framework choice matters far less than you think. The success of an RL project hinges on the environment, the reward, evaluation, and hyperparameters — not on which library you picked (the "simplicity first" rule in RL design principles says exactly this). What really determines your productivity is understanding the fundamentals — and the best vehicle for understanding is a clean, readable reference implementation. So this page's conclusion leans toward "CleanRL for learning, SB3 for production, RLlib/Brax for scale," rather than handing you a single "one true answer."

1. The Three Dimensions of Framework Choice ​

Every framework trades off along these three dimensions — figure out what you actually need first:

Dimension 1  Maturity    (How complete are the algorithm implementations? Are the
                         pits patched? How large is the community?)
Dimension 2  Flexibility (Can you modify internals? Roll a custom algorithm?)
Dimension 3  Parallelism (Single machine, single GPU → multi-machine distributed —
                         how far can it scale?)

Choosing = finding your spot on the Pareto frontier of these three
DimensionQuestions to askImplication
MaturityHow many production deployments? Is the issue tracker active? Have the algorithm implementations been vetted by the community?More mature means more peace of mind, but more "black box"
FlexibilityDo I need to change network architectures, define custom losses, invent new algorithms?More flexible means more effort, but more malleable
ParallelismIs a single process not enough? How many environments in parallel? Do I need GPU/multi-machine?More parallel means more complex

For the vast majority of personal projects and small-to-medium business projects, the parallelism requirement is "one GPU is plenty" — which eliminates most options, leaving essentially just "SB3 or CleanRL."

2. The Six Options, One by One ​

1. Stable-Baselines3 (SB3) — production-grade, works out of the box ​

  • In one sentence: The most mature, best-documented, most "foolproof" RL library in the PyTorch ecosystem.
  • Strengths: Unified API (PPO(env).learn()), broad built-in algorithm coverage (A2C/DQN/PPO/SAC/TD3/DDPG), high-quality docs and tutorials, seamless Gymnasium integration.
  • Good for: A business team that needs a reliable algorithm that "just runs"; teaching and comparison; your default starting point for every project.
  • The price: Poor flexibility — changing networks or losses means hacking internals; parallelism relies on SubprocVecEnv, which caps out at modest scale.
python
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

env = make_vec_env("CartPole-v1", n_envs=4)
model = PPO("MlpPolicy", env, learning_rate=3e-4, seed=0, verbose=1)
model.learn(total_timesteps=50_000)
model.save("ppo_cartpole")          # save the model

2. Ray RLlib — the distributed veteran ​

  • In one sentence: An industrial-strength framework built for "large-scale distributed training."
  • Strengths: Distributed scheduling on top of Ray, multi-machine multi-GPU support, massive environment parallelism, Tune for automated hyperparameter search; broad algorithm coverage.
  • Good for: Large-scale environment sampling, multi-machine clusters, teams already invested in the Ray ecosystem (Tune, Serve, Datasets).
  • The price: A steep learning curve, deep abstraction layers, and massive overkill for small projects; the API evolves fast, so online tutorials go stale quickly.

3. Tianshou — the academic's flexible "Swiss Army knife" ​

  • In one sentence: A high-level PyTorch RL library led by the Chinese research community, with flexibility approaching hand-written code.
  • Strengths: Modular design (policies / collectors / environments / trainers compose freely), extremely friendly to research newcomers, complete built-in algorithm coverage (including newer ones like QRDQN, IQN).
  • Good for: Research, ablations, and rapidly validating new ideas in academic/lab settings.
  • The price: Flexible = you must understand its abstractions; docs and community are smaller than SB3's.

4. CleanRL — readability-first teaching material and engineering blueprint ​

  • In one sentence: Compresses a complete PPO implementation into a single file plus a matching write-up — code as paper.
  • Strengths: Single Python file, zero abstraction, argparse for hyperparameters; every algorithm ships with an explanation doc aligned to the paper; full coverage of PPO/SAC/DQN/DDPG and more.
  • Good for: Learning the fundamentals (highly recommended), a self-research starting point that is "readable, modifiable, auditable," and paper reproduction.
  • The price: It's a "reference implementation," not a "general-purpose library" — no distributed training, no scheduling, you wire up logging yourself.
python
# CleanRL's PPO: single file, richly commented (skeleton excerpt)
# python cleanrl/ppo.py --env-id CartPole-v1 --total-timesteps 50000
class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        self.critic = nn.Sequential(...)
        self.actor = nn.Sequential(...)
    def get_value(self, x): ...
    def get_action_and_value(self, x, action=None): ...
# Main loop: rollout → GAE → update → logging, all laid flat in one file

5. Brax / PureJaxRL — GPU parallelism and the JAX ecosystem ​

  • In one sentence: Write your environments in differentiable JAX, run entire batches of them directly on GPU/TPU, and leave CPU sampling in the dust.
  • Strengths: Brax provides JAX versions of physics environments (Ant, Humanoid, HalfCheetah, etc.); PureJaxRL builds minimal pure-JAX PPO/SAC implementations on top, finishing massively parallel training within a few thousand lines.
  • Good for: Tight sample budgets, extreme-throughput ambitions, research/frontier engineering that wants "environment + algorithm all on GPU," and anyone seriously pushing training scale.
  • The price: JAX has a steep learning curve, a poor debugging experience (abstract error messages), and an ecosystem still evolving at high speed.

6. Rolling your own — the last resort, and the best teacher ​

  • In one sentence: Write the training loop, the algorithm, the replay buffer, and the logging entirely yourself.
  • Good for: Learning (every RL practitioner should hand-write it once — see the progressive tutorial); or requirements no existing framework can meet (custom environment protocols, exotic parallelism, differentiable environments).
  • The price: You step into every pit yourself — debugging costs are enormous, and you'll likely end up re-implementing an unmaintained SB3.

The right way to "roll your own"

Rolling your own ≠ reinventing every wheel from zero. Copy CleanRL first, then modify: take a readable reference implementation as the skeleton, and gradually replace parts with what you need. Writing everything yourself = spending two or three months on engineering details someone else already solved. This advice is the most important output of this page — see section 5.

3. The Six Options, Side by Side ​

DimensionSB3RLlibTianshouCleanRLBrax/PureJaxRLRoll your own
Language/backendPyTorchPyTorch/TFPyTorchPyTorchJAXAnything
Algorithm coverageBroad and stableBroadestBroad and currentMainstream algorithmsMostly PPO/SACUp to you
Works out of the box★★★★★★★★★★★★★★★★★★★★
Flexibility★★★★★★★★★★★★★★★★★★★★★★★
Distributed/parallelSingle-machine vectorMulti-machine, large scaleSingle-machineSingle-machineGPU-scaleUp to you
Learning costLowHighMediumLowHighHighest
Docs/communityBestGoodMediumExcellent (one doc per file)Medium—
Best forBusiness/beginnersLarge-scale teamsResearchersLearners/self-research starting pointScaling researchersGeeks/special needs

4. Practical Advice: Choose by Scenario ​

ScenarioRecommendationWhy
Learning the fundamentals, writing RL for the first timeCleanRL + SB3CleanRL for understanding every number, SB3 as a quick cross-check
The business team wants a policy that "runs, reliably"SB3Complete docs, few pits, low maintenance
Need to modify networks / reward shaping on top of SB3SB3 + custom MlpExtractor / callbacksMost changes fit SB3's extension points
Need large-scale parallelism / multi-machine trainingRLlib (CPU cluster) or Brax (GPU)Pick by hardware
Doing research, ablations, trying new ideasTianshouThe most ergonomic modularity
Deep customization (differentiable envs, custom protocols)Roll your own + copy CleanRLCopy first, modify after
Reproducing a specific paperThat paper's official implementationOfficial code is closer to the paper than any framework

An often-overlooked option: just use the paper's official implementation

To reproduce PPO/DQN/RLHF, the first choice is always the paper's repo (CleanRL often maintains a corresponding replica as well). Framework implementations are wrapped in abstraction, and many details (reward normalization, the exact advantage computation) may differ from the paper — this is exactly the origin of "check implementation differences before reproducing" in tuning and hyperparameter optimization.

5. The Boundary with Rolling Your Own: "Copy CleanRL First, Then Modify" ​

This advice deserves its own section. It has three progressive layers:

text
Layer 1  Copy:   Clone CleanRL's ppo.py and understand every line of it
Layer 2  Modify: Swap in your environment, your network, your reward, your logging
Layer 3  Grow:   Bring in components from SB3/RLlib (callbacks, vectorization,
                  logging) or migrate wholesale — only when needed

When should you graduate from "copying CleanRL" to "a bigger framework"?

Trigger signalAction
Environment sampling becomes the bottleneck (single-process CPU too slow)Add gymnasium.vector / SubprocVecEnv first; RLlib/Brax only if that's not enough
Need multi-machine trainingGo straight to RLlib or Brax — don't struggle on your own stack
Need integration with the company's scheduling/storage systemsStart with SB3 (easy ecosystem integration), or evaluate RLlib directly
The team will maintain this code for 2+ yearsPrefer SB3 (the community maintains it for you)
Just validating an idea, no long-term maintenanceCleanRL / your own fork is fine

The true cost of rolling your own

The biggest hidden cost of rolling your own isn't development — it's maintenance and reproducibility: six months later your "quick prototype" has grown into technical debt nobody understands, while the SB3 community has fixed hundreds of edge-case bugs for you. Unless you have a clear reason (learning, special protocols, performance), "copy CleanRL first, then modify; switch to SB3 the moment you can."

Extension-point examples: "controlled self-research" on top of SB3 ​

Many people think choosing SB3 means losing control. In fact, SB3's extension points cover most "I just want to change one thing" needs. Four levels, from lightest to heaviest:

LevelNeedSB3's entry point
1Change network architectureCustom MlpExtractor / features_extractor_class
2Change training-loop behaviorcallbacks (saving, evaluation, LR scheduling, custom logging)
3Change reward/observationGymnasium wrappers (RewardWrapper/ObservationWrapper)
4Change core algorithm logicSubclass PPO/SAC and override methods (at this point you're near "copy and modify")

A typical level-2 example — anneal the learning rate dynamically during training and evaluate periodically:

python
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback

class TrainCallback(BaseCallback):
    """After each rollout: linearly decay the learning rate + print eval return periodically."""
    def _on_step(self):
        if self.n_calls % self.model.n_steps == 0:
            # Linearly decay the learning rate to 1/10
            progress = self.num_timesteps / self.model.total_timesteps
            new_lr = 3e-4 * (1.0 - 0.9 * progress)
            self.model.lr_schedule = lambda _: new_lr
        return True

model = PPO("MlpPolicy", "CartPole-v1", learning_rate=3e-4, verbose=0)
model.learn(total_timesteps=100_000, callback=TrainCallback())

A typical level-1 example — give the policy a feature extractor (pass observations through a custom network first):

python
import torch.nn as nn
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor

class ObsNormExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space, features_dim=128):
        super().__init__(observation_space, features_dim)
        self.net = nn.Sequential(
            nn.Linear(observation_space.shape[0], 128), nn.ReLU(),
        )
    def forward(self, observations):
        return self.net(observations)

model = PPO("MlpPolicy", "CartPole-v1",
            policy_kwargs={"features_extractor_class": ObsNormExtractor})

When "extending SB3" turns into "must roll your own"

SB3's extension points solve 90% of needs; the remaining 10% concentrate on: needing a custom training loop (non-standard update rhythms), needing a differentiable environment (the environment participates in backprop), and needing custom multi-machine scheduling. When those show up, it's not too late to start "copying CleanRL and modifying."

6. Frameworks in Concert with the Toolchain ​

A framework is always just the middle layer. The full toolchain looks like this:

text
┌──────────────────────────────────────────────────────────────┐
│  Environment layer   Gymnasium (unified interface)           │
│                      ├─ MuJoCo (continuous control)          │
│                      ├─ Atari (pixel observations)           │
│                      ├─ Brax (JAX environments)              │
│                      └─ Custom environments                  │
├──────────────────────────────────────────────────────────────┤
│  Sampling layer      SubprocVecEnv / gymnasium.vector /      │
│                      multi-machine                           │
├──────────────────────────────────────────────────────────────┤
│  Algorithm layer     SB3 / RLlib / Tianshou / CleanRL / DIY  │
├──────────────────────────────────────────────────────────────┤
│  Config layer        Hydra (config management)               │
│  Logging layer       W&B / TensorBoard / CSV                 │
│  Tuning layer        Optuna (Bayesian optimization)          │
│  Evaluation layer    rliable / custom eval scripts           │
└──────────────────────────────────────────────────────────────┘

Choose the framework with the layers above and below in mind:

  1. Environment compatibility: Is your environment Gymnasium-interfaced? Brax environments are available via gymnasium.make (brax provides Gymnasium wrappers); Atari/MuJoCo both have Gymnasium integrations. Non-Gymnasium environments (some industrial simulators) instantly eliminate most framework options.
  2. Observation preprocessing: Do pixel observations (Atari) need frame_stack/grayscale? SB3 has AtariPreprocessing built in; CleanRL has ready-made flags — both are painless on image tasks.
  3. Toolchain integration: SB3 has official EvalCallback and tensorboard integration; CleanRL uses argparse + manual logging and can plug into W&B; RLlib has Tune for tuning and Serve for deployment built in.

Quick Selection Flowchart ​

text
Start
  ├─ Writing RL for the first time? --------------------→ CleanRL tutorials + SB3 cross-check
  ├─ Business needs a stable policy? ------------------→ SB3
  ├─ Single-machine parallelism not enough? -----------→ RLlib or Brax (by CPU/GPU)
  ├─ Doing research / ablations? ----------------------→ Tianshou
  └─ Want full control of everything? -----------------→ Copy CleanRL first, then decide whether to roll your own

7. One-Sentence Summary for Your Team ​

CleanRL for teaching, SB3 for production, RLlib/Brax for scale, and for special needs — copy CleanRL first, then modify.

The framework is scaffolding, not the work. Your work is the whole experimental system of "environment + reward + evaluation + tuning" — no framework can help you with that part; you have to build it yourself following the pipeline in build an RL project from scratch.

Further Reading ​

References ​