Q Learning From Scratch Teaching An Agent To Navig

You’ve seen the toddler reach for the cookie jar. Now imagine dropping that toddler into a frozen lake — a 4x4 grid of ice tiles, some solid, some ready to crack and drop them into freezing water. The toddler has no map, no instructions, and no one telling them which way to go. All they can do is try moving up, down, left, or right, and feel the consequences. If they reach the far corner, they get a cookie (+1). If they fall through a hole, they get dunked (game over, 0 reward). If they just shuffle around on safe ice, nothing happens (0 reward, keep going).

That’s exactly the problem we’re going to solve in this tutorial. We’ll teach a computer agent to navigate the FrozenLake environment from the Gymnasium library — a classic test bed for reinforcement learning. By the end, you’ll have built a complete Q-learning agent from scratch, watched it fail, watched it learn, and understood every line of code along the way.

Let’s start by seeing just how bad a random agent is at this task.

1. The Problem: Your Agent is Lost in a Frozen Lake

Here’s the setup. You have a 4x4 grid:

• S: Start position (top-left corner)
• F: Frozen surface (safe to walk on)
• H: Hole (fall in and you lose)
• G: Goal (the treasure, bottom-right corner)

The agent starts at S and can move in four directions: up, down, left, right. Each step moves it one tile. If it reaches G, it gets +1 reward and the episode ends. If it falls into an H, it gets 0 reward and the episode ends. Otherwise, it gets 0 reward and keeps going.

Let’s see what happens when the agent just picks random actions. We’ll use the gymnasium library (the modern successor to OpenAI Gym) to create the environment.

import gymnasium as gym
import numpy as np

# Create the FrozenLake environment (non-slippery for simplicity)
env = gym.make("FrozenLake-v1", is_slippery=False, render_mode="ansi")

# Run a random agent for 10 episodes
num_episodes = 10
successes = 0

for episode in range(num_episodes):
    state, _ = env.reset()
    done = False
    total_reward = 0
    steps = 0

    while not done:
        # Random action: 0=left, 1=down, 2=right, 3=up
        action = env.action_space.sample()
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        total_reward += reward
        steps += 1
        state = next_state

    if total_reward > 0:
        successes += 1
    print(f"Episode {episode + 1}: Reward = {total_reward}, Steps = {steps}, "
          f"Outcome = {'WIN' if total_reward > 0 else 'LOSE'}")

print(f"\nSuccess rate: {successes}/{num_episodes} ({successes/num_episodes*100:.0f}%)")

When you run this, you’ll see something like:

Episode 1: Reward = 0, Steps = 24, Outcome = LOSE
Episode 2: Reward = 0, Steps = 26, Outcome = LOSE
...
Success rate: 0/10 (0%)

The random agent almost never reaches the goal. It wanders around, hits walls, falls into holes, and occasionally gets stuck in loops. The problem is clear: the agent has no memory of what worked before. It’s like the toddler who never learns from the hot mug — every attempt is a fresh start.

So how can the agent learn from its failures and successes?

2. The Intuition: Learning a Cheat Sheet for Every Situation

Imagine you’re at a new restaurant with a huge menu. You’ve never been there before, so you have no idea what’s good. But you decide to keep a little notebook. Every time you order something, you write down: “In this mood (state), I ordered this dish (action), and it was delicious (reward).” Over time, your notebook becomes a cheat sheet: “When I’m hungry for comfort food, get the mac and cheese. When I want something light, get the salad.”

That’s exactly what a Q-table is. It’s a lookup table where:

• Rows = every possible state the agent can be in (16 tiles on the FrozenLake grid)
• Columns = every possible action the agent can take (4 directions: left, down, right, up)
• Each cell = a number called a Q-value, which estimates “how good” that action is in that state

A high Q-value means “taking this action in this state tends to lead to good things.” A low Q-value means “this action usually ends badly.”

The agent’s policy — its strategy for making decisions — becomes simple: in any state, look at the Q-values for all four actions, and pick the one with the highest number. That’s called a greedy policy because it always grabs the best option.

Let’s create an empty Q-table and see what it looks like.

import numpy as np

# FrozenLake has 16 states (4x4 grid) and 4 actions (left, down, right, up)
n_states = 16
n_actions = 4

# Initialize Q-table with zeros
q_table = np.zeros((n_states, n_actions))

print("Empty Q-table (all zeros):")
print(q_table)
print(f"\nShape: {q_table.shape}")
print(f"Rows (states): {n_states}")
print(f"Columns (actions): {n_actions}")

Right now, every cell is 0. The agent has no idea what’s good. But as it explores, it will fill in these numbers.

Here’s the greedy policy function:

def greedy_policy(state, q_table):
    """Return the action with the highest Q-value for this state.
    If there's a tie, pick randomly among the tied actions."""
    q_values = q_table[state]
    max_value = np.max(q_values)
    # Find all actions that have the max value (handles ties)
    best_actions = np.where(q_values == max_value)[0]
    return np.random.choice(best_actions)

# Test it on state 0 (the start position)
state = 0
action = greedy_policy(state, q_table)
print(f"In state {state}, greedy policy picks action {action}")
print(f"Action meanings: 0=left, 1=down, 2=right, 3=up")

Since all Q-values are 0, the policy just picks randomly among the four actions. That’s no better than the random agent. But soon, the Q-table will start to reflect real experience.

3. The Learning Rule: How to Update the Cheat Sheet (Bellman Equation in Plain English)

Now we need a way to update the Q-table based on experience. Every time the agent takes a step, it learns something new. Here’s what it observes:

1. Old state (S): Where it was before moving
2. Action (A): What it did
3. Reward (R): What it got for that action
4. New state (S’): Where it ended up

The question is: how do we combine these four pieces of information to update the Q-value for (S, A)?

Here’s the update rule, also known as the Bellman equation (don’t let the name scare you — it’s just a formula for updating estimates based on new evidence):

New Q(S, A) = Old Q(S, A) + learning_rate * (R + discount * max Q(S', :) - Old Q(S, A))

Let’s translate each term into plain English:

• learning_rate (often called alpha, α): How much to trust the new information vs. the old estimate. A value of 0 means “never learn anything new.” A value of 1 means “completely replace the old estimate with the new one.” We typically use something in between, like 0.1 or 0.5.

• discount (often called gamma, γ): How much to care about future rewards vs. immediate rewards. A value of 0 means “only care about the reward right now.” A value of 1 means “care equally about all future rewards.”

• max Q(S’, :): The best possible future reward from the new state. This is the agent’s estimate of “how good is it to be in this new state?”

• R + discount * max Q(S’, :): This is the target — the actual reward plus the best future reward we can expect. It’s the “true” value we’re trying to estimate.

• Old Q(S, A): Our current estimate.

• The difference: The gap between the target and our current estimate. We nudge our estimate toward the target.

Think of it like this: you’re trying to guess how long it takes to drive to work. Your current guess is 30 minutes. Today it took 35 minutes (the new information). If your learning rate is 0.5, your new guess becomes: 30 + 0.5 * (35 - 30) = 32.5 minutes. You’re adjusting toward the new data, but not completely replacing your old estimate.

Let’s implement this update rule as a function.

import numpy as np

def update_q_table(q_table, state, action, reward, next_state, learning_rate, discount):
    """
    Update the Q-table using the Bellman equation.
    
    Parameters:
    - q_table: the current Q-table (numpy array)
    - state: the current state (integer)
    - action: the action taken (integer)
    - reward: the reward received (float)
    - next_state: the state we transitioned to (integer)
    - learning_rate: how much to trust new info (float, 0-1)
    - discount: how much to care about future rewards (float, 0-1)
    
    Returns:
    - q_table: the updated Q-table
    - old_q: the Q-value before the update
    - new_q: the Q-value after the update
    """
    old_q = q_table[state, action]
    
    # Best possible future reward from the next state
    best_future_q = np.max(q_table[next_state])
    
    # The target: actual reward + discounted best future reward
    target = reward + discount * best_future_q
    
    # The update: nudge old Q toward the target
    new_q = old_q + learning_rate * (target - old_q)
    
    q_table[state, action] = new_q
    
    return q_table, old_q, new_q

# Let's test it with a concrete example
# Create a small Q-table (just for testing)
q_table = np.zeros((16, 4))

# Suppose we're in state 0, we took action 2 (right), got reward 0, and ended up in state 1
state = 0
action = 2
reward = 0
next_state = 1
learning_rate = 0.1
discount = 0.9

print("Before update:")
print(f"Q(0, 2) = {q_table[0, 2]:.4f}")
print(f"Q-values for state 1: {q_table[1]}")

q_table, old_q, new_q = update_q_table(
    q_table, state, action, reward, next_state, learning_rate, discount
)

print("\nAfter update:")
print(f"Old Q(0, 2) = {old_q:.4f}")
print(f"Target (R + discount * max Q(S', :)) = {reward} + {discount} * {np.max(q_table[next_state]):.4f} = {reward + discount * np.max(q_table[next_state]):.4f}")
print(f"New Q(0, 2) = {new_q:.4f}")
print(f"Change: {new_q - old_q:.4f}")

In this example, since the next state’s Q-values are all zero, the target is just 0 + 0.9 * 0 = 0. So the new Q-value is still 0 — no learning happened because there’s no information yet. But once the agent reaches the goal and gets a +1 reward, that value will propagate backward through the Q-table, teaching the agent which actions lead to success.

4. The Exploration-Exploitation Dilemma: Why the Agent Must Try New Things

Here’s a problem. If the agent always follows the greedy policy (always picks the action with the highest Q-value), it will never discover new paths. Imagine the agent accidentally stumbles into a hole early on. The Q-value for that action becomes negative (or at least, lower than alternatives). From then on, the greedy policy will avoid that action. But what if there was a better path that required going near that hole first?

The agent faces a fundamental trade-off:

• Exploitation: Use what you already know to get good rewards. Safe, but you might miss better options.
• Exploration: Try new things to discover better strategies. Risky in the short term, but potentially better in the long run.

This is called the exploration-exploitation dilemma, and it’s one of the hardest problems in reinforcement learning.

The simplest solution is epsilon-greedy (ε-greedy). Here’s how it works:

• With probability epsilon (ε), pick a random action (explore)
• With probability 1 - ε, pick the best action according to the Q-table (exploit)

We start with a high epsilon (like 1.0 — pure exploration) and gradually decay it over time, so the agent explores a lot early on and exploits more as it learns.

import numpy as np
import matplotlib.pyplot as plt

def epsilon_greedy_policy(state, q_table, epsilon):
    """
    Choose an action using epsilon-greedy.
    
    Parameters:
    - state: current state (integer)
    - q_table: the Q-table
    - epsilon: exploration rate (0-1)
    
    Returns:
    - action: chosen action (integer)
    """
    if np.random.random() < epsilon:
        # Explore: pick a random action
        return np.random.randint(0, q_table.shape[1])
    else:
        # Exploit: pick the best action
        return greedy_policy(state, q_table)

def greedy_policy(state, q_table):
    """Return the action with the highest Q-value."""
    q_values = q_table[state]
    max_value = np.max(q_values)
    best_actions = np.where(q_values == max_value)[0]
    return np.random.choice(best_actions)

# Let's visualize epsilon decay over episodes
n_episodes = 1000
epsilon_start = 1.0
epsilon_end = 0.01
decay_rate = 0.995  # exponential decay

epsilons = []
epsilon = epsilon_start
for episode in range(n_episodes):
    epsilons.append(epsilon)
    epsilon = max(epsilon_end, epsilon * decay_rate)

# Plot the decay
plt.figure(figsize=(10, 4))
plt.plot(epsilons)
plt.xlabel("Episode")
plt.ylabel("Epsilon (exploration rate)")
plt.title("Epsilon Decay Over Episodes")
plt.grid(True, alpha=0.3)
plt.show()

You’ll see a curve that starts at 1.0 and gradually drops toward 0.01. Early on, the agent is almost always exploring. By episode 500, it’s mostly exploiting what it’s learned.

5. The Training Loop: Putting It All Together

Now we have all the pieces:

• A Q-table to store knowledge
• An epsilon-greedy policy to balance exploration and exploitation
• A Bellman update rule to learn from each step

Let’s combine them into a full training loop. This is the moment where the agent goes from random stumbling to purposeful navigation.

import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt

# Create environment
env = gym.make("FrozenLake-v1", is_slippery=False, render_mode="ansi")

# Hyperparameters
n_episodes = 2000
learning_rate = 0.1
discount = 0.99
epsilon_start = 1.0
epsilon_end = 0.01
decay_rate = 0.995

# Initialize Q-table
n_states = env.observation_space.n
n_actions = env.action_space.n
q_table = np.zeros((n_states, n_actions))

# Tracking metrics
success_rates = []
success_window = 100  # sliding window for success rate
recent_successes = []

def greedy_policy(state, q_table):
    q_values = q_table[state]
    max_value = np.max(q_values)
    best_actions = np.where(q_values == max_value)[0]
    return np.random.choice(best_actions)

def epsilon_greedy_policy(state, q_table, epsilon):
    if np.random.random() < epsilon:
        return np.random.randint(0, n_actions)
    else:
        return greedy_policy(state, q_table)

def update_q_table(q_table, state, action, reward, next_state, learning_rate, discount):
    old_q = q_table[state, action]
    best_future_q = np.max(q_table[next_state])
    target = reward + discount * best_future_q
    new_q = old_q + learning_rate * (target - old_q)
    q_table[state, action] = new_q
    return q_table

# Training loop
epsilon = epsilon_start
for episode in range(n_episodes):
    state, _ = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        action = epsilon_greedy_policy(state, q_table, epsilon)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        
        # Update Q-table
        q_table = update_q_table(q_table, state, action, reward, next_state, learning_rate, discount)
        
        total_reward += reward
        state = next_state
    
    # Track success
    recent_successes.append(1 if total_reward > 0 else 0)
    if len(recent_successes) > success_window:
        recent_successes.pop(0)
    
    # Decay epsilon
    epsilon = max(epsilon_end, epsilon * decay_rate)
    
    # Print progress every 200 episodes
    if (episode + 1) % 200 == 0:
        success_rate = np.mean(recent_successes) * 100
        print(f"Episode {episode + 1}: Success rate (last {success_window} eps) = {success_rate:.1f}%, Epsilon = {epsilon:.3f}")

print("\nTraining complete!")
print(f"Final Q-table shape: {q_table.shape}")

# Plot learning curve
plt.figure(figsize=(10, 5))
# Smooth the success rate with a rolling average
window = 50
smoothed_successes = np.convolve(recent_successes, np.ones(window)/window, mode='valid')
plt.plot(smoothed_successes * 100)
plt.xlabel("Episode (last 200)")
plt.ylabel("Success rate (%)")
plt.title("Learning Curve: Success Rate Over Time")
plt.grid(True, alpha=0.3)
plt.show()

When you run this, you’ll see something amazing happen. Early on, the success rate is near zero. But as the agent explores and updates its Q-table, the success rate climbs. By episode 1000 or so, it should be consistently reaching the goal 80-100% of the time.

What’s actually happening here? The agent is learning a path from start to goal. It discovers that moving right from state 0 leads to state 1, which eventually leads to the goal. It also learns which tiles are holes and avoids them. The Q-values for good actions become positive, while actions that lead to holes stay at zero (or become negative in more sophisticated setups).

6. What the Agent Learned: Visualizing the Q-Table

After training, the Q-table contains real knowledge. Let’s inspect it and visualize the policy the agent learned.

import numpy as np

# We'll reuse the Q-table from the training loop above
# For this standalone block, let's recreate a trained Q-table
# (In practice, you'd save and load it)

# Let's assume we have a trained Q-table. We'll create one manually
# to demonstrate the visualization (since we can't share state between blocks)

# Create a mock trained Q-table for demonstration
# In reality, you'd use the q_table from the training loop
q_table = np.zeros((16, 4))

# Manually set some realistic Q-values for a trained agent
# State 0 (start): best action is right (action 2)
q_table[0] = [0.0, 0.0, 0.8, 0.0]
# State 1: best action is right (action 2) or down (action 1)
q_table[1] = [0.0, 0.7, 0.8, 0.0]
# State 2: best action is down (action 1)
q_table[2] = [0.0, 0.9, 0.0, 0.0]
# State 3: hole (action doesn't matter, all low)
q_table[3] = [0.1, 0.1, 0.1, 0.1]
# State 4: best action is right (action 2)
q_table[4] = [0.0, 0.0, 0.7, 0.0]
# State 5: hole
q_table[5] = [0.1, 0.1, 0.1, 0.1]
# State 6: best action is down (action 1)
q_table[6] = [0.0, 0.8, 0.0, 0.0]
# State 7: hole
q_table[7] = [0.1, 0.1, 0.1, 0.1]
# State 8: best action is right (action 2)
q_table[8] = [0.0, 0.0, 0.6, 0.0]
# State 9: best action is right (action 2) or down (action 1)
q_table[9] = [0.0, 0.7, 0.6, 0.0]
# State 10: best action is down (action 1)
q_table[10] = [0.0, 0.9, 0.0, 0.0]
# State 11: hole
q_table[11] = [0.1, 0.1, 0.1, 0.1]
# State 12: hole
q_table[12] = [0.1, 0.1, 0.1, 0.1]
# State 13: best action is right (action 2)
q_table[13] = [0.0, 0.0, 0.7, 0.0]
# State 14: best action is right (action 2)
q_table[14] = [0.0, 0.0, 0.8, 0.0]
# State 15: goal (all high because once you're here, you win)
q_table[15] = [1.0, 1.0, 1.0, 1.0]

def visualize_policy(q_table):
    """Print the grid with arrows showing the best action in each state."""
    # Arrow symbols for actions: 0=left, 1=down, 2=right, 3=up
    arrows = ['←', '↓', '→', '↑']
    
    print("Learned Policy (arrows show best action):")
    print("=" * 25)
    
    for row in range(4):
        row_str = ""
        for col in range(4):
            state = row * 4 + col
            best_action = np.argmax(q_table[state])
            
            # Special markers for start, goal, and holes
            if state == 0:
                marker = "S"
            elif state == 15:
                marker = "G"
            elif state in [3, 5, 7, 11, 12]:  # hole positions
                marker = "H"
            else:
                marker = arrows[best_action]
            
            row_str += f"  {marker}  "
        print(row_str)
        print()  # blank line between rows
    
    print("Legend: S=Start, G=Goal, H=Hole, arrows=best direction")

# Visualize the policy
visualize_policy(q_table)

# Also show Q-values for a few key states
print("\nQ-values for key states:")
print("-" * 50)
for state in [0, 2, 6, 10, 14, 15]:
    print(f"State {state} (row {state // 4}, col {state % 4}):")
    print(f"  Left: {q_table[state][0]:.3f}, Down: {q_table[state][1]:.3f}, "
          f"Right: {q_table[state][2]:.3f}, Up: {q_table[state][3]:.3f}")
    print(f"  Best action: {['Left', 'Down', 'Right', 'Up'][np.argmax(q_table[state])]}")
    print()

You’ll see a grid where arrows point the way from start to goal. The agent has learned to navigate around the holes and find the shortest path. Notice how states near the goal have higher Q-values — the agent knows it’s close to the treasure.

7. The Hard Part: Why Q-Learning is ‘Off-Policy’ (and Why That Matters)

This is the hardest part of Q-learning conceptually, so let’s take it slow.

Remember our update rule:

New Q(S, A) = Old Q(S, A) + learning_rate * (R + discount * max Q(S', :) - Old Q(S, A))

The key part is max Q(S’, :): we’re using the best possible action from the next state, regardless of what action the agent actually took next.

This means Q-learning is off-policy. Here’s what that means:

• The agent follows an exploratory policy (epsilon-greedy) to choose actions
• But it updates its Q-values assuming a greedy policy (always taking the best action)

It’s like a student who sometimes gets distracted in class (exploring), but when taking notes, writes down what the best student would do (the greedy policy). The student can learn the optimal strategy even while behaving suboptimally.

Contrast this with SARSA (State-Action-Reward-State-Action), which is on-policy:

New Q(S, A) = Old Q(S, A) + learning_rate * (R + discount * Q(S', A') - Old Q(S, A))

Notice the difference: SARSA uses Q(S’, A’) — the Q-value for the action actually taken in the next state, not the best possible action. This means SARSA learns about the policy it’s actually following, including its exploration.

Why does this matter? Q-learning can learn the optimal policy even while exploring randomly. It’s more sample-efficient and can converge to the optimal policy faster. But it can also be more unstable because it’s learning about a policy it isn’t actually following.

Let’s see this in action with a quick comparison.

import gymnasium as gym
import numpy as np

def train_q_learning(env, n_episodes=1000, learning_rate=0.1, discount=0.99,
                     epsilon_start=1.0, epsilon_end=0.01, decay_rate=0.995):
    """Train using Q-learning (off-policy)."""
    n_states = env.observation_space.n
    n_actions = env.action_space.n
    q_table = np.zeros((n_states, n_actions))
    
    epsilon = epsilon_start
    success_count = 0
    
    for episode in range(n_episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0
        
        while not done:
            # Epsilon-greedy action selection
            if np.random.random() < epsilon:
                action = np.random.randint(0, n_actions)
            else:
                action = np.argmax(q_table[state])
            
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            
            # Q-learning update (off-policy: uses max Q of next state)
            best_next_q = np.max(q_table[next_state])
            target = reward + discount * best_next_q
            q_table[state, action] += learning_rate * (target - q_table[state, action])
            
            total_reward += reward
            state = next_state
        
        if total_reward > 0:
            success_count += 1
        
        epsilon = max(epsilon_end, epsilon * decay_rate)
    
    return success_count / n_episodes

def train_sarsa(env, n_episodes=1000, learning_rate=0.1, discount=0.99,
                epsilon_start=1.0, epsilon_end=0.01, decay_rate=0.995):
    """Train using SARSA (on-policy)."""
    n_states = env.observation_space.n
    n_actions = env.action_space.n
    q_table = np.zeros((n_states, n_actions))
    
    epsilon = epsilon_start
    success_count = 0
    
    for episode in range(n_episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0
        
        # Choose first action
        if np.random.random() < epsilon:
            action = np.random.randint(0, n_actions)
        else:
            action = np.argmax(q_table[state])
        
        while not done:
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            
            # Choose next action (for SARSA update)
            if np.random.random() < epsilon:
                next_action = np.random.randint(0, n_actions)
            else:
                next_action = np.argmax(q_table[next_state])
            
            # SARSA update (on-policy: uses the action actually taken)
            target = reward + discount * q_table[next_state, next_action]
            q_table[state, action] += learning_rate * (target - q_table[state, action])
            
            total_reward += reward
            state = next_state
            action = next_action
        
        if total_reward > 0:
            success_count += 1
        
        epsilon = max(epsilon_end, epsilon * decay_rate)
    
    return success_count / n_episodes

# Compare the two algorithms
env = gym.make("FrozenLake-v1", is_slippery=False)

n_trials = 10
ql_successes = []
sarsa_successes = []

for trial in range(n_trials):
    ql_rate = train_q_learning(env, n_episodes=500)
    sarsa_rate = train_sarsa(env, n_episodes=500)
    ql_successes.append(ql_rate)
    sarsa_successes.append(sarsa_rate)
    print(f"Trial {trial + 1}: Q-learning = {ql_rate:.2%}, SARSA = {sarsa_rate:.2%}")

print(f"\nAverage success rate over {n_trials} trials:")
print(f"Q-learning: {np.mean(ql_successes):.2%} (off-policy)")
print(f"SARSA: {np.mean(sarsa_successes):.2%} (on-policy)")

You’ll likely see that Q-learning outperforms SARSA on this deterministic FrozenLake environment. That’s because Q-learning can learn the optimal path even while exploring randomly, while SARSA’s learning is constrained by the exploration policy.

8. Hyperparameter Tuning: What Happens When You Change the Knobs

Q-learning has three key knobs you can turn: learning rate, discount factor, and epsilon decay. Let’s see what happens when we change them.

import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt

def train_q_learning(env, n_episodes=500, learning_rate=0.1, discount=0.99,
                     epsilon_start=1.0, epsilon_end=0.01, decay_rate=0.995):
    """Train Q-learning and return success rate."""
    n_states = env.observation_space.n
    n_actions = env.action_space.n
    q_table = np.zeros((n_states, n_actions))
    
    epsilon = epsilon_start
    successes = 0
    
    for episode in range(n_episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0
        
        while not done:
            if np.random.random() < epsilon:
                action = np.random.randint(0, n_actions)
            else:
                action = np.argmax(q_table[state])
            
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            
            best_next_q = np.max(q_table[next_state])
            target = reward + discount * best_next_q
            q_table[state, action] += learning_rate * (target - q_table[state, action])
            
            total_reward += reward
            state = next_state
        
        if total_reward > 0:
            successes += 1
        
        epsilon = max(epsilon_end, epsilon * decay_rate)
    
    return successes / n_episodes, q_table

# Grid search over hyperparameters
env = gym.make("FrozenLake-v1", is_slippery=False)

learning_rates = [0.1, 0.5, 0.9]
discounts = [0.5, 0.9, 0.99]

results = {}

for lr in learning_rates:
    for disc in discounts:
        # Run 5 trials for each combination
        trial_rates = []
        for _ in range(5):
            rate, _ = train_q_learning(env, n_episodes=500, learning_rate=lr, discount=disc)
            trial_rates.append(rate)
        
        avg_rate = np.mean(trial_rates)
        results[(lr, disc)] = avg_rate
        print(f"LR={lr:.1f}, Discount={disc:.2f}: Success rate = {avg_rate:.2%}")

# Visualize results
print("\n" + "=" * 50)
print("Hyperparameter Tuning Results")
print("=" * 50)
print("\nSuccess rates (higher is better):")
print(f"{'Learning Rate':<15} {'Discount':<10} {'Success Rate':<15}")
print("-" * 40)
for (lr, disc), rate in sorted(results.items()):
    print(f"{lr:<15.1f} {disc:<10.2f} {rate:<15.2%}")

You’ll see clear patterns:

• Learning rate too high (0.9): The Q-values bounce around too much. The agent overreacts to each new experience and never settles on a stable policy.
• Learning rate too low (0.1): Learning is slow but stable. It might need more episodes to converge.
• Discount too low (0.5): The agent is myopic. It doesn’t care enough about future rewards, so it might get stuck in local optima (like staying near the start where it’s safe).
• Discount too high (0.99): The agent plans for the long term and learns to reach the goal efficiently.

For FrozenLake, the sweet spot is usually a low learning rate (0.1-0.3) and a high discount factor (0.9-0.99).

9. From FrozenLake to Any Grid: Generalizing the Code

The beauty of Q-learning is that it works for any grid world, not just FrozenLake. Let’s create a custom grid world and train an agent on it.

import numpy as np
import gymnasium as gym
from gymnasium import spaces

class CustomGridWorld(gym.Env):
    """
    A custom 5x5 grid world.
    - Start at (0, 0)
    - Goal at (4, 4)
    - Reward of -1 per step (to encourage shorter paths)
    - +10 for reaching the goal
    - Walls at certain positions
    """
    
    def __init__(self):
        super().__init__()
        self.grid_size = 5
        self.n_states = self.grid_size * self.grid_size
        self.n_actions = 4  # 0=left, 1=down, 2=right, 3=up
        
        # Define observation and action spaces
        self.observation_space = spaces.Discrete(self.n_states)
        self.action_space = spaces.Discrete(self.n_actions)
        
        # Define walls (positions the agent can't enter)
        self.walls = {(1, 1), (2, 2), (3, 1)}
        
        # Start and goal positions
        self.start_pos = (0, 0)
        self.goal_pos = (4, 4)
        
        self.current_pos = self.start_pos
    
    def _pos_to_state(self, pos):
        """Convert (row, col) position to state index."""
        return pos[0] * self.grid_size + pos[1]
    
    def _state_to_pos(self, state):
        """Convert state index to (row, col) position."""
        return (state // self.grid_size, state % self.grid_size)
    
    def reset(self, seed=None):
        super().reset(seed=seed)
        self.current_pos = self.start_pos
        return self._pos_to_state(self.current_pos), {}
    
    def step(self, action):
        """Take an action and return (next_state, reward, terminated, truncated, info)."""
        row, col = self.current_pos
        
        # Calculate new position based on action
        if action == 0:  # left
            new_pos = (row, max(0, col - 1))
        elif action == 1:  # down
            new_pos = (min(self.grid_size - 1, row + 1), col)
        elif action == 2:  # right
            new_pos = (row, min(self.grid_size - 1, col + 1))
        elif action == 3:  # up
            new_pos = (max(0, row - 1), col)
        else:
            raise ValueError(f"Invalid action: {action}")
        
        # Check if the new position is a wall (if so, stay put)
        if new_pos in self.walls:
            new_pos = self.current_pos
        
        self.current_pos = new_pos
        
        # Calculate reward
        if self.current_pos == self.goal_pos:
            reward = 10
            terminated = True
        else:
            reward = -1  # penalty for each step to encourage efficiency
            terminated = False
        
        truncated = False
        
        return self._pos_to_state(self.current_pos), reward, terminated, truncated, {}
    
    def render(self):
        """Print the grid."""
        grid = []
        for row in range(self.grid_size):
            row_str = ""
            for col in range(self.grid_size):
                pos = (row, col)
                if pos == self.current_pos:
                    row_str += " A "  # Agent
                elif pos == self.goal_pos:
                    row_str += " G "  # Goal
                elif pos in self.walls:
                    row_str += " W "  # Wall
                else:
                    row_str += " . "  # Empty
            grid.append(row_str)
        return "\n".join(grid)

# Train Q-learning on the custom grid world
env = CustomGridWorld()

# Hyperparameters
n_episodes = 2000
learning_rate = 0.1
discount = 0.99
epsilon_start = 1.0
epsilon_end = 0.01
decay_rate = 0.995

# Initialize Q-table
q_table = np.zeros((env.n_states, env.n_actions))

# Training
epsilon = epsilon_start
for episode in range(n_episodes):
    state, _ = env.reset()
    done = False
    total_reward = 0
    steps = 0
    
    while not done:
        # Epsilon-greedy action selection
        if np.random.random() < epsilon:
            action = np.random.randint(0, env.n_actions)
        else:
            action = np.argmax(q_table[state])
        
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        
        # Q-learning update
        best_next_q = np.max(q_table[next_state])
        target = reward + discount * best_next_q
        q_table[state, action] += learning_rate * (target - q_table[state, action])
        
        total_reward += reward
        steps += 1
        state = next_state
    
    epsilon = max(epsilon_end, epsilon * decay_rate)
    
    if (episode + 1) % 200 == 0:
        print(f"Episode {episode + 1}: Total reward = {total_reward}, Steps = {steps}")

print("\nTraining complete!")

# Visualize the learned policy
def visualize_custom_policy(q_table, env):
    """Visualize the learned policy on the custom grid."""
    arrows = ['←', '↓', '→', '↑']
    
    print("\nLearned Policy:")
    print("=" * 25)
    
    for row in range(env.grid_size):
        row_str = ""
        for col in range(env.grid_size):
            state = row * env.grid_size + col
            pos = (row, col)
            
            if pos == env.goal_pos:
                marker = "G"
            elif pos in env.walls:
                marker = "W"
            else:
                best_action = np.argmax(q_table[state])
                marker = arrows[best_action]
            
            row_str += f"  {marker}  "
        print(row_str)
        print()

visualize_custom_policy(q_table, env)

# Test the learned policy (no exploration)
print("\nTesting the learned policy (greedy):")
state, _ = env.reset()
done = False
total_reward = 0
steps = 0

while not done:
    action = np.argmax(q_table[state])
    next_state, reward, terminated, truncated, _ = env.step(action)
    done = terminated or truncated
    total_reward += reward
    steps += 1
    state = next_state
    
    print(f"Step {steps}: State {state}, Action {['Left', 'Down', 'Right', 'Up'][action]}, "
          f"Reward {reward}, Position {env._state_to_pos(state)}")

print(f"\nReached goal in {steps} steps with total reward {total_reward}")

You’ll see the agent learn to navigate around the walls and find the shortest path to the goal. The -1 per step penalty encourages it to take the most efficient route.

Recap: What You Learned

Let’s recap the key ideas:

1. Q-table: A lookup table that stores the expected future reward for each (state, action) pair. It’s the agent’s “cheat sheet” for decision-making.

2. Bellman equation: The update rule that nudges Q-values toward the true expected return. It combines the immediate reward with the best future reward from the next state.

3. Epsilon-greedy: A simple strategy for balancing exploration (trying new things) and exploitation (using what you know).

4. Off-policy learning: Q-learning learns about the optimal policy even while following an exploratory policy. This is both a strength (it can learn faster) and a complexity (it’s learning about a policy it isn’t following).

5. Hyperparameter tuning: The learning rate, discount factor, and epsilon decay all affect how the agent learns. There’s no one-size-fits-all setting — you need to experiment.

Check Your Understanding

Remember: What does the Q-table store? What do the rows and columns represent?

Understand: Explain in your own words why Q-learning is called “off-policy.”

Apply: Given a Q-table for a 3x3 grid world, what action would a greedy policy choose in state 4 if the Q-values are [0.2, 0.8, 0.5, 0.1]?

Analyze: Compare the learning curves of Q-learning with a high learning rate (0.9) vs. a low learning rate (0.1). Why might one be better than the other?

Evaluate: In the FrozenLake environment, would a discount factor of 0.1 ever be useful? Why or why not?

Create: Design a custom grid world with a different reward structure (e.g., positive rewards for visiting certain tiles, negative rewards for others). How would you modify the Q-learning algorithm to handle this?

Related articles

• Part 1: What Is Reinforcement Learning, Really? — If you haven’t read it yet, this introduces the core concepts of agents, environments, and rewards that we built on here.