Reinforcement Learning

Reinforcement Learning Best Practices

Overview

This skill provides comprehensive guidance for implementing reinforcement learning in Python using the modern ecosystem (2024-2025). Gymnasium has replaced OpenAI Gym as the standard environment interface. Stable-Baselines3 (SB3) is recommended for prototyping, RLlib for production/distributed training, and CleanRL for research.

When to Use

• Building RL agents for discrete or continuous control tasks
• Creating custom simulation environments
• Tuning hyperparameters for RL algorithms
• Debugging training issues (reward curves, policy collapse, numerical instability)
• Deploying trained policies to production

Library Selection

Algorithm Decision Tree

Start
  |
  v
Action space type?
  |
  +-- Discrete --> Sample efficiency critical?
  |                  |
  |                  +-- Yes --> DQN (or Double/Dueling DQN)
  |                  +-- No  --> Stability critical?
  |                               |
  |                               +-- Yes --> PPO
  |                               +-- No  --> A2C (faster iterations)
  |
  +-- Continuous --> Sample efficiency critical?
                       |
                       +-- Yes --> SAC (auto entropy) or TD3
                       +-- No  --> PPO (more stable, less efficient)

Quick Selection Table:

Quick Start with Stable-Baselines3

Basic Training

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Create vectorized environment (4 parallel envs)
env = make_vec_env("CartPole-v1", n_envs=4)

# Initialize and train
model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=100_000)

# Save and load
model.save("ppo_cartpole")
loaded_model = PPO.load("ppo_cartpole")

# Evaluate
obs = env.reset()
for _ in range(1000):
    action, _ = loaded_model.predict(obs, deterministic=True)
    obs, reward, done, info = env.step(action)

Custom Environment Template

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

class CustomEnv(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"]}

    def __init__(self, render_mode=None):
        super().__init__()
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32
        )
        self.action_space = spaces.Discrete(2)
        self.render_mode = render_mode

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.state = self.np_random.uniform(low=-0.05, high=0.05, size=(4,))
        return self.state.astype(np.float32), {}

    def step(self, action):
        # Implement environment dynamics here
        observation = self.state.astype(np.float32)
        reward = 1.0
        terminated = False  # Episode ended due to task completion/failure
        truncated = False   # Episode ended due to time limit
        info = {}
        return observation, reward, terminated, truncated, info

    def render(self):
        pass

Hyperparameter Tuning with Optuna

import optuna
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy

def objective(trial):
    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-3, log=True)
    n_steps = trial.suggest_categorical("n_steps", [256, 512, 1024, 2048])
    gamma = trial.suggest_float("gamma", 0.9, 0.9999)

    model = PPO(
        "MlpPolicy", "CartPole-v1",
        learning_rate=learning_rate,
        n_steps=n_steps,
        gamma=gamma,
        verbose=0
    )
    model.learn(total_timesteps=50_000)

    mean_reward, _ = evaluate_policy(model, model.get_env(), n_eval_episodes=10)
    return mean_reward

study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=50)
print(f"Best params: {study.best_params}")

Core Workflow

1. Define the environment - Use Gymnasium API, validate spaces 2. Select algorithm - Based on action space and requirements 3. Start simple - Default hyperparameters, short training 4. Monitor training - TensorBoard, check reward curves 5. Debug issues - Use the debugging playbook 6. Tune hyperparameters - Optuna for systematic search 7. Evaluate properly - Separate eval env, multiple seeds 8. Deploy - Export to ONNX/TorchScript

Reference Files

• algorithms.md - Deep dive on DQN, PPO, SAC, A2C, TD3
• environments.md - Gymnasium setup, custom envs, wrappers
• training.md - Hyperparameters, reward engineering, normalization
• debugging.md - Failure modes, diagnostics, sanity checks
• evaluation.md - Metrics, logging, reproducibility
• deployment.md - ONNX export, inference optimization, safety

Essential Dependencies

pip install gymnasium stable-baselines3 tensorboard optuna
# For Atari environments
pip install gymnasium[atari] gymnasium[accept-rom-license]
# For MuJoCo
pip install gymnasium[mujoco]

Common Pitfalls to Avoid

1. Not normalizing observations - Use VecNormalize wrapper 2. Wrong action space handling - Check discrete vs continuous 3. Ignoring seed management - Set seeds for reproducibility 4. Training and eval on same env - Use separate eval environment 5. Not monitoring entropy - Low entropy = policy collapse 6. Sparse rewards without shaping - Add intermediate rewards 7. Too large/small learning rate - Start with 3e-4 for most algorithms

RL Algorithms Deep Dive

Overview

This reference covers the most important RL algorithms, when to use each, and key hyperparameters.

Algorithm Categories

On-Policy vs Off-Policy

• On-Policy (PPO, A2C): Learn from current policy's experience only. More stable but less sample efficient.
• Off-Policy (DQN, SAC, TD3): Learn from replay buffer of past experience. More sample efficient but can be unstable.

Value-Based vs Policy-Based

• Value-Based (DQN): Learn Q-values, derive policy by taking argmax. Only works for discrete actions.
• Policy-Based (PPO, SAC): Directly optimize the policy. Works for both discrete and continuous.

---

DQN (Deep Q-Network)

When to Use

• Discrete action spaces
• Sample efficiency is important
• Can tolerate some instability

How It Works

1. Maintain Q-network that estimates Q(s, a) 2. Store transitions in replay buffer 3. Sample mini-batches and minimize TD error 4. Use target network for stability (soft or hard updates)

Key Variants

Double DQN: Addresses overestimation bias by decoupling action selection and evaluation.

# Standard DQN (overestimates)
target = r + gamma * max(Q_target(s'))

# Double DQN (less bias)
a_best = argmax(Q_online(s'))
target = r + gamma * Q_target(s', a_best)

Dueling DQN: Separates value and advantage streams.

Q(s, a) = V(s) + A(s, a) - mean(A(s, :))

Key Hyperparameters

SB3 Example

from stable_baselines3 import DQN

model = DQN(
    "MlpPolicy",
    "CartPole-v1",
    learning_rate=1e-4,
    buffer_size=100_000,
    learning_starts=1000,
    batch_size=32,
    gamma=0.99,
    target_update_interval=1000,
    exploration_fraction=0.1,
    exploration_final_eps=0.05,
    verbose=1
)
model.learn(total_timesteps=100_000)

---

PPO (Proximal Policy Optimization)

When to Use

• General purpose, good default choice
• Stability is important
• Discrete or continuous actions
• Parallel environments available

How It Works

1. Collect rollouts using current policy 2. Compute advantages (GAE) 3. Optimize clipped surrogate objective 4. Repeat

Key Innovation

Clips the policy ratio to prevent too large updates:

ratio = pi_new(a|s) / pi_old(a|s)
clipped_ratio = clip(ratio, 1 - epsilon, 1 + epsilon)
loss = -min(ratio * A, clipped_ratio * A)

Key Hyperparameters

SB3 Example

from stable_baselines3 import PPO

model = PPO(
    "MlpPolicy",
    "CartPole-v1",
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    gae_lambda=0.95,
    clip_range=0.2,
    ent_coef=0.01,  # Add entropy bonus for exploration
    verbose=1,
    tensorboard_log="./ppo_logs/"
)
model.learn(total_timesteps=100_000)

---

SAC (Soft Actor-Critic)

When to Use

• Continuous action spaces
• Sample efficiency is critical
• Can handle entropy tuning automatically

How It Works

1. Maximize expected return + entropy bonus 2. Learn Q-function, V-function, and policy 3. Automatic entropy coefficient tuning 4. Off-policy with replay buffer

Key Innovation

Maximizes entropy-augmented reward:

J(pi) = E[sum(r + alpha * H(pi(.|s)))]

This encourages exploration and robustness.

Key Hyperparameters

SB3 Example

from stable_baselines3 import SAC

model = SAC(
    "MlpPolicy",
    "Pendulum-v1",
    learning_rate=3e-4,
    buffer_size=1_000_000,
    batch_size=256,
    gamma=0.99,
    tau=0.005,
    ent_coef="auto",  # Automatic entropy tuning
    verbose=1,
    tensorboard_log="./sac_logs/"
)
model.learn(total_timesteps=100_000)

---

TD3 (Twin Delayed DDPG)

When to Use

• Continuous action spaces
• Alternative to SAC
• When SAC's entropy bonus causes issues

How It Works

1. Twin Q-networks (take minimum to reduce overestimation) 2. Delayed policy updates (update policy less frequently than Q) 3. Target policy smoothing (add noise to target actions)

Key Improvements Over DDPG

# Twin Q-networks
target_Q = min(Q1_target(s', a'), Q2_target(s', a'))

# Target policy smoothing
noise = clip(N(0, sigma), -c, c)
a' = clip(pi_target(s') + noise, a_low, a_high)

# Delayed policy updates
if step % policy_delay == 0:
    update_policy()

Key Hyperparameters

SB3 Example

from stable_baselines3 import TD3
from stable_baselines3.common.noise import NormalActionNoise
import numpy as np

env = gym.make("Pendulum-v1")
n_actions = env.action_space.shape[-1]
action_noise = NormalActionNoise(
    mean=np.zeros(n_actions),
    sigma=0.1 * np.ones(n_actions)
)

model = TD3(
    "MlpPolicy",
    env,
    learning_rate=1e-3,
    buffer_size=1_000_000,
    batch_size=100,
    gamma=0.99,
    tau=0.005,
    policy_delay=2,
    action_noise=action_noise,
    verbose=1
)
model.learn(total_timesteps=100_000)

---

A2C (Advantage Actor-Critic)

When to Use

• Simpler alternative to PPO
• Fast iteration needed
• Many parallel environments available
• Lower sample efficiency is acceptable

How It Works

1. Synchronous version of A3C 2. Multiple parallel environments 3. Actor-critic with advantage estimation 4. No replay buffer (on-policy)

Key Hyperparameters

SB3 Example

from stable_baselines3 import A2C

model = A2C(
    "MlpPolicy",
    "CartPole-v1",
    learning_rate=7e-4,
    n_steps=5,
    gamma=0.99,
    gae_lambda=1.0,
    ent_coef=0.01,
    vf_coef=0.5,
    verbose=1
)
model.learn(total_timesteps=100_000)

---

Algorithm Comparison Summary

Decision Flowchart

Need to train an RL agent?
│
├─ Discrete actions?
│  ├─ Sample efficiency critical? → DQN (Double/Dueling)
│  ├─ Stability critical? → PPO
│  └─ Fast prototyping? → A2C
│
└─ Continuous actions?
   ├─ Sample efficiency critical?
   │  ├─ Want automatic entropy? → SAC
   │  └─ Prefer deterministic policy? → TD3
   └─ Stability critical? → PPO

RL Debugging Playbook

Overview

Debugging RL is notoriously difficult due to high variance, delayed feedback, and complex interactions between components. This guide provides systematic approaches to diagnosing and fixing common issues.

Pre-Training Sanity Checks

Before Training: The Checklist

def pre_training_checklist(env, model):
    """Run before any serious training."""

    print("=" * 50)
    print("PRE-TRAINING SANITY CHECKS")
    print("=" * 50)

    # 1. Environment validation
    print("\n1. Environment Validation")
    from gymnasium.utils.env_checker import check_env
    try:
        check_env(env.envs[0] if hasattr(env, 'envs') else env)
        print("   [PASS] Environment API check")
    except Exception as e:
        print(f"   [FAIL] Environment API: {e}")

    # 2. Random policy baseline
    print("\n2. Random Policy Baseline")
    obs = env.reset()
    total_rewards = []
    for _ in range(10):
        episode_reward = 0
        done = False
        while not done:
            action = env.action_space.sample()
            obs, reward, done, info = env.step(action)
            episode_reward += reward
        total_rewards.append(episode_reward)
        obs = env.reset()
    print(f"   Random policy mean reward: {np.mean(total_rewards):.2f} +/- {np.std(total_rewards):.2f}")

    # 3. Observation statistics
    print("\n3. Observation Statistics")
    obs = env.reset()
    obs_samples = [obs]
    for _ in range(100):
        obs, _, done, _ = env.step(env.action_space.sample())
        obs_samples.append(obs)
        if done:
            obs = env.reset()
    obs_array = np.array(obs_samples)
    print(f"   Obs shape: {obs_array.shape}")
    print(f"   Obs range: [{obs_array.min():.2f}, {obs_array.max():.2f}]")
    print(f"   Obs mean: {obs_array.mean():.2f}, std: {obs_array.std():.2f}")
    if np.abs(obs_array).max() > 100:
        print("   [WARN] Large observation values - consider normalization")

    # 4. Reward statistics
    print("\n4. Reward Statistics")
    rewards = []
    obs = env.reset()
    for _ in range(1000):
        obs, reward, done, _ = env.step(env.action_space.sample())
        rewards.append(reward)
        if done:
            obs = env.reset()
    rewards = np.array(rewards)
    print(f"   Reward range: [{rewards.min():.2f}, {rewards.max():.2f}]")
    print(f"   Reward mean: {rewards.mean():.4f}, std: {rewards.std():.4f}")
    print(f"   Nonzero rewards: {(rewards != 0).sum() / len(rewards) * 100:.1f}%")
    if (rewards != 0).sum() < 10:
        print("   [WARN] Very sparse rewards - learning may be difficult")

    # 5. Action space check
    print("\n5. Action Space")
    print(f"   Type: {type(env.action_space).__name__}")
    print(f"   Shape/Size: {env.action_space.shape if hasattr(env.action_space, 'shape') else env.action_space.n}")

    # 6. Model forward pass
    print("\n6. Model Forward Pass")
    try:
        obs = env.reset()
        action, _ = model.predict(obs, deterministic=True)
        print(f"   [PASS] Model produces valid actions: {action}")
    except Exception as e:
        print(f"   [FAIL] Model forward pass: {e}")

    print("\n" + "=" * 50)

# Usage
env = make_vec_env("CartPole-v1", n_envs=1)
model = PPO("MlpPolicy", env, verbose=0)
pre_training_checklist(env, model)

---

Common Failure Modes

1. Reward Not Improving

Symptoms:

• Flat reward curve
• No learning progress after many timesteps

Diagnostics:

# Check if policy is updating
# In TensorBoard, look for:
# - policy_gradient_loss changing
# - value_loss changing
# - learning_rate (if scheduled)

Causes and Solutions:

2. Reward Improves Then Collapses

Symptoms:

• Initial improvement
• Sudden drop in performance
• Never recovers

Diagnostics:

# Monitor these in TensorBoard:
# - entropy: dropping to near 0 = policy collapse
# - value_loss: spiking = value function issues
# - clip_fraction (PPO): high = updates too aggressive

Causes and Solutions:

3. High Variance / Unstable Training

Symptoms:

• Reward oscillates wildly
• Different seeds give very different results
• Hard to reproduce results

Diagnostics:

# Run multiple seeds
results = []
for seed in [0, 1, 2, 3, 4]:
    model = PPO("MlpPolicy", env, seed=seed, verbose=0)
    model.learn(total_timesteps=100_000)
    mean_reward, _ = evaluate_policy(model, env, n_eval_episodes=10)
    results.append(mean_reward)

print(f"Mean: {np.mean(results):.2f}, Std: {np.std(results):.2f}")
# Std > 50% of mean indicates high variance

Causes and Solutions:

4. NaN or Numerical Instability

Symptoms:

• Loss becomes NaN
• Actions become NaN/Inf
• Training crashes

Diagnostics:

# Add NaN checks
import torch

def check_for_nan(model, name=""):
    for param_name, param in model.policy.named_parameters():
        if torch.isnan(param).any():
            print(f"NaN in {name} {param_name}")
        if torch.isinf(param).any():
            print(f"Inf in {name} {param_name}")

Causes and Solutions:

---

Reward Curve Interpretation

Healthy Training Curves

Good PPO curve:
    ^
    |           ____----
    |      ____/
reward|  __/
    | /
    |/
    +-----------------------> timesteps

Characteristics:
- Gradual improvement
- May have plateaus
- Some variance is normal

Problem Patterns

Policy Collapse:
    ^
    |    ___
    |   /   \
reward| /     \___
    |/           \_______
    +-----------------------> timesteps

Diagnosis: Entropy dropping, increase ent_coef


Reward Hacking:
    ^
    |                 ____
    |            ____/
reward|       ____/
    |  ____/
    |_/
    +-----------------------> timesteps

But evaluation performance is poor!
Diagnosis: Agent exploiting reward, fix reward function


No Learning:
    ^
    |
    |  ~~~~~~~~~~~~~~~~~~~~~~~~
reward|
    |
    +-----------------------> timesteps

Diagnosis: Check LR, exploration, reward sparsity

---

Policy Entropy Monitoring

Why Entropy Matters

• High entropy: Policy is stochastic, exploring
• Low entropy: Policy is deterministic, exploiting
• Zero entropy: Policy collapse, stuck on single action

Monitoring Code

from stable_baselines3.common.callbacks import BaseCallback

class EntropyMonitorCallback(BaseCallback):
    """Monitor policy entropy during training."""

    def __init__(self, warning_threshold=0.1, verbose=0):
        super().__init__(verbose)
        self.warning_threshold = warning_threshold
        self.entropies = []

    def _on_step(self):
        # For PPO, entropy is logged automatically
        # Access via self.logger
        if len(self.model.ep_info_buffer) > 0:
            # Custom entropy calculation for inspection
            if hasattr(self.model, 'policy'):
                obs = self.training_env.reset()
                with torch.no_grad():
                    dist = self.model.policy.get_distribution(
                        torch.tensor(obs).float()
                    )
                    entropy = dist.entropy().mean().item()
                    self.entropies.append(entropy)

                    if entropy < self.warning_threshold:
                        print(f"WARNING: Low entropy ({entropy:.4f}) at step {self.num_timesteps}")

        return True

# Usage
callback = EntropyMonitorCallback(warning_threshold=0.1)
model.learn(total_timesteps=100_000, callback=callback)

Entropy Guidelines by Algorithm

---

Value Function Diagnostics

Explained Variance

Measures how well value function predicts returns.

# SB3 logs this as "train/explained_variance"
# In TensorBoard, look for values:
# - Near 1.0: Value function very accurate (good)
# - Near 0.0: Value function no better than mean (needs tuning)
# - Negative: Value function worse than mean (problem!)

Manual Value Function Check

def diagnose_value_function(model, env, n_episodes=10):
    """Check if value function predictions are reasonable."""

    all_values = []
    all_returns = []

    for _ in range(n_episodes):
        obs = env.reset()
        episode_rewards = []
        episode_values = []

        done = False
        while not done:
            # Get value prediction
            obs_tensor = torch.tensor(obs).float().unsqueeze(0)
            with torch.no_grad():
                value = model.policy.predict_values(obs_tensor).item()
            episode_values.append(value)

            # Take step
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, done, info = env.step(action)
            episode_rewards.append(reward)

        # Compute actual returns (discounted)
        gamma = model.gamma
        returns = []
        G = 0
        for r in reversed(episode_rewards):
            G = r + gamma * G
            returns.insert(0, G)

        all_values.extend(episode_values)
        all_returns.extend(returns)

    values = np.array(all_values)
    returns = np.array(all_returns)

    # Compute correlation
    correlation = np.corrcoef(values, returns)[0, 1]

    # Compute explained variance
    var_returns = np.var(returns)
    var_unexplained = np.var(returns - values)
    explained_var = 1 - var_unexplained / var_returns if var_returns > 0 else 0

    print(f"Value-Return Correlation: {correlation:.3f}")
    print(f"Explained Variance: {explained_var:.3f}")
    print(f"Value range: [{values.min():.2f}, {values.max():.2f}]")
    print(f"Return range: [{returns.min():.2f}, {returns.max():.2f}]")

    if explained_var < 0:
        print("WARNING: Negative explained variance - value function is harmful")
    elif explained_var < 0.5:
        print("WARNING: Low explained variance - consider tuning vf_coef or network")

    return explained_var, correlation

---

The Deadly Triad

The deadly triad in RL refers to the combination that can cause instability: 1. Function approximation (neural networks) 2. Bootstrapping (using value estimates to update value estimates) 3. Off-policy learning (learning from old experience)

Mitigation Strategies

# 1. Target networks (for off-policy)
# DQN, SAC, TD3 use target networks by default
model = DQN("MlpPolicy", env, target_update_interval=1000)

# 2. Gradient clipping
model = PPO("MlpPolicy", env, max_grad_norm=0.5)

# 3. Conservative value updates
model = PPO("MlpPolicy", env, vf_coef=0.5)  # Weight value loss less

# 4. Double Q-learning (DQN)
# Decouples action selection from value estimation

# 5. Soft updates (SAC, TD3)
model = SAC("MlpPolicy", env, tau=0.005)  # Slow target updates

---

Debugging Workflow

Step-by-Step Process

1. Verify environment works
   └─ Run random policy, check rewards

2. Check observations
   └─ Appropriate scale? Contains needed info?

3. Test with simple baseline
   └─ Does DQN/PPO work on CartPole?

4. Start with defaults
   └─ Only tune after confirming learning

5. Monitor training
   └─ TensorBoard: reward, losses, entropy

6. If no learning:
   ├─ Check learning rate (try 10x higher/lower)
   ├─ Check exploration (entropy, epsilon)
   ├─ Check reward sparsity
   └─ Verify environment determinism

7. If unstable:
   ├─ Reduce learning rate
   ├─ Increase batch size
   ├─ Add normalization
   └─ Try different seeds

8. If NaN:
   ├─ Check observation/reward scale
   ├─ Reduce learning rate
   └─ Increase gradient clipping

Minimal Reproducible Example

When debugging, reduce to simplest case:

import gymnasium as gym
from stable_baselines3 import PPO

# 1. Start with known-working environment
env = gym.make("CartPole-v1")

# 2. Default parameters
model = PPO("MlpPolicy", env, verbose=1)

# 3. Short training
model.learn(total_timesteps=10_000)

# 4. Quick evaluation
obs, _ = env.reset()
for _ in range(100):
    action, _ = model.predict(obs)
    obs, reward, done, truncated, info = env.step(action)
    if done or truncated:
        break

# If this works, incrementally add complexity
# If this fails, environment issue or installation problem

---

Logging for Debugging

from stable_baselines3.common.logger import configure

# Set up detailed logging
logger = configure("./logs", ["stdout", "tensorboard", "csv"])

model = PPO("MlpPolicy", env, verbose=1)
model.set_logger(logger)

# Custom logging in callbacks
class DebugCallback(BaseCallback):
    def _on_step(self):
        # Log custom metrics
        self.logger.record("debug/custom_metric", some_value)

        # Log histograms (TensorBoard)
        self.logger.record("debug/actions", self.locals["actions"], exclude="stdout")

        return True

Key TensorBoard Metrics

RL Deployment Guide

Overview

Deploying RL models to production requires careful attention to model export, inference optimization, monitoring, and safety. This guide covers best practices for production-ready RL systems.

Model Export

ONNX Export

ONNX (Open Neural Network Exchange) enables deployment across different frameworks and platforms.

import torch
import numpy as np
from stable_baselines3 import PPO

# Train model
model = PPO("MlpPolicy", "CartPole-v1")
model.learn(total_timesteps=50_000)

# Export to ONNX
def export_to_onnx(model, onnx_path, env):
    """Export SB3 model to ONNX format."""

    # Get a sample observation
    obs = env.observation_space.sample()
    obs_tensor = torch.tensor(obs).float().unsqueeze(0)

    # Export the policy network
    torch.onnx.export(
        model.policy,
        obs_tensor,
        onnx_path,
        export_params=True,
        opset_version=11,
        do_constant_folding=True,
        input_names=['observation'],
        output_names=['action', 'value', 'log_prob'],
        dynamic_axes={
            'observation': {0: 'batch_size'},
            'action': {0: 'batch_size'},
            'value': {0: 'batch_size'},
            'log_prob': {0: 'batch_size'}
        }
    )
    print(f"Model exported to {onnx_path}")

export_to_onnx(model, "policy.onnx", model.get_env())

# Verify ONNX model
import onnx
onnx_model = onnx.load("policy.onnx")
onnx.checker.check_model(onnx_model)

ONNX Inference

import onnxruntime as ort
import numpy as np

class ONNXPolicy:
    """Wrapper for ONNX model inference."""

    def __init__(self, onnx_path):
        self.session = ort.InferenceSession(onnx_path)
        self.input_name = self.session.get_inputs()[0].name

    def predict(self, observation):
        """Get action from observation."""
        if isinstance(observation, np.ndarray):
            obs = observation.astype(np.float32)
        else:
            obs = np.array(observation, dtype=np.float32)

        if obs.ndim == 1:
            obs = obs.reshape(1, -1)

        outputs = self.session.run(None, {self.input_name: obs})
        action = outputs[0]  # First output is action

        return action[0]  # Remove batch dimension

# Usage
policy = ONNXPolicy("policy.onnx")
obs = env.reset()[0]
action = policy.predict(obs)

TorchScript Export

For PyTorch-native deployment:

import torch
from stable_baselines3 import PPO

model = PPO.load("trained_model")

class TorchScriptPolicy(torch.nn.Module):
    """Wrapper for TorchScript export."""

    def __init__(self, policy):
        super().__init__()
        self.features_extractor = policy.features_extractor
        self.mlp_extractor = policy.mlp_extractor
        self.action_net = policy.action_net

    def forward(self, obs):
        features = self.features_extractor(obs)
        latent_pi, _ = self.mlp_extractor(features)
        return self.action_net(latent_pi)

# Create wrapper and trace
wrapper = TorchScriptPolicy(model.policy)
wrapper.eval()

# Trace with example input
example_obs = torch.randn(1, model.observation_space.shape[0])
traced = torch.jit.trace(wrapper, example_obs)

# Save
traced.save("policy_traced.pt")

# Load and use
loaded = torch.jit.load("policy_traced.pt")
loaded.eval()

with torch.no_grad():
    action_logits = loaded(example_obs)
    action = torch.argmax(action_logits, dim=1)

---

Inference Optimization

Batched Inference

import numpy as np
import torch
from collections import deque

class BatchedInference:
    """Accumulate observations and run batched inference."""

    def __init__(self, model, batch_size=32, max_wait_ms=10):
        self.model = model
        self.batch_size = batch_size
        self.max_wait_ms = max_wait_ms
        self.pending = deque()
        self.last_batch_time = time.time()

    def add_observation(self, obs, callback):
        """Add observation to batch queue."""
        self.pending.append((obs, callback))

        # Check if should process batch
        if len(self.pending) >= self.batch_size:
            self._process_batch()
        elif (time.time() - self.last_batch_time) * 1000 > self.max_wait_ms:
            self._process_batch()

    def _process_batch(self):
        """Process accumulated observations."""
        if not self.pending:
            return

        # Collect observations
        obs_list = []
        callbacks = []
        while self.pending and len(obs_list) < self.batch_size:
            obs, callback = self.pending.popleft()
            obs_list.append(obs)
            callbacks.append(callback)

        # Batched inference
        obs_batch = np.stack(obs_list)
        with torch.no_grad():
            actions, _, _ = self.model.policy(torch.tensor(obs_batch).float())
            actions = actions.numpy()

        # Return results
        for action, callback in zip(actions, callbacks):
            callback(action)

        self.last_batch_time = time.time()

GPU Optimization

import torch

class GPUPolicy:
    """GPU-optimized policy inference."""

    def __init__(self, model_path, device="cuda"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")

        # Load model to GPU
        model = PPO.load(model_path, device=self.device)
        self.policy = model.policy.to(self.device)
        self.policy.eval()

        # Use half precision for faster inference
        if self.device.type == "cuda":
            self.policy = self.policy.half()

    @torch.no_grad()
    def predict(self, observations):
        """Batched prediction on GPU."""
        obs_tensor = torch.tensor(observations, device=self.device)

        if self.device.type == "cuda":
            obs_tensor = obs_tensor.half()

        actions, _, _ = self.policy(obs_tensor)
        return actions.cpu().numpy()

    @torch.no_grad()
    def predict_single(self, observation):
        """Single observation prediction."""
        obs_tensor = torch.tensor(observation, device=self.device).unsqueeze(0)

        if self.device.type == "cuda":
            obs_tensor = obs_tensor.half()

        action, _, _ = self.policy(obs_tensor)
        return action.cpu().numpy()[0]

Quantization

import torch

def quantize_policy(model_path, output_path):
    """Quantize model for faster CPU inference."""

    model = PPO.load(model_path)
    policy = model.policy
    policy.eval()

    # Dynamic quantization (for CPU)
    quantized_policy = torch.quantization.quantize_dynamic(
        policy,
        {torch.nn.Linear},  # Quantize linear layers
        dtype=torch.qint8
    )

    # Save quantized model
    torch.save(quantized_policy.state_dict(), output_path)

    # Compare size
    import os
    original_size = os.path.getsize(model_path + ".zip")
    quantized_size = os.path.getsize(output_path)
    print(f"Original size: {original_size / 1024:.1f} KB")
    print(f"Quantized size: {quantized_size / 1024:.1f} KB")
    print(f"Compression: {original_size / quantized_size:.1f}x")

    return quantized_policy

---

Production Monitoring

Metrics to Track

from dataclasses import dataclass
from typing import List, Optional
import time
import logging

@dataclass
class InferenceMetrics:
    """Metrics for production monitoring."""
    latency_ms: float
    action: int
    observation_norm: float
    timestamp: float
    episode_id: str

class ProductionPolicy:
    """Policy wrapper with monitoring."""

    def __init__(self, model_path, metrics_logger=None):
        self.model = PPO.load(model_path)
        self.model.policy.eval()
        self.metrics_logger = metrics_logger or logging.getLogger(__name__)

        # Running statistics
        self.inference_count = 0
        self.total_latency = 0
        self.action_distribution = {}
        self.episode_rewards = []
        self.current_episode_reward = 0
        self.current_episode_id = None

    def predict(self, observation, episode_id=None):
        """Predict with monitoring."""
        start_time = time.time()

        # Get action
        action, _ = self.model.predict(observation, deterministic=True)

        # Calculate latency
        latency_ms = (time.time() - start_time) * 1000

        # Track metrics
        self.inference_count += 1
        self.total_latency += latency_ms

        # Track action distribution
        action_key = int(action) if isinstance(action, (int, np.integer)) else tuple(action)
        self.action_distribution[action_key] = self.action_distribution.get(action_key, 0) + 1

        # Log metrics
        metrics = InferenceMetrics(
            latency_ms=latency_ms,
            action=action_key,
            observation_norm=np.linalg.norm(observation),
            timestamp=time.time(),
            episode_id=episode_id or "unknown"
        )

        self._log_metrics(metrics)

        return action

    def record_reward(self, reward, done=False):
        """Track episode rewards."""
        self.current_episode_reward += reward

        if done:
            self.episode_rewards.append(self.current_episode_reward)
            self.current_episode_reward = 0

    def _log_metrics(self, metrics: InferenceMetrics):
        """Log metrics to monitoring system."""
        self.metrics_logger.info(
            f"inference latency_ms={metrics.latency_ms:.2f} "
            f"action={metrics.action} "
            f"obs_norm={metrics.observation_norm:.2f}"
        )

    def get_statistics(self):
        """Get aggregated statistics."""
        return {
            "total_inferences": self.inference_count,
            "avg_latency_ms": self.total_latency / max(1, self.inference_count),
            "action_distribution": self.action_distribution,
            "episodes_completed": len(self.episode_rewards),
            "avg_episode_reward": np.mean(self.episode_rewards) if self.episode_rewards else 0,
        }

Alerting

class PolicyMonitor:
    """Monitor policy for anomalies."""

    def __init__(self, policy, alert_callback):
        self.policy = policy
        self.alert_callback = alert_callback

        # Thresholds
        self.max_latency_ms = 100
        self.min_entropy = 0.1
        self.obs_range = (-10, 10)

        # Baselines (set from training)
        self.baseline_action_dist = None

    def check_latency(self, latency_ms):
        """Alert on high latency."""
        if latency_ms > self.max_latency_ms:
            self.alert_callback(
                "HIGH_LATENCY",
                f"Inference latency {latency_ms:.1f}ms exceeds threshold {self.max_latency_ms}ms"
            )

    def check_observation(self, observation):
        """Alert on out-of-distribution observations."""
        obs_min, obs_max = observation.min(), observation.max()

        if obs_min < self.obs_range[0] or obs_max > self.obs_range[1]:
            self.alert_callback(
                "OOD_OBSERVATION",
                f"Observation out of expected range: [{obs_min:.2f}, {obs_max:.2f}]"
            )

        if np.isnan(observation).any():
            self.alert_callback("NAN_OBSERVATION", "Observation contains NaN values")

    def check_action_distribution(self, recent_actions, window=1000):
        """Alert on distribution shift."""
        if len(recent_actions) < window:
            return

        if self.baseline_action_dist is None:
            return

        # Calculate current distribution
        current_dist = {}
        for a in recent_actions[-window:]:
            current_dist[a] = current_dist.get(a, 0) + 1

        # Chi-squared test for distribution shift
        # ... implementation

    def check_reward_degradation(self, recent_rewards, window=100):
        """Alert on performance degradation."""
        if len(recent_rewards) < window:
            return

        recent_mean = np.mean(recent_rewards[-window:])
        baseline_mean = np.mean(recent_rewards[:-window]) if len(recent_rewards) > window else recent_mean

        # Alert if performance drops significantly
        if recent_mean < baseline_mean * 0.8:  # 20% degradation
            self.alert_callback(
                "PERFORMANCE_DEGRADATION",
                f"Recent reward {recent_mean:.2f} is {(1 - recent_mean/baseline_mean)*100:.1f}% below baseline"
            )

Prometheus/Grafana Integration

from prometheus_client import Counter, Histogram, Gauge, start_http_server

# Define metrics
INFERENCE_COUNTER = Counter('rl_inference_total', 'Total inference calls')
INFERENCE_LATENCY = Histogram('rl_inference_latency_seconds', 'Inference latency')
EPISODE_REWARD = Gauge('rl_episode_reward', 'Latest episode reward')
ACTION_COUNTER = Counter('rl_action_total', 'Actions taken', ['action'])

class PrometheusPolicy:
    """Policy with Prometheus metrics."""

    def __init__(self, model_path, port=8000):
        self.model = PPO.load(model_path)
        start_http_server(port)  # Start metrics endpoint

    def predict(self, observation):
        INFERENCE_COUNTER.inc()

        with INFERENCE_LATENCY.time():
            action, _ = self.model.predict(observation, deterministic=True)

        ACTION_COUNTER.labels(action=str(action)).inc()

        return action

    def record_episode_end(self, reward):
        EPISODE_REWARD.set(reward)

---

Safety Considerations

Action Constraints

import numpy as np

class SafePolicy:
    """Policy wrapper with safety constraints."""

    def __init__(self, model, action_bounds=None, safety_margin=0.1):
        self.model = model
        self.action_bounds = action_bounds
        self.safety_margin = safety_margin

    def predict(self, observation, state_constraints=None):
        """Predict with safety checks."""

        # Get raw action
        action, _ = self.model.predict(observation, deterministic=True)

        # Apply action bounds
        if self.action_bounds is not None:
            low, high = self.action_bounds
            action = np.clip(action, low, high)

        # Apply safety constraints
        if state_constraints is not None:
            action = self._apply_safety_constraints(observation, action, state_constraints)

        return action

    def _apply_safety_constraints(self, obs, action, constraints):
        """Modify action to satisfy safety constraints."""
        # Example: limit velocity changes
        if 'max_acceleration' in constraints:
            current_velocity = obs[1]  # Assuming velocity is in observation
            max_accel = constraints['max_acceleration']

            # Predict next velocity
            predicted_velocity = current_velocity + action * 0.1  # dt = 0.1

            # Clamp acceleration
            if abs(predicted_velocity - current_velocity) > max_accel:
                action = np.sign(action) * max_accel / 0.1

        return action

Fallback Policies

class FallbackPolicy:
    """Primary policy with fallback on failure."""

    def __init__(self, primary_model_path, fallback_model_path=None):
        self.primary = PPO.load(primary_model_path)
        self.fallback = PPO.load(fallback_model_path) if fallback_model_path else None
        self.use_fallback = False
        self.error_count = 0
        self.max_errors = 5

    def predict(self, observation):
        """Predict with fallback on error."""
        try:
            if self.use_fallback and self.fallback:
                return self._fallback_predict(observation)

            action, _ = self.primary.predict(observation, deterministic=True)

            # Validate action
            if np.isnan(action).any() or np.isinf(action).any():
                raise ValueError("Invalid action from primary policy")

            # Reset error count on success
            self.error_count = 0
            return action

        except Exception as e:
            self.error_count += 1
            logging.error(f"Primary policy error: {e}")

            if self.error_count >= self.max_errors:
                self.use_fallback = True
                logging.warning("Switching to fallback policy")

            return self._fallback_predict(observation)

    def _fallback_predict(self, observation):
        """Safe fallback action."""
        if self.fallback:
            return self.fallback.predict(observation, deterministic=True)[0]
        else:
            # Return neutral/safe action
            return np.zeros(self.primary.action_space.shape)

Human Override

class HumanOverridablePolicy:
    """Policy that allows human override."""

    def __init__(self, model, override_callback=None):
        self.model = model
        self.override_callback = override_callback
        self.override_active = False
        self.human_action = None

    def predict(self, observation):
        """Get action, allowing human override."""

        # Check for override
        if self.override_callback:
            override = self.override_callback(observation)
            if override is not None:
                return override

        if self.override_active and self.human_action is not None:
            return self.human_action

        return self.model.predict(observation, deterministic=True)[0]

    def set_human_action(self, action):
        """Set human override action."""
        self.human_action = action
        self.override_active = True

    def release_override(self):
        """Return control to AI."""
        self.override_active = False
        self.human_action = None

---

Online Learning Patterns

Continuous Learning

from collections import deque
import numpy as np

class OnlineLearningPolicy:
    """Policy that continues learning in production."""

    def __init__(self, model_path, buffer_size=10000, update_freq=1000):
        self.model = PPO.load(model_path)
        self.buffer = deque(maxlen=buffer_size)
        self.update_freq = update_freq
        self.step_count = 0

    def predict(self, observation):
        """Predict action."""
        action, _ = self.model.predict(observation, deterministic=False)  # Stochastic for exploration
        return action

    def record_transition(self, obs, action, reward, next_obs, done):
        """Record transition for learning."""
        self.buffer.append((obs, action, reward, next_obs, done))
        self.step_count += 1

        if self.step_count % self.update_freq == 0:
            self._update()

    def _update(self):
        """Update policy from buffer."""
        if len(self.buffer) < self.update_freq:
            return

        # Sample batch
        batch = list(self.buffer)[-self.update_freq:]

        # Convert to training format
        # ... implementation depends on algorithm

        # Fine-tune model
        # self.model.learn(...)

        logging.info(f"Policy updated at step {self.step_count}")

A/B Testing

import random
import hashlib

class ABTestPolicy:
    """A/B testing between policies."""

    def __init__(self, policy_a, policy_b, traffic_split=0.5):
        self.policy_a = policy_a
        self.policy_b = policy_b
        self.traffic_split = traffic_split

        self.metrics_a = {"rewards": [], "count": 0}
        self.metrics_b = {"rewards": [], "count": 0}

    def predict(self, observation, user_id=None):
        """Route to policy based on user_id or random."""

        # Deterministic routing for same user
        if user_id:
            hash_val = int(hashlib.md5(user_id.encode()).hexdigest(), 16)
            use_a = (hash_val % 100) < (self.traffic_split * 100)
        else:
            use_a = random.random() < self.traffic_split

        if use_a:
            self.metrics_a["count"] += 1
            return self.policy_a.predict(observation), "A"
        else:
            self.metrics_b["count"] += 1
            return self.policy_b.predict(observation), "B"

    def record_reward(self, reward, variant):
        """Record reward for variant."""
        if variant == "A":
            self.metrics_a["rewards"].append(reward)
        else:
            self.metrics_b["rewards"].append(reward)

    def get_results(self):
        """Get A/B test results."""
        mean_a = np.mean(self.metrics_a["rewards"]) if self.metrics_a["rewards"] else 0
        mean_b = np.mean(self.metrics_b["rewards"]) if self.metrics_b["rewards"] else 0

        return {
            "policy_a": {
                "count": self.metrics_a["count"],
                "mean_reward": mean_a,
                "total_rewards": len(self.metrics_a["rewards"])
            },
            "policy_b": {
                "count": self.metrics_b["count"],
                "mean_reward": mean_b,
                "total_rewards": len(self.metrics_b["rewards"])
            },
            "winner": "A" if mean_a > mean_b else "B",
            "improvement": abs(mean_a - mean_b) / max(mean_a, mean_b, 1) * 100
        }

---

Deployment Checklist

Pre-Deployment

• [ ] Model exported to ONNX/TorchScript
• [ ] Inference latency benchmarked
• [ ] Memory usage profiled
• [ ] Input validation implemented
• [ ] Action bounds enforced
• [ ] Fallback policy configured
• [ ] Monitoring metrics defined
• [ ] Alerting thresholds set
• [ ] Logging configured
• [ ] A/B test plan ready

Post-Deployment

• [ ] Latency within SLA
• [ ] No NaN/Inf in outputs
• [ ] Action distribution stable
• [ ] Reward not degrading
• [ ] Error rate acceptable
• [ ] Resource usage stable
• [ ] Alerts working
• [ ] Rollback tested

Gymnasium Environments Guide

Overview

Gymnasium (the maintained fork of OpenAI Gym) is the standard API for RL environments. This guide covers setup, custom environments, vectorization, and essential wrappers.

Installation

# Core gymnasium
pip install gymnasium

# Atari environments
pip install gymnasium[atari] gymnasium[accept-rom-license]

# MuJoCo environments
pip install gymnasium[mujoco]

# Box2D environments (LunarLander, etc.)
pip install gymnasium[box2d]

# All classic control
pip install gymnasium[classic-control]

Migration from OpenAI Gym

Key Changes

# Old (gym)
import gym
env = gym.make("CartPole-v1")
obs = env.reset()
obs, reward, done, info = env.step(action)

# New (gymnasium)
import gymnasium as gym
env = gym.make("CartPole-v1")
obs, info = env.reset()  # Returns tuple now
obs, reward, terminated, truncated, info = env.step(action)  # 5 values
done = terminated or truncated

Terminated vs Truncated

• terminated: Episode ended due to task completion or failure (e.g., pole fell, goal reached)
• truncated: Episode ended due to time limit or external condition (e.g., max steps)

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

if terminated:
    print("Task ended naturally")
if truncated:
    print("Episode cut short (time limit)")
if terminated or truncated:
    obs, info = env.reset()

---

Custom Environment Template

Complete Example

import gymnasium as gym
from gymnasium import spaces
import numpy as np
from typing import Optional, Tuple, Dict, Any

class CustomEnv(gym.Env):
    """
    Custom Environment following Gymnasium API.

    Description:
        A simple environment where an agent must reach a goal position.

    Observation Space:
        Box(4,) - [agent_x, agent_y, goal_x, goal_y]

    Action Space:
        Discrete(4) - [up, down, left, right]

    Rewards:
        -1 per step, +10 for reaching goal
    """

    metadata = {
        "render_modes": ["human", "rgb_array"],
        "render_fps": 30
    }

    def __init__(self, render_mode: Optional[str] = None, grid_size: int = 10):
        super().__init__()

        self.grid_size = grid_size
        self.render_mode = render_mode

        # Define observation space
        self.observation_space = spaces.Box(
            low=0,
            high=grid_size,
            shape=(4,),
            dtype=np.float32
        )

        # Define action space
        self.action_space = spaces.Discrete(4)

        # Action mappings
        self._action_to_direction = {
            0: np.array([0, 1]),   # up
            1: np.array([0, -1]),  # down
            2: np.array([-1, 0]),  # left
            3: np.array([1, 0])    # right
        }

        # State
        self.agent_pos = None
        self.goal_pos = None
        self.steps = 0
        self.max_steps = 100

    def reset(
        self,
        seed: Optional[int] = None,
        options: Optional[Dict[str, Any]] = None
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """Reset the environment to initial state."""
        super().reset(seed=seed)

        # Initialize positions randomly
        self.agent_pos = self.np_random.integers(0, self.grid_size, size=2).astype(np.float32)
        self.goal_pos = self.np_random.integers(0, self.grid_size, size=2).astype(np.float32)

        # Ensure agent and goal are not at same position
        while np.array_equal(self.agent_pos, self.goal_pos):
            self.goal_pos = self.np_random.integers(0, self.grid_size, size=2).astype(np.float32)

        self.steps = 0

        observation = self._get_obs()
        info = self._get_info()

        return observation, info

    def step(self, action: int) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
        """Execute one step in the environment."""
        # Move agent
        direction = self._action_to_direction[action]
        self.agent_pos = np.clip(
            self.agent_pos + direction,
            0,
            self.grid_size - 1
        ).astype(np.float32)

        self.steps += 1

        # Check termination
        reached_goal = np.array_equal(self.agent_pos, self.goal_pos)
        terminated = reached_goal
        truncated = self.steps >= self.max_steps

        # Compute reward
        if reached_goal:
            reward = 10.0
        else:
            reward = -1.0

        observation = self._get_obs()
        info = self._get_info()

        return observation, reward, terminated, truncated, info

    def _get_obs(self) -> np.ndarray:
        """Construct observation from state."""
        return np.concatenate([self.agent_pos, self.goal_pos]).astype(np.float32)

    def _get_info(self) -> Dict[str, Any]:
        """Return auxiliary info."""
        return {
            "distance": np.linalg.norm(self.agent_pos - self.goal_pos),
            "steps": self.steps
        }

    def render(self):
        """Render the environment."""
        if self.render_mode == "human":
            self._render_human()
        elif self.render_mode == "rgb_array":
            return self._render_rgb_array()

    def _render_human(self):
        """Render to screen."""
        print(f"Agent: {self.agent_pos}, Goal: {self.goal_pos}")

    def _render_rgb_array(self) -> np.ndarray:
        """Return RGB array of current state."""
        # Create simple grid visualization
        img = np.zeros((self.grid_size * 10, self.grid_size * 10, 3), dtype=np.uint8)
        # Add agent (blue) and goal (green)
        ax, ay = (self.agent_pos * 10).astype(int)
        gx, gy = (self.goal_pos * 10).astype(int)
        img[ay:ay+10, ax:ax+10] = [0, 0, 255]  # Agent blue
        img[gy:gy+10, gx:gx+10] = [0, 255, 0]  # Goal green
        return img

    def close(self):
        """Clean up resources."""
        pass

Registering Custom Environment

from gymnasium.envs.registration import register

register(
    id="CustomEnv-v0",
    entry_point="my_module:CustomEnv",
    max_episode_steps=100,
)

# Now can use:
env = gym.make("CustomEnv-v0")

---

Environment Validation Checklist

Before training, validate your custom environment:

from gymnasium.utils.env_checker import check_env

env = CustomEnv()
check_env(env, warn=True)  # Raises errors if API violated

Manual Checks

def validate_environment(env_class):
    """Comprehensive environment validation."""
    env = env_class()

    # 1. Check spaces are defined
    assert env.observation_space is not None, "observation_space not defined"
    assert env.action_space is not None, "action_space not defined"

    # 2. Check reset returns correct format
    obs, info = env.reset(seed=42)
    assert env.observation_space.contains(obs), f"Reset obs not in space: {obs}"
    assert isinstance(info, dict), "Info must be dict"

    # 3. Check step returns correct format
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    assert env.observation_space.contains(obs), f"Step obs not in space: {obs}"
    assert isinstance(reward, (int, float)), "Reward must be numeric"
    assert isinstance(terminated, bool), "Terminated must be bool"
    assert isinstance(truncated, bool), "Truncated must be bool"
    assert isinstance(info, dict), "Info must be dict"

    # 4. Check determinism with seed
    env.reset(seed=42)
    actions = [env.action_space.sample() for _ in range(10)]
    env.reset(seed=42)
    obs1, _ = env.reset(seed=123)
    for a in actions[:5]:
        obs1, _, _, _, _ = env.step(a)

    env.reset(seed=123)
    obs2, _ = env.reset(seed=123)
    for a in actions[:5]:
        obs2, _, _, _, _ = env.step(a)

    assert np.allclose(obs1, obs2), "Environment not deterministic with same seed"

    # 5. Check episode termination
    env.reset()
    for _ in range(10000):
        _, _, terminated, truncated, _ = env.step(env.action_space.sample())
        if terminated or truncated:
            break
    else:
        print("Warning: Episode did not terminate in 10000 steps")

    print("All checks passed!")
    env.close()

validate_environment(CustomEnv)

---

Vectorized Environments

Vectorized environments run multiple instances in parallel for faster training.

Types

from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
from stable_baselines3.common.env_util import make_vec_env

# DummyVecEnv - Sequential (single process, good for debugging)
env = DummyVecEnv([lambda: gym.make("CartPole-v1") for _ in range(4)])

# SubprocVecEnv - Parallel (separate processes, faster)
env = SubprocVecEnv([lambda: gym.make("CartPole-v1") for _ in range(4)])

# Convenience function
env = make_vec_env("CartPole-v1", n_envs=4, vec_env_cls=SubprocVecEnv)

When to Use Each

Custom Environment with Vectorization

def make_env(env_id, rank, seed=0):
    """Create a wrapped, monitored environment."""
    def _init():
        env = gym.make(env_id)
        env.reset(seed=seed + rank)
        return env
    return _init

# Create vectorized environment
n_envs = 4
env = SubprocVecEnv([make_env("CartPole-v1", i) for i in range(n_envs)])

---

Essential Wrappers

Observation Wrappers

from gymnasium.wrappers import (
    NormalizeObservation,
    TransformObservation,
    FrameStack,
    GrayScaleObservation,
    ResizeObservation
)

# Normalize observations (running mean/std)
env = NormalizeObservation(env)

# Custom transformation
env = TransformObservation(env, lambda obs: obs / 255.0)

# Stack frames (for temporal info)
env = FrameStack(env, num_stack=4)

# Image preprocessing
env = GrayScaleObservation(env)
env = ResizeObservation(env, shape=(84, 84))

Reward Wrappers

from gymnasium.wrappers import (
    NormalizeReward,
    ClipReward,
    TransformReward
)

# Normalize rewards (running stats)
env = NormalizeReward(env, gamma=0.99)

# Clip rewards to range
env = ClipReward(env, min_reward=-1, max_reward=1)

# Custom reward transformation
env = TransformReward(env, lambda r: np.sign(r))

Action Wrappers

from gymnasium.wrappers import (
    ClipAction,
    RescaleAction
)

# Clip actions to valid range
env = ClipAction(env)

# Rescale actions
env = RescaleAction(env, min_action=-1.0, max_action=1.0)

Monitoring Wrappers

from gymnasium.wrappers import RecordVideo, RecordEpisodeStatistics

# Record videos
env = RecordVideo(env, video_folder="./videos", episode_trigger=lambda x: x % 100 == 0)

# Track episode statistics
env = RecordEpisodeStatistics(env)
# Access via info["episode"]["r"], info["episode"]["l"], info["episode"]["t"]

SB3 Wrappers

from stable_baselines3.common.vec_env import VecNormalize, VecFrameStack

# Normalize observations and rewards for vectorized envs
env = make_vec_env("Pendulum-v1", n_envs=4)
env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0)

# Stack frames
env = VecFrameStack(env, n_stack=4)

# Save and load normalization stats
env.save("vec_normalize.pkl")
env = VecNormalize.load("vec_normalize.pkl", env)

---

Common Space Types

from gymnasium import spaces
import numpy as np

# Discrete - single integer action
action_space = spaces.Discrete(4)  # 0, 1, 2, or 3

# MultiDiscrete - multiple discrete values
action_space = spaces.MultiDiscrete([3, 2, 4])  # [0-2, 0-1, 0-3]

# Box - continuous values
obs_space = spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32)
obs_space = spaces.Box(low=np.array([0, -np.inf]), high=np.array([1, np.inf]), dtype=np.float32)

# Dict - structured observations
obs_space = spaces.Dict({
    "position": spaces.Box(low=-10, high=10, shape=(2,)),
    "velocity": spaces.Box(low=-1, high=1, shape=(2,)),
    "target": spaces.Discrete(5)
})

# Tuple - multiple spaces
obs_space = spaces.Tuple([
    spaces.Box(low=0, high=255, shape=(84, 84, 3), dtype=np.uint8),
    spaces.Discrete(10)
])

# MultiBinary - binary flags
obs_space = spaces.MultiBinary(8)  # 8 binary values

---

Environment Testing Script

import gymnasium as gym
import numpy as np

def test_environment(env_or_id, n_episodes=5, render=False):
    """Test environment with random policy."""
    if isinstance(env_or_id, str):
        env = gym.make(env_or_id, render_mode="human" if render else None)
    else:
        env = env_or_id

    episode_rewards = []
    episode_lengths = []

    for ep in range(n_episodes):
        obs, info = env.reset()
        total_reward = 0
        steps = 0

        while True:
            action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            steps += 1

            if render:
                env.render()

            if terminated or truncated:
                break

        episode_rewards.append(total_reward)
        episode_lengths.append(steps)
        print(f"Episode {ep+1}: reward={total_reward:.2f}, length={steps}")

    print(f"\nSummary over {n_episodes} episodes:")
    print(f"  Mean reward: {np.mean(episode_rewards):.2f} +/- {np.std(episode_rewards):.2f}")
    print(f"  Mean length: {np.mean(episode_lengths):.1f} +/- {np.std(episode_lengths):.1f}")

    env.close()
    return episode_rewards, episode_lengths

# Usage
test_environment("CartPole-v1", n_episodes=10)

RL Evaluation and Reproducibility

Overview

Proper evaluation is critical in RL due to high variance. This guide covers metrics, evaluation procedures, logging, reproducibility, and statistical comparisons.

Essential Metrics to Track

During Training

During Evaluation

from stable_baselines3.common.evaluation import evaluate_policy
import numpy as np

def comprehensive_evaluation(model, env, n_episodes=100):
    """Comprehensive evaluation with multiple metrics."""

    episode_rewards = []
    episode_lengths = []
    success_count = 0

    for _ in range(n_episodes):
        obs, _ = env.reset()
        episode_reward = 0
        episode_length = 0
        done = False

        while not done:
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            episode_reward += reward
            episode_length += 1
            done = terminated or truncated

            # Track success if defined
            if 'is_success' in info and info['is_success']:
                success_count += 1

        episode_rewards.append(episode_reward)
        episode_lengths.append(episode_length)

    rewards = np.array(episode_rewards)
    lengths = np.array(episode_lengths)

    metrics = {
        "mean_reward": np.mean(rewards),
        "std_reward": np.std(rewards),
        "min_reward": np.min(rewards),
        "max_reward": np.max(rewards),
        "median_reward": np.median(rewards),
        "mean_length": np.mean(lengths),
        "std_length": np.std(lengths),
        "success_rate": success_count / n_episodes if 'is_success' in info else None,
        "n_episodes": n_episodes
    }

    # Confidence interval (95%)
    ci = 1.96 * metrics["std_reward"] / np.sqrt(n_episodes)
    metrics["ci_95"] = ci
    metrics["ci_lower"] = metrics["mean_reward"] - ci
    metrics["ci_upper"] = metrics["mean_reward"] + ci

    return metrics

# Usage
metrics = comprehensive_evaluation(model, eval_env, n_episodes=100)
print(f"Mean Reward: {metrics['mean_reward']:.2f} +/- {metrics['std_reward']:.2f}")
print(f"95% CI: [{metrics['ci_lower']:.2f}, {metrics['ci_upper']:.2f}]")

---

Separate Evaluation from Training

Why Separate?

1. Training uses exploration - Stochastic policy 2. Evaluation should be deterministic - Test true performance 3. Training env may have wrappers - Normalization stats differ 4. Prevents data leakage - Don't tune on test performance

Implementation

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import VecNormalize
from stable_baselines3.common.callbacks import EvalCallback

# Training environment (with normalization)
train_env = make_vec_env("Pendulum-v1", n_envs=4)
train_env = VecNormalize(train_env, norm_obs=True, norm_reward=True)

# Evaluation environment (separate, no reward normalization)
eval_env = make_vec_env("Pendulum-v1", n_envs=1)
eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False, training=False)

# Sync normalization stats (important!)
eval_env.obs_rms = train_env.obs_rms

# Evaluation callback
eval_callback = EvalCallback(
    eval_env,
    best_model_save_path="./logs/best_model",
    log_path="./logs/eval",
    eval_freq=10000,
    n_eval_episodes=20,
    deterministic=True,
    render=False
)

# Train
model = PPO("MlpPolicy", train_env, verbose=1)
model.learn(total_timesteps=100_000, callback=eval_callback)

Evaluation with Normalized Observations

# After training, load model and normalization stats
model = PPO.load("best_model")
eval_env = make_vec_env("Pendulum-v1", n_envs=1)

# Load normalization stats
eval_env = VecNormalize.load("vec_normalize.pkl", eval_env)
eval_env.training = False  # Don't update running stats
eval_env.norm_reward = False  # Don't normalize rewards

# Evaluate
mean_reward, std_reward = evaluate_policy(
    model, eval_env, n_eval_episodes=50, deterministic=True
)

---

Logging Frameworks

TensorBoard

from stable_baselines3 import PPO

# Enable TensorBoard logging
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    tensorboard_log="./tb_logs/"
)

model.learn(total_timesteps=100_000)

# View logs
# tensorboard --logdir ./tb_logs/

Weights & Biases

import wandb
from wandb.integration.sb3 import WandbCallback

# Initialize W&B
wandb.init(
    project="rl-experiments",
    config={
        "algorithm": "PPO",
        "env": "CartPole-v1",
        "learning_rate": 3e-4,
    },
    sync_tensorboard=True
)

model = PPO("MlpPolicy", env, verbose=1, tensorboard_log=f"runs/{wandb.run.id}")

model.learn(
    total_timesteps=100_000,
    callback=WandbCallback(
        gradient_save_freq=1000,
        model_save_path=f"models/{wandb.run.id}",
        verbose=2
    )
)

wandb.finish()

Custom Logging Callback

from stable_baselines3.common.callbacks import BaseCallback
import json
from datetime import datetime

class DetailedLoggingCallback(BaseCallback):
    """Log detailed metrics to JSON file."""

    def __init__(self, log_path, verbose=0):
        super().__init__(verbose)
        self.log_path = log_path
        self.logs = []

    def _on_step(self):
        # Log every 1000 steps
        if self.n_calls % 1000 == 0:
            log_entry = {
                "timestep": self.num_timesteps,
                "time": datetime.now().isoformat(),
            }

            # Episode info
            if len(self.model.ep_info_buffer) > 0:
                ep_rewards = [ep['r'] for ep in self.model.ep_info_buffer]
                ep_lengths = [ep['l'] for ep in self.model.ep_info_buffer]
                log_entry["mean_reward"] = np.mean(ep_rewards)
                log_entry["std_reward"] = np.std(ep_rewards)
                log_entry["mean_length"] = np.mean(ep_lengths)

            self.logs.append(log_entry)

        return True

    def _on_training_end(self):
        # Save all logs
        with open(self.log_path, 'w') as f:
            json.dump(self.logs, f, indent=2)
        print(f"Logs saved to {self.log_path}")

MLflow

import mlflow
from stable_baselines3 import PPO

# Start MLflow run
mlflow.set_experiment("rl-experiments")

with mlflow.start_run():
    # Log parameters
    mlflow.log_params({
        "algorithm": "PPO",
        "learning_rate": 3e-4,
        "n_steps": 2048,
        "env": "CartPole-v1"
    })

    # Train
    model = PPO("MlpPolicy", env, learning_rate=3e-4, n_steps=2048)
    model.learn(total_timesteps=100_000)

    # Evaluate and log metrics
    mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=100)
    mlflow.log_metrics({
        "mean_reward": mean_reward,
        "std_reward": std_reward
    })

    # Save model as artifact
    model.save("ppo_model")
    mlflow.log_artifact("ppo_model.zip")

---

Seed Management and Determinism

Setting Seeds Properly

import random
import numpy as np
import torch
import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.utils import set_random_seed

def make_deterministic(seed):
    """Set all seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # PyTorch deterministic operations
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def train_with_seed(env_id, seed, total_timesteps):
    """Train model with specific seed."""
    make_deterministic(seed)

    # Create environment with seed
    env = gym.make(env_id)
    env.reset(seed=seed)

    # For vectorized envs
    # env = make_vec_env(env_id, n_envs=4, seed=seed)

    # Create model with seed
    model = PPO("MlpPolicy", env, seed=seed, verbose=0)

    # Train
    model.learn(total_timesteps=total_timesteps)

    return model

# Run multiple seeds
results = {}
for seed in [0, 42, 123, 456, 789]:
    model = train_with_seed("CartPole-v1", seed, 50_000)
    mean_reward, _ = evaluate_policy(model, model.get_env(), n_eval_episodes=50)
    results[seed] = mean_reward
    print(f"Seed {seed}: {mean_reward:.2f}")

print(f"\nMean across seeds: {np.mean(list(results.values())):.2f}")
print(f"Std across seeds: {np.std(list(results.values())):.2f}")

Reproducibility Checklist

• [ ] Set random seed for Python's random
• [ ] Set numpy seed
• [ ] Set PyTorch seed (CPU and CUDA)
• [ ] Set environment seed in reset()
• [ ] Pass seed to model constructor
• [ ] Use deterministic PyTorch operations
• [ ] Document all library versions
• [ ] Log hardware information

Saving Experiment Configuration

import json
import torch
import stable_baselines3

def save_experiment_config(filepath, model, env, seed, hyperparams):
    """Save complete experiment configuration."""
    config = {
        "seed": seed,
        "environment": {
            "id": env.spec.id if hasattr(env, 'spec') else str(type(env)),
            "observation_space": str(env.observation_space),
            "action_space": str(env.action_space),
        },
        "hyperparameters": hyperparams,
        "versions": {
            "python": sys.version,
            "torch": torch.__version__,
            "stable_baselines3": stable_baselines3.__version__,
            "numpy": np.__version__,
            "gymnasium": gym.__version__,
        },
        "hardware": {
            "cuda_available": torch.cuda.is_available(),
            "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        }
    }

    with open(filepath, 'w') as f:
        json.dump(config, f, indent=2)

    return config

---

Statistical Comparison Methods

Comparing Two Algorithms

from scipy import stats
import numpy as np

def compare_algorithms(rewards_a, rewards_b, alpha=0.05):
    """
    Statistical comparison of two algorithms.

    Args:
        rewards_a: List of episode rewards for algorithm A
        rewards_b: List of episode rewards for algorithm B
        alpha: Significance level

    Returns:
        Dictionary with comparison results
    """
    rewards_a = np.array(rewards_a)
    rewards_b = np.array(rewards_b)

    # Basic statistics
    results = {
        "algo_a": {
            "mean": np.mean(rewards_a),
            "std": np.std(rewards_a),
            "median": np.median(rewards_a),
            "n": len(rewards_a)
        },
        "algo_b": {
            "mean": np.mean(rewards_b),
            "std": np.std(rewards_b),
            "median": np.median(rewards_b),
            "n": len(rewards_b)
        }
    }

    # Welch's t-test (doesn't assume equal variance)
    t_stat, p_value = stats.ttest_ind(rewards_a, rewards_b, equal_var=False)
    results["welch_ttest"] = {
        "t_statistic": t_stat,
        "p_value": p_value,
        "significant": p_value < alpha
    }

    # Mann-Whitney U test (non-parametric)
    u_stat, p_value_mw = stats.mannwhitneyu(rewards_a, rewards_b, alternative='two-sided')
    results["mann_whitney"] = {
        "u_statistic": u_stat,
        "p_value": p_value_mw,
        "significant": p_value_mw < alpha
    }

    # Effect size (Cohen's d)
    pooled_std = np.sqrt((rewards_a.std()**2 + rewards_b.std()**2) / 2)
    cohens_d = (rewards_a.mean() - rewards_b.mean()) / pooled_std
    results["effect_size"] = {
        "cohens_d": cohens_d,
        "interpretation": interpret_cohens_d(cohens_d)
    }

    # 95% confidence interval for difference
    diff_mean = rewards_a.mean() - rewards_b.mean()
    diff_se = np.sqrt(rewards_a.var()/len(rewards_a) + rewards_b.var()/len(rewards_b))
    results["difference"] = {
        "mean": diff_mean,
        "ci_lower": diff_mean - 1.96 * diff_se,
        "ci_upper": diff_mean + 1.96 * diff_se
    }

    return results

def interpret_cohens_d(d):
    """Interpret Cohen's d effect size."""
    d = abs(d)
    if d < 0.2:
        return "negligible"
    elif d < 0.5:
        return "small"
    elif d < 0.8:
        return "medium"
    else:
        return "large"

# Example usage
ppo_rewards = [195, 200, 198, 188, 200, 195, 192, 200, 197, 199]  # 10 eval episodes
dqn_rewards = [180, 175, 190, 185, 178, 182, 188, 179, 183, 186]

comparison = compare_algorithms(ppo_rewards, dqn_rewards)
print(f"PPO mean: {comparison['algo_a']['mean']:.2f} +/- {comparison['algo_a']['std']:.2f}")
print(f"DQN mean: {comparison['algo_b']['mean']:.2f} +/- {comparison['algo_b']['std']:.2f}")
print(f"Significant difference: {comparison['welch_ttest']['significant']}")
print(f"Effect size: {comparison['effect_size']['interpretation']}")

Multiple Seeds Comparison

def multi_seed_comparison(algo_results):
    """
    Compare algorithms across multiple seeds.

    Args:
        algo_results: Dict mapping algorithm name to list of (seed, reward) tuples

    Returns:
        Comparison results
    """
    # Aggregate by algorithm
    algo_rewards = {}
    for algo, results in algo_results.items():
        algo_rewards[algo] = [r for _, r in results]

    # Kruskal-Wallis test (non-parametric ANOVA)
    groups = list(algo_rewards.values())
    h_stat, p_value = stats.kruskal(*groups)

    print("=" * 50)
    print("Multi-Seed Algorithm Comparison")
    print("=" * 50)

    for algo, rewards in algo_rewards.items():
        print(f"\n{algo}:")
        print(f"  Seeds: {len(rewards)}")
        print(f"  Mean: {np.mean(rewards):.2f} +/- {np.std(rewards):.2f}")
        print(f"  Min/Max: {np.min(rewards):.2f} / {np.max(rewards):.2f}")

    print(f"\nKruskal-Wallis test:")
    print(f"  H-statistic: {h_stat:.4f}")
    print(f"  p-value: {p_value:.4f}")
    print(f"  Significant (p<0.05): {p_value < 0.05}")

    return algo_rewards, (h_stat, p_value)

# Example
results = {
    "PPO": [(0, 195), (1, 198), (2, 192), (3, 200), (4, 197)],
    "DQN": [(0, 180), (1, 175), (2, 185), (3, 178), (4, 182)],
    "A2C": [(0, 188), (1, 190), (2, 185), (3, 192), (4, 187)]
}

multi_seed_comparison(results)

---

Checkpointing Strategies

Basic Checkpointing

from stable_baselines3.common.callbacks import CheckpointCallback

checkpoint_callback = CheckpointCallback(
    save_freq=10000,  # Save every 10k steps
    save_path="./checkpoints/",
    name_prefix="rl_model",
    save_replay_buffer=True,  # For off-policy
    save_vecnormalize=True    # Save normalization stats
)

model.learn(total_timesteps=100_000, callback=checkpoint_callback)

Best Model Saving

from stable_baselines3.common.callbacks import EvalCallback

eval_callback = EvalCallback(
    eval_env,
    best_model_save_path="./best_model/",
    log_path="./logs/",
    eval_freq=5000,
    n_eval_episodes=10,
    deterministic=True
)

model.learn(total_timesteps=100_000, callback=eval_callback)

# Load best model
best_model = PPO.load("./best_model/best_model")

Complete Checkpoint with Normalization

def save_complete_checkpoint(model, env, path):
    """Save model, normalization stats, and config."""
    import os
    os.makedirs(path, exist_ok=True)

    # Save model
    model.save(os.path.join(path, "model"))

    # Save normalization stats if using VecNormalize
    if isinstance(env, VecNormalize):
        env.save(os.path.join(path, "vec_normalize.pkl"))

    # Save config
    config = {
        "algorithm": type(model).__name__,
        "policy": model.policy_class.__name__,
        "n_envs": model.n_envs,
        "gamma": model.gamma,
    }
    with open(os.path.join(path, "config.json"), 'w') as f:
        json.dump(config, f)

def load_complete_checkpoint(path, env_fn):
    """Load model with normalization stats."""
    import os

    # Load config
    with open(os.path.join(path, "config.json"), 'r') as f:
        config = json.load(f)

    # Create environment
    env = env_fn()

    # Load normalization stats if they exist
    norm_path = os.path.join(path, "vec_normalize.pkl")
    if os.path.exists(norm_path):
        env = VecNormalize.load(norm_path, env)
        env.training = False

    # Load model
    model_path = os.path.join(path, "model")
    model = PPO.load(model_path, env=env)

    return model, env

---

Evaluation Report Template

def generate_evaluation_report(model, env, n_episodes=100, output_path="eval_report.md"):
    """Generate comprehensive evaluation report."""

    # Run evaluation
    metrics = comprehensive_evaluation(model, env, n_episodes)

    # Generate report
    report = f"""# Reinforcement Learning Evaluation Report

## Summary
- **Algorithm**: {type(model).__name__}
- **Environment**: {env.spec.id if hasattr(env, 'spec') else 'Custom'}
- **Evaluation Episodes**: {n_episodes}

## Performance Metrics

| Metric | Value |
|--------|-------|
| Mean Reward | {metrics['mean_reward']:.2f} |
| Std Reward | {metrics['std_reward']:.2f} |
| 95% CI | [{metrics['ci_lower']:.2f}, {metrics['ci_upper']:.2f}] |
| Min Reward | {metrics['min_reward']:.2f} |
| Max Reward | {metrics['max_reward']:.2f} |
| Median Reward | {metrics['median_reward']:.2f} |
| Mean Episode Length | {metrics['mean_length']:.1f} |

## Interpretation

The agent achieves a mean reward of **{metrics['mean_reward']:.2f}** with a standard deviation of {metrics['std_reward']:.2f}.

The 95% confidence interval for the true mean performance is [{metrics['ci_lower']:.2f}, {metrics['ci_upper']:.2f}].

## Notes

- Evaluation was performed with deterministic policy
- All episodes were run to completion (terminated or truncated)
"""

    with open(output_path, 'w') as f:
        f.write(report)

    print(f"Report saved to {output_path}")
    return report

RL Training Best Practices

Overview

Training RL agents effectively requires careful attention to hyperparameters, reward engineering, exploration, and normalization. This guide covers practical techniques for successful training.

Hyperparameter Tuning

Starting Point Strategy

1. Start with defaults - SB3 defaults are well-tuned for common cases 2. Scale learning rate - Adjust based on environment complexity 3. Adjust rollout length - Match to environment horizon 4. Tune exploration - Critical for hard exploration problems

Optuna Integration

import optuna
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.env_util import make_vec_env

def objective(trial):
    """Optuna objective for PPO hyperparameter tuning."""

    # Suggest hyperparameters
    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-3, log=True)
    n_steps = trial.suggest_categorical("n_steps", [256, 512, 1024, 2048])
    batch_size = trial.suggest_categorical("batch_size", [32, 64, 128, 256])
    n_epochs = trial.suggest_int("n_epochs", 3, 30)
    gamma = trial.suggest_float("gamma", 0.9, 0.9999)
    gae_lambda = trial.suggest_float("gae_lambda", 0.8, 1.0)
    clip_range = trial.suggest_float("clip_range", 0.1, 0.4)
    ent_coef = trial.suggest_float("ent_coef", 1e-8, 0.1, log=True)

    # Ensure batch_size <= n_steps
    batch_size = min(batch_size, n_steps)

    # Create environment
    env = make_vec_env("CartPole-v1", n_envs=4)

    # Create model
    model = PPO(
        "MlpPolicy",
        env,
        learning_rate=learning_rate,
        n_steps=n_steps,
        batch_size=batch_size,
        n_epochs=n_epochs,
        gamma=gamma,
        gae_lambda=gae_lambda,
        clip_range=clip_range,
        ent_coef=ent_coef,
        verbose=0
    )

    # Train
    model.learn(total_timesteps=50_000)

    # Evaluate
    mean_reward, std_reward = evaluate_policy(
        model, model.get_env(), n_eval_episodes=10, deterministic=True
    )

    # Report intermediate value (for pruning)
    trial.report(mean_reward, step=50_000)

    return mean_reward

# Create study with pruning
study = optuna.create_study(
    direction="maximize",
    pruner=optuna.pruners.MedianPruner(n_warmup_steps=5)
)

# Optimize
study.optimize(objective, n_trials=100, timeout=3600)

# Results
print(f"Best trial: {study.best_trial.value}")
print(f"Best params: {study.best_params}")

SAC Hyperparameter Tuning

def sac_objective(trial):
    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-3, log=True)
    buffer_size = trial.suggest_categorical("buffer_size", [10000, 50000, 100000, 500000])
    batch_size = trial.suggest_categorical("batch_size", [64, 128, 256, 512])
    gamma = trial.suggest_float("gamma", 0.9, 0.9999)
    tau = trial.suggest_float("tau", 0.001, 0.05)

    model = SAC(
        "MlpPolicy", "Pendulum-v1",
        learning_rate=learning_rate,
        buffer_size=buffer_size,
        batch_size=batch_size,
        gamma=gamma,
        tau=tau,
        verbose=0
    )
    model.learn(total_timesteps=20_000)

    mean_reward, _ = evaluate_policy(model, model.get_env(), n_eval_episodes=10)
    return mean_reward

Key Hyperparameters by Algorithm

PPO

SAC

---

Reward Engineering

Principles

1. Dense > Sparse - More frequent feedback accelerates learning 2. Shaped rewards - Guide agent toward goal without changing optimal policy 3. Potential-based shaping - Guarantees policy invariance 4. Normalize rewards - Keep rewards in reasonable range

Sparse Rewards Problem

# Bad: Sparse reward (agent may never see positive reward)
def sparse_reward(achieved_goal, desired_goal):
    return 1.0 if np.allclose(achieved_goal, desired_goal) else 0.0

# Better: Distance-based shaping
def shaped_reward(achieved_goal, desired_goal, prev_distance=None):
    current_distance = np.linalg.norm(achieved_goal - desired_goal)

    # Terminal reward
    if current_distance < 0.1:
        return 10.0

    # Shaping reward (progress toward goal)
    if prev_distance is not None:
        progress = prev_distance - current_distance
        return progress  # Positive if getting closer

    return -0.1  # Small negative to encourage speed

Potential-Based Shaping

def potential(state):
    """Potential function based on distance to goal."""
    goal = np.array([0, 0])
    return -np.linalg.norm(state - goal)

def shaped_reward(state, next_state, base_reward, gamma=0.99):
    """Add potential-based shaping that preserves optimal policy."""
    shaping = gamma * potential(next_state) - potential(state)
    return base_reward + shaping

Reward Normalization

from stable_baselines3.common.vec_env import VecNormalize

# Automatic reward normalization
env = make_vec_env("Pendulum-v1", n_envs=4)
env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_reward=10.0)

# Manual running average
class RewardNormalizer:
    def __init__(self, gamma=0.99):
        self.return_rms = RunningMeanStd()
        self.returns = 0
        self.gamma = gamma

    def normalize(self, reward, done):
        self.returns = self.returns * self.gamma + reward
        self.return_rms.update(np.array([self.returns]))
        normalized = reward / (np.sqrt(self.return_rms.var) + 1e-8)
        if done:
            self.returns = 0
        return normalized

Common Reward Mistakes

---

Exploration Strategies

Entropy Bonus

Encourages policy to maintain exploration by penalizing deterministic policies.

# PPO with entropy bonus
model = PPO(
    "MlpPolicy", "CartPole-v1",
    ent_coef=0.01,  # Entropy coefficient
    verbose=1
)

# Monitor entropy during training (via TensorBoard)
# Look for "train/entropy_loss" - should not drop to 0

Epsilon-Greedy (DQN)

# DQN exploration schedule
model = DQN(
    "MlpPolicy", "CartPole-v1",
    exploration_fraction=0.2,      # Fraction of training for decay
    exploration_initial_eps=1.0,   # Start with full random
    exploration_final_eps=0.05,    # End with 5% random
    verbose=1
)

Action Noise (Continuous)

from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
import numpy as np

n_actions = env.action_space.shape[-1]

# Gaussian noise
action_noise = NormalActionNoise(
    mean=np.zeros(n_actions),
    sigma=0.1 * np.ones(n_actions)
)

# Ornstein-Uhlenbeck noise (temporally correlated)
action_noise = OrnsteinUhlenbeckActionNoise(
    mean=np.zeros(n_actions),
    sigma=0.1 * np.ones(n_actions),
    theta=0.15,  # Rate of mean reversion
    dt=1e-2
)

model = TD3("MlpPolicy", env, action_noise=action_noise)

Curiosity-Driven Exploration

For environments with very sparse rewards:

# Using stable-baselines3-contrib
from sb3_contrib import RecurrentPPO

# Or implement intrinsic motivation
class ICMReward:
    """Intrinsic Curiosity Module reward."""

    def __init__(self, feature_dim=64, lr=1e-3):
        self.forward_model = ...  # Predicts next state encoding
        self.inverse_model = ...  # Predicts action from states

    def compute_intrinsic_reward(self, obs, action, next_obs):
        # Reward based on prediction error
        predicted_next = self.forward_model(obs, action)
        actual_next = self.encode(next_obs)
        intrinsic_reward = torch.mean((predicted_next - actual_next) ** 2)
        return intrinsic_reward.item()

---

Normalization Techniques

Observation Normalization

from stable_baselines3.common.vec_env import VecNormalize

env = make_vec_env("Pendulum-v1", n_envs=4)

# Normalize observations with running statistics
env = VecNormalize(
    env,
    norm_obs=True,
    norm_reward=False,
    clip_obs=10.0,       # Clip normalized obs to [-10, 10]
    gamma=0.99
)

# Train
model = PPO("MlpPolicy", env)
model.learn(total_timesteps=100_000)

# Important: Save normalization stats with model
env.save("vec_normalize.pkl")
model.save("ppo_model")

# Load for evaluation
eval_env = make_vec_env("Pendulum-v1", n_envs=1)
eval_env = VecNormalize.load("vec_normalize.pkl", eval_env)
eval_env.training = False  # Don't update stats during eval
eval_env.norm_reward = False  # Don't normalize rewards during eval

Advantage Normalization

PPO and A2C automatically normalize advantages:

# In PPO, advantages are normalized by default
# (advantage - mean) / (std + eps)

# This happens inside the algorithm, but you can control via:
model = PPO(
    "MlpPolicy", env,
    normalize_advantage=True  # Default
)

Manual Normalization

class RunningMeanStd:
    """Running mean and standard deviation."""

    def __init__(self, epsilon=1e-4, shape=()):
        self.mean = np.zeros(shape, dtype=np.float64)
        self.var = np.ones(shape, dtype=np.float64)
        self.count = epsilon

    def update(self, batch):
        batch_mean = np.mean(batch, axis=0)
        batch_var = np.var(batch, axis=0)
        batch_count = batch.shape[0]
        self._update_from_moments(batch_mean, batch_var, batch_count)

    def _update_from_moments(self, batch_mean, batch_var, batch_count):
        delta = batch_mean - self.mean
        tot_count = self.count + batch_count

        new_mean = self.mean + delta * batch_count / tot_count
        m_a = self.var * self.count
        m_b = batch_var * batch_count
        m2 = m_a + m_b + np.square(delta) * self.count * batch_count / tot_count
        new_var = m2 / tot_count

        self.mean = new_mean
        self.var = new_var
        self.count = tot_count

    def normalize(self, x):
        return (x - self.mean) / (np.sqrt(self.var) + 1e-8)

---

Learning Rate Schedules

Linear Decay

from stable_baselines3.common.callbacks import BaseCallback

def linear_schedule(initial_value):
    """Linear learning rate schedule."""
    def func(progress_remaining):
        return progress_remaining * initial_value
    return func

model = PPO(
    "MlpPolicy", "CartPole-v1",
    learning_rate=linear_schedule(3e-4),  # Will decay from 3e-4 to 0
    verbose=1
)

Custom Schedules

def exponential_schedule(initial_value, decay_rate=0.99):
    """Exponential decay."""
    def func(progress_remaining):
        return initial_value * (decay_rate ** (1 - progress_remaining))
    return func

def warmup_schedule(initial_value, warmup_fraction=0.1):
    """Linear warmup then constant."""
    def func(progress_remaining):
        # progress_remaining goes from 1 to 0
        progress = 1 - progress_remaining
        if progress < warmup_fraction:
            return initial_value * (progress / warmup_fraction)
        return initial_value
    return func

def cosine_schedule(initial_value, min_value=1e-6):
    """Cosine annealing."""
    def func(progress_remaining):
        return min_value + 0.5 * (initial_value - min_value) * (1 + np.cos(np.pi * (1 - progress_remaining)))
    return func

---

Batch Size and Buffer Guidelines

On-Policy (PPO, A2C)

Total timesteps per update = n_steps * n_envs

Off-Policy (DQN, SAC, TD3)

Guidelines:

• Buffer should hold at least 10x the episode length
• Larger buffers improve stability but use more memory
• learning_starts should allow for some exploration first

---

Training Loop Best Practices

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import EvalCallback, CheckpointCallback
from stable_baselines3.common.env_util import make_vec_env

# Separate training and evaluation environments
train_env = make_vec_env("CartPole-v1", n_envs=4)
eval_env = make_vec_env("CartPole-v1", n_envs=1)

# Callbacks
eval_callback = EvalCallback(
    eval_env,
    best_model_save_path="./logs/best_model",
    log_path="./logs/eval",
    eval_freq=10000,
    n_eval_episodes=10,
    deterministic=True
)

checkpoint_callback = CheckpointCallback(
    save_freq=50000,
    save_path="./logs/checkpoints",
    name_prefix="ppo_model"
)

# Create model
model = PPO(
    "MlpPolicy",
    train_env,
    verbose=1,
    tensorboard_log="./logs/tensorboard"
)

# Train with callbacks
model.learn(
    total_timesteps=500_000,
    callback=[eval_callback, checkpoint_callback],
    progress_bar=True
)

# Final save
model.save("final_model")