Train a Robot with Computer Vision + RL

Build a simulated robot that learns to navigate a floor plan from pixels.

Introduction

30 Second Summary

Think about how a robot vacuum learns to navigate your home. It can only see a small area around itself, yet it figures out how to move between rooms and avoid furniture. That same challenge sits at the heart of modern robotics research.

In this project, you will train a virtual robot to navigate a house floor plan using only its limited local camera view. You will combine computer vision with reinforcement learning to teach the robot spatial reasoning from raw pixel observations.

What You'll Build

A complete visual navigation pipeline where a simulated robot learns to navigate a 2D house floor plan from a top-down camera perspective using a CNN policy trained with PPO.

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

  • A floor plan occupancy grid processed from an image using OpenCV morphological operations.
  • A custom Gymnasium environment with egocentric visual observations that simulate a robot's limited camera view.
  • A trained reinforcement learning agent that learns efficient navigation purely from pixels, with training curves and path visualizations showing the robot's learned behavior.
  • Secret Mission: Add multiple goal rooms and train the agent to navigate to any of them using goal-conditioned observations.

Are there any prerequisites?

You should be comfortable writing Python. A basic understanding of machine learning concepts (training, loss, neural networks) is helpful but not required.

Everything runs for free on Google Colab. No local GPU or paid services needed.

Before We Start

Training an AI agent to navigate a floor plan requires serious computing power. A convolutional neural network processes thousands of image observations during training, so a GPU dramatically speeds things up.

Google Colab gives us free access to an NVIDIA T4 GPU. In this step, you will create a Colab notebook, enable the GPU, and install all the libraries needed for computer vision and reinforcement learning.

In this step, get ready to:
  • Create a Google Colab notebook with a T4 GPU runtime.
  • Install the required packages for reinforcement learning and computer vision.
  • Import all libraries and verify GPU availability.
Create a Colab notebook and enable GPU

Google Colab runs Python notebooks in the cloud with free GPU access. You need to create a new notebook and switch it to a GPU runtime so training runs fast.

  • Open Google Colab in your browser and sign in with your Google account.
  • Click New notebook in the bottom-right corner of the pop-up dialog.
  • In the top menu bar, click Runtime.
  • Click Change runtime type.
  • Under Hardware accelerator, select T4 GPU.
  • Click Save.

Why do we need a GPU?

Our agent uses a convolutional neural network (CNN) to process image observations. During training, the CNN processes hundreds of thousands of images. A GPU handles this massively parallel math far faster than a CPU.

The free T4 GPU in Colab is more than enough for this project. Training 100,000 timesteps takes about 2-5 minutes on a T4.

Install the required packages

The project uses Stable-Baselines3 for reinforcement learning, Gymnasium for the custom environment, and OpenCV for image processing. You need to install them with specific version pins so everything stays compatible.

  • Install all required packages by pasting this command into the first cell and running it:
!pip install stable-baselines3==2.8.0 gymnasium==1.2.3 opencv-python-headless numpy matplotlib

What does this command do?

  • stable-baselines3==2.8.0 provides the PPO reinforcement learning algorithm and CNN policy tools.
  • gymnasium==1.2.3 is the framework for building custom RL environments with standard interfaces.
  • opencv-python-headless gives us image processing functions (no GUI needed in Colab).
  • numpy and matplotlib handle numerical arrays and plotting.
  • Run the cell by clicking the play button on the left, or pressing Shift+Enter.
  • Wait for the installation to finish. You will see a series of "Successfully installed" messages in the output.

Seeing dependency errors?

If pip reports a dependency conflict, restart your runtime by clicking Runtime then Restart runtime. Then re-run the install cell. Colab sometimes has pre-installed packages that conflict with the pinned versions.

help me fix a pip dependency conflict in Google Colab when installing stable-baselines3 and gymnasium

Import all libraries and verify GPU access

With the packages installed, you need to import all the modules your project will use and confirm that PyTorch can see the GPU.

  • Create a new code cell below the install cell by clicking + Code in the toolbar.
  • Paste the following imports and GPU check into the new cell:
import numpy as np
import cv2
import matplotlib.pyplot as plt
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
import torch
import torch.nn as nn

# Check if GPU is available
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

What are these imports for?

  • numpy and cv2 handle the floor plan as a grid of numbers and apply image processing.
  • gymnasium provides the base class for our custom navigation environment.
  • PPO, BaseFeaturesExtractor, DummyVecEnv, and Monitor are the Stable-Baselines3 tools for training and evaluating the agent.
  • torch and torch.nn let us build the custom CNN that the agent uses to understand its visual observations.
  • Run the cell with Shift+Enter.
  • You should see output showing the PyTorch version and CUDA available: True followed by GPU: Tesla T4 (or similar).

Seeing CUDA available: False?

This means the GPU runtime is not active. Go back to Runtime then Change runtime type and confirm T4 GPU is selected. Click Save and re-run both cells.

help me enable GPU in Google Colab - torch.cuda.is_available() returns False

✔️ Awesome, I've got everything!

Great. Double check you have both cells saved in your notebook (the install cell and the imports cell).

ⓧ I'd like to double check the full code

Here is everything your notebook should contain so far across two cells:

Cell 1 (installation):

!pip install stable-baselines3==2.8.0 gymnasium==1.2.3 opencv-python-headless numpy matplotlib

Cell 2 (imports and GPU check):

import numpy as np
import cv2
import matplotlib.pyplot as plt
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
import torch
import torch.nn as nn

# Check if GPU is available
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

Your environment is ready. Next up, you will process a floor plan image into a navigable grid using computer vision.

Set Up the Colab Environment

Our robot needs a world to navigate. In reinforcement learning, the "world" is called an environment. Before building the full environment, we need the physical space itself: a floor plan that defines where the robot can and cannot move.

In this step, you'll generate a procedural floor plan using NumPy. The floor plan simulates a small house with four rooms connected by hallways, represented as a 64x64 grid.

In this step, get ready to:
  • Define a function that generates a binary floor plan grid with rooms and hallways.
  • Visualize the floor plan to confirm the layout is correct.
Define the floor plan generator function

A floor plan for reinforcement learning is stored as an occupancy grid. Each cell in the grid is either a wall (0) or free space (1) that the robot can walk through.

  • In your Colab notebook, add a new code cell below your imports cell by clicking + Code in the toolbar.
  • Paste the following function definition into the new cell:
# --- Cell 3: Generate a Procedural Floor Plan ---
def create_floor_plan(size=64):
    """Generate a simple house floor plan as a binary occupancy grid.

    Args:
        size: The width and height of the grid in pixels.

    Returns:
        A 2D numpy array where 0 = wall (obstacle) and 1 = free space.
    """
    # Start with all walls (zeros)
    grid = np.zeros((size, size), dtype=np.uint8)

    # Carve out rooms (set to 1 = free space)
    # Room 1: Top-left (living room)
    grid[2:28, 2:28] = 1
    # Room 2: Top-right (kitchen)
    grid[2:28, 36:62] = 1
    # Room 3: Bottom-left (bedroom)
    grid[36:62, 2:28] = 1
    # Room 4: Bottom-right (bathroom)
    grid[36:62, 36:62] = 1

    # Carve out hallways connecting rooms (doorways)
    # Horizontal hallway connecting top rooms
    grid[12:18, 28:36] = 1
    # Vertical hallway connecting left rooms
    grid[28:36, 10:20] = 1
    # Horizontal hallway connecting bottom rooms
    grid[46:52, 28:36] = 1
    # Vertical hallway connecting right rooms
    grid[28:36, 44:54] = 1

    return grid

What does this code do?

The function starts with a grid of all zeros (walls everywhere), then "carves out" rooms and hallways by setting regions to 1.

  • Four rooms occupy the corners of the 64x64 grid, each roughly 26x26 cells.
  • Four hallways connect adjacent rooms so the robot has a valid path between any two rooms.
  • The dtype=np.uint8 keeps memory usage low, which matters when we later stack thousands of observations.

Why generate the floor plan instead of loading an image?

Generating the floor plan procedurally with NumPy guarantees a valid layout with rooms that are always connected. If we loaded a real image, we'd need extra validation to ensure the robot can actually reach every room.

This approach also makes the project self-contained. No file downloads or uploads required.

You won't see any output yet. The next chunk calls this function and displays the result.

Generate and visualize the floor plan
  • In the same cell, below the function definition, add the following code to create and display the floor plan:
# Create and visualize the floor plan
floor_plan = create_floor_plan(size=64)

plt.figure(figsize=(6, 6))
plt.imshow(floor_plan, cmap="gray", origin="upper")
plt.title("Generated Floor Plan\n(White = Free Space, Black = Walls)")
plt.axis("off")
plt.show()

print(f"Floor plan shape: {floor_plan.shape}")
print(f"Free cells: {floor_plan.sum()} / {floor_plan.size}")

What does this code do?

This calls your create_floor_plan function and stores the result in floor_plan. Then it uses Matplotlib to render the grid as an image.

  • White pixels (value 1) represent free space the robot can walk on.
  • Black pixels (value 0) represent walls the robot cannot pass through.
  • The print statements confirm the grid shape and how much of the grid is walkable.
  • Run the cell by pressing Shift+Enter.
  • You should see a black-and-white image with four white squares (rooms) connected by narrow white rectangles (hallways).
  • Below the image, confirm you see output similar to Floor plan shape: (64, 64) and a free-cells count.

Seeing a blank or all-black image?

  • Make sure the function and visualization code are in the same cell. If they are in separate cells, run both cells in order.
  • Check that your imports cell (Cell 2) ran successfully first. Matplotlib must be imported before calling plt.show().
  • If you see a NameError: name 'np' is not defined error, re-run your imports cell above.

Help me debug my floor plan visualization

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook.

ⓧ I'd like to double check the full code

Here is the full notebook code up to this point. Make sure your cells match:

# --- Cell 1: Installation ---
# !pip install stable-baselines3==2.8.0 gymnasium==1.2.3 opencv-python-headless numpy matplotlib

# --- Cell 2: Imports and GPU Check ---
import numpy as np
import cv2
import matplotlib.pyplot as plt
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
import torch
import torch.nn as nn

# Check if GPU is available
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# --- Cell 3: Generate a Procedural Floor Plan ---
def create_floor_plan(size=64):
    grid = np.zeros((size, size), dtype=np.uint8)
    grid[2:28, 2:28] = 1
    grid[2:28, 36:62] = 1
    grid[36:62, 2:28] = 1
    grid[36:62, 36:62] = 1
    grid[12:18, 28:36] = 1
    grid[28:36, 10:20] = 1
    grid[46:52, 28:36] = 1
    grid[28:36, 44:54] = 1
    return grid

floor_plan = create_floor_plan(size=64)

plt.figure(figsize=(6, 6))
plt.imshow(floor_plan, cmap="gray", origin="upper")
plt.title("Generated Floor Plan\n(White = Free Space, Black = Walls)")
plt.axis("off")
plt.show()

print(f"Floor plan shape: {floor_plan.shape}")
print(f"Free cells: {floor_plan.sum()} / {floor_plan.size}")

Your floor plan is ready. Next up, you'll process this grid with OpenCV and define the robot's limited camera view.

Process the Floor Plan with Computer Vision

You have a raw floor plan grid, but real-world images are noisy. Small gaps in walls could let the robot "cheat" through barriers that should be solid.

In this step, you will clean the floor plan using OpenCV morphological operations, define where the robot starts and where it needs to go, and build the egocentric observation function that gives the robot a limited camera view of its surroundings.

In this step, get ready to:
  • Apply computer vision operations to clean the floor plan into a usable occupancy grid.
  • Define start and goal positions and visualize them on the map.
  • Implement the egocentric view function that simulates the robot's limited camera.
Apply morphological operations to clean the floor plan

Morphological closing is a classic computer vision operation that fills small gaps and smooths edges in binary images. It works by dilating the image first (expanding white regions), then eroding it back (shrinking them). The net effect is that tiny holes in walls get sealed.

  • Create a new code cell in your Colab notebook and add the floor plan processing function by pasting this code:
# --- Cell 4: Process Floor Plan and Define Positions ---
def process_floor_plan(grid):
    """Apply morphological operations to clean up the floor plan.

    Args:
        grid: Binary occupancy grid (0=wall, 1=free).

    Returns:
        Cleaned binary grid with smoother walls.
    """
    # Convert to uint8 image format for OpenCV (0 or 255)
    img = (grid * 255).astype(np.uint8)

    # Apply morphological closing to fill small gaps in walls
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    cleaned = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)

    # Convert back to binary (0 or 1)
    return (cleaned > 127).astype(np.uint8)


# Create and visualize the floor plan
processed_plan = process_floor_plan(floor_plan)

What does this code do?

The function converts the binary grid (0s and 1s) to a 0-255 image that OpenCV can work with. It then creates a 3x3 rectangular kernel and uses cv2.morphologyEx with MORPH_CLOSE to fill any small gaps in the walls.

Finally it converts back to binary (0 or 1) by thresholding at 127. This gives us a clean occupancy grid where 1 means "the robot can walk here" and 0 means "wall."

What is an occupancy grid?

An occupancy grid is a grid-based representation of space where each cell stores whether that location is occupied (wall) or free (navigable). Real robots use these grids built from LIDAR or camera data to plan paths through their environment.

  • Do not run this cell yet. You will add more code to it in the next substep.
Define start and goal positions

The robot needs a starting position and a destination. You will place the start in the top-left room and the goal in the bottom-right room, forcing the robot to navigate through hallways to reach it.

  • Add the following code to the same cell, below the processed_plan line:
# Define start and goal positions (row, col)
start_pos = np.array([5, 5], dtype=np.int32)       # Top-left room
goal_pos = np.array([49, 49], dtype=np.int32)       # Bottom-right room

# Visualize with start and goal
plt.figure(figsize=(6, 6))
plt.imshow(processed_plan, cmap="gray", origin="upper")
plt.plot(start_pos[1], start_pos[0], "go", markersize=12, label="Start")
plt.plot(goal_pos[1], goal_pos[0], "r*", markersize=15, label="Goal")
plt.legend(loc="upper right", fontsize=10)
plt.title("Floor Plan with Start and Goal Positions")
plt.axis("off")
plt.show()

What does this code do?

The start position [5, 5] places the robot in row 5, column 5 of the grid. That is inside the top-left room. The goal position [49, 49] is inside the bottom-right room.

The Matplotlib plot shows the cleaned floor plan with a green circle at the start and a red star at the goal. Notice that plt.plot uses (col, row) order because Matplotlib's x-axis corresponds to columns.

  • Do not run this cell yet. You will add one more block of code below.
Implement the egocentric observation function

A real robot does not see the entire floor plan from above. It only sees what is immediately around it through its camera. To simulate this, you will build a function that extracts a small window of the grid centered on the robot's current position.

  • Add the following code to the bottom of the same cell:
def get_egocentric_view(grid, position, view_size=7):
    """Extract a local observation window centered on the agent's position.

    This simulates a top-down camera with limited field of view.
    Areas outside the grid boundaries are treated as walls (0).

    Args:
        grid: The full floor plan occupancy grid.
        position: Agent's current (row, col) position.
        view_size: The width/height of the observation window (must be odd).

    Returns:
        A (view_size, view_size) numpy array representing the local view.
    """
    half = view_size // 2
    row, col = position

    # Create a padded version of the grid (pad with walls = 0)
    padded = np.pad(grid, pad_width=half, mode="constant", constant_values=0)

    # Extract the view window (adjust indices for padding offset)
    view = padded[row:row + view_size, col:col + view_size].copy()

    return view


# Test the egocentric view
test_view = get_egocentric_view(processed_plan, start_pos, view_size=7)
print(f"Egocentric view shape: {test_view.shape}")

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].imshow(processed_plan, cmap="gray", origin="upper")
axes[0].plot(start_pos[1], start_pos[0], "go", markersize=10)
axes[0].set_title("Full Floor Plan (Agent at Start)")
axes[0].axis("off")

axes[1].imshow(test_view, cmap="gray", origin="upper", vmin=0, vmax=1)
axes[1].set_title("Agent's Egocentric View (7x7)")
axes[1].axis("off")
plt.tight_layout()
plt.show()

What does this code do?

The get_egocentric_view function pads the entire grid with zeros (walls) around the edges. This ensures that when the robot is near a boundary, the view window still has the correct size. It just sees walls beyond the edge.

The function then slices a 7x7 window centered on the robot's position. This 7x7 patch is all the robot will "see" when learning to navigate. The side-by-side plot compares the full map to the tiny local view the agent receives.

Why only a 7x7 view?

Partial observability is what makes this problem interesting. If the robot could see the entire map, navigation would be trivial. By limiting it to a 7x7 window, the agent must learn spatial memory and wall-following strategies from experience, similar to how a real robot navigates with a limited camera field of view.

  • Run the cell by pressing Shift+Enter.
  • You should see two plots appear. The first shows the floor plan with a green circle (start) and red star (goal). The second shows the full floor plan next to a small 7x7 grayscale patch representing the robot's local view.

Not seeing the plots?

Make sure you ran the previous cells first (Cell 2 for imports and Cell 3 for the floor plan). If you see a NameError mentioning floor_plan, go back and run Cell 3.

If the plot appears blank or all-black, check that processed_plan.sum() returns a number greater than zero. A zero sum means the morphological operation eliminated all free space, which should not happen with a 3x3 kernel.

Help me debug why my floor plan plots are not showing.

✔️ Awesome, I've got everything!

Great. Make sure you have saved your notebook before moving on.

ⓧ I'd like to double check the full code

Here is the complete Cell 4 code. Make sure yours matches:

# --- Cell 4: Process Floor Plan and Define Positions ---
def process_floor_plan(grid):
    img = (grid * 255).astype(np.uint8)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    cleaned = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
    return (cleaned > 127).astype(np.uint8)

processed_plan = process_floor_plan(floor_plan)

start_pos = np.array([5, 5], dtype=np.int32)
goal_pos = np.array([49, 49], dtype=np.int32)

plt.figure(figsize=(6, 6))
plt.imshow(processed_plan, cmap="gray", origin="upper")
plt.plot(start_pos[1], start_pos[0], "go", markersize=12, label="Start")
plt.plot(goal_pos[1], goal_pos[0], "r*", markersize=15, label="Goal")
plt.legend(loc="upper right", fontsize=10)
plt.title("Floor Plan with Start and Goal Positions")
plt.axis("off")
plt.show()

def get_egocentric_view(grid, position, view_size=7):
    half = view_size // 2
    row, col = position
    padded = np.pad(grid, pad_width=half, mode="constant", constant_values=0)
    view = padded[row:row + view_size, col:col + view_size].copy()
    return view

test_view = get_egocentric_view(processed_plan, start_pos, view_size=7)
print(f"Egocentric view shape: {test_view.shape}")

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].imshow(processed_plan, cmap="gray", origin="upper")
axes[0].plot(start_pos[1], start_pos[0], "go", markersize=10)
axes[0].set_title("Full Floor Plan (Agent at Start)")
axes[0].axis("off")
axes[1].imshow(test_view, cmap="gray", origin="upper", vmin=0, vmax=1)
axes[1].set_title("Agent's Egocentric View (7x7)")
axes[1].axis("off")
plt.tight_layout()
plt.show()

Your floor plan is processed, your positions are set, and your robot has its "camera" function ready. Next up, you will build the full custom Gymnasium environment that wraps all of this into a proper reinforcement learning interface.

Build the Custom Gymnasium Environment

You have a processed floor plan, a start position, a goal position, and a function that extracts what the robot sees. Now you need to wrap all of that into a format that a reinforcement learning algorithm can interact with.

Gymnasium is the standard interface for RL environments. It defines a contract: the environment must accept actions, return observations, and signal when an episode is done. In this step, you will build a custom FloorPlanEnv class that follows this contract and teaches the agent through reward shaping.

In this step, get ready to:
  • Create a custom Gymnasium environment class with image-based observations and discrete actions.
  • Implement reward shaping that guides the agent toward efficient navigation.
  • Test the environment with random actions to verify it works correctly.
Create the FloorPlanEnv class with observation and action spaces

A Gymnasium environment needs two things defined upfront: what the agent can see (the observation space) and what it can do (the action space). Your robot sees a small image and can move in four directions.

  • Create a new code cell in your notebook and define the FloorPlanEnv class with its __init__ method by pasting this code:
# --- Cell 5: Build the Custom Gymnasium Environment ---
class FloorPlanEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array"]}

    def __init__(self, floor_plan, start_pos, goal_pos, view_size=7, max_steps=500):
        super().__init__()
        self.floor_plan = floor_plan.copy()
        self.start_pos = np.array(start_pos, dtype=np.int32)
        self.goal_pos = np.array(goal_pos, dtype=np.int32)
        self.view_size = view_size
        self.max_steps = max_steps
        # 4 discrete actions: up, down, left, right
        self.action_space = spaces.Discrete(4)
        # Single-channel image observation (channel-first for SB3)
        self.observation_space = spaces.Box(
            low=0, high=255, shape=(1, view_size, view_size), dtype=np.uint8
        )
        # Map action indices to movement deltas (row, col)
        self._action_to_direction = {
            0: np.array([-1, 0]),  # Up
            1: np.array([1, 0]),   # Down
            2: np.array([0, -1]),  # Left
            3: np.array([0, 1]),   # Right
        }
        self._agent_pos = None
        self._step_count = 0

What does this code do?

This defines the skeleton of your RL environment. The key parts are:

  • spaces.Discrete(4) tells the agent it has exactly four possible actions (up, down, left, right).
  • spaces.Box(low=0, high=255, shape=(1, view_size, view_size)) defines the observation as a single-channel 7x7 image with pixel values 0-255. The channel-first format is what Stable-Baselines3 expects.
  • _action_to_direction maps each action number to a row/column movement delta on the grid.

Why channel-first format?

Stable-Baselines3 expects image observations in (C, H, W) order (channel-first). If you provide (H, W, C) instead, SB3 silently wraps the environment to transpose it. Using channel-first from the start avoids that overhead.

The environment also needs a method that generates what the robot sees at its current position. This is the observation the neural network will learn from.

  • Add the _get_obs method inside the FloorPlanEnv class, below the __init__ method:
    def _get_obs(self):
        # Get the local view around the agent's current position
        raw_view = get_egocentric_view(self.floor_plan, self._agent_pos, self.view_size)
        # Scale binary grid (0/1) to image range (0/255)
        obs = (raw_view * 255).astype(np.uint8)
        # Mark the goal with a distinct value if it's within the view
        half = self.view_size // 2
        goal_relative = self.goal_pos - self._agent_pos + half
        gr, gc = goal_relative
        if 0 <= gr < self.view_size and 0 <= gc < self.view_size:
            obs[gr, gc] = 128
        # Add channel dimension for (1, H, W) format
        return obs[np.newaxis, :, :]

What does _get_obs do?

This method generates the observation the agent receives at each timestep. It uses the get_egocentric_view function from the previous step to grab the 7x7 window around the agent.

The observation encodes three types of information through pixel values: 0 = wall, 255 = free space, 128 = goal (only visible when the goal is within the 7x7 window). This gives the CNN three visually distinct signals to learn from.

Implement reset() and step() with reward shaping

The reset() method starts a new episode. The step() method processes an action, moves the agent, and returns a reward signal that teaches the agent what behaviors are good or bad.

  • Add the reset() method below _get_obs:
    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # Place the agent back at the start
        self._agent_pos = self.start_pos.copy()
        self._step_count = 0
        observation = self._get_obs()
        info = {"distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos)}
        return observation, info

What does reset() do?

Every time an episode ends (the agent reaches the goal or runs out of steps), the environment resets. The agent goes back to the start position, the step counter resets, and a fresh observation is returned. The info dict includes the Euclidean distance to the goal for tracking progress.

  • Add the step() method below reset():
    def step(self, action):
        self._step_count += 1
        # Calculate new position from action
        direction = self._action_to_direction[action]
        new_pos = self._agent_pos + direction
        row, col = new_pos
        grid_h, grid_w = self.floor_plan.shape
        # Check if the move is valid (in bounds and not a wall)
        if (0 <= row < grid_h and 0 <= col < grid_w
                and self.floor_plan[row, col] == 1):
            self._agent_pos = new_pos
            reward = -0.01  # Small step penalty encourages efficiency
        else:
            reward = -0.1  # Wall penalty discourages bumping into walls
        # Check if the agent reached the goal
        terminated = np.array_equal(self._agent_pos, self.goal_pos)
        if terminated:
            reward = 10.0  # Large reward for reaching the goal
        # Truncate if max steps exceeded
        truncated = self._step_count >= self.max_steps
        observation = self._get_obs()
        info = {
            "distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos),
            "steps": self._step_count,
        }
        return observation, reward, terminated, truncated, info

How does reward shaping work here?

Reward shaping teaches the agent what behaviors are desirable through three signals:

  • +10.0 for reaching the goal. This is the primary objective.
  • -0.01 per step. Without this, the agent would wander aimlessly even after finding a path. The penalty encourages shorter routes.
  • -0.1 for hitting a wall. This is 10x the step penalty, so the agent quickly learns to avoid walls rather than repeatedly bumping into them.

What is the difference between terminated and truncated?

Gymnasium distinguishes two reasons an episode can end. terminated means the task completed naturally (agent reached the goal). truncated means the episode was cut short (hit the 500-step limit). This distinction matters for training because the agent should not learn that "running out of time" equals failure.

Test the environment with random actions

Before training an agent, you should always verify that your custom environment produces valid observations, returns rewards in the expected range, and terminates correctly.

  • Create a new code cell and test the environment by running:
# --- Cell 6: Test the Environment ---
env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)

# Verify the spaces are defined correctly
print(f"Observation space: {env.observation_space}")
print(f"Action space: {env.action_space}")

# Run a reset and check the initial observation
obs, info = env.reset()
print(f"Initial obs shape: {obs.shape}, dtype: {obs.dtype}")
print(f"Initial info: {info}")

# Take 10 random steps to verify step() works
total_reward = 0
for i in range(10):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    if terminated or truncated:
        break

print(f"After 10 random steps - Reward: {total_reward:.3f}, Distance: {info['distance_to_goal']:.1f}")

What should I see?

You should see output confirming the observation space is Box(0, 255, (1, 7, 7), uint8) and the action space is Discrete(4). The initial observation shape should be (1, 7, 7) with dtype uint8.

The total reward after 10 random steps should be slightly negative (around -0.1 to -1.0), since random actions earn step penalties and wall penalties. The distance should still be large because random movement rarely reaches the goal.

  • Run the cell and verify the output matches the expected shapes and values.

Seeing a shape mismatch or attribute error?

  • Make sure you ran the cell containing the class definition BEFORE this test cell. Colab cells run independently.
  • If you see an error about get_egocentric_view not being defined, go back and run Cell 4 from the previous step.
  • Check that your _get_obs method is indented inside the class body, not at the module level.

Why is my FloorPlanEnv giving unexpected observation shapes or errors?

✔️ Awesome, I've got everything!

Great. Double check that your test cell printed valid shapes and a negative total reward, then move on.

ⓧ I'd like to double check the full code

Here is the complete code for Cells 5 and 6 combined:

# --- Cell 5: Build the Custom Gymnasium Environment ---
class FloorPlanEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array"]}

    def __init__(self, floor_plan, start_pos, goal_pos, view_size=7, max_steps=500):
        super().__init__()
        self.floor_plan = floor_plan.copy()
        self.start_pos = np.array(start_pos, dtype=np.int32)
        self.goal_pos = np.array(goal_pos, dtype=np.int32)
        self.view_size = view_size
        self.max_steps = max_steps
        self.action_space = spaces.Discrete(4)
        self.observation_space = spaces.Box(
            low=0, high=255, shape=(1, view_size, view_size), dtype=np.uint8
        )
        self._action_to_direction = {
            0: np.array([-1, 0]),
            1: np.array([1, 0]),
            2: np.array([0, -1]),
            3: np.array([0, 1]),
        }
        self._agent_pos = None
        self._step_count = 0

    def _get_obs(self):
        raw_view = get_egocentric_view(self.floor_plan, self._agent_pos, self.view_size)
        obs = (raw_view * 255).astype(np.uint8)
        half = self.view_size // 2
        goal_relative = self.goal_pos - self._agent_pos + half
        gr, gc = goal_relative
        if 0 <= gr < self.view_size and 0 <= gc < self.view_size:
            obs[gr, gc] = 128
        return obs[np.newaxis, :, :]

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self._agent_pos = self.start_pos.copy()
        self._step_count = 0
        observation = self._get_obs()
        info = {"distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos)}
        return observation, info

    def step(self, action):
        self._step_count += 1
        direction = self._action_to_direction[action]
        new_pos = self._agent_pos + direction
        row, col = new_pos
        grid_h, grid_w = self.floor_plan.shape
        if (0 <= row < grid_h and 0 <= col < grid_w
                and self.floor_plan[row, col] == 1):
            self._agent_pos = new_pos
            reward = -0.01
        else:
            reward = -0.1
        terminated = np.array_equal(self._agent_pos, self.goal_pos)
        if terminated:
            reward = 10.0
        truncated = self._step_count >= self.max_steps
        observation = self._get_obs()
        info = {
            "distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos),
            "steps": self._step_count,
        }
        return observation, reward, terminated, truncated, info

# --- Cell 6: Test the Environment ---
env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
print(f"Observation space: {env.observation_space}")
print(f"Action space: {env.action_space}")
obs, info = env.reset()
print(f"Initial obs shape: {obs.shape}, dtype: {obs.dtype}")
print(f"Initial info: {info}")
total_reward = 0
for i in range(10):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    if terminated or truncated:
        break
print(f"After 10 random steps - Reward: {total_reward:.3f}, Distance: {info['distance_to_goal']:.1f}")

Your custom environment is working. The random agent performs poorly (as expected), but the observation and reward mechanics are solid. Next, you will build a custom CNN that can actually learn from these small 7x7 observations and train a PPO agent to navigate the floor plan.

Build the CNN Extractor and Train the Agent

Your environment is ready and producing observations. The next challenge is building a brain that can actually learn from those tiny 7x7 pixel views.

The default CNN in Stable-Baselines3 uses large 8x8 and 4x4 filters designed for 84x84 Atari game frames. Those filters would reduce our 7x7 observation to nothing useful. In this step, you will build a custom CNN with smaller filters, then train a PPO agent that learns to navigate from start to goal.

In this step, get ready to:
  • Define a custom CNN feature extractor sized for 7x7 observations.
  • Configure and train a PPO agent with 100,000 timesteps.
  • Plot the training reward curve to verify learning progress.
Define the custom CNN feature extractor

The default Nature CNN in Stable-Baselines3 expects input images of at least 36x36 pixels. Our agent sees a 7x7 window, so we need a custom network with 3x3 convolutional filters that preserve spatial information at this small scale.

  • Create a new code cell in your notebook and add the custom CNN class:
# --- Cell 7: Custom CNN Feature Extractor ---
class SmallCNN(BaseFeaturesExtractor):
    def __init__(self, observation_space: spaces.Box, features_dim: int = 64):
        super().__init__(observation_space, features_dim)
        n_input_channels = observation_space.shape[0]
        self.cnn = nn.Sequential(
            nn.Conv2d(n_input_channels, 16, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten(),
        )
        with torch.no_grad():
            sample = torch.as_tensor(observation_space.sample()[None]).float()
            n_flatten = self.cnn(sample).shape[1]
        self.linear = nn.Sequential(
            nn.Linear(n_flatten, features_dim),
            nn.ReLU()
        )

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        return self.linear(self.cnn(observations))

What does this code do?

SmallCNN subclasses BaseFeaturesExtractor from Stable-Baselines3. It takes the raw pixel observation and compresses it into a 64-dimensional feature vector that the PPO policy network can use for decision-making.

  • Two Conv2d layers with 3x3 kernels progressively extract spatial features from the 7x7 input. The first layer produces 16 feature maps, the second produces 32.
  • The torch.no_grad() block runs a dummy observation through the CNN to calculate the flattened output size automatically. This avoids manual calculation when architecture changes.
  • A final linear layer maps the flattened CNN output to a compact 64-dimensional vector.

Why not use the default Nature CNN?

The Nature CNN was designed for Atari games with 84x84 pixel frames. Its first layer alone uses an 8x8 kernel with stride 4. On a 7x7 input, that kernel is literally larger than the entire image. Our custom CNN uses 3x3 kernels that can meaningfully detect spatial patterns (walls, free space, goal) within the small observation window.

  • Run the cell. You should see no output (the class is just defined, not instantiated yet).

Seeing a NameError?

If you see an error about BaseFeaturesExtractor or nn not being defined, make sure you ran Cell 2 (Imports) earlier in this session. Colab does not persist variables between sessions.

Help me fix a NameError in my SmallCNN class definition

Configure and train the PPO agent

Before training, you need a way to track how well the agent performs over time. A callback will record the total reward from each completed episode so you can plot the learning curve later.

  • Create a new cell and define the reward logging callback:
# --- Cell 8: Training Callback for Logging ---
class RewardLoggerCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
        self.episode_rewards = []

    def _on_step(self) -> bool:
        infos = self.locals.get("infos", [])
        for info in infos:
            if "episode" in info:
                self.episode_rewards.append(info["episode"]["r"])
        return True

What does this code do?

RewardLoggerCallback hooks into the training loop. Every time an episode finishes, the Monitor wrapper adds an "episode" key to the info dict with the total reward ("r"). The callback collects these into a list for plotting.

  • Run the cell. No output is expected.

Now you are ready to wire everything together and launch training.

  • Create a new cell and add the training configuration and execution code:
# --- Cell 9: Train the PPO Agent ---
def make_env():
    env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
    env = Monitor(env)
    return env

vec_env = DummyVecEnv([make_env])

policy_kwargs = dict(
    features_extractor_class=SmallCNN,
    features_extractor_kwargs=dict(features_dim=64),
)

model = PPO(
    "CnnPolicy",
    vec_env,
    policy_kwargs=policy_kwargs,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    verbose=1,
    device="auto",
)

print("Training the navigation agent...")
callback = RewardLoggerCallback()
model.learn(total_timesteps=100_000, callback=callback)
print("Training complete!")

What does this code do?

  • make_env() is a factory function that creates a fresh FloorPlanEnv wrapped with Monitor (which tracks episode statistics like total reward).
  • DummyVecEnv wraps the environment in the vectorized interface that Stable-Baselines3 requires.
  • policy_kwargs tells PPO to use your custom SmallCNN instead of the default Nature CNN.
  • The PPO model collects 2048 steps per rollout, trains for 10 epochs on each batch of 64, and uses a discount factor of 0.99 (the agent values future rewards highly).
  • model.learn(total_timesteps=100_000) runs the full training loop. The agent interacts with the environment, collects experience, and updates its policy 100,000 times.
  • Run the cell. Training will take approximately 2-5 minutes on a T4 GPU.
  • You should see periodic output lines showing timesteps, FPS, and other training statistics. When it finishes, you will see Training complete! printed.

Training is very slow or erroring out?

  • Check that your runtime is set to T4 GPU. Go to **Runtime** then **Change runtime type** and confirm **T4 GPU** is selected.
  • If you see a CUDA out-of-memory error, restart the runtime (**Runtime** then **Restart runtime**) and re-run all cells from the top.
  • If no GPU is available, training will still work on CPU. It will just be slower (5-10 minutes instead of 2-5).

Help me debug my PPO training setup in Google Colab

Plot the training reward curve

A reward curve shows whether your agent is actually improving over time. If the curve trends upward, the agent is learning to reach the goal more often and with fewer wasted steps.

  • Create a new cell and add the plotting code:
# --- Cell 10: Plot Training Curve ---
if callback.episode_rewards:
    rewards = np.array(callback.episode_rewards)
    window = min(50, len(rewards) // 4) if len(rewards) > 4 else 1
    smoothed = np.convolve(rewards, np.ones(window) / window, mode="valid")
    plt.figure(figsize=(10, 4))
    plt.plot(rewards, alpha=0.3, color="blue", label="Raw episode reward")
    plt.plot(
        np.arange(window - 1, len(rewards)),
        smoothed,
        color="red",
        linewidth=2,
        label=f"Smoothed (window={window})"
    )
    plt.xlabel("Episode")
    plt.ylabel("Total Reward")
    plt.title("Training Progress: Episode Reward over Time")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()
else:
    print("No episode rewards logged yet - training may need more timesteps.")

What does this code do?

  • The blue line shows the raw reward for every single episode. It is noisy because individual episodes vary in how the agent explores.
  • The red line is a rolling average that smooths out the noise, making the overall learning trend visible.
  • An upward trend in the smoothed line means the agent is earning higher rewards over time. This means it is reaching the goal more frequently or taking shorter paths.
  • Run the cell. You should see a chart with a noisy blue line and a smoother red trend line.
  • Look for the red smoothed line trending upward from left to right. Early episodes will have low (negative) rewards, and later episodes should show higher rewards as the agent improves.

What if the curve is flat or going down?

With 100,000 timesteps, you should see clear improvement. The agent may not be perfect yet, but the trend should be upward. If the curve is completely flat, the agent might be stuck hitting walls repeatedly. This is normal to see in the early portion of training before the agent discovers a path to the goal.

Seeing the else branch message instead of a chart?

If you see "No episode rewards logged yet" it means the callback did not record any completed episodes. Make sure you ran Cell 9 (training) fully before this cell. If training was interrupted, re-run Cell 9 to completion first.

Help me figure out why my training callback has no episode rewards

✔️ Awesome, I've got everything!

Great. Double check you have all four cells (7 through 10) saved and that training completed successfully.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point for comparison:

# --- Cell 1: Installation ---
# !pip install stable-baselines3==2.8.0 gymnasium==1.2.3 opencv-python-headless numpy matplotlib

# --- Cell 2: Imports and GPU Check ---
import numpy as np
import cv2
import matplotlib.pyplot as plt
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
import torch
import torch.nn as nn

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# --- Cell 3: Generate a Procedural Floor Plan ---
def create_floor_plan(size=64):
    grid = np.zeros((size, size), dtype=np.uint8)
    grid[2:28, 2:28] = 1
    grid[2:28, 36:62] = 1
    grid[36:62, 2:28] = 1
    grid[36:62, 36:62] = 1
    grid[12:18, 28:36] = 1
    grid[28:36, 10:20] = 1
    grid[46:52, 28:36] = 1
    grid[28:36, 44:54] = 1
    return grid

floor_plan = create_floor_plan(size=64)

plt.figure(figsize=(6, 6))
plt.imshow(floor_plan, cmap="gray", origin="upper")
plt.title("Generated Floor Plan\n(White = Free Space, Black = Walls)")
plt.axis("off")
plt.show()

print(f"Floor plan shape: {floor_plan.shape}")
print(f"Free cells: {floor_plan.sum()} / {floor_plan.size}")

# --- Cell 4: Process Floor Plan and Define Positions ---
def process_floor_plan(grid):
    img = (grid * 255).astype(np.uint8)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    cleaned = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
    return (cleaned > 127).astype(np.uint8)

processed_plan = process_floor_plan(floor_plan)

start_pos = np.array([5, 5], dtype=np.int32)
goal_pos = np.array([49, 49], dtype=np.int32)

plt.figure(figsize=(6, 6))
plt.imshow(processed_plan, cmap="gray", origin="upper")
plt.plot(start_pos[1], start_pos[0], "go", markersize=12, label="Start")
plt.plot(goal_pos[1], goal_pos[0], "r*", markersize=15, label="Goal")
plt.legend(loc="upper right", fontsize=10)
plt.title("Floor Plan with Start and Goal Positions")
plt.axis("off")
plt.show()

def get_egocentric_view(grid, position, view_size=7):
    half = view_size // 2
    row, col = position
    padded = np.pad(grid, pad_width=half, mode="constant", constant_values=0)
    view = padded[row:row + view_size, col:col + view_size].copy()
    return view

test_view = get_egocentric_view(processed_plan, start_pos, view_size=7)
print(f"Egocentric view shape: {test_view.shape}")

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].imshow(processed_plan, cmap="gray", origin="upper")
axes[0].plot(start_pos[1], start_pos[0], "go", markersize=10)
axes[0].set_title("Full Floor Plan (Agent at Start)")
axes[0].axis("off")
axes[1].imshow(test_view, cmap="gray", origin="upper", vmin=0, vmax=1)
axes[1].set_title("Agent's Egocentric View (7x7)")
axes[1].axis("off")
plt.tight_layout()
plt.show()

# --- Cell 5: Build the Custom Gymnasium Environment ---
class FloorPlanEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array"]}

    def __init__(self, floor_plan, start_pos, goal_pos, view_size=7, max_steps=500):
        super().__init__()
        self.floor_plan = floor_plan.copy()
        self.start_pos = np.array(start_pos, dtype=np.int32)
        self.goal_pos = np.array(goal_pos, dtype=np.int32)
        self.view_size = view_size
        self.max_steps = max_steps
        self.action_space = spaces.Discrete(4)
        self.observation_space = spaces.Box(
            low=0, high=255, shape=(1, view_size, view_size), dtype=np.uint8
        )
        self._action_to_direction = {
            0: np.array([-1, 0]),
            1: np.array([1, 0]),
            2: np.array([0, -1]),
            3: np.array([0, 1]),
        }
        self._agent_pos = None
        self._step_count = 0

    def _get_obs(self):
        raw_view = get_egocentric_view(self.floor_plan, self._agent_pos, self.view_size)
        obs = (raw_view * 255).astype(np.uint8)
        half = self.view_size // 2
        goal_relative = self.goal_pos - self._agent_pos + half
        gr, gc = goal_relative
        if 0 <= gr < self.view_size and 0 <= gc < self.view_size:
            obs[gr, gc] = 128
        return obs[np.newaxis, :, :]

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self._agent_pos = self.start_pos.copy()
        self._step_count = 0
        observation = self._get_obs()
        info = {"distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos)}
        return observation, info

    def step(self, action):
        self._step_count += 1
        direction = self._action_to_direction[action]
        new_pos = self._agent_pos + direction
        row, col = new_pos
        grid_h, grid_w = self.floor_plan.shape
        if (0 <= row < grid_h and 0 <= col < grid_w
                and self.floor_plan[row, col] == 1):
            self._agent_pos = new_pos
            reward = -0.01
        else:
            reward = -0.1
        terminated = np.array_equal(self._agent_pos, self.goal_pos)
        if terminated:
            reward = 10.0
        truncated = self._step_count >= self.max_steps
        observation = self._get_obs()
        info = {
            "distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos),
            "steps": self._step_count,
        }
        return observation, reward, terminated, truncated, info

# --- Cell 6: Test the Environment ---
env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
print(f"Observation space: {env.observation_space}")
print(f"Action space: {env.action_space}")
obs, info = env.reset()
print(f"Initial obs shape: {obs.shape}, dtype: {obs.dtype}")
print(f"Initial info: {info}")
total_reward = 0
for i in range(10):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    if terminated or truncated:
        break
print(f"After 10 random steps - Reward: {total_reward:.3f}, Distance: {info['distance_to_goal']:.1f}")

# --- Cell 7: Custom CNN Feature Extractor ---
class SmallCNN(BaseFeaturesExtractor):
    def __init__(self, observation_space: spaces.Box, features_dim: int = 64):
        super().__init__(observation_space, features_dim)
        n_input_channels = observation_space.shape[0]
        self.cnn = nn.Sequential(
            nn.Conv2d(n_input_channels, 16, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten(),
        )
        with torch.no_grad():
            sample = torch.as_tensor(observation_space.sample()[None]).float()
            n_flatten = self.cnn(sample).shape[1]
        self.linear = nn.Sequential(
            nn.Linear(n_flatten, features_dim),
            nn.ReLU()
        )

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        return self.linear(self.cnn(observations))

# --- Cell 8: Training Callback for Logging ---
class RewardLoggerCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
        self.episode_rewards = []

    def _on_step(self) -> bool:
        infos = self.locals.get("infos", [])
        for info in infos:
            if "episode" in info:
                self.episode_rewards.append(info["episode"]["r"])
        return True

# --- Cell 9: Train the PPO Agent ---
def make_env():
    env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
    env = Monitor(env)
    return env

vec_env = DummyVecEnv([make_env])

policy_kwargs = dict(
    features_extractor_class=SmallCNN,
    features_extractor_kwargs=dict(features_dim=64),
)

model = PPO(
    "CnnPolicy",
    vec_env,
    policy_kwargs=policy_kwargs,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    verbose=1,
    device="auto",
)

print("Training the navigation agent...")
callback = RewardLoggerCallback()
model.learn(total_timesteps=100_000, callback=callback)
print("Training complete!")

# --- Cell 10: Plot Training Curve ---
if callback.episode_rewards:
    rewards = np.array(callback.episode_rewards)
    window = min(50, len(rewards) // 4) if len(rewards) > 4 else 1
    smoothed = np.convolve(rewards, np.ones(window) / window, mode="valid")
    plt.figure(figsize=(10, 4))
    plt.plot(rewards, alpha=0.3, color="blue", label="Raw episode reward")
    plt.plot(
        np.arange(window - 1, len(rewards)),
        smoothed,
        color="red",
        linewidth=2,
        label=f"Smoothed (window={window})"
    )
    plt.xlabel("Episode")
    plt.ylabel("Total Reward")
    plt.title("Training Progress: Episode Reward over Time")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()
else:
    print("No episode rewards logged yet - training may need more timesteps.")

Your agent has now learned a navigation policy from 100,000 interactions with the floor plan. Next up, you will evaluate how well it performs and visualize the path it learned to take from start to goal.

Evaluate and Visualize the Trained Agent

Your agent has been training for 100,000 timesteps, learning from thousands of episodes of trial and error. The training curve gave you a high-level view of improvement, but it does not tell you how the agent navigates.

In this step, you will run the trained agent through the floor plan in evaluation mode and record its exact path. Then you will visualize what the agent learned: its navigation route on the full map, and the egocentric camera view it used to make decisions at each point along the way.

In this step, get ready to:
  • Run the trained agent deterministically and record its trajectory.
  • Visualize the learned navigation path overlaid on the floor plan.
  • Create a frame sequence showing the agent's camera view at key moments.
Run the trained agent and record its trajectory

During training, the agent explored randomly at first and gradually improved. For evaluation, you want the agent to use its best learned behavior without any randomness. This is called deterministic inference.

  • Create a new code cell in your notebook and add the evaluation loop by pasting this code:
# --- Cell 11: Evaluate and Visualize the Trained Agent ---
eval_env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
obs, info = eval_env.reset()
path = [eval_env._agent_pos.copy()]
views = [obs.copy()]
done = False
total_reward = 0

while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = eval_env.step(action)
    path.append(eval_env._agent_pos.copy())
    views.append(obs.copy())
    total_reward += reward
    done = terminated or truncated

path = np.array(path)
print(f"Evaluation: {len(path)} steps, Total reward: {total_reward:.2f}")
print(f"Reached goal: {terminated}")

What does this code do?

This creates a fresh environment instance for evaluation, separate from the training environment.

  • model.predict(obs, deterministic=True) tells the agent to always pick its best action (no exploration noise).
  • path records every position the agent visits, so you can plot the route later.
  • views stores the egocentric observation at each timestep for the frame sequence visualization.
  • The loop ends when the agent either reaches the goal (terminated) or runs out of steps (truncated).
  • Run the cell.
  • You should see output like Evaluation: <N> steps, Total reward: <value> followed by Reached goal: True (or False if the agent did not reach the goal within 500 steps).

Reached goal says False?

This can happen if the agent did not learn enough during 100,000 timesteps. The agent may still show improvement over random (fewer steps, higher reward) even if it did not reach the goal perfectly. If your training curve showed improvement, the visualization in the next substep will still reveal learned navigation behavior.

Help me understand why my agent didn't reach the goal

Visualize the navigation path on the floor plan

Now you will overlay the agent's recorded trajectory on the full floor plan. This gives you an immediate visual understanding of whether the agent learned a sensible route through the hallways.

  • Add the path visualization code below the existing code in the same cell:
plt.figure(figsize=(8, 8))
plt.imshow(processed_plan, cmap="gray", origin="upper")
plt.plot(path[:, 1], path[:, 0], "b-", linewidth=1.5, alpha=0.7, label="Agent path")
plt.plot(path[0, 1], path[0, 0], "go", markersize=12, label="Start")
plt.plot(goal_pos[1], goal_pos[0], "r*", markersize=15, label="Goal")
plt.plot(path[-1, 1], path[-1, 0], "bs", markersize=8, label="Final position")
plt.legend(loc="upper right", fontsize=10)
plt.title(f"Learned Navigation Path ({len(path)} steps, Reward: {total_reward:.1f})")
plt.axis("off")
plt.show()

What does this visualization show?

  • The blue line traces the agent's path from start to its final position.
  • The green circle marks the start position (top-left room).
  • The red star marks the goal (bottom-right room).
  • The blue square marks where the agent stopped. If it overlaps the red star, the agent reached the goal.
  • Run the cell again.
  • You should see a floor plan with a blue path line tracing the route the agent took. A well-trained agent's path will follow hallways between rooms rather than bumping into walls repeatedly.

Path looks like random noise?

If the path zig-zags wildly without clear direction, the agent may not have trained enough. Check your training curve from the previous step. If rewards were still climbing at 100,000 timesteps, try retraining with more timesteps (e.g. 200,000).

Help me improve my agent's navigation path

Create a frame sequence of the agent's egocentric view

The path plot shows where the agent went, but it does not show what the agent saw at each moment. This final visualization pairs the agent's map position with its 7x7 egocentric camera view at evenly spaced timesteps throughout the episode.

  • Create a new code cell and add the frame sequence visualization by pasting this code:
# --- Cell 12: Show Egocentric View Sequence ---
n_frames = min(8, len(views))
indices = np.linspace(0, len(views) - 1, n_frames, dtype=int)

fig, axes = plt.subplots(2, n_frames, figsize=(2.5 * n_frames, 5))

for i, idx in enumerate(indices):
    axes[0, i].imshow(processed_plan, cmap="gray", origin="upper")
    axes[0, i].plot(path[idx, 1], path[idx, 0], "ro", markersize=6)
    axes[0, i].plot(goal_pos[1], goal_pos[0], "r*", markersize=8)
    axes[0, i].set_title(f"Step {idx}", fontsize=9)
    axes[0, i].axis("off")
    view_img = views[idx][0]
    axes[1, i].imshow(view_img, cmap="gray", origin="upper", vmin=0, vmax=255)
    axes[1, i].set_title("Agent view", fontsize=9)
    axes[1, i].axis("off")

plt.suptitle("Agent's Navigation: Map Position (top) vs Egocentric View (bottom)", fontsize=11)
plt.tight_layout()
plt.show()

What does this visualization show?

This creates a 2-row grid of up to 8 frames sampled evenly from the agent's episode.

  • The top row shows the agent's position (red dot) on the full floor plan at each sampled timestep.
  • The bottom row shows the corresponding 7x7 egocentric view. This is the only information the CNN had to make its decision.
  • views[idx][0] removes the channel dimension so Matplotlib can display it as a 2D grayscale image.
  • Run the cell.
  • You should see a wide figure with two rows: the top row showing red dots progressing across the floor plan, and the bottom row showing small grayscale squares representing the agent's limited camera view at each point.

Seeing an index error or blank plot?

Make sure you ran Cell 11 first (the evaluation loop that populates path and views). If those variables are empty or undefined, this cell cannot plot anything.

Help me fix my frame sequence visualization

Why is this visualization useful?

In real robotics, engineers debug navigation failures by reviewing what the robot's sensors saw at the moment things went wrong. This frame sequence gives you the same capability. You can see whether the agent learned to turn at hallway junctions or got stuck bouncing off walls.

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all cells.

ⓧ I'd like to double check the full code

Here is the complete notebook code from Cell 1 through Cell 12. Compare your cells against this:

# --- Cell 1: Installation ---
# !pip install stable-baselines3==2.8.0 gymnasium==1.2.3 opencv-python-headless numpy matplotlib

# --- Cell 2: Imports and GPU Check ---
import numpy as np
import cv2
import matplotlib.pyplot as plt
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
import torch
import torch.nn as nn

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# --- Cell 3: Generate a Procedural Floor Plan ---
def create_floor_plan(size=64):
    grid = np.zeros((size, size), dtype=np.uint8)
    grid[2:28, 2:28] = 1
    grid[2:28, 36:62] = 1
    grid[36:62, 2:28] = 1
    grid[36:62, 36:62] = 1
    grid[12:18, 28:36] = 1
    grid[28:36, 10:20] = 1
    grid[46:52, 28:36] = 1
    grid[28:36, 44:54] = 1
    return grid

floor_plan = create_floor_plan(size=64)

plt.figure(figsize=(6, 6))
plt.imshow(floor_plan, cmap="gray", origin="upper")
plt.title("Generated Floor Plan\n(White = Free Space, Black = Walls)")
plt.axis("off")
plt.show()

print(f"Floor plan shape: {floor_plan.shape}")
print(f"Free cells: {floor_plan.sum()} / {floor_plan.size}")

# --- Cell 4: Process Floor Plan and Define Positions ---
def process_floor_plan(grid):
    img = (grid * 255).astype(np.uint8)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    cleaned = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
    return (cleaned > 127).astype(np.uint8)

processed_plan = process_floor_plan(floor_plan)

start_pos = np.array([5, 5], dtype=np.int32)
goal_pos = np.array([49, 49], dtype=np.int32)

plt.figure(figsize=(6, 6))
plt.imshow(processed_plan, cmap="gray", origin="upper")
plt.plot(start_pos[1], start_pos[0], "go", markersize=12, label="Start")
plt.plot(goal_pos[1], goal_pos[0], "r*", markersize=15, label="Goal")
plt.legend(loc="upper right", fontsize=10)
plt.title("Floor Plan with Start and Goal Positions")
plt.axis("off")
plt.show()

def get_egocentric_view(grid, position, view_size=7):
    half = view_size // 2
    row, col = position
    padded = np.pad(grid, pad_width=half, mode="constant", constant_values=0)
    view = padded[row:row + view_size, col:col + view_size].copy()
    return view

test_view = get_egocentric_view(processed_plan, start_pos, view_size=7)
print(f"Egocentric view shape: {test_view.shape}")

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].imshow(processed_plan, cmap="gray", origin="upper")
axes[0].plot(start_pos[1], start_pos[0], "go", markersize=10)
axes[0].set_title("Full Floor Plan (Agent at Start)")
axes[0].axis("off")
axes[1].imshow(test_view, cmap="gray", origin="upper", vmin=0, vmax=1)
axes[1].set_title("Agent's Egocentric View (7x7)")
axes[1].axis("off")
plt.tight_layout()
plt.show()

# --- Cell 5: Build the Custom Gymnasium Environment ---
class FloorPlanEnv(gym.Env):
    metadata = {"render_modes": ["rgb_array"]}

    def __init__(self, floor_plan, start_pos, goal_pos, view_size=7, max_steps=500):
        super().__init__()
        self.floor_plan = floor_plan.copy()
        self.start_pos = np.array(start_pos, dtype=np.int32)
        self.goal_pos = np.array(goal_pos, dtype=np.int32)
        self.view_size = view_size
        self.max_steps = max_steps
        self.action_space = spaces.Discrete(4)
        self.observation_space = spaces.Box(
            low=0, high=255, shape=(1, view_size, view_size), dtype=np.uint8
        )
        self._action_to_direction = {
            0: np.array([-1, 0]),
            1: np.array([1, 0]),
            2: np.array([0, -1]),
            3: np.array([0, 1]),
        }
        self._agent_pos = None
        self._step_count = 0

    def _get_obs(self):
        raw_view = get_egocentric_view(self.floor_plan, self._agent_pos, self.view_size)
        obs = (raw_view * 255).astype(np.uint8)
        half = self.view_size // 2
        goal_relative = self.goal_pos - self._agent_pos + half
        gr, gc = goal_relative
        if 0 <= gr < self.view_size and 0 <= gc < self.view_size:
            obs[gr, gc] = 128
        return obs[np.newaxis, :, :]

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self._agent_pos = self.start_pos.copy()
        self._step_count = 0
        observation = self._get_obs()
        info = {"distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos)}
        return observation, info

    def step(self, action):
        self._step_count += 1
        direction = self._action_to_direction[action]
        new_pos = self._agent_pos + direction
        row, col = new_pos
        grid_h, grid_w = self.floor_plan.shape
        if (0 <= row < grid_h and 0 <= col < grid_w
                and self.floor_plan[row, col] == 1):
            self._agent_pos = new_pos
            reward = -0.01
        else:
            reward = -0.1
        terminated = np.array_equal(self._agent_pos, self.goal_pos)
        if terminated:
            reward = 10.0
        truncated = self._step_count >= self.max_steps
        observation = self._get_obs()
        info = {
            "distance_to_goal": np.linalg.norm(self._agent_pos - self.goal_pos),
            "steps": self._step_count,
        }
        return observation, reward, terminated, truncated, info

# --- Cell 6: Test the Environment ---
env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
print(f"Observation space: {env.observation_space}")
print(f"Action space: {env.action_space}")
obs, info = env.reset()
print(f"Initial obs shape: {obs.shape}, dtype: {obs.dtype}")
print(f"Initial info: {info}")
total_reward = 0
for i in range(10):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    total_reward += reward
    if terminated or truncated:
        break
print(f"After 10 random steps - Reward: {total_reward:.3f}, Distance: {info['distance_to_goal']:.1f}")

# --- Cell 7: Custom CNN Feature Extractor ---
class SmallCNN(BaseFeaturesExtractor):
    def __init__(self, observation_space: spaces.Box, features_dim: int = 64):
        super().__init__(observation_space, features_dim)
        n_input_channels = observation_space.shape[0]
        self.cnn = nn.Sequential(
            nn.Conv2d(n_input_channels, 16, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten(),
        )
        with torch.no_grad():
            sample = torch.as_tensor(observation_space.sample()[None]).float()
            n_flatten = self.cnn(sample).shape[1]
        self.linear = nn.Sequential(
            nn.Linear(n_flatten, features_dim),
            nn.ReLU()
        )

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        return self.linear(self.cnn(observations))

# --- Cell 8: Training Callback for Logging ---
class RewardLoggerCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
        self.episode_rewards = []

    def _on_step(self) -> bool:
        infos = self.locals.get("infos", [])
        for info in infos:
            if "episode" in info:
                self.episode_rewards.append(info["episode"]["r"])
        return True

# --- Cell 9: Train the PPO Agent ---
def make_env():
    env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
    env = Monitor(env)
    return env

vec_env = DummyVecEnv([make_env])

policy_kwargs = dict(
    features_extractor_class=SmallCNN,
    features_extractor_kwargs=dict(features_dim=64),
)

model = PPO(
    "CnnPolicy",
    vec_env,
    policy_kwargs=policy_kwargs,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=64,
    n_epochs=10,
    gamma=0.99,
    verbose=1,
    device="auto",
)

print("Training the navigation agent...")
callback = RewardLoggerCallback()
model.learn(total_timesteps=100_000, callback=callback)
print("Training complete!")

# --- Cell 10: Plot Training Curve ---
if callback.episode_rewards:
    rewards = np.array(callback.episode_rewards)
    window = min(50, len(rewards) // 4) if len(rewards) > 4 else 1
    smoothed = np.convolve(rewards, np.ones(window) / window, mode="valid")
    plt.figure(figsize=(10, 4))
    plt.plot(rewards, alpha=0.3, color="blue", label="Raw episode reward")
    plt.plot(
        np.arange(window - 1, len(rewards)),
        smoothed,
        color="red",
        linewidth=2,
        label=f"Smoothed (window={window})"
    )
    plt.xlabel("Episode")
    plt.ylabel("Total Reward")
    plt.title("Training Progress: Episode Reward over Time")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()
else:
    print("No episode rewards logged yet - training may need more timesteps.")

# --- Cell 11: Evaluate and Visualize the Trained Agent ---
eval_env = FloorPlanEnv(processed_plan, start_pos, goal_pos, view_size=7, max_steps=500)
obs, info = eval_env.reset()
path = [eval_env._agent_pos.copy()]
views = [obs.copy()]
done = False
total_reward = 0

while not done:
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = eval_env.step(action)
    path.append(eval_env._agent_pos.copy())
    views.append(obs.copy())
    total_reward += reward
    done = terminated or truncated

path = np.array(path)
print(f"Evaluation: {len(path)} steps, Total reward: {total_reward:.2f}")
print(f"Reached goal: {terminated}")

plt.figure(figsize=(8, 8))
plt.imshow(processed_plan, cmap="gray", origin="upper")
plt.plot(path[:, 1], path[:, 0], "b-", linewidth=1.5, alpha=0.7, label="Agent path")
plt.plot(path[0, 1], path[0, 0], "go", markersize=12, label="Start")
plt.plot(goal_pos[1], goal_pos[0], "r*", markersize=15, label="Goal")
plt.plot(path[-1, 1], path[-1, 0], "bs", markersize=8, label="Final position")
plt.legend(loc="upper right", fontsize=10)
plt.title(f"Learned Navigation Path ({len(path)} steps, Reward: {total_reward:.1f})")
plt.axis("off")
plt.show()

# --- Cell 12: Show Egocentric View Sequence ---
n_frames = min(8, len(views))
indices = np.linspace(0, len(views) - 1, n_frames, dtype=int)

fig, axes = plt.subplots(2, n_frames, figsize=(2.5 * n_frames, 5))

for i, idx in enumerate(indices):
    axes[0, i].imshow(processed_plan, cmap="gray", origin="upper")
    axes[0, i].plot(path[idx, 1], path[idx, 0], "ro", markersize=6)
    axes[0, i].plot(goal_pos[1], goal_pos[0], "r*", markersize=8)
    axes[0, i].set_title(f"Step {idx}", fontsize=9)
    axes[0, i].axis("off")
    view_img = views[idx][0]
    axes[1, i].imshow(view_img, cmap="gray", origin="upper", vmin=0, vmax=255)
    axes[1, i].set_title("Agent view", fontsize=9)
    axes[1, i].axis("off")

plt.suptitle("Agent's Navigation: Map Position (top) vs Egocentric View (bottom)", fontsize=11)
plt.tight_layout()
plt.show()

Secret mission

Goal-Conditioned Navigation

Your agent currently navigates to a single fixed goal. Real robots need to reach different destinations on command. Can you teach your agent to navigate to any room in the house based on a goal signal in its observation?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on Google Colab's free tier, so there are no ongoing costs. Decide whether to keep your notebook available, disconnect the runtime, or delete everything.

Resources you used:

  • Google Colab notebook (stored in Google Drive).
  • Colab GPU runtime (active T4 session).

Keep everything running

No action needed. Choose this if you want to keep experimenting with hyperparameters, training for more timesteps, or tweaking the floor plan layout.

  • Your notebook stays in Google Drive and you can reopen it anytime.
  • The GPU runtime will automatically disconnect after a period of inactivity (typically 90 minutes idle or 12 hours total).
  • Pip packages are installed in the session only. They reset when the runtime disconnects, so you will need to re-run Cell 1 next time.

Pause - I'll come back to this later

Disconnect the GPU runtime to free up resources for other Colab users. Your notebook and code remain saved in Google Drive.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • Your trained model and variables will be cleared from memory, but your code cells remain intact.
  • When you return, click Runtime > Run all to re-install packages and retrain from scratch.

Delete - I don't want to use this again

Remove the notebook from your Google Drive entirely.

  • First, disconnect the runtime: click Runtime > Disconnect and delete runtime.
  • Open Google Drive at drive.google.com.
  • Navigate to the Colab Notebooks folder (or wherever you saved your notebook).
  • Right-click the notebook file and select Move to trash.
  • Empty your trash to permanently delete it: click Trash in the left sidebar, then click Empty trash.

Nice Work!

Nice Work!

You just built an AI robot that learned to navigate a house purely by looking at its surroundings. From raw image processing to a trained neural network that solves a maze, you now have hands-on experience bridging computer vision and reinforcement learning.

You've learned how to:

  • Process a floor plan into a navigable occupancy grid using OpenCV and build a custom Gymnasium environment with egocentric visual observations.
  • Implement a custom CNN feature extractor and train a PPO agent that learns spatial navigation purely from pixel observations using Stable-Baselines3.
  • Visualize the agent's learned navigation behavior with path plots and egocentric view sequences using Matplotlib.
  • Secret Mission: Extended to goal-conditioned navigation where the agent navigates to different rooms on command using a multi-channel observation.

Ready to quiz yourself?