Hello and welcome to the fourth and final lesson of "Game On: Integrating RL Agents with Environments"! So far in our journey, we've covered the fundamentals of integrating agents with environments, explored the crucial balance between exploration and exploitation, and learned how to track and visualize training statistics to monitor agent performance.

Today, we'll dive into one of the most illuminating aspects of reinforcement learning: visualizing policies and value functions. While our previous lesson helped us understand how well our agent is learning through performance metrics, today we'll look at what our agent has learned by visualizing its decision-making strategy and value estimations.

By the end of this lesson, we'll have implemented two powerful visualization techniques that will allow us to see our agent's learned policy as directional arrows and its value function as a colorful heatmap. These visualizations will transform abstract numbers in our Q-table into intuitive representations that reveal the agent's understanding of the environment. Let's begin!

Before we can visualize our agent's policy, we need a method to extract the best action for each state according to our Q-table. The policy is derived from the Q-table by selecting the action with the highest Q-value for each state:

class QLearningAgent:
    # ...
    def get_policy(self, state):
        """
        Return the best action for a state according to current Q-values.
        
        Args:
            state: The state to get policy for.
            
        Returns:
            int: The best action.
        """
        return int(np.argmax(self.Q[state]))  # Find action with highest Q-value

This method simply returns the action index with the highest Q-value for the given state using np.argmax(), which finds the index of the maximum value in an array. This gives us the agent's current policy – what it believes is the best action to take in each state.

Now that we can retrieve the policy, let's implement a function to visualize it in a grid. We'll represent each action as a directional arrow, making it easy to interpret at a glance:

def visualize_policy(agent, env_size):
    """
    Display the learned policy as arrows in a grid.
    
    Args:
        agent (QLearningAgent): The agent with a learned policy.
        env_size (int): Size of the environment grid.
    """
    action_symbols = ["↑", "↓", "←", "→"]  # Symbols for up, down, left, right
    
    print("Learned Policy:")
    for row in range(env_size):
        line = []
        for col in range(env_size):
            state = (row, col)
            if state == (env_size - 1, env_size - 1):
                line.append("G")  # Mark goal state
            else:
                best_action = agent.get_policy(state)
                line.append(action_symbols[best_action])
        print(" ".join(line))
    print()

This function takes our agent and the environment size as inputs and prints a grid of arrows. For each state in our grid world, we call the agent's get_policy(state) method to get the best action, then use our arrow symbols (↑, ↓, ←, →) to visualize that action. The goal state is marked with "G" to make it stand out.

After training our agent successfully, the policy visualization might look something like this:

→ → → → ↓
← → → ↓ ↓
→ → → → ↓
→ ↑ → ← ↓
→ ↑ ↑ → G

By examining this policy grid, we can easily trace the agent's strategy from any starting position. For example, if the agent starts at the top-left corner (0,0), it will follow the arrows: right, right, right, down, down, down, down, reaching the goal in 7 steps – the optimal path length for that starting position, showing us that the agent has learned an optimal policy.

Now let's turn our attention to the value function. As mentioned earlier, the value function tells us how good it is to be in each state, estimated as the expected future reward.

For a Q-learning agent, the value function is directly related to the Q-table. Specifically, the value of a state is the maximum Q-value for that state across all possible actions. This makes sense because a rational agent would always choose the action that promises the highest expected return.

Here's the agent method that extracts the state value from the Q-table:

class QLearningAgent:
    # ...
    def get_state_value(self, state):
        """
        Return the value of a state (maximum Q-value).
        
        Args:
            state: The state to get value for.
            
        Returns:
            float: The state value.
        """
        return np.max(self.Q[state])  # Return highest Q-value for this state

This simple method takes a state and returns the maximum Q-value for that state, representing the state's value according to our agent's current estimates.

In our grid world environment:

• States close to the goal should have higher values because they can lead to rewards sooner.
• States far from the goal typically have lower values.
• The goal state itself usually has the highest value.

The distribution of these values across the grid creates a "value landscape" that guides the agent's decision-making. Visualizing this landscape can reveal how the agent perceives different states and their potential for future reward.

Now let's implement a function to visualize the value function using a colorful heatmap:

def visualize_value_function(agent, env_size):
    """
    Display the state values as a heatmap.
    
    Args:
        agent (QLearningAgent): The agent with learned Q-values.
        env_size (int): Size of the environment grid.
    """
    plt.figure(figsize=(8, 6))
    value_grid = np.zeros((env_size, env_size))
    
    # Populate the value grid with state values
    for row in range(env_size):
        for col in range(env_size):
            state = (row, col)
            value_grid[row, col] = agent.get_state_value(state)
    
    # Create the heatmap visualization
    plt.imshow(value_grid, cmap='viridis')
    plt.colorbar(label='State Value')
    plt.title('Learned State Values')
    
    # Add value text in each cell for precision
    for row in range(env_size):
        for col in range(env_size):
            value = value_grid[row, col]
            text_color = "white" if value < 0.5 * np.max(value_grid) else "black"
            plt.text(col, row, f"{value:.2f}",
                     ha="center", va="center", color=text_color)
    
    plt.savefig("value_function.png")
    plt.close()
    print("Value function visualization saved to 'value_function.png'")

This function creates a matrix to hold the value of each state in our grid world. We iterate through each cell, get its value using our agent's get_state_value() method, and store it in the corresponding position in our value_grid matrix.

The visualization uses Matplotlib's imshow() to create a heatmap with the 'viridis' colormap, which transitions from blue (low values) to yellow (high values). We add a colorbar for interpretation and overlay the actual numerical values on each cell for precise reading.

The resulting heatmap makes it immediately apparent where the agent believes the valuable states are, with bright yellow areas representing high-value states (typically near the goal) and dark blue areas representing low-value states (typically far from the goal).

With our value function heatmap generated, let's examine the learned state values in more detail. In the image above, each cell corresponds to a state in our 5×5 grid, and the color indicates the magnitude of that state’s value — darker blues represent lower values, and brighter yellows represent higher values:

At a glance, notice how the bottom-right corner — our goal location — is assigned the highest values (1.00). This is expected: since the agent receives a reward for reaching this terminal state, the algorithm learns that being anywhere close to the goal is advantageous, as the agent can collect a high reward quickly. Consequently, cells adjacent to the goal also display high values (0.98–0.99), reflecting the agent’s confidence that from those states, it can secure the reward with only a few moves.

In contrast, states farther away from the goal tend to have lower values (around 0.90–0.94). These cells are still part of a viable path to the goal, but the agent anticipates a longer journey or a bit more uncertainty in receiving the reward.

Overall, this gradient from lower values in the upper-left to higher values in the lower-right reveals that the agent has successfully learned which states are more valuable for collecting future rewards. When combined with the policy visualization, you can trace how the agent uses these state-value estimates to decide the next action, continuously moving toward states with higher expected returns.

Fantastic job! You've now mastered two powerful visualization techniques that give you deeper insights into your reinforcement learning agents. By visualizing both the policy and value function, we've transformed abstract Q-table numbers into intuitive visual representations that reveal what the agent has actually learned. We've covered the relationship between Q-values and the state-value function, implemented text-based policy visualization using directional arrows, created heatmap visualizations of the value function, and learned how to interpret these visualizations to gain insights into agent behavior.

In the upcoming practice exercises, you'll get hands-on experience implementing these visualization techniques and analyzing the results. You'll be able to see firsthand how different training parameters affect the learned policy and value function, and how visualizations can help you debug and understand agent behavior. These visualization skills are invaluable tools in your reinforcement learning toolkit as you tackle increasingly complex environments and agents.