Reinforcement Learning with Gymnasium and Stable-Baselines3

Table of Contents

Introduction

Prerequisites

Core RL Concepts

Getting Started with Gymnasium

Deep Q-Network (DQN)

Proximal Policy Optimization (PPO)

Advantage Actor-Critic (A2C)

Creating a Custom Environment

Training Loop and Callback Monitoring

Model Evaluation

TensorBoard Integration

Best Practices

Conclusion

---

Introduction

Reinforcement Learning (RL) is a paradigm of machine learning where an agent learns to make decisions by interacting with an environment. Unlike supervised learning, RL does not require labeled data. Instead, the agent discovers optimal behavior through trial and error, guided by reward signals.

This tutorial provides a hands-on guide to building RL agents using Gymnasium (the maintained fork of OpenAI Gym) for environment simulation and Stable-Baselines3 (SB3) for state-of-the-art RL algorithm implementations. You will learn to train agents with DQN, PPO, and A2C, create custom environments, monitor training, and evaluate model performance.

---

Prerequisites

• Python 3.9+
• Basic understanding of neural networks
• Familiarity with PyTorch (helpful but not required)

Install the required packages:

pip install gymnasium stable-baselines3[extra] tensorboard numpy matplotlib

The [extra] flag installs additional dependencies including TensorBoard support and extra Gymnasium environments.

---

Core RL Concepts

Before writing any code, it is essential to understand the fundamental components of an RL system.

The RL Framework

Agent <---> Environment
  |              |
  |-- Action --> |
  |              |-- State, Reward -->|
  |<------------------------------ ---|

• Agent: The learner and decision-maker that selects actions.
• Environment: The world the agent interacts with. It receives actions and returns states and rewards.
• State (Observation): A representation of the current situation in the environment.
• Action: A decision made by the agent that affects the environment.
• Reward: A scalar feedback signal indicating how good the action was.
• Policy: A mapping from states to actions. This is what the agent learns.
• Episode: A complete sequence from the initial state to a terminal state.
• Discount Factor (gamma): A value between 0 and 1 that determines how much future rewards are valued relative to immediate ones.

Value Functions

• V(s): The expected cumulative reward starting from state s, following policy pi.
• Q(s, a): The expected cumulative reward starting from state s, taking action a, then following policy pi.

Exploration vs. Exploitation

The agent must balance exploring new actions (to discover potentially better strategies) with exploiting known good actions (to maximize reward). This is managed through strategies like epsilon-greedy or entropy regularization.

---

Getting Started with Gymnasium

Gymnasium provides standardized interfaces for RL environments.

Exploring an Environment

import gymnasium as gym
import numpy as np

Create the CartPole environment
env = gym.make("CartPole-v1", rendermode="human")


print(f"Observation space: {env.observationspace}")
print(f"Action space: {env.actionspace}")

print(f"Observation shape: {env.observationspace.shape}")
print(f"Number of actions: {env.actionspace.n}")


Run a random agent
obs, info = env.reset(seed=42)
totalreward = 0

for step in range(500):
    action = env.actionspace.sample()  # random action

    obs, reward, terminated, truncated, info = env.step(action)
    totalreward += reward

    if terminated or truncated:
        print(f"Episode ended at step {step + 1} with total reward: {totalreward}")

        obs, info = env.reset()
        totalreward = 0

env.close()

Common Gymnasium Environments

# Discrete action spaces
envcartpole = gym.make("CartPole-v1")          # Balance a pole on a cart

envlunar = gym.make("LunarLander-v3")           # Land a spacecraft
envmountain = gym.make("MountainCar-v0")         # Drive up a mountain


Continuous action spaces
envpendulum = gym.make("Pendulum-v1")            # Swing a pendulum upright
envbipedal = gym.make("BipedalWalker-v3")         # Walk with a 2D biped


Inspect any environment
for name, env in [("CartPole", envcartpole), ("LunarLander", envlunar)]:

    print(f"\n{name}:")
    print(f"  Obs space:    {env.observationspace}")
    print(f"  Action space: {env.actionspace}")

    env.close()

---

Deep Q-Network (DQN)

DQN approximates the Q-function using a neural network. It is designed for environments with discrete action spaces.

Training a DQN Agent

from stablebaselines3 import DQN
from stablebaselines3.common.evaluation import evaluatepolicy

Create the environment
env = gym.make("CartPole-v1")

Initialize the DQN agent
modeldqn = DQN(

    policy="MlpPolicy",
    env=env,
    learningrate=1e-3,
    buffersize=100000,        # replay buffer size
    learningstarts=1000,       # steps before training starts

    batchsize=64,
    tau=1.0,                    # target network update coefficient
    gamma=0.99,                 # discount factor
    trainfreq=4,               # update the model every 4 steps

    targetupdateinterval=1000,# update target network every 1000 steps

    explorationfraction=0.1,   # fraction of training for exploration decay
    explorationinitialeps=1.0,
    explorationfinaleps=0.05,
    verbose=1,
    seed=42,
)

Train the agent
modeldqn.learn(totaltimesteps=50000)


Evaluate the trained agent
meanreward, stdreward = evaluatepolicy(modeldqn, env, nevalepisodes=20)

print(f"DQN - Mean reward: {meanreward:.2f} +/- {stdreward:.2f}")


Save the model
modeldqn.save("dqncartpole")

env.close()

Loading and Using a Trained Model

# Load the saved model
loadedmodel = DQN.load("dqncartpole")


env = gym.make("CartPole-v1", rendermode="human")
obs, info = env.reset()

for  in range(1000):

    action, states = loadedmodel.predict(obs, deterministic=True)

    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        obs, info = env.reset()

env.close()

---

Proximal Policy Optimization (PPO)

PPO is a policy gradient method that is widely considered the default algorithm for RL. It works with both discrete and continuous action spaces and is known for its stability and sample efficiency.

Training a PPO Agent

from stablebaselines3 import PPO

env = gym.make("LunarLander-v3")

modelppo = PPO(

    policy="MlpPolicy",
    env=env,
    learningrate=3e-4,
    nsteps=2048,               # steps per rollout

    batchsize=64,
    nepochs=10,                # optimization epochs per update

    gamma=0.999,
    gaelambda=0.98,            # GAE parameter
    cliprange=0.2,             # PPO clipping parameter

    entcoef=0.01,              # entropy coefficient for exploration
    vfcoef=0.5,                # value function coefficient

    maxgradnorm=0.5,

    verbose=1,
    seed=42,
)

modelppo.learn(totaltimesteps=200000)

meanreward, stdreward = evaluatepolicy(modelppo, env, nevalepisodes=20)
print(f"PPO - Mean reward: {meanreward:.2f} +/- {stdreward:.2f}")

modelppo.save("ppolunarlander")
env.close()

PPO with Continuous Actions

envcontinuous = gym.make("Pendulum-v1")

modelppocont = PPO(

    policy="MlpPolicy",
    env=envcontinuous,
    learningrate=3e-4,

    nsteps=1024,
    batchsize=64,

    nepochs=10,
    gamma=0.98,
    gaelambda=0.95,

    cliprange=0.2,
    verbose=1,
)

modelppocont.learn(totaltimesteps=100000)
meanreward, stdreward = evaluatepolicy(modelppocont, envcontinuous, nevalepisodes=20)
print(f"PPO Continuous - Mean reward: {meanreward:.2f} +/- {stdreward:.2f}")
envcontinuous.close()

---

Advantage Actor-Critic (A2C)

A2C is a synchronous variant of the Asynchronous Advantage Actor-Critic (A3C). It uses an actor network (policy) and a critic network (value function) trained simultaneously.

Training an A2C Agent

from stablebaselines3 import A2C

env = gym.make("CartPole-v1")

modela2c = A2C(

    policy="MlpPolicy",
    env=env,
    learningrate=7e-4,
    nsteps=5,                  # steps per environment update

    gamma=0.99,
    gaelambda=1.0,
    entcoef=0.0,

    vfcoef=0.5,
    maxgradnorm=0.5,
    usermsprop=True,
    normalizeadvantage=False,

    verbose=1,
    seed=42,
)

modela2c.learn(totaltimesteps=50000)

meanreward, stdreward = evaluatepolicy(modela2c, env, nevalepisodes=20)
print(f"A2C - Mean reward: {meanreward:.2f} +/- {stdreward:.2f}")

modela2c.save("a2ccartpole")
env.close()

---

Creating a Custom Environment

Real-world RL problems often require custom environments. Gymnasium provides a clear interface for building them.

Custom Grid World Environment

import gymnasium as gym
from gymnasium import spaces
import numpy as np

class GridWorldEnv(gym.Env):
    """A simple grid world where the agent must reach a goal position."""

    metadata = {"rendermodes": ["human", "ansi"], "renderfps": 4}

    def init(self, rendermode=None, gridsize=5):
        super().init()
        self.gridsize = gridsize
        self.rendermode = rendermode

        # Action space: 0=up, 1=down, 2=left, 3=right
        self.actionspace = spaces.Discrete(4)


        # Observation space: agent position (row, col)
        self.observationspace = spaces.Box(
            low=0, high=gridsize - 1, shape=(2,), dtype=np.float32

        )

        self.goalposition = np.array([gridsize - 1, gridsize - 1], dtype=np.float32)
        self.agentposition = None

        self.maxsteps = gridsize  gridsize  2
        self.currentstep = 0


    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.agentposition = np.array([0, 0], dtype=np.float32)
        self.currentstep = 0

        return self.agentposition.copy(), {}

    def step(self, action):
        self.currentstep += 1


        # Apply movement
        if action == 0 and self.agentposition[0] > 0:
            self.agentposition[0] -= 1  # up

        elif action == 1 and self.agentposition[0] < self.gridsize - 1:

            self.agentposition[0] += 1  # down
        elif action == 2 and self.agentposition[1] > 0:

            self.agentposition[1] -= 1  # left
        elif action == 3 and self.agentposition[1] < self.gridsize - 1:
            self.agentposition[1] += 1  # right


        # Calculate reward
        distance = np.linalg.norm(self.agentposition - self.goalposition)

        reachedgoal = distance < 0.5

        if reachedgoal:

            reward = 100.0
        else:
            reward = -1.0  # step penalty to encourage efficiency

        terminated = reachedgoal
        truncated = self.currentstep >= self.maxsteps

        return self.agentposition.copy(), reward, terminated, truncated, {}


    def render(self):
        if self.rendermode == "ansi":
            grid = [["." for  in range(self.gridsize)] for  in range(self.gridsize)]
            r, c = int(self.agentposition[0]), int(self.agentposition[1])
            grid[r][c] = "A"
            gr, gc = int(self.goalposition[0]), int(self.goalposition[1])
            grid[gr][gc] = "G" if (r, c) != (gr, gc) else "W"
            return "\n".join(" ".join(row) for row in grid)

Registering and Using the Custom Environment

from gymnasium.envs.registration import register

register(
    id="GridWorld-v0",
    entrypoint="main:GridWorldEnv",

    maxepisodesteps=200,

)

Verify the environment
from stablebaselines3.common.envchecker import checkenv

env = GridWorldEnv(gridsize=5)

checkenv(env, warn=True)
print("Environment check passed.")

Train PPO on the custom environment
env = GridWorldEnv(gridsize=5)

model = PPO("MlpPolicy", env, verbose=1, learningrate=3e-4, nsteps=512)

model.learn(totaltimesteps=50000)


meanreward, stdreward = evaluatepolicy(model, env, nevalepisodes=50)
print(f"Custom Env - Mean reward: {meanreward:.2f} +/- {stdreward:.2f}")

---

Training Loop and Callback Monitoring

Callbacks let you monitor training progress, save checkpoints, and implement early stopping.

Essential Callbacks

from stablebaselines3.common.callbacks import (
    EvalCallback,
    CheckpointCallback,
    CallbackList,
    StopTrainingOnRewardThreshold,
)

env = gym.make("LunarLander-v3")
evalenv = gym.make("LunarLander-v3")

Stop training when mean reward reaches 200
stopcallback = StopTrainingOnRewardThreshold(rewardthreshold=200, verbose=1)

Evaluate every 10,000 steps and save the best model
evalcallback = EvalCallback(

    evalenv,
    bestmodelsavepath="./bestmodel/",
    logpath="./evallogs/",
    evalfreq=10000,
    nevalepisodes=10,
    deterministic=True,
    callbackaftereval=stopcallback,

    verbose=1,
)

Save checkpoints every 20,000 steps
checkpointcallback = CheckpointCallback(
    savefreq=20000,
    savepath="./checkpoints/",

    nameprefix="ppolunar",

)

Combine callbacks
callback = CallbackList([evalcallback, checkpointcallback])


model = PPO("MlpPolicy", env, verbose=1)
model.learn(totaltimesteps=500000, callback=callback)


env.close()
evalenv.close()

Custom Callback

from stablebaselines3.common.callbacks import BaseCallback

class TrainingMetricsCallback(BaseCallback):
    """Logs custom metrics during training."""

    def init(self, verbose=0):
        super().init(verbose)
        self.episoderewards = []
        self.episodelengths = []


    def onstep(self) -> bool:

        # Check if an episode has ended
        if self.locals.get("dones") is not None:
            for i, done in enumerate(self.locals["dones"]):
                if done:
                    info = self.locals["infos"][i]
                    if "episode" in info:
                        self.episoderewards.append(info["episode"]["r"])
                        self.episodelengths.append(info["episode"]["l"])


                        if len(self.episoderewards) % 50 == 0:
                            recent = self.episoderewards[-50:]

                            print(f"  Last 50 episodes - Mean: {np.mean(recent):.1f}, "
                                  f"Min: {np.min(recent):.1f}, Max: {np.max(recent):.1f}")
        return True  # return False to stop training

    def ontrainingend(self):
        print(f"\nTraining complete. Total episodes: {len(self.episoderewards)}")

---

Model Evaluation

Comprehensive Evaluation

from stablebaselines3.common.evaluation import evaluatepolicy
import matplotlib.pyplot as plt

def comprehensiveevaluation(model, envid, nepisodes=100):
    """Run a thorough evaluation of a trained RL model."""
    env = gym.make(envid)

    episoderewards = []
    episodelengths = []


    for ep in range(nepisodes):
        obs, info = env.reset()
        totalreward = 0

        steps = 0

        while True:
            action,  = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            totalreward += reward

            steps += 1

            if terminated or truncated:
                break

        episoderewards.append(totalreward)

        episodelengths.append(steps)

    env.close()

    rewards = np.array(episoderewards)

    lengths = np.array(episodelengths)

    print(f"=== Evaluation over {nepisodes} episodes ===")

    print(f"Reward  - Mean: {rewards.mean():.2f}, Std: {rewards.std():.2f}, "
          f"Min: {rewards.min():.2f}, Max: {rewards.max():.2f}")
    print(f"Length  - Mean: {lengths.mean():.1f}, Std: {lengths.std():.1f}")

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    axes[0].hist(rewards, bins=30, color="steelblue", edgecolor="black")
    axes[0].axvline(rewards.mean(), color="red", linestyle="--", label=f"Mean: {rewards.mean():.1f}")
    axes[0].settitle("Reward Distribution")
    axes[0].setxlabel("Total Reward")

    axes[0].legend()

    axes[1].plot(rewards, alpha=0.6)
    axes[1].axhline(rewards.mean(), color="red", linestyle="--")
    axes[1].settitle("Reward per Episode")
    axes[1].setxlabel("Episode")

    axes[1].setylabel("Total Reward")

    plt.tightlayout()

    plt.savefig("rlevaluation.png", dpi=150)
    plt.show()

    return rewards, lengths

Usage
modelppo = PPO.load("ppolunarlander")
rewards, lengths = comprehensiveevaluation(modelppo, "LunarLander-v3", nepisodes=100)

---

TensorBoard Integration

Stable-Baselines3 has built-in TensorBoard logging.

Setting Up TensorBoard Logging

from stablebaselines3 import PPO

env = gym.make("LunarLander-v3")

model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    tensorboardlog="./tblogs/",
)

Each call to learn creates a new run
model.learn(totaltimesteps=100000, tblogname="pporun1")
model.learn(totaltimesteps=100000, tblogname="pporun2", resetnumtimesteps=False)

env.close()

Launching TensorBoard

# Run this in your terminal
tensorboard --logdir ./tblogs/ --port 6006

Then open http://localhost:6006 in your browser. You will see metrics including:

• rollout/eprewmean - Mean episode reward over time
• rollout/eplenmean - Mean episode length
• train/loss - Training loss
• train/entropyloss - Policy entropy (exploration)
• train/policygradientloss - Policy gradient loss

Comparing Multiple Algorithms

env = gym.make("CartPole-v1")

Train all three algorithms with TensorBoard logging
for AlgoCls, name in [(DQN, "dqn"), (PPO, "ppo"), (A2C, "a2c")]:
    model = AlgoCls("MlpPolicy", env, verbose=0, tensorboardlog="./tbcomparison/")
    model.learn(totaltimesteps=50000, tblogname=name)
    model.save(f"{name}cartpole")


env.close()
Then run: tensorboard --logdir ./tbcomparison/

---

Best Practices

Normalize observations. Use VecNormalize wrapper from SB3 to normalize observations and rewards, especially for continuous environments.

from stablebaselines3.common.vecenv import DummyVecEnv, VecNormalize

env = DummyVecEnv([lambda: gym.make("Pendulum-v1")])
env = VecNormalize(env, normobs=True, normreward=True, clipobs=10.0)

Use vectorized environments for faster training. Run multiple environments in parallel to collect more data per update.

from stablebaselines3.common.vecenv import SubprocVecEnv

def makeenv(envid, rank, seed=0):

    def init():
        env = gym.make(envid)

        env.reset(seed=seed + rank)
        return env
    return init

numenvs = 4

env = SubprocVecEnv([makeenv("LunarLander-v3", i) for i in range(numenvs)])

Start with PPO. It is the most stable and versatile algorithm. Only switch to DQN for discrete problems where sample efficiency matters, or to SAC/TD3 for continuous control.

Tune hyperparameters systematically. Use the RL Zoo3 hyperparameter optimization or Optuna integration.

Monitor training curves. Always use TensorBoard. Watch for reward plateau, entropy collapse, or diverging loss.

Validate environment correctness. Use checkenv() before training. A buggy environment will silently produce poor agents.

Use deterministic evaluation. Set deterministic=True during evaluation to get reproducible results that reflect the learned policy.

Save and version models. Save models at regular checkpoints and always save the final model alongside the normalization statistics.

Reward shaping matters. A well-designed reward function is often more impactful than algorithm choice. Make rewards dense, informative, and aligned with the desired behavior.

Set seeds for reproducibility. Set seeds in the environment, the algorithm, and NumPy/PyTorch for reproducible experiments.

---

Conclusion

In this tutorial, we covered the foundations of reinforcement learning and demonstrated how to apply it using modern tooling:

• Gymnasium provides standardized, well-maintained environments with a clean API for both built-in and custom environments.
• DQN is effective for discrete action spaces and leverages experience replay and target networks.
• PPO is the go-to general-purpose algorithm, offering stability and good performance across diverse tasks.
• A2C provides a simpler actor-critic approach that can be effective with proper tuning.
• Custom environments allow you to apply RL to domain-specific problems.
• Callbacks and TensorBoard are essential for monitoring and managing training runs.

The key to success in RL is iterative experimentation: start simple, monitor everything, and gradually increase complexity. With the tools covered in this tutorial, you have a solid foundation for tackling both standard benchmarks and real-world RL challenges.