Visualizing Policies and Value Functions in Reinforcement Learning

Visualizing Policies and Value Functions in Reinforcement Learning | CodeSignal Learn

Introduction

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!

Understanding Policies and Value Functions

V(S)

Join the 1M+ learners on CodeSignal

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

Before we start coding our visualizations, let's understand what we're actually trying to visualize.

A policy in reinforcement learning is essentially the agent's strategy — it defines what action the agent will take in each state. For our Q-learning agent, the policy is derived from the Q-table by selecting the action with the highest Q-value for each state.

The value function represents how good it is to be in a particular state. It estimates the expected future reward the agent can obtain starting from that state and following its policy thereafter. There are two types of value functions:

State-value function $V(S)$ : Estimates the expected return when starting in a state and following the current policy.
Action-value function $Q(S, A)$ : Estimates the expected return when taking a specific action in a state and following the current policy afterward.

In our Q-learning agent, we directly learn the Q-values (using the Q-table), but we can derive the state-value function from them by simply taking the maximum Q-value for each state:

V(s) = max(Q(s,a))

This relationship is crucial — our Q-table contains action-values for each state-action pair, and the maximum Q-value for a state gives us its state-value. Visualizing these functions helps us understand which states the agent considers valuable, see the direction of the agent's movement strategy, identify potential misconceptions, and confirm if the agent has successfully learned the optimal path to the goal.

Retrieving the Policy from the Q-table

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: Python 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 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.

Visualizing the Policy with Directional Arrows

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: Python 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 = ["↑", "↓", "\leftarrow", "\to"] # 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() 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 = ["↑", "↓", "\leftarrow", "\to"] # 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 (↑, ↓, \leftarrow, \to) to visualize that action. The goal state is marked with "G" to make it stand out.

Interpreting the Policy Visualization

After training our agent successfully, the policy visualization might look something like this: text \to \to \to \to ↓ \leftarrow \to \to ↓ ↓ \to \to \to \to ↓ \to ↑ \to \leftarrow ↓ \to ↑ ↑ \to G \to \to \to \to ↓ \leftarrow \to \to ↓ ↓ \to \to \to \to ↓ \to ↑ \to \leftarrow ↓ \to ↑ ↑ \to 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.

The Value Function and State Assessment

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: Python 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 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.

Creating Value Function Heatmaps

Now let's implement a function to visualize the value function using a colorful heatmap: Pythondef 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'")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).