Train an RL Agent with MuJoCo

Train a neural network to balance a pole using PPO and MuJoCo physics.

Introduction

30 Second Summary

Think about balancing a broomstick on your palm. Your brain constantly makes tiny corrections based on what it sees and feels. Teaching a computer to solve the same problem through trial and error is the core challenge of reinforcement learning.

In this project, you will train a neural network agent to balance a pole on a moving cart using reinforcement learning and the MuJoCo physics simulator. You'll watch the agent go from immediately failing to perfectly balancing, then capture the entire learning journey as portfolio-ready visualizations.

What You'll Build

You'll run your trained agent and watch it hold the pole perfectly upright for the full episode, then pull up the training curve that tells the story of how it learned, ready to show in a portfolio or explain in an interview.

By the end of this project, you'll have:

  • A random agent baseline that immediately drops the pole, demonstrating why the balance problem is hard for an untrained system.
  • A trained PPO agent that keeps the pole balanced for the full 1000-timestep episode through learned control.
  • Training curve plots and a comparison chart showing reward climbing from near-zero to maximum as the agent masters the task.
  • Secret Mission: Transfer your trained workflow to the HalfCheetah locomotion environment and train a simulated robot to run.

Are there any prerequisites?

You only need basic Python familiarity (loops, functions, imports). No reinforcement learning or MuJoCo experience is required.

Before We Start

Before we dive into code, let's lock in what you're about to build and why it matters.

You're going to train a neural network to solve a classic mechatronics control problem: balancing a pole on a moving cart. This is the exact workflow used in Physical AI labs before deploying policies to real robots.

Set Up Your RL Environment

Reinforcement learning experiments rely on a physics simulator to generate training data and a library to run the learning algorithms. Without these installed, you cannot run a single training episode.

In this step, you will create an isolated Python environment and install the packages your project needs: MuJoCo for physics simulation, Gymnasium for the RL environment interface, Stable Baselines3 for the PPO algorithm, and Matplotlib for visualizations.

In this step, get ready to:
  • Create and activate a Python virtual environment.
  • Install MuJoCo, Gymnasium, Stable Baselines3, and Matplotlib.
  • Verify all packages import correctly and the InvertedPendulum environment loads.
Create a virtual environment and activate it

A virtual environment keeps your project's packages isolated from your system Python. This prevents version conflicts between projects.

  • In VS Code, open a terminal by pressing Ctrl+` (backtick key, top-left of your keyboard).
  • Confirm your Python version is 3.10 or higher by running:
python3 --version

✔️ I see version 3.10 or higher

You are ready to continue. Your Python version supports MuJoCo and Stable Baselines3.

ⓧ I see a version lower than 3.10

MuJoCo and Stable Baselines3 both require Python 3.10 or higher. Download the latest Python from the official site:

https://www.python.org/downloads/

After installing, close and reopen your terminal, then re-run the version check. On macOS, you may need to use python3.12 instead of python3 if the system default still points to an older version.

ⓧ Command not found

Python is not installed or not on your PATH. Download Python 3.10 or higher from the official site:

https://www.python.org/downloads/

During installation on Windows, check the box that says Add python.exe to PATH. After installing, close and reopen your terminal, then re-run the version check.

  • Create a virtual environment by running this command:
python3 -m venv rl_env

What does this command do?

python3 -m venv rl_env tells Python to create a new virtual environment in a folder called rl_env. This folder contains its own copy of python and pip so that every package you install stays inside this project.

  • Activate the virtual environment by running the command for your operating system:

macOS

source rl_env/bin/activate

This command also works on Linux.

Windows

rl_env\Scripts\activate
  • Confirm activation succeeded by checking your terminal prompt. You should see (rl_env) at the beginning of your prompt line.

Don't see (rl_env) in your prompt?

Make sure you ran the activation command for your specific operating system. macOS and Linux use source while Windows runs the script directly.

If you see a permissions error on Windows, try running Set-ExecutionPolicy -Scope CurrentUser RemoteSigned in PowerShell first, then retry the activation command.

help me activate my Python virtual environment

Install the required packages
  • Install all four packages in one command by running:
pip install mujoco "gymnasium[mujoco]" "stable-baselines3[extra]" matplotlib

What does this install?

  • mujoco is the physics engine that simulates the pendulum and cart. A copy of the MuJoCo library is bundled inside the Python package, so there is nothing else to download.
  • gymnasium[mujoco] provides the standard RL environment interface and registers MuJoCo-specific environments like InvertedPendulum.
  • stable-baselines3[extra] includes the PPO algorithm plus extra dependencies for evaluation, logging, and TensorBoard.
  • matplotlib handles plotting your training curves and comparison charts.

This installation takes 1-2 minutes depending on your internet speed. You will see download progress bars and a final summary listing all successfully installed packages.

Installation failed or timed out?

Check that (rl_env) still appears in your terminal prompt. If it disappeared, re-run the activation command before retrying the install.

If you see a network timeout, retry the same pip install command. Pip resumes from cached downloads so it will be faster on a second attempt.

help me fix my pip install errors

Verify all imports and load the environment

Before you run the next command, what do you think will happen if one of the packages failed to install silently?

  • Test that all three core packages import correctly by running:
python -c "import mujoco; import gymnasium; import stable_baselines3; print('All imports OK')"

You should see All imports OK printed in your terminal. If any package failed to install, Python would throw an ImportError naming the missing module.

  • Now confirm the InvertedPendulum environment loads by running:
python -c "import gymnasium as gym; env = gym.make('InvertedPendulum-v5'); print(f'Environment loaded! Observation space: {env.observation_space}'); env.close()"

You should see Environment loaded! Observation space: Box(-inf, inf, (4,), float64) confirming MuJoCo and Gymnasium are working together.

What does this output mean?

The observation space tells you what the agent will see at each timestep: 4 floating-point numbers representing cart position, pole angle, cart velocity, and pole angular velocity. You will explore these in detail in the next step.

Seeing an import error?

If import mujoco fails, make sure you are running python from inside the activated virtual environment. Look for (rl_env) in your prompt.

If the InvertedPendulum environment fails to load, check that you installed gymnasium[mujoco] with the mujoco extra in brackets, not just gymnasium alone.

help me fix MuJoCo import errors in my Python environment

Your RL environment is ready. Next up, you will explore the InvertedPendulum environment to understand what the agent sees and does before training it.

Explore the Environment with a Random Agent

Your RL environment is installed and working. Now you need to understand what your agent actually interacts with before you can train it.

Every reinforcement learning problem has three parts: what the agent observes, what actions it can take, and how it gets rewarded. Running a random agent first establishes a baseline for comparison. Without any intelligence, how quickly does the pole fall?

In this step, get ready to:
  • Inspect the InvertedPendulum-v5 observation space and action space.
  • Run a random agent for 10 episodes and measure its performance.
  • Confirm the pole falls almost immediately without a learned policy.
Create the environment and inspect its spaces

Before an agent can learn, you need to know what information it receives (observations) and what levers it can pull (actions). The InvertedPendulum-v5 environment provides both in a compact form.

  • In VS Code, create a new file called 01_explore_environment.py in your project folder.
  • Add the following code to create the environment and print its spaces:
import gymnasium as gym
import numpy as np

# Create the Inverted Pendulum environment
env = gym.make("InvertedPendulum-v5", render_mode="human")

# Print environment details
print("=== Inverted Pendulum Environment ===")
print(f"Observation Space: {env.observation_space}")
print(f"  Shape: {env.observation_space.shape}")
print(f"  Observations: [cart_position, pole_angle, cart_velocity, pole_angular_velocity]")
print(f"\nAction Space: {env.action_space}")
print(f"  Shape: {env.action_space.shape}")
print(f"  Range: [{env.action_space.low[0]}, {env.action_space.high[0]}]")
print(f"  Action: force applied to cart in Newtons")

What does this code do?

This script creates a MuJoCo physics simulation of an inverted pendulum (a pole balanced on a cart).

  • The observation space has shape (4,) meaning the agent receives 4 numbers each timestep: cart position, pole angle, cart velocity, and pole angular velocity.
  • The action space has shape (1,) with range [-3, 3], meaning the agent applies a single force (in Newtons) to push the cart left or right.
  • The render_mode="human" setting tells Gymnasium to open a visual simulation window when the environment runs.
  • Save 01_explore_environment.py.
  • Run the script in your terminal:
python 01_explore_environment.py

You should see the observation and action space details printed in your terminal. The observation space is Box(-inf, inf, (4,), float64) and the action space is Box(-3.0, 3.0, (1,), float32).

Seeing an import or rendering error?

  • Make sure your virtual environment is still activated. You should see the environment name in your terminal prompt.
  • If you see a GL or rendering error, try changing render_mode="human" to render_mode="rgb_array" temporarily to skip the visual window.

Can't figure out the error? Help me debug my MuJoCo environment creation error

Run random episodes to establish a baseline

Now you know what the agent sees and does. The next question is: how badly does a random agent perform? The episode ends when the pole tilts past 0.2 radians, and a perfect agent would earn 1000 reward (one point per timestep for the full episode).

  • Add the following code below the print statements in 01_explore_environment.py:
# Run episodes with a random agent
num_episodes = 10
episode_rewards = []

for episode in range(num_episodes):
    observation, info = env.reset()
    total_reward = 0
    done = False

    while not done:
        action = env.action_space.sample()  # Random action
        observation, reward, terminated, truncated, info = env.step(action)
        total_reward += reward
        done = terminated or truncated

    episode_rewards.append(total_reward)
    print(f"Episode {episode + 1}: Reward = {total_reward:.0f}")

env.close()

print(f"\n=== Random Agent Results ===")
print(f"Mean Reward: {np.mean(episode_rewards):.1f} / 1000")
print(f"The pole falls almost immediately without a learned policy!")

What does this code do?

This loop runs 10 complete episodes where the agent takes random actions each timestep.

  • env.action_space.sample() picks a random force between -3 and 3 Newtons each step.
  • The episode ends when terminated becomes True (pole angle exceeds 0.2 radians) or truncated becomes True (1000 timesteps reached).
  • env.close() cleanly shuts down the simulation window after all episodes finish.
  • Save 01_explore_environment.py.

Before you run this, what do you think the average reward will be out of 1000? Remember, a random agent is just flailing the cart left and right with no strategy.

  • Run the script again:
python 01_explore_environment.py

A MuJoCo simulation window opens and you can watch the pole topple over almost instantly each episode. The mean reward prints around 20-40 out of a possible 1000.

Why is the reward so low?

The agent earns +1 for every timestep the pole stays upright. A perfect run lasts 1000 timesteps for a reward of 1000.

A random agent pushes the cart in arbitrary directions, so the pole crosses the 0.2 radian threshold within about 20-40 timesteps. This is the baseline your trained agent will dramatically beat in the next step.

Not seeing the simulation window?

  • The window may appear behind your VS Code window. Check your taskbar for a new application.
  • If the script runs but you see no window and no errors, it is still computing episodes. Wait for the terminal output to finish.
  • On some systems, the window closes very quickly because episodes are short. The terminal output confirms everything worked.

Still stuck? Help me troubleshoot my MuJoCo rendering window not appearing

✔️ Awesome, I've got everything!

Great. Double check you have saved 01_explore_environment.py and your terminal shows the random agent's mean reward around 20-40.

ⓧ I'd like to double check the full code

Your 01_explore_environment.py should have two sections in order:

  • The imports and environment creation chunk (import gymnasium through the action space print statements).
  • The random episodes chunk (num_episodes = 10 through the final print statement about the pole falling).

Scroll up and confirm both chunks are present in that order, with no extra code between them.

Your baseline is set. A random agent scores around 20-40 out of 1000, proving this control problem is impossible to solve without learning. Next up, you will train a PPO agent to do dramatically better.

Train a PPO Agent (Intentionally Undertrained)

Your random agent proved that brute force doesn't work. The pole falls in milliseconds without intelligence behind the controls. Now you'll bring in reinforcement learning to teach a neural network how to balance.

PPO (Proximal Policy Optimization) learns by collecting experience and updating its policy. In this step, you'll intentionally train for too few timesteps to watch the agent struggle. This builds intuition for why training duration matters in RL.

In this step, get ready to:
  • Write a training script that creates a PPO agent and trains it for 10,000 timesteps.
  • Evaluate the trained agent and render an episode to see how it performs.
Write the training script

PPO works by running the agent in the environment, collecting batches of experience, then updating the neural network's weights to make successful actions more likely. The "MlpPolicy" means the network is a simple multi-layer perceptron, perfect for a low-dimensional problem like balancing a pole.

  • Create a new file called 02_train_undertrained.py in your project folder.
  • Add the training code by pasting the following:
import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy

# Create environment
env = gym.make("InvertedPendulum-v5")

# Train PPO for only 10,000 timesteps (intentionally too few)
print("Training PPO for 10,000 timesteps (intentionally undertrained)...")
model = PPO("MlpPolicy", env, verbose=1, device="auto")
model.learn(total_timesteps=10_000)

# Save the undertrained model
model.save("ppo_pendulum_undertrained")
print("Model saved as 'ppo_pendulum_undertrained.zip'")

What does this code do?

This script sets up the RL training loop in just a few lines thanks to Stable Baselines3's simple API.

  • "MlpPolicy" tells PPO to use a multi-layer perceptron neural network as the policy. It maps the 4 observations directly to the 1 action value.
  • device="auto" automatically uses your GPU if available, falling back to CPU otherwise. Your RTX 5060 will speed up training.
  • verbose=1 prints training progress so you can watch the agent learn in real time.
  • 10,000 timesteps is intentionally low. Each timestep is one action in the environment. A full episode is up to 1000 timesteps, so 10,000 means the agent only gets about 10 full episodes of practice.
Add evaluation and rendering

Training produces a model file, but you need to see how the agent actually performs. The evaluation code runs 10 test episodes and renders one so you can watch the agent attempt to balance.

  • Add the following code below the save line in 02_train_undertrained.py:
# Evaluate the undertrained agent
print("\n=== Evaluating Undertrained Agent ===")
eval_env = gym.make("InvertedPendulum-v5")
mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10)
print(f"Mean Reward: {mean_reward:.1f} +/- {std_reward:.1f}")
print(f"The agent is learning but not yet good enough to balance reliably.")

eval_env.close()
env.close()

# Render one episode to visualize
render_env = gym.make("InvertedPendulum-v5", render_mode="human")
obs, info = render_env.reset()
total_reward = 0
done = False

while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)
    total_reward += reward
    done = terminated or truncated

print(f"Rendered episode reward: {total_reward:.0f}")
render_env.close()

What does this code do?

  • evaluate_policy() runs the trained agent for 10 episodes and calculates the mean reward. This gives you a reliable number instead of judging from a single run.
  • deterministic=True means the agent always picks its best action (no random exploration). This shows the agent's true ability.
  • The render section opens a MuJoCo window so you can visually watch the agent try to balance the pole.

Before you run this, think about what you expect. The random agent scored about 20-40 out of 1000. Do you think 10,000 timesteps of training is enough to master balancing?

  • Save 02_train_undertrained.py.
  • Run the script by executing this command in your terminal:
python 02_train_undertrained.py

Training takes about 30 seconds. You'll see PPO's training logs scroll by, then the evaluation results print.

The mean reward lands somewhere around 100-300. That's better than the random agent's 20-40, so the network IS learning. But it's far from the maximum 1000. When the render window opens, you'll see the pole wobble back and forth before eventually falling over.

Why does the agent still fail?

10,000 timesteps is only about 10 full episodes of practice. The PPO algorithm needs far more experience to converge on a good policy. The agent has learned that pushing the cart matters, but hasn't refined the timing and force needed for sustained balance.

This is the core challenge of RL: finding the right training duration. Too little, and the policy is half-baked. You'll fix this in the next step with proper training.

Script crashes or no render window?

  • If you see an import error, make sure your virtual environment is activated. You should see the environment name in your terminal prompt.
  • If the render window doesn't appear, check that your display server is running. On Windows or macOS this should work out of the box.
  • If training is extremely slow, confirm your GPU is detected. The training logs will show the device being used.

Help me debug my PPO training script for InvertedPendulum

✔️ Awesome, I've got everything!

Great. Double check you've saved your file, then move on.

ⓧ I'd like to double check the full code

Here's the complete 02_train_undertrained.py file:

import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy

# Create environment
env = gym.make("InvertedPendulum-v5")

# Train PPO for only 10,000 timesteps (intentionally too few)
print("Training PPO for 10,000 timesteps (intentionally undertrained)...")
model = PPO("MlpPolicy", env, verbose=1, device="auto")
model.learn(total_timesteps=10_000)

# Save the undertrained model
model.save("ppo_pendulum_undertrained")
print("Model saved as 'ppo_pendulum_undertrained.zip'")

# Evaluate the undertrained agent
print("\n=== Evaluating Undertrained Agent ===")
eval_env = gym.make("InvertedPendulum-v5")
mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10)
print(f"Mean Reward: {mean_reward:.1f} +/- {std_reward:.1f}")
print(f"The agent is learning but not yet good enough to balance reliably.")

eval_env.close()
env.close()

# Render one episode to visualize
render_env = gym.make("InvertedPendulum-v5", render_mode="human")
obs, info = render_env.reset()
total_reward = 0
done = False

while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)
    total_reward += reward
    done = terminated or truncated

print(f"Rendered episode reward: {total_reward:.0f}")
render_env.close()

Your agent is learning, but it's clearly not there yet. The pole wobbles and eventually gives up. In the next step, you'll diagnose exactly why this happens using training curves and fix it with proper training duration.

Diagnose and Fix with Full Training

In the previous step, your PPO agent scored between 100-300 out of a possible 1000. That's clearly better than the random agent's ~30, which proves the reinforcement learning algorithm IS working.

The problem is convergence. The agent didn't get enough practice episodes to fully learn the balancing strategy. In this step, you'll diagnose the gap using training curves and fix it by giving PPO 10x more experience.

In this step, get ready to:
  • Train a new PPO agent for 100,000 timesteps with evaluation callbacks that track learning progress.
  • Plot the training curve to visualize how reward climbs over time.
  • Evaluate the fully trained model and confirm a perfect score of 1000.
Train a PPO agent for 100,000 timesteps

The undertrained agent from the last step only got 10,000 timesteps of experience. That's like giving someone 10 minutes to learn to ride a bicycle. This time you'll provide 100,000 timesteps and track evaluation metrics at regular intervals so you can watch the learning happen.

  • Create a new file called 03_train_full.py in your project folder.
  • Add the imports, environment setup, and evaluation callback by pasting this code:
import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.callbacks import EvalCallback

# Create environment
env = gym.make("InvertedPendulum-v5")
eval_env = gym.make("InvertedPendulum-v5")

# Set up evaluation callback to track progress
eval_callback = EvalCallback(
    eval_env,
    best_model_save_path="./best_model/",
    log_path="./eval_logs/",
    eval_freq=5000,
    n_eval_episodes=10,
    deterministic=True,
)

What does this code do?

  • EvalCallback automatically evaluates the agent every 5,000 timesteps during training and saves the best-performing model.
  • eval_freq=5000 means you'll get 20 evaluation checkpoints across 100,000 timesteps. These checkpoints become the data points on your training curve.
  • The callback saves results to ./eval_logs/ so you can plot them after training finishes.
  • Now add the PPO model with full hyperparameters and the training call below the callback. Paste this code at the bottom of the file:
# Train PPO for 100,000 timesteps
print("Training PPO for 100,000 timesteps...")
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    device="auto",
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    gae_lambda=0.95,
    clip_range=0.2,
    tensorboard_log="./tensorboard_logs/",
)
model.learn(total_timesteps=100_000, callback=eval_callback)

# Save the trained model
model.save("ppo_pendulum_trained")
print("Model saved as 'ppo_pendulum_trained.zip'")

What do these hyperparameters mean?

  • learning_rate=3e-4 controls how much the neural network adjusts on each update. This is PPO's default, which works well for most tasks.
  • n_steps=2048 means the agent collects 2048 timesteps of experience before each policy update.
  • gamma=0.99 is the discount factor. High values like 0.99 make the agent care about future rewards, which is important for balancing.
  • clip_range=0.2 prevents PPO from making too-large policy changes at once, keeping training stable.
  • Save 03_train_full.py.
  • Run the training script by running this command:
python 03_train_full.py

Training takes about 2-3 minutes with GPU acceleration. You'll see rolling statistics printed as the agent collects experience and updates its policy.

When it finishes, you should see Model saved as 'ppo_pendulum_trained.zip' printed in your terminal.

Training seems stuck or very slow?

  • If you see CUDA warnings, the training will still work on CPU. It just takes slightly longer (about 3-5 minutes instead of 2-3).
  • If you get a memory error, close other applications and try again. MuJoCo environments are lightweight but PyTorch needs memory for the neural network.

Why is my PPO training running slowly or throwing errors?

Evaluate and visualize the trained agent

Your model is saved, but you haven't yet seen proof that it actually learned to balance. Now you'll add evaluation code that measures the trained agent's reward and plots the entire learning journey as a training curve.

  • In 03_train_full.py, add the evaluation and plotting code below the model save line. Paste this at the bottom of the file:
# Evaluate the fully trained agent
print("\n=== Evaluating Fully Trained Agent ===")
mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=10)
print(f"Mean Reward: {mean_reward:.1f} +/- {std_reward:.1f}")

# Plot training curve from eval logs
eval_results = np.load("./eval_logs/evaluations.npz")
timesteps = eval_results["timesteps"]
results = eval_results["results"]
mean_rewards = np.mean(results, axis=1)

plt.figure(figsize=(10, 6))
plt.plot(timesteps, mean_rewards, linewidth=2, color="blue")
plt.fill_between(
    timesteps,
    mean_rewards - np.std(results, axis=1),
    mean_rewards + np.std(results, axis=1),
    alpha=0.2,
    color="blue",
)
plt.xlabel("Training Timesteps", fontsize=12)
plt.ylabel("Mean Episode Reward", fontsize=12)
plt.title("PPO Training Progress: Inverted Pendulum", fontsize=14)
plt.axhline(y=1000, color="green", linestyle="--", label="Maximum Reward (1000)")
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("training_curve.png", dpi=150)
plt.show()
print("Training curve saved as 'training_curve.png'")

What does this code do?

  • evaluate_policy runs the trained agent for 10 full episodes and reports the mean reward. A perfect score of 1000 means the pole never fell.
  • The plotting section loads the evaluation checkpoints that EvalCallback saved every 5,000 timesteps, then draws the reward curve with a shaded standard deviation band.
  • The green dashed line at 1000 marks the theoretical maximum. You should see the blue curve climbing toward it.
  • Finally, add the rendering code and environment cleanup at the bottom of the file:
eval_env.close()
env.close()

# Render the trained agent
render_env = gym.make("InvertedPendulum-v5", render_mode="human")
obs, info = render_env.reset()
total_reward = 0
done = False

while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)
    total_reward += reward
    done = terminated or truncated

print(f"Rendered episode reward: {total_reward:.0f}")
render_env.close()

What does the render code do?

This opens a MuJoCo visualization window and runs one full episode using the trained model. You'll see the cart making tiny adjustments to keep the pole perfectly upright for the entire 1000 timesteps.

  • Save 03_train_full.py.

Before you run this, what do you think the mean reward will be after 100,000 timesteps of training? Will it hit the maximum 1000?

  • Run the complete training script by running this command:
python 03_train_full.py

The full training takes about 2-3 minutes with GPU acceleration. After training completes, you'll see the evaluation results, a matplotlib plot window, and the MuJoCo render window.

Your terminal should print Mean Reward: 1000.0 +/- 0.0. That's a perfect score. The pole never falls during the entire 1000-timestep episode.

What should I see in the training curve?

The plot shows reward starting low on the left (where the agent is still learning) and climbing toward the green dashed line at 1000 on the right. The shaded band shows variance across evaluation episodes. As the agent masters balancing, the band narrows because performance becomes consistent.

Why did 100,000 timesteps work when 10,000 failed?

PPO learns by collecting experience and making small policy updates. With only 10,000 timesteps, the agent only got about 5 policy updates. With 100,000 timesteps, it got roughly 49 updates, enough to converge on the optimal balancing strategy.

This is the core diagnostic skill in RL. When an agent underperforms, the first question is always: did it get enough training? The training curve answers that question visually.

Mean reward isn't reaching 1000?

  • If you see a reward around 800-900, the agent may need slightly more training. Try increasing total_timesteps to 150,000.
  • If the plot doesn't appear, close the MuJoCo render window first. Matplotlib may wait for focus.
  • Check that your file has callback=eval_callback in the model.learn() call. Without it, the eval logs won't exist and the plot will fail.

Why isn't my trained PPO agent reaching a perfect score on InvertedPendulum?

✔️ Awesome, I've got everything!

Great. Double check you've saved 03_train_full.py and confirm you see training_curve.png in your project folder.

ⓧ I'd like to double check the full code

Your 03_train_full.py file is made up of four sections you built above. Scroll up and confirm each is present in this order:

  • Imports, environment creation, and EvalCallback setup (ends with deterministic=True,).
  • PPO model creation with hyperparameters, model.learn() call, and model save (ends with the zip print statement).
  • Evaluation, training curve plot with plt.savefig("training_curve.png", dpi=150), and the saved-confirmation print.
  • Environment close calls, then the render loop ending with render_env.close().

Your agent went from random failure (~30 reward) to wobbly attempts (~200 reward) to perfect balancing (1000 reward). That training curve tells the whole story. Next up, you'll package everything into portfolio-ready artifacts.

Create Portfolio Artifacts

Your agent can balance a pole perfectly. But a trained model file sitting in a folder doesn't tell anyone what you accomplished.

In this step, you'll generate visual proof of your RL pipeline: a comparison chart that shows the dramatic improvement from random to trained, and a recorded video of your agent in action.

In this step, get ready to:
  • Generate a comparison bar chart showing random vs undertrained vs fully trained performance.
  • Confirm the training curve plot from the previous step is ready.
  • Record a video of your trained agent balancing the pole.
Evaluate all three agent types

To build a comparison chart, you first need the performance numbers for each stage of training: the random agent, the undertrained model, and the fully trained model.

  • Create a new file called 04_portfolio_artifacts.py in VS Code.
  • Add the imports and random agent evaluation loop by pasting this code:
import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy
from gymnasium.wrappers import RecordVideo

# Load models
model_undertrained = PPO.load("ppo_pendulum_undertrained")
model_trained = PPO.load("ppo_pendulum_trained")

# Evaluate all three approaches
eval_env = gym.make("InvertedPendulum-v5")

# Random agent evaluation
random_rewards = []
for _ in range(10):
    obs, info = eval_env.reset()
    total_reward = 0
    done = False
    while not done:
        action = eval_env.action_space.sample()
        obs, reward, terminated, truncated, info = eval_env.step(action)
        total_reward += reward
        done = terminated or truncated
    random_rewards.append(total_reward)

What does this code do?

  • PPO.load() loads the two saved models from earlier steps so you can evaluate them without retraining.
  • The random evaluation loop runs 10 episodes using env.action_space.sample() (completely random forces) to establish the baseline.
  • Below the random evaluation loop, add the model evaluations by pasting this code:
# Undertrained agent evaluation
undertrained_reward, _ = evaluate_policy(model_undertrained, eval_env, n_eval_episodes=10)

# Trained agent evaluation
trained_reward, _ = evaluate_policy(model_trained, eval_env, n_eval_episodes=10)

eval_env.close()

What does this code do?

evaluate_policy() is a Stable Baselines3 utility that runs the model for a set number of episodes and returns the mean reward. It handles the predict-step-done loop internally, so you get a clean number for comparison.

  • Save 04_portfolio_artifacts.py and run it to confirm the models load and evaluations complete:
python 04_portfolio_artifacts.py

The script should run for a few seconds (loading models and running episodes) then exit cleanly with no errors.

Seeing a FileNotFoundError?

  • Make sure ppo_pendulum_undertrained.zip and ppo_pendulum_trained.zip exist in the same folder as your script. These were saved in Steps 3 and 4.
  • Confirm your terminal is in the project folder by running ls *.zip to list the model files.

Why can't my script find the model files?

Add the comparison bar chart

Now you'll visualize the performance gap between all three agent stages as a bar chart. This is the single most impactful portfolio image because it tells the full story at a glance.

  • In 04_portfolio_artifacts.py, add the chart code below eval_env.close():
# Create comparison bar chart
labels = ["Random Agent", "Undertrained\n(10K steps)", "Fully Trained\n(100K steps)"]
rewards = [np.mean(random_rewards), undertrained_reward, trained_reward]
colors = ["#e74c3c", "#f39c12", "#27ae60"]

plt.figure(figsize=(8, 6))
bars = plt.bar(labels, rewards, color=colors, edgecolor="black", linewidth=0.8)
plt.ylabel("Mean Episode Reward", fontsize=12)
plt.title("RL Agent Performance Comparison: Inverted Pendulum", fontsize=14)
plt.axhline(y=1000, color="green", linestyle="--", alpha=0.5, label="Maximum (1000)")
plt.legend(fontsize=11)
plt.ylim(0, 1100)

for bar, reward in zip(bars, rewards):
    plt.text(
        bar.get_x() + bar.get_width() / 2.0,
        bar.get_height() + 20,
        f"{reward:.0f}",
        ha="center",
        fontsize=12,
        fontweight="bold",
    )

plt.tight_layout()
plt.savefig("performance_comparison.png", dpi=150)
plt.show()
print("Comparison chart saved as 'performance_comparison.png'")

What does this code do?

  • Three bars (red, orange, green) represent random, undertrained, and fully trained performance side by side.
  • A dashed green line at 1000 marks the maximum possible reward (the pole never falls for the entire episode).
  • The plt.text() loop labels each bar with its exact numeric reward value.
  • Save 04_portfolio_artifacts.py and run it:
python 04_portfolio_artifacts.py

A chart window pops up showing three bars. You should see the random agent bar barely visible near the bottom, the undertrained bar in the middle range, and the fully trained bar reaching the 1000 maximum line.

Close the chart window. You should see Comparison chart saved as 'performance_comparison.png' printed in the terminal.

Chart window not appearing?

  • If you're running over SSH or in a headless environment, plt.show() won't display a window. The PNG is still saved regardless. Check your project folder for performance_comparison.png.
  • If you see a Matplotlib backend error, the chart still saved to disk. Open the PNG directly to verify it looks correct.

Why isn't my matplotlib chart displaying?

Confirm your training curve from the previous step

You already generated training_curve.png in Step 4 when you ran 03_train_full.py. This plot shows your agent's reward climbing from near-zero to the 1000 maximum over 100,000 timesteps.

  • Confirm the file exists by running:
ls training_curve.png

You should see training_curve.png listed. This is your second portfolio artifact alongside the comparison chart.

Record a video of the trained agent

A static chart proves numbers. A video proves the agent actually works in real time. Gymnasium's RecordVideo wrapper captures frames during an episode and saves them as an mp4.

  • In 04_portfolio_artifacts.py, add the video recording code at the bottom of the file:
# Record video of trained agent
video_env = RecordVideo(
    gym.make("InvertedPendulum-v5", render_mode="rgb_array"),
    video_folder="./videos/",
    name_prefix="trained_agent",
    episode_trigger=lambda x: True,
)

obs, info = video_env.reset()
done = False
while not done:
    action, _ = model_trained.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = video_env.step(action)
    done = terminated or truncated

video_env.close()
print("Video saved in ./videos/ folder")

What does this code do?

  • RecordVideo wraps the environment so every frame gets captured. The render_mode="rgb_array" tells MuJoCo to render pixel data instead of opening a window.
  • episode_trigger=lambda x: True means every episode gets recorded (we only run one).
  • The loop uses the fully trained model to control the cart, and the wrapper saves the result as an mp4 in the ./videos/ folder.
  • Save 04_portfolio_artifacts.py.

Before you run the full script, what do you expect to see in the video? Will the pole stay up for the full episode or fall partway through?

  • Run the complete script:
python 04_portfolio_artifacts.py

After the chart window closes, you should see Video saved in ./videos/ folder printed. The trained agent achieves a perfect 1000 reward, meaning the pole never falls during the entire recording.

  • Confirm all your portfolio artifacts exist by running:
ls performance_comparison.png training_curve.png videos/

You should see both PNG files listed and a video file (prefixed trained_agent) inside the videos/ folder.

No video file in the videos folder?

  • RecordVideo requires the moviepy or ffmpeg package to encode video. If you see an encoding error, install it with pip install moviepy and re-run the script.
  • Make sure the environment uses render_mode="rgb_array" (not "human"). The video wrapper captures pixel arrays, not window frames.

Why isn't RecordVideo saving my video file?

✔️ Awesome, I've got everything!

Great. Double check you've saved 04_portfolio_artifacts.py.

ⓧ I'd like to double check the full code

Here is the complete 04_portfolio_artifacts.py. Confirm your file matches, scrolling through each section in order:

  • Imports and model loading (the first chunk you pasted).
  • Random agent evaluation loop.
  • Undertrained and trained agent evaluations with evaluate_policy().
  • Comparison bar chart with plt.savefig("performance_comparison.png").
  • RecordVideo section saving to ./videos/.

You now have three portfolio-ready artifacts: a training curve showing your agent's learning journey, a comparison chart proving the RL pipeline works, and a video of perfect pole balancing. Next up, the Secret Mission challenges you to apply this same workflow to a much harder locomotion problem.

Secret mission

Train a Running Robot

You proved you can balance a pole. Now train a simulated cheetah robot to run using the exact same PPO workflow, scaling from 4 observations and 1 action to 17 observations and 6 actuators.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally with free, open-source tools. There are no ongoing costs or cloud resources to worry about.

Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • Project folder on your Desktop (contains all Python scripts, trained model .zip files, training plots, videos, and log directories).
  • Python virtual environment inside the project folder (contains installed packages: mujoco, gymnasium, stable-baselines3, matplotlib, and torch).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with your trained agents or want to try different hyperparameters.

  • Your trained models (ppo_pendulum_undertrained.zip, ppo_pendulum_trained.zip, ppo_halfcheetah.zip) are saved locally and use minimal disk space.
  • Your plots (training_curve.png, performance_comparison.png, cheetah_training_curve.png) and videos are ready for your portfolio.
  • The virtual environment stays on disk but only activates when you explicitly run the activate command.

Pause - I'll come back to this later

Deactivate your virtual environment to free up your terminal session. Your project files and trained models stay intact on disk.

  • Deactivate the virtual environment by running this in your terminal:
deactivate
  • Close VS Code when you are done.
  • To resume later, open your project folder in VS Code, activate the virtual environment again, and run your scripts as before.

Delete - I don't want to use this again

Remove all project files, trained models, and the virtual environment. This frees up disk space and leaves no trace of the project.

  • Deactivate the virtual environment first by running:
deactivate
  • Close VS Code.
  • Delete the project folder from your Desktop. On macOS or Linux, run:
cd ~/Desktop
rm -rf [[PROJECT_FOLDER="your project folder name"]]
  • On Windows, run:
cd ~/Desktop
rmdir /s /q [[PROJECT_FOLDER="your project folder name"]]
  • This removes all Python scripts, trained model files, plots, videos, log directories, and the virtual environment in one step.

That's a Wrap!

That's a Wrap!

Nice work! You just built a complete reinforcement learning pipeline that trained a neural network to solve a classic mechatronics control problem using the same physics engine powering cutting-edge robotics research.

You've learned how to:

  • Trained a PPO agent from scratch that went from random failure (reward ~30) to perfect pole balancing (reward 1000) using Stable Baselines3 and MuJoCo physics simulation.
  • Diagnosed an undertrained model using training curves and fixed it by extending training from 10,000 to 100,000 timesteps, building intuition for how RL agents learn through trial and error.
  • Created portfolio-ready artifacts including performance comparison charts, training curve visualizations, and rendered videos of your trained agent in action.
  • Secret Mission: Applied the same RL workflow to the HalfCheetah-v5 locomotion environment and trained a multi-joint simulated robot to run.

Ready to quiz yourself?