Solve a Grid World with Value Iteration

Build an MDP solver that finds optimal strategies using the Bellman equation.

Introduction

30 Second Summary

Every time you use a GPS that reroutes around traffic or watch a game AI pick a path through a maze, you're seeing an algorithm that weighs risk against reward to choose the best action in an uncertain world.

In this project, you will build a Markov Decision Process solver from scratch that finds the optimal strategy for navigating a grid world with rewards, penalties, and noisy actions.

What You'll Build

Picture running your script and watching a colorful heatmap appear, with arrows showing your agent the safest, most rewarding path through a grid. States near the goal glow green, states near danger glow red, and the arrows reveal exactly how uncertainty reshapes the agent's decisions.

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

  • A working value iteration algorithm that converges on the optimal value of every state, proving you can translate the Bellman equation into running code.
  • A side-by-side policy comparison showing how your agent reroutes away from danger when its actions become unreliable.
  • A Matplotlib heatmap with policy arrows that makes abstract math visually tangible at a glance.
  • Secret Mission: Implement Policy Iteration and race it against Value Iteration to compare convergence speed.

Are there any prerequisites?

You should be comfortable with basic Python (loops, dictionaries, functions) and have a basic understanding of probability (what it means for an event to happen with a certain chance). No prior reinforcement learning experience is needed.

Before We Start

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

Set Up Your Environment

To solve a Markov Decision Process computationally, you need tools for efficient array math and for visualizing the solution.

In this step, you'll create the project folder, install NumPy (for array operations) and Matplotlib (for plotting), and verify they work before writing any algorithm code.

In this step, get ready to:
  • Create the project folder and Python file.
  • Install NumPy and Matplotlib.
  • Verify the dependencies are working.
Create your project folder and file
  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it.
  • Open the integrated terminal by pressing Ctrl+` (Windows) or Cmd+` (macOS).

You need a dedicated folder on your Desktop so you can find it easily later.

  • Move to your Desktop and create the project folder by running these commands:
cd ~/Desktop
mkdir mdp-gridworld
cd mdp-gridworld
  • Create the Python file you'll work in by running:
touch mdp_gridworld.py
  • Confirm the file exists by running:
ls

You should see mdp_gridworld.py listed in the output.

  • Open the project folder in VS Code by clicking File in the top menu bar.
  • Click Open Folder.
  • Select the mdp-gridworld folder on your Desktop, then click Open.

You should see mdp_gridworld.py in the Explorer sidebar on the left.

Don't see the file in the Explorer?

Make sure you selected the mdp-gridworld folder itself, not the Desktop. The folder should appear as the top-level item in the Explorer sidebar.

If you accidentally opened the wrong folder, click File then Open Folder again and navigate to Desktop/mdp-gridworld.

Help me find my project folder in VS Code

Install NumPy and Matplotlib

NumPy handles the array math for storing state values. Matplotlib renders the heatmap visualization you'll build in Step 5.

  • Open the integrated terminal again (it may have reset after opening the folder). If you see a terminal panel at the bottom of VS Code, you're good.
  • Install both libraries by running:
pip install numpy matplotlib

What does this command do?

pip is Python's package manager. This command downloads and installs both NumPy and Matplotlib (plus their dependencies) into your Python environment.

Seeing a 'pip not found' error?

Try running pip3 install numpy matplotlib instead. On some systems, Python 3's package manager is named pip3.

If you see a permissions error, try adding --user at the end of the command to install into your user directory.

Help me fix my pip installation error

Verify the installation

Before you write any algorithm code, confirm that both libraries import correctly and print their version numbers.

What do you think will happen when you run the verification command below? If both libraries installed correctly, you should see two version numbers printed.

  • Verify the installation by running:
python -c "import numpy; import matplotlib; print(numpy.__version__, matplotlib.__version__)"

You should see version numbers printed without errors (e.g., 2.5.1 3.11.1). Your exact numbers may differ depending on when you installed.

Seeing a ModuleNotFoundError?

This means the library did not install into the Python version you're currently running. Try running python3 instead of python in the verification command.

If that works, use python3 for all commands in this project.

Help me fix my import error for numpy or matplotlib

Your environment is ready. Next up, you'll build the grid world itself, encoding the MDP's states, actions, rewards, and transitions into a Python class.

Build the Grid World

Before you can solve a Markov Decision Process, you need one to solve. Building the environment yourself forces you to encode the five components of an MDP (states, actions, transitions, rewards, discount) into code.

In this step, you will define a 4x3 grid world where each cell is a state, give the agent four directional actions, and set up a goal, a pit, and a wall. This is where the math becomes concrete.

In this step, get ready to:
  • Define the GridWorld class with terminal states, walls, actions, and boundary checking.
  • Implement a transition model that returns where the agent ends up after an action.
  • Print the grid to the terminal showing the goal, pit, and wall positions.
Define the GridWorld class

An MDP needs five components: a set of states, a set of actions, a transition function, a reward function, and a discount factor. The GridWorld class packages all of these into a single object.

  • In mdp_gridworld.py, add the imports and the class with its __init__ method by pasting the following code:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches


class GridWorld:
    """A 4x3 Grid World MDP."""

    def __init__(self, living_reward=-0.04, noise=0.0):
        self.rows = 3
        self.cols = 4
        self.living_reward = living_reward
        self.noise = noise  # action uncertainty (0=deterministic)

        self.terminals = {(3, 2): 1.0, (3, 1): -1.0}  # goal and pit
        self.walls = [(1, 1)]

        self.actions = {  # each action is a (dx, dy) vector
            "UP": (0, 1),
            "DOWN": (0, -1),
            "LEFT": (-1, 0),
            "RIGHT": (1, 0),
        }

        # Build list of all valid (non-wall) states
        self.states = []
        for x in range(self.cols):
            for y in range(self.rows):
                if (x, y) not in self.walls:
                    self.states.append((x, y))

What does this code do?

This sets up the MDP's structure:

  • The grid is 4 columns by 3 rows. Position (3,2) is the goal with reward +1. Position (3,1) is the pit with reward -1.
  • The living_reward of -0.04 penalizes each step, motivating the agent to reach the goal quickly instead of wandering.
  • Each action maps to a direction vector. UP is (0,1) because y increases upward in this coordinate system.
  • The states list holds every cell that is not a wall. This gives the agent 11 valid positions on the 4x3 grid.

The agent also needs to know what happens when it tries to move off the edge or into a wall. It should simply stay in place.

  • Add the following three helper methods directly below the __init__ method, inside the GridWorld class:
    def is_terminal(self, state):
        return state in self.terminals

    def get_reward(self, state):
        # Terminal states give their stored reward, others give living_reward
        if state in self.terminals:
            return self.terminals[state]
        return self.living_reward

    def get_next_state(self, state, direction):
        # Try moving, but stay in place if hitting a wall or edge
        dx, dy = direction
        next_state = (state[0] + dx, state[1] + dy)
        if (
            next_state[0] < 0
            or next_state[0] >= self.cols
            or next_state[1] < 0
            or next_state[1] >= self.rows
            or next_state in self.walls
        ):
            return state
        return next_state

What does this code do?

  • is_terminal checks whether a state is the goal or the pit. Terminal states end the episode.
  • get_reward returns +1 for the goal, -1 for the pit, and -0.04 for every other cell.
  • get_next_state computes where the agent lands after moving. If the destination is off-grid or a wall, the agent stays put.
  • Save mdp_gridworld.py.
  • Verify the class can be instantiated by running this command in your terminal:
python -c "from mdp_gridworld import GridWorld; gw = GridWorld(); print(len(gw.states), 'states')"

You should see 11 states printed. The 4x3 grid has 12 cells minus 1 wall.

Seeing an ImportError or SyntaxError?

  • Make sure your terminal is in the mdp-gridworld folder where mdp_gridworld.py lives.
  • Check that the helper methods are indented inside the class (4 spaces in from class GridWorld:).
  • Confirm you saved the file before running the command.

Help me debug my GridWorld class import error.

Implement the transition model

The transition function is the heart of an MDP. It answers: given a state and an action, where does the agent end up? For now you will make it deterministic. Every action succeeds with probability 1.0.

  • Add the get_transitions method below get_next_state, still inside the GridWorld class:
    def get_transitions(self, state, action_name):
        # Terminal states have no transitions
        if self.is_terminal(state):
            return []

        intended = self.actions[action_name]

        # Deterministic: 100% chance of intended direction
        if self.noise == 0.0:
            return [(1.0, self.get_next_state(state, intended))]

        # Stochastic: slip perpendicular with probability noise
        transitions = []
        if intended in [(0, 1), (0, -1)]:
            perpendiculars = [(-1, 0), (1, 0)]
        else:
            perpendiculars = [(0, 1), (0, -1)]

        transitions.append(
            (1.0 - 2 * self.noise, self.get_next_state(state, intended))
        )
        for perp in perpendiculars:
            transitions.append((self.noise, self.get_next_state(state, perp)))

        return transitions

What does this code do?

This method returns a list of (probability, next_state) tuples representing every possible outcome of taking an action.

  • Terminal states return an empty list because the episode is over.
  • When noise is 0, the agent always lands where it intends. The method returns a single tuple: (1.0, next_state).
  • The stochastic branch (used in Step 4) adds a chance of slipping perpendicular to the intended direction.
  • Save mdp_gridworld.py.

Indentation issues?

Make sure get_transitions is at the same indentation level as get_next_state above it. Both should be one indent level inside the class body.

Help me fix my get_transitions indentation.

Print the grid and run the script

Numbers in a dictionary are hard to interpret. A printed grid showing the goal, pit, and wall makes the environment tangible.

  • Add the print_grid method below get_transitions, then add the __main__ block at the bottom of the file (outside the class):
    def print_grid(self, values=None, policy=None):
        # Map action names to arrow symbols for display
        action_symbols = {"UP": "^", "DOWN": "v", "LEFT": "<", "RIGHT": ">"}
        print()
        print("+" + "------+" * self.cols)
        for y in range(self.rows - 1, -1, -1):
            row_str = "|"
            for x in range(self.cols):
                if (x, y) in self.walls:
                    row_str += " #### |"
                elif (x, y) in self.terminals:
                    row_str += f" {self.terminals[(x, y)]:+.1f} |"
                elif policy and (x, y) in policy:
                    row_str += f"   {action_symbols[policy[(x, y)]]}  |"
                elif values and (x, y) in values:
                    row_str += f"{values[(x, y)]:+5.2f} |"
                else:
                    row_str += "   .  |"
            print(row_str)
            print("+" + "------+" * self.cols)


if __name__ == "__main__":
    gw = GridWorld(living_reward=-0.04, noise=0.0)
    print("Grid layout (+1=goal, -1=pit, ####=wall):")
    gw.print_grid()

What does this code do?

  • print_grid loops row by row (top to bottom visually) and fills each cell with the appropriate label: #### for walls, +1.0/-1.0 for terminals, arrows for policy, values for state values, or a dot for empty.
  • The __main__ block creates a deterministic grid world and prints it. This is what runs when you execute the script directly.
  • Save mdp_gridworld.py.

Before you run this, what do you expect the printed grid to look like? Think about which cells should show values and which should show symbols.

  • Run the script by entering this command in your terminal:
python mdp_gridworld.py

You should see a 4x3 grid printed to the terminal with +1.0 for the goal (top right), -1.0 for the pit (below the goal), #### for the wall, and dots for all other cells.

Not seeing the grid?

  • Make sure the if __name__ == "__main__": block is at the very bottom of the file with no indentation (it is outside the class).
  • Check that print_grid is indented as a method inside GridWorld, not at the module level.
  • Confirm you are running the command from the mdp-gridworld folder.

Help me debug why my grid world script isn't printing output.

✔️ Awesome, I've got everything!

Great. Double check you have saved mdp_gridworld.py and your grid prints cleanly.

ⓧ I'd like to double check the full code

Scroll up and confirm your file contains each of these sections in order:

  • The three imports (numpy, matplotlib.pyplot, matplotlib.patches).
  • The GridWorld class with __init__ defining rows, cols, terminals, walls, actions, and states.
  • The is_terminal, get_reward, and get_next_state helper methods.
  • The get_transitions method returning (probability, next_state) tuples.
  • The print_grid method.
  • The if __name__ == "__main__": block at the bottom, outside the class.

Your grid world is alive. You have the states, actions, rewards, and transitions that define a complete MDP. Next up, you will write the algorithm that solves it.

Implement Value Iteration

You have a grid world with states, actions, and rewards. Now you need an algorithm that can figure out the best action to take at every single state.

Value iteration solves this by repeatedly asking one question at every state: "If I pick the best action here, what's the total reward I can expect from now on?" It keeps updating its answer until the numbers stop changing. That's the Bellman equation in action.

In this step, get ready to:
  • Implement value iteration to compute the optimal value of every state.
  • Extract the optimal policy from the converged value function.
  • Run the solver and see the optimal path to the goal printed as arrows.
Write the value iteration function

Value iteration starts every state at value 0, then sweeps through the grid over and over. At each state, it tries every action and keeps the one that gives the highest expected reward plus discounted future value. When the biggest change in any state drops below a tiny threshold, the values have converged.

  • In mdp_gridworld.py, add the following function below the GridWorld class (after the class ends, before the if __name__ block). Start with the function signature, initialization, and loop structure:
def value_iteration(grid_world, gamma=0.9, epsilon=0.001):
    values = {state: 0.0 for state in grid_world.states}
    iterations = 0

    while True:
        delta = 0
        new_values = values.copy()

        for state in grid_world.states:
            if grid_world.is_terminal(state):
                new_values[state] = grid_world.get_reward(state)
                continue
  • Directly below the continue line (still inside the for state loop), add the Bellman update and convergence check:
            max_value = float("-inf")
            for action_name in grid_world.actions:
                action_value = 0
                for prob, next_state in grid_world.get_transitions(
                    state, action_name
                ):
                    reward = grid_world.get_reward(state)
                    action_value += prob * (reward + gamma * values[next_state])
                max_value = max(max_value, action_value)

            new_values[state] = max_value
            delta = max(delta, abs(new_values[state] - values[state]))

        values = new_values
        iterations += 1

        if delta < epsilon:
            break

    return values, iterations

What does this code do?

  • values is a dictionary mapping every state to its estimated long-term value. It starts at 0 everywhere.
  • For each non-terminal state, it tries all four actions and computes the expected value using the transition probabilities, immediate reward, and discounted future value.
  • delta tracks the largest value change in a single sweep. When it drops below epsilon (0.001), the values have converged and the loop stops.
  • gamma=0.9 is the discount factor. It tells the agent that a reward received sooner is worth more than the same reward later.
  • Now replace your existing if __name__ == "__main__": block at the bottom of the file with the following to test value iteration:
if __name__ == "__main__":
    print("=== DETERMINISTIC GRID WORLD ===")
    gw_det = GridWorld(living_reward=-0.04, noise=0.0)
    print("\nGrid layout (+1=goal, -1=pit, ####=wall):")
    gw_det.print_grid()

    values_det, iters_det = value_iteration(gw_det, gamma=0.9)
    print(f"\nValue Iteration converged in {iters_det} iterations")
    print("\nState Values:")
    gw_det.print_grid(values=values_det)
  • Save mdp_gridworld.py.
  • Run the script to see value iteration in action:
python mdp_gridworld.py

You should see the grid layout followed by a line reporting how many iterations it took to converge (fewer than 20), then a grid of state values showing positive numbers near the goal and negative numbers near the pit.

Seeing a NameError or unexpected values?

  • If you see NameError: name 'value_iteration' is not defined, make sure the function is placed outside the GridWorld class (at the module level, not indented inside the class body).
  • If all values are 0.0, check that the while True loop uses values[next_state] (the old values) not new_values[next_state] in the Bellman update.
  • help me debug my value iteration function
Extract the optimal policy

The value function tells you how good each state is, but not which direction to go. The policy extractor reads those converged values and picks the action that leads to the highest-value next state at every position.

  • Add the extract_policy function below value_iteration (above the if __name__ block):
def extract_policy(grid_world, values, gamma=0.9):
    policy = {}

    for state in grid_world.states:
        if grid_world.is_terminal(state):
            continue

        best_action = None
        best_value = float("-inf")

        for action_name in grid_world.actions:
            action_value = 0
            for prob, next_state in grid_world.get_transitions(
                state, action_name
            ):
                reward = grid_world.get_reward(state)
                action_value += prob * (reward + gamma * values[next_state])

            if action_value > best_value:
                best_value = action_value
                best_action = action_name

        policy[state] = best_action

    return policy

What does this code do?

  • It loops through every non-terminal state and calculates the expected value of each action (the same math as value iteration).
  • Instead of taking the max to update a value, it records which action achieved the max. That action becomes the optimal policy for that state.
  • The result is a dictionary mapping each state to its best action (UP, DOWN, LEFT, or RIGHT).
  • Replace your if __name__ == "__main__": block with the final version that calls both functions and prints the policy as arrows:
if __name__ == "__main__":
    print("=== DETERMINISTIC GRID WORLD ===")
    gw_det = GridWorld(living_reward=-0.04, noise=0.0)
    print("\nGrid layout (+1=goal, -1=pit, ####=wall):")
    gw_det.print_grid()

    values_det, iters_det = value_iteration(gw_det, gamma=0.9)
    policy_det = extract_policy(gw_det, values_det, gamma=0.9)

    print(f"\nValue Iteration converged in {iters_det} iterations")
    print("\nOptimal Policy (deterministic):")
    gw_det.print_grid(policy=policy_det)
    print("\nState Values:")
    gw_det.print_grid(values=values_det)

What does this code do?

  • gw_det creates a deterministic grid world (noise=0.0) so actions always go where intended.
  • value_iteration computes the optimal value of every state, returning the values dictionary and how many sweeps it took.
  • extract_policy reads those values and picks the best direction at each state.
  • print_grid(policy=policy_det) displays the policy using arrow symbols: ^, v, <, >.
  • Save mdp_gridworld.py.

Before you run this, think about what the optimal policy should look like. In a deterministic world with a goal at the top-right and a pit below it, which direction would each cell point?

  • Run the script:
python mdp_gridworld.py

You should see the optimal policy printed as arrows (^, v, <, >) showing the shortest path from any state to the goal. The convergence message should report fewer than 20 iterations.

What should I see?

The policy grid shows arrows pointing toward the goal at (3,2). States in the bottom row point right and then up. States in the top row point right toward the goal. The state values are positive near the goal and slightly negative in far-away states (because of the -0.04 living reward).

Arrows pointing in random directions?

  • Make sure extract_policy uses the same gamma=0.9 value that value_iteration used. Mismatched discount factors produce incorrect policies.
  • Check that extract_policy skips terminal states with continue. Terminal states have no actions to choose.
  • help me figure out why my policy arrows look wrong

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

Scroll up and confirm your file contains these sections in order:

  • The three imports (numpy, matplotlib.pyplot, matplotlib.patches) at the top.
  • The GridWorld class with all its methods (__init__, is_terminal, get_reward, get_next_state, get_transitions, print_grid).
  • The value_iteration function (initializes values, loops with Bellman updates, stops when delta < epsilon).
  • The extract_policy function (picks the best action per state from converged values).
  • The if __name__ == "__main__": block that creates gw_det, runs value iteration, extracts the policy, and prints everything.

Your solver finds the optimal path through a world where actions always work perfectly. Next up, you'll introduce uncertainty and watch the optimal strategy change when the agent can't trust its own movements.

Add Stochastic Transitions

Your deterministic solver finds the shortest path to the goal, but that is just fancy pathfinding. Real decision-making involves uncertainty.

In this step, you will make the agent's actions noisy. When the agent tries to move in one direction, there is now a chance it slips sideways. This forces the algorithm to weigh risk against reward, and you will see the optimal policy physically change near the dangerous pit.

In this step, get ready to:
  • Understand how stochastic transitions work with the 0.8/0.1/0.1 probability split.
  • Run value iteration on a noisy grid world and extract its policy.
  • Compare the deterministic and stochastic policies to find where uncertainty changes the strategy.
Understand the stochastic transition model

The get_transitions method you wrote in Step 2 already handles non-zero noise. You built it that way, but until now you have only used it with noise=0.0 (deterministic). Here is the relevant section of your existing code:

if self.noise == 0.0:
    return [(1.0, self.get_next_state(state, intended))]

transitions = []
if intended in [(0, 1), (0, -1)]:
    perpendiculars = [(-1, 0), (1, 0)]
else:
    perpendiculars = [(0, 1), (0, -1)]

transitions.append(
    (1.0 - 2 * self.noise, self.get_next_state(state, intended))
)
for perp in perpendiculars:
    transitions.append((self.noise, self.get_next_state(state, perp)))

What does this code do?

When noise=0.1, the probability of moving in the intended direction is 1 - 2*0.1 = 0.8. Each perpendicular direction gets probability 0.1.

If the agent tries to go UP, there is an 80% chance it goes UP, a 10% chance it slips LEFT, and a 10% chance it slips RIGHT. The agent cannot slip backward. This is the classic stochastic transition model from reinforcement learning textbooks.

Why does this change the optimal policy?

With deterministic actions, the agent can safely walk alongside the pit because it will never accidentally fall in. With stochastic actions, walking next to the pit means a 10% chance of slipping into it and receiving -1 reward.

The agent must now weigh the shorter path against the risk of catastrophic failure. This is what separates real decision-making from simple pathfinding.

Create a stochastic grid world and solve it

Your existing value_iteration and extract_policy functions already handle stochastic transitions. They call get_transitions which returns multiple (probability, next_state) tuples when noise is present. No changes needed to those functions.

  • In mdp_gridworld.py, scroll to the bottom of your if __name__ == "__main__": block.
  • Find the last line that currently reads:
    gw_det.print_grid(values=values_det)
  • Add the following code below that line to create and solve a stochastic grid world:
    print("\n\n=== STOCHASTIC GRID WORLD (noise=0.1) ===")
    gw_stoch = GridWorld(living_reward=-0.04, noise=0.1)

    values_stoch, iters_stoch = value_iteration(gw_stoch, gamma=0.9)
    policy_stoch = extract_policy(gw_stoch, values_stoch, gamma=0.9)

    print(f"\nValue Iteration converged in {iters_stoch} iterations")
    print("\nOptimal Policy (stochastic):")
    gw_stoch.print_grid(policy=policy_stoch)
    print("\nState Values:")
    gw_stoch.print_grid(values=values_stoch)

What does this code do?

  • Creates a new GridWorld with noise=0.1, giving the classic 80/10/10 probability split on every action.
  • Runs the same value_iteration algorithm on this noisy world. The algorithm automatically accounts for uncertainty because it sums over all possible next states weighted by their probabilities.
  • Extracts and prints the new optimal policy, which may now differ from the deterministic one.
  • Save mdp_gridworld.py.
  • Run the script to see the stochastic policy by running:
python mdp_gridworld.py

You should see both the deterministic section from before and a new stochastic section. Notice the stochastic policy arrows may already look different in some cells.

Not seeing the stochastic section?

  • Make sure your new code is indented under the if __name__ == "__main__": block. Python requires consistent indentation (4 spaces).
  • Check that you added the code after the last gw_det.print_grid(values=values_det) line, not inside a function definition.

Help me debug why the stochastic section isn't printing.

Compare the deterministic and stochastic policies

The two policies are printed separately, but you need a direct comparison to see exactly which states changed strategy due to uncertainty.

  • In mdp_gridworld.py, add the following comparison loop at the bottom of the __main__ block, below the stochastic section you just added:
    print("\n\n=== POLICY COMPARISON ===")
    print("States where the policy CHANGED due to stochasticity:")
    changes_found = False
    for state in gw_det.states:
        if state in policy_det and state in policy_stoch:
            if policy_det[state] != policy_stoch[state]:
                print(
                    f"  State {state}: "
                    f"{policy_det[state]} -> {policy_stoch[state]}"
                )
                changes_found = True
    if not changes_found:
        print("  (no changes)")

What does this code do?

  • Loops over every state in the grid and compares the deterministic policy to the stochastic policy.
  • Prints any state where the optimal action changed. These are the states where the risk of slipping into the pit outweighs the benefit of the shorter path.
  • Includes a fallback message in case no changes are found (which would indicate a bug, since the pit should cause at least one policy change).
  • Save mdp_gridworld.py.

Before you run the script, think about this: which states near the pit do you expect to change their policy when there is a 10% chance of slipping sideways?

  • Run the script to see the full comparison by running:
python mdp_gridworld.py

In the POLICY COMPARISON section, you should see at least one state where the policy changed. A state near the pit that previously pointed RIGHT (alongside the danger) now points UP or another direction to avoid the risk of slipping in.

Why did the policy change near the pit?

With deterministic actions, walking alongside the pit at (3,1) is perfectly safe. The agent can hug the bottom row and turn up at the last moment.

With stochastic actions, every step next to the pit carries a 10% chance of slipping into it and receiving -1 reward. The Bellman equation weighs this risk automatically. The expected value of walking alongside the pit drops, so the agent reroutes to approach the goal from a safer direction.

Seeing "(no changes)" in the comparison?

  • Verify that gw_stoch uses noise=0.1 and not noise=0.0.
  • Check that the comparison loop references policy_det and policy_stoch (not the same variable twice).
  • Make sure you are iterating over gw_det.states which includes all non-wall states.

Help me figure out why my policies are identical.

✔️ Awesome, I've got everything!

Great. Double check that you have saved mdp_gridworld.py and that your terminal shows the policy comparison output.

ⓧ I'd like to double check the full code

Your mdp_gridworld.py should contain these sections in order. Scroll through your file and confirm each is present:

  • Imports (numpy, matplotlib.pyplot, matplotlib.patches).
  • The GridWorld class with __init__, is_terminal, get_reward, get_next_state, get_transitions, and print_grid.
  • The value_iteration function.
  • The extract_policy function.
  • The __main__ block with three sections: DETERMINISTIC GRID WORLD, STOCHASTIC GRID WORLD (noise=0.1), and POLICY COMPARISON.

You have just proven that uncertainty fundamentally changes how an agent makes decisions. Next up, you will turn these numbers and arrows into a visual heatmap that makes the structure of the solution obvious at a glance.

Visualize the Value Function and Policy

Your algorithm works. The terminal prints correct values and policies. But rows of numbers and arrow symbols are hard to interpret at a glance.

A heatmap makes the structure of the solution instantly visible. Colors reveal which states are valuable (green, near the goal) and which are dangerous (red, near the pit). Arrows show the flow of optimal decisions. Together they tell a story no table of numbers can.

In this step, get ready to:
  • Render the value function as a color-coded heatmap using Matplotlib.
  • Draw policy arrows and label special cells (goal, pit, wall).
  • Save the finished figure as a PNG file.
Build the heatmap

The value dictionary maps (x, y) tuples to floats. To use plt.imshow, you need a regular 2D array. You will convert the dictionary into a NumPy grid, then render it with a red-yellow-green colormap.

  • In mdp_gridworld.py, add the following visualize function above the if __name__ == "__main__": block:
def visualize(grid_world, values, policy, title="MDP Solution"):
    # Convert value dictionary to 2D array for heatmap
    value_grid = np.full((grid_world.rows, grid_world.cols), np.nan)
    for state, value in values.items():
        x, y = state
        value_grid[y, x] = value

    # Map actions to arrow offsets for drawing later
    arrow_map = {
        "UP": (0, 0.3),
        "DOWN": (0, -0.3),
        "LEFT": (-0.3, 0),
        "RIGHT": (0.3, 0),
    }

    fig, ax = plt.subplots(1, 1, figsize=(8, 6))

    # Render heatmap: green = high value, red = low value
    im = ax.imshow(
        value_grid, cmap="RdYlGn", origin="lower", vmin=-1, vmax=1
    )
    plt.colorbar(im, ax=ax, label="State Value")

    plt.tight_layout()
    plt.savefig("mdp_solution.png", dpi=150, bbox_inches="tight")
    plt.show()
    print("\nVisualization saved to 'mdp_solution.png'")

What does this code do?

np.full(..., np.nan) creates a grid pre-filled with NaN so walls (which have no value) render as blank.

origin="lower" places row 0 at the bottom of the plot, matching the grid world's coordinate system where y=0 is the bottom row.

cmap="RdYlGn" maps values from -1 (red) through 0 (yellow) to +1 (green), so dangerous states glow red and the goal glows green.

  • Now scroll to the very end of your if __name__ == "__main__": block. After the policy comparison code, add this call:
    visualize(
        gw_stoch, values_stoch, policy_stoch,
        title="MDP Grid World - Stochastic (noise=0.1)",
    )
  • Save mdp_gridworld.py.
  • Run the script to see your first heatmap:
python mdp_gridworld.py

A Matplotlib window opens showing a colored grid. Green cells appear near the goal at (3,2), red near the pit at (3,1), and a blank white cell where the wall lives. Close the window to continue.

No window appeared?

  • If you see a warning about a non-interactive backend, you may need to set the backend. Add import matplotlib; matplotlib.use('TkAgg') at the very top of the file, above the other imports.
  • Make sure the visualize call is indented inside the if __name__ block (4 spaces).
  • Help me debug my matplotlib visualization.
Draw policy arrows and mark special cells

The heatmap shows which states are valuable, but not what the agent should do there. Arrows make the policy visible. You will also label the wall and terminal states so the grid reads like a map.

  • In your visualize function, find the line that says plt.tight_layout(). Add the following code above it to draw value labels and policy arrows:
    # Draw value text and policy arrows for each state
    for state in grid_world.states:
        x, y = state
        value = values[state]
        text_color = "white" if abs(value) > 0.7 else "black"
        ax.text(
            x, y - 0.35, f"{value:.2f}", ha="center", va="center",
            fontsize=8, color=text_color,
        )

        if state in policy:
            dx, dy = arrow_map[policy[state]]
            ax.arrow(
                x - dx / 2, y - dy / 2, dx, dy,
                head_width=0.12, head_length=0.06, fc="black", ec="black",
            )

What does this code do?

  • For each state, it prints the numeric value below center and picks a text color (white on dark cells, black on light cells) for readability.
  • ax.arrow() draws from slightly behind center to slightly ahead, using the arrow_map offsets to point UP, DOWN, LEFT, or RIGHT.
  • Save mdp_gridworld.py and run the script again:
python mdp_gridworld.py

The heatmap now shows black arrows in each non-terminal cell, pointing toward the goal. You can see the flow of optimal decisions across the grid.

  • Close the matplotlib window. Now add the wall and terminal labels. Still above plt.tight_layout() (and below the arrow loop you just added), paste this code:
    # Mark walls with a gray rectangle
    for wall in grid_world.walls:
        x, y = wall
        rect = mpatches.FancyBboxPatch(
            (x - 0.5, y - 0.5), 1, 1,
            boxstyle="round,pad=0", facecolor="gray", edgecolor="black",
        )
        ax.add_patch(rect)
        ax.text(
            x, y, "WALL", ha="center", va="center",
            fontsize=8, fontweight="bold", color="white",
        )

    # Label terminal states (Goal and Pit)
    for state, reward in grid_world.terminals.items():
        x, y = state
        label = f"Goal\n+{reward:.0f}" if reward > 0 else f"Pit\n{reward:.0f}"
        ax.text(
            x, y + 0.2, label, ha="center", va="center",
            fontsize=9, fontweight="bold",
        )

What does this code do?

  • FancyBboxPatch draws a filled gray rectangle over the wall cell so it stands out from the heatmap.
  • The terminal-label loop prints "Goal" or "Pit" with the reward value so the learner can immediately identify the endpoints without checking coordinates.
  • Save mdp_gridworld.py and run the script again:
python mdp_gridworld.py

The wall cell is now a solid gray rectangle labeled "WALL". The goal and pit cells show their names and reward values.

Arrows not visible or overlapping labels?

  • Check that the arrow loop comes before the wall/terminal loops in your function. The wall patch draws over anything beneath it.
  • Make sure you have a blank line between the arrow loop's closing and the wall loop's start. Python's indentation must stay at 4 spaces throughout.
  • Help me fix my matplotlib arrows or labels.
Add grid lines and save the final figure

The heatmap looks good, but without grid lines it is hard to tell where one cell ends and the next begins. Adding lines and axis labels finishes the visualization.

  • In the visualize function, find plt.tight_layout(). Add the following code above it (after the terminal-label loop):
    # Add grid lines between cells and label axes
    ax.set_xticks(np.arange(-0.5, grid_world.cols, 1), minor=True)
    ax.set_yticks(np.arange(-0.5, grid_world.rows, 1), minor=True)
    ax.grid(which="minor", color="black", linewidth=1)
    ax.set_xticks(range(grid_world.cols))
    ax.set_yticks(range(grid_world.rows))
    ax.set_xticklabels([f"Col {i}" for i in range(grid_world.cols)])
    ax.set_yticklabels([f"Row {i}" for i in range(grid_world.rows)])
    ax.set_title(title, fontsize=14, fontweight="bold")

What does this code do?

  • Minor ticks at -0.5, 0.5, 1.5... place grid lines on cell boundaries instead of through cell centers.
  • Major ticks at 0, 1, 2... provide the axis labels ("Col 0", "Row 1", etc.) so you can read coordinates directly from the plot.
  • Save mdp_gridworld.py.

Before you run the script for the last time, predict what you expect to see. Will the arrows near the pit point toward it or away from it?

  • Run the script:
python mdp_gridworld.py

A Matplotlib window opens with the finished visualization. The grid cells have clear black borders, labeled axes, and a bold title. Green cells cluster near the goal at (3,2), red near the pit at (3,1). Arrows near the pit route away from it because the stochastic agent accounts for the risk of slipping.

  • Close the matplotlib window. Confirm that mdp_solution.png was saved to your project folder by running:
ls

You should see mdp_solution.png listed alongside mdp_gridworld.py.

Don't see mdp_solution.png?

  • Make sure plt.savefig("mdp_solution.png", dpi=150, bbox_inches="tight") appears before plt.show() in your function. If savefig comes after show(), the figure is already cleared and the PNG will be blank.
  • Check that your terminal is inside the mdp-gridworld folder. The file saves relative to where you run the script.
  • Help me find my saved PNG file.

✔️ Awesome, I've got everything!

Great. Double check you have saved mdp_gridworld.py and confirmed mdp_solution.png exists in your project folder.

ⓧ I'd like to double check the full code

Your visualize function should contain these sections in order. Scroll up through your file and confirm each is present:

  • Function signature: def visualize(grid_world, values, policy, title="MDP Solution"):
  • The value_grid array creation and arrow_map dictionary.
  • The ax.imshow() call with cmap="RdYlGn" and the colorbar.
  • The loop over grid_world.states that draws value text and policy arrows with ax.arrow().
  • The wall loop with FancyBboxPatch and the terminal-label loop.
  • Grid lines (ax.grid(which="minor", ...)), axis labels, and the title.
  • plt.tight_layout(), plt.savefig("mdp_solution.png", dpi=150, bbox_inches="tight"), plt.show(), and the print statement at the very end.
  • At the bottom of __main__: the visualize(gw_stoch, values_stoch, policy_stoch, ...) call.

Your MDP solver now produces a publication-quality visualization. Colors, arrows, and labels work together to tell the story of how uncertainty reshapes optimal decisions. Next up, you will tackle a bonus challenge that implements an entirely different algorithm for solving the same MDP.

Secret mission

Implement Policy Iteration

Policy Iteration is the other classical MDP algorithm. It converges in fewer improvement steps but does more work per step. Can you implement it and race it against Value Iteration to compare convergence speed?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • The mdp-gridworld project folder on your Desktop (contains mdp_gridworld.py and mdp_solution.png).
  • NumPy and Matplotlib Python packages installed via pip.

Keep everything running

No action needed. Choose this if you plan to experiment with different grid sizes, reward structures, or discount factors to see how they change the policy.

  • Your mdp-gridworld folder stays on your Desktop with all files intact.
  • NumPy and Matplotlib remain installed and available for other Python projects.

Pause - I'll come back to this later

Since this project runs entirely locally with no background processes, there is nothing to pause. Your files will remain exactly as you left them.

  • Your mdp-gridworld folder and all generated files stay in place.
  • When you return, open the folder in VS Code and run python mdp_gridworld.py to pick up where you left off.

Delete - I don't want to use this again

Remove all project files and optionally uninstall the Python packages.

  • Delete the project folder by running this in your terminal:

macOS

rm -rf ~/Desktop/mdp-gridworld

Windows

rmdir /s /q %USERPROFILE%\Desktop\mdp-gridworld
  • Optionally uninstall the Python packages by running:
pip uninstall numpy matplotlib -y

Keep NumPy and Matplotlib?

NumPy and Matplotlib are widely used Python libraries. If you plan to do other data science or visualization work, you may want to keep them installed.

Nice Work!

Nice Work!

You just built a complete Markov Decision Process solver from scratch and watched it find optimal strategies under uncertainty.

You've learned how to:

  • Build a Grid World environment encoding the five components of an MDP and implement value iteration using the Bellman equation to compute optimal state values until convergence.
  • Extract an optimal policy from the converged value function and demonstrate how stochastic transitions reshape the agent's strategy near dangerous states.
  • Visualize the value function as a color-coded heatmap with policy arrows using Matplotlib, turning abstract math into a readable picture of optimal decision-making.
  • Secret Mission: Implemented Policy Iteration and compared its convergence speed against Value Iteration on the same stochastic grid world.

Ready to quiz yourself?