Visualizing Training Statistics in Reinforcement Learning

Introduction

Hello and welcome to the third lesson of " Game On: Integrating RL Agents with Environments "! In our previous lessons, we learned how to integrate a reinforcement learning agent with an environment and explored the crucial exploration-exploitation trade-off that affects learning performance. Today, we're diving into an equally important aspect of reinforcement learning: visualizing training statistics . When building RL systems, it's essential to track and visualize how our agents are performing over time. Without proper visualization, it can be challenging to understand whether our agent is actually learning, how quickly it's improving, or if it's getting stuck in suboptimal behaviors. By the end of this lesson, we'll be able to track, analyze, and visualize key performance metrics of our reinforcement learning agent. These visualization tools will help us gain insights into the learning process, identify issues, and make informed decisions about how to improve our agents. Let's get started!

Why Visualization Matters in Reinforcement Learning

When training RL agents, we often run hundreds or thousands of episodes over many hours. Without proper visualization, it's nearly impossible to understand what's happening during this training process. Thus, visualization serves several critical purposes in RL: Monitoring progress : Visualizations provide an immediate sense of whether your agent is learning or not. Debugging : Abnormal patterns in learning curves can alert you to bugs in your implementation. Hyperparameter tuning : Comparing visualization results helps identify optimal hyperparameter settings. Insight generation : Patterns in visualizations can reveal insights about the nature of your environment and agent. Communication : Graphs and charts communicate your results effectively to others. Think of visualization as your window into the black box of RL. Without it, you're essentially flying blind, hoping your agent is learning properly but having no way to confirm. In research papers and industry projects, learning curves are standard practice for a reason — they tell the story of an agent's learning journey in a clear, interpretable format. Let's build the tools to create these visualizations for our own agents.

Key Metrics to Track in RL Training

Before we start implementing visualization functions, we need to understand the key metrics that help us evaluate an agent's performance. We are focusing on three simple, yet fundamental, metrics: Total rewards per episode : The sum of rewards received during each episode. This is the primary indicator of how well the agent is performing its task. Steps per episode : The number of actions needed to complete an episode. For goal-oriented tasks like our grid world, fewer steps generally indicate more efficient behavior. Success rate : The percentage of episodes where the agent successfully achieved the goal. This is especially relevant for tasks with clear success/failure outcomes. But specifically, what should we be expecting to see in these metrics as learning progresses? Rewards : Should generally increase over time as the agent learns better strategies. Steps : Should decrease over time as the agent finds more efficient paths. Success rate : Should increase from initial random performance to consistently achieving the goal.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Collecting Training Statistics

To create meaningful visualizations, we first need to collect relevant statistics during the training process. We'll modify our training function to track the three key metrics we just discussed, as these metrics will provide a comprehensive view of our agent's learning progress:

Python

def train_agent(env, agent, num_episodes=200, eval_interval=20):
    # Initialize tracking variables
    rewards_per_episode = []
    steps_per_episode = []
    success_rates = []
    window_size = 10  # For calculating moving average
    
    for episode in range(num_episodes):
        # Run episode (environment interaction code omitted for brevity)
        # ...
        
        # After episode completes, record statistics
        success = not info.get("timeout", False)
        success_rates.append(float(success))
        rewards_per_episode.append(total_reward)
        steps_per_episode.append(env.steps_count)
        
        # Periodically print progress using moving averages
        if (episode + 1) % eval_interval == 0:
            print_progress_stats(episode, num_episodes, rewards_per_episode, 
                                steps_per_episode, success_rates, window_size)
    
    # Return collected statistics for visualization
    stats = {
        'rewards': rewards_per_episode,
        'steps': steps_per_episode,
        'success_rates': success_rates
    }
    
    return stats

def train_agent(env, agent, num_episodes=200, eval_interval=20):
    # Initialize tracking variables
    rewards_per_episode = []
    steps_per_episode = []
    success_rates = []
    window_size = 10  # For calculating moving average
    
    for episode in range(num_episodes):
        # Run episode (environment interaction code omitted for brevity)
        # ...
        
        # After episode completes, record statistics
        success = not info.get("timeout", False)
        success_rates.append(float(success))
        rewards_per_episode.append(total_reward)
        steps_per_episode.append(env.steps_count)
        
        # Periodically print progress using moving averages
        if (episode + 1) % eval_interval == 0:
            print_progress_stats(episode, num_episodes, rewards_per_episode, 
                                steps_per_episode, success_rates, window_size)
    
    # Return collected statistics for visualization
    stats = {
        'rewards': rewards_per_episode,
        'steps': steps_per_episode,
        'success_rates': success_rates
    }
    
    return stats

This training function incorporates several important elements for effective visualization:

Storage of episodic data in separate arrays for each metric
Calculation of moving averages to smooth out noise and identify trends
Periodic progress reporting to monitor training in real-time
Return of all statistics in a structured format that will be easy to pass to our visualization functions

Adding Steps and Success Rate Plots

Now let's complete our visualization by adding plots for the other two metrics:

Python

    # Plot steps
    plt.subplot(1, 3, 2)
    plt.plot(steps, alpha=0.5)
    
    if len(steps) >= window:
        moving_avg = np.convolve(steps, np.ones(window)/window, mode='valid')
        plt.plot(np.arange(window-1, window-1+len(moving_avg)), moving_avg)
    
    plt.title('Steps per Episode')
    plt.xlabel('Episode')
    plt.ylabel('Steps')
    
    # Plot success rate
    plt.subplot(1, 3, 3)
    if len(success_rates) >= window:
        moving_avg = np.convolve(success_rates, np.ones(window)/window, mode='valid')
        plt.plot(np.arange(window-1, window-1+len(moving_avg)), moving_avg)
    
    plt.title('Success Rate')
    plt.xlabel('Episode')
    plt.ylabel('Success Rate')
    plt.ylim([-0.05, 1.05])  # Set y-axis limits for percentage display

    # Plot steps
    plt.subplot(1, 3, 2)
    plt.plot(steps, alpha=0.5)
    
    if len(steps) >= window:
        moving_avg = np.convolve(steps, np.ones(window)/window, mode='valid')
        plt.plot(np.arange(window-1, window-1+len(moving_avg)), moving_avg)
    
    plt.title('Steps per Episode')
    plt.xlabel('Episode')
    plt.ylabel('Steps')
    
    # Plot success rate
    plt.subplot(1, 3, 3)
    if len(success_rates) >= window:
        moving_avg = np.convolve(success_rates, np.ones(window)/window, mode='valid')
        plt.plot(np.arange(window-1, window-1+len(moving_avg)), moving_avg)
    
    plt.title('Success Rate')
    plt.xlabel('Episode')
    plt.ylabel('Success Rate')
    plt.ylim([-0.05, 1.05])  # Set y-axis limits for percentage display

For the steps plot, we follow the same approach as with rewards — showing both raw data and a moving average. However, for the success rate, we only display the moving average because the raw data would be binary (0 or 1) and would create a visually cluttered plot.

Notice the plt.ylim([-0.05, 1.05]) line — this ensures our success rate plot has a fixed y-axis range from 0 to 1 (with a small margin). This makes the plot more interpretable as a percentage.

Interpreting Learning Curves

Now that we've set up our plots, let's walk through how to interpret the results and what they tell us about the agent's learning process. Below is a typical outcome for a 200-episode training run, with three side-by-side plots:

Training Stats Plots

Let's breakdown these plots:

Rewards per Episode
- Raw vs. Moving Average: The raw reward values (blue line) fluctuate from episode to episode, especially early on when the agent has not yet learned a stable policy. The moving average (orange line) smooths these fluctuations, allowing you to see clearer trends.
- Early Episodes: In the initial phase, the agent often receives low or zero reward because it acts randomly and does not consistently reach the goal. You may notice a handful of “spikes” in reward if it happens to stumble into the correct behavior by chance.
- Later Episodes: Over time, as learning progresses, you can see the moving average of the rewards rising, indicating that the agent is more consistently achieving the goal.
Steps per Episode
- Efficiency Measure: In many goal-oriented tasks (like a GridWorld or Maze), fewer steps per episode generally indicates more efficient behavior. Early in training, the agent typically takes many steps—often the environment's maximum step limit—because it does not know how to reach the goal quickly.
- Downward Trend: As the agent learns a better strategy, you can see the steps-per-episode curve trend downward. This reflects that the agent is discovering shorter (and thus more efficient) paths to the goal.
- Variability: Even well-performing agents can show occasional spikes in the number of steps if they explore or get temporarily “lost.” However, the overall moving average should descend if learning is effective.
Success Rate
- Binary Outcomes: The success rate tracks whether the agent succeeded or failed in each episode. The raw data for a single episode would be either 0 or 1, which can look very scattered in a plot. Consequently, we often plot only the moving average to see a smooth trend.
- Initial Phase: Typically, the success rate starts near 0 when the agent is learning from scratch. You might see a slow, incremental climb as it discovers partial strategies.
- Breakthrough and Convergence: Once the agent finds a good policy (often after some key exploratory episodes), the success rate can jump significantly. In successful training, the moving average may approach or reach 1.0, indicating the agent consistently achieves the goal.

Occasional dips or plateaus are part of the normal learning process—especially in stochastic environments or when an agent's exploration temporarily overrides exploitation. By carefully analyzing these plots, you can diagnose whether training is on track, tune hyperparameters (e.g., learning rate or epsilon in (\epsilon)-greedy strategies), or spot potential bugs if the curves behave unexpectedly.