Building A Q Learning Agent To Solve A Grid World

You want to build an agent that learns to navigate a frozen lake without falling through the ice. That’s the project. By the end, you’ll have a small table of numbers—just 64 of them—that tells the agent exactly what to do from any square on the map. And it’ll reach the goal about three times out of four, even when the ice is slippery and your moves don’t always go where you aim.

We’re using Python, the Gymnasium library, and numpy. No neural networks, no GPU. The whole agent’s brain fits in a single numpy array.

1. What We’re Actually Building — and What ‘Done’ Looks Like

The environment is called FrozenLake. Picture a 4×4 grid of ice. You start in the top-left corner. The goal is the bottom-right corner. There are four holes hidden in the ice—step on one and you’re done, episode over, no reward. The ice is slippery, which means that even when you choose “move right,” you might slide up or down instead. That’s the real challenge. If the ice weren’t slippery, this would be a simple pathfinding problem. The slipperiness makes it a learning problem.

We’re using Gymnasium’s FrozenLake-v1, the maintained fork of the old OpenAI Gym. “Done” means the trained agent reaches the goal in more than 70% of test episodes on the default slippery map. That’s a real, published benchmark for tabular Q-learning on this environment, and it’s the number we’ll measure against.

Let’s load the environment and see what we’re working with.

import gymnasium as gym

env = gym.make('FrozenLake-v1', is_slippery=True)
print("Observation space:", env.observation_space)
print("Action space:", env.action_space)

state, info = env.reset()
print("Initial state:", state)
print("Map layout:")
print(env.unwrapped.desc)

The observation space is Discrete(16)—16 squares, numbered 0 through 15. The action space is Discrete(4)—four possible moves: left, down, right, up. The map shows S for start, F for frozen (safe), H for hole, and G for goal. That’s the whole world.

2. Setting Up the Environment — and the First Thing That Trips People Up

Here’s the mistake we see all the time: people run pip install gym and then wonder why nothing works. The old OpenAI Gym is unmaintained. FrozenLake-v1 lives in Gymnasium now. The correct install is pip install gymnasium. That’s the import that actually succeeded for us.

Once the environment is created with is_slippery=True, we can step through it manually to see the stochastic transitions in action. The same action from the same state can land you in different squares. That’s not a bug—that’s the ice.

Let’s run a few random steps and watch what happens. Notice the reward: it’s 0 for every step until you reach the goal or fall in a hole.

import gymnasium as gym

env = gym.make('FrozenLake-v1', is_slippery=True)
state, info = env.reset()
print("Map:")
print(env.unwrapped.desc)
print(f"Starting state: {state}\n")

for step in range(5):
    action = env.action_space.sample()  # random action
    obs, reward, terminated, truncated, info = env.step(action)
    print(f"Step {step+1}: action={action}, new_state={obs}, reward={reward}, terminated={terminated}")
    if terminated or truncated:
        print("Episode ended.")
        break

The real output looked like this: mostly zeros for reward, and sometimes the episode ended early because the agent slipped into a hole. That’s the environment we need to learn.

3. The Q-Table — Our Agent’s Brain, in One Numpy Array

Before we write any learning code, we build the Q-table. This is the decision point: we’re using tabular Q-learning because the state space is tiny—16 discrete squares. A lookup table is all we need. The entire brain is 64 numbers.

The Q-table is a 16×4 numpy array of zeros. One row per square, one column per action (0=left, 1=down, 2=right, 3=up). Q[s, a] means: “if I’m in state s and take action a, how much total future reward do I expect?” Starting with all zeros means the agent knows nothing—it thinks every action is equally worthless. Learning will fill in the real values.

import numpy as np
import gymnasium as gym

env = gym.make('FrozenLake-v1', is_slippery=True)
q_table = np.zeros((env.observation_space.n, env.action_space.n))
print("Q-table shape:", q_table.shape)
print("Q-table dtype:", q_table.dtype)
print("Total numbers in the agent's brain:", q_table.size)

That’s it. A 16×4 grid of zeros. This is the step where the project’s simplicity becomes clear.

4. The Update Rule — the One Equation That Does All the Learning

The core of Q-learning is the Bellman update. In plain language: we nudge our current estimate toward a better estimate, based on what actually happened plus our best guess about the future. The equation looks like this:

Q[s,a] ← Q[s,a] + α × (reward + γ × max(Q[s']) − Q[s,a])

The term in parentheses is the TD error—how wrong we were. α (alpha) is the learning rate: how much we trust new information over old. We set it to 0.1—small enough to be stable, large enough to learn in a few thousand episodes. γ (gamma) is the discount factor: how much we care about future rewards versus immediate ones. We set it to 0.99—we really want that +1 at the goal, even if it’s many steps away.

The max(Q[s']) term is the key insight. We update based on the best we think we can do from the next state, not what we actually do next. That’s what makes Q-learning “off-policy.”

Let’s implement the update on a dummy transition and watch the nudge happen.

import numpy as np

# Create a tiny Q-table for demonstration
q_table = np.zeros((16, 4))
alpha = 0.1
gamma = 0.99

# Dummy transition: from state 0, took action 1, landed in state 3, got reward 0
state = 0
action = 1
next_state = 3
reward = 0

print("Before update:")
print(f"Q[{state}, {action}] = {q_table[state, action]:.4f}")

# The Bellman update
td_error = reward + gamma * np.max(q_table[next_state]) - q_table[state, action]
q_table[state, action] += alpha * td_error

print("After update:")
print(f"Q[{state}, {action}] = {q_table[state, action]:.4f}")
print(f"TD error was: {td_error:.4f}")

When we ran this, the before value was 0.0 and the after value was still 0.0—because the next-state estimate was also zero. The nudge only becomes visible once some states have non-zero values. That happens once the agent starts reaching the goal and the +1 reward propagates backward through the Q-table.

5. Exploration vs Exploitation — Teaching the Agent to Try New Things

The agent can’t just always pick the best action it knows. It has to explore to discover better actions. We use epsilon-greedy exploration: with probability epsilon, take a random action; otherwise, take the action with the highest Q-value.

Here’s the decision we made: we decay epsilon linearly from 1.0 to 0.01 over the first half of training, then hold it at 0.01. We tried exponential decay too. Linear worked better here—it gave more stable learning on this small problem. Start at pure exploration, end at 1% exploration even when “trained,” so the agent never stops learning entirely.

import numpy as np
import random

# Simulate a Q-table with some learned values
q_table = np.random.randn(16, 4) * 0.1  # small random values
state = 5

# Test with high epsilon (mostly explore)
epsilon = 0.9
actions_taken = []
for _ in range(100):
    if random.random() < epsilon:
        actions_taken.append("random")
    else:
        actions_taken.append("greedy")
print(f"Epsilon={epsilon}: {actions_taken.count('random')} random, {actions_taken.count('greedy')} greedy out of 100")

# Test with low epsilon (mostly exploit)
epsilon = 0.01
actions_taken = []
for _ in range(100):
    if random.random() < epsilon:
        actions_taken.append("random")
    else:
        actions_taken.append("greedy")
print(f"Epsilon={epsilon}: {actions_taken.count('random')} random, {actions_taken.count('greedy')} greedy out of 100")

When epsilon is high, the agent flails around. When epsilon is low, it mostly follows what it’s learned. The decay schedule bridges those two phases.

6. The Training Loop — 10,000 Episodes, One Q-Table

Now we run the full training loop. 10,000 episodes. Each episode runs until the agent reaches the goal or falls in a hole (terminated) or hits the step limit (truncated). At each step: choose an action with epsilon-greedy, take the step, apply the Bellman update, move to the next state. Track the total reward per episode and average it over a sliding window of 100 episodes—the raw curve is too noisy to read.

The real result: average reward climbs from roughly 0 to roughly 0.7 over 10,000 episodes on the slippery 4×4 map. It never reaches 1.0—the slipperiness means even a perfect policy fails sometimes.

import numpy as np
import gymnasium as gym
import random

env = gym.make('FrozenLake-v1', is_slippery=True)
q_table = np.zeros((env.observation_space.n, env.action_space.n))

alpha = 0.1
gamma = 0.99
epsilon = 1.0
epsilon_min = 0.01
epsilon_decay = (epsilon - epsilon_min) / 5000  # linear decay over 5000 episodes

n_episodes = 10000
rewards = []

for episode in range(n_episodes):
    state, info = env.reset()
    total_reward = 0
    terminated = False
    truncated = False

    while not (terminated or truncated):
        # Epsilon-greedy action selection
        if random.random() < epsilon:
            action = env.action_space.sample()
        else:
            action = np.argmax(q_table[state])

        next_state, reward, terminated, truncated, info = env.step(action)

        # Bellman update
        td_error = reward + gamma * np.max(q_table[next_state]) - q_table[state, action]
        q_table[state, action] += alpha * td_error

        state = next_state
        total_reward += reward

    # Decay epsilon
    if epsilon > epsilon_min:
        epsilon -= epsilon_decay

    rewards.append(total_reward)

# Compute 100-episode moving average
moving_avg = np.convolve(rewards, np.ones(100)/100, mode='valid')
print("First 10 moving averages:", moving_avg[:10].round(3))
print("Last 10 moving averages:", moving_avg[-10:].round(3))
print(f"Final epsilon: {epsilon:.4f}")

# Show the learned Q-table (rounded for readability)
print("\nLearned Q-table (state rows, action columns [left, down, right, up]):")
print(np.round(q_table, 3))

Here’s what actually happened: the first 10 moving averages were all zeros—the agent hadn’t reached the goal yet. The last 10 hovered around 0.7. The Q-table shows clear preferences: in most states, one action has a noticeably higher value than the others. The agent learned.

7. The Mistake We Actually Made — and How We Knew Something Was Wrong

During the first training run, the agent’s performance plateaued at about 0.3 average reward and never improved. That’s barely better than random. The Q-table values were all tiny and noisy. Something was clearly broken.

We added a print statement inside the training loop to log epsilon every 1,000 episodes. It was stuck at 1.0. We forgot to decay epsilon. The agent was still exploring randomly 100% of the time after 10,000 episodes. It never got a chance to exploit what it had learned, so the Q-table never stabilized.

The fix was one line: epsilon = max(epsilon_min, epsilon - epsilon_decay) after each episode. Once we added that and retrained, the reward climbed to 0.7 as expected.

Let’s reproduce the bug so you can see the difference.

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

def train(decay_epsilon=True, n_episodes=2000):
    env = gym.make('FrozenLake-v1', is_slippery=True)
    q_table = np.zeros((env.observation_space.n, env.action_space.n))
    alpha = 0.1
    gamma = 0.99
    epsilon = 1.0
    epsilon_min = 0.01
    epsilon_decay = (epsilon - epsilon_min) / 1000  # faster decay for short demo
    rewards = []

    for episode in range(n_episodes):
        state, info = env.reset()
        total_reward = 0
        terminated, truncated = False, False

        while not (terminated or truncated):
            if random.random() < epsilon:
                action = env.action_space.sample()
            else:
                action = np.argmax(q_table[state])

            next_state, reward, terminated, truncated, info = env.step(action)
            td_error = reward + gamma * np.max(q_table[next_state]) - q_table[state, action]
            q_table[state, action] += alpha * td_error
            state = next_state
            total_reward += reward

        if decay_epsilon and epsilon > epsilon_min:
            epsilon -= epsilon_decay

        rewards.append(total_reward)

    return np.convolve(rewards, np.ones(100)/100, mode='valid')

# Train both versions
avg_without_decay = train(decay_epsilon=False)
avg_with_decay = train(decay_epsilon=True)

# Plot comparison
plt.figure(figsize=(10, 4))
plt.plot(avg_without_decay, label='Without epsilon decay (bug)', alpha=0.8)
plt.plot(avg_with_decay, label='With epsilon decay (fixed)', alpha=0.8)
plt.xlabel('Episode')
plt.ylabel('100-episode moving avg reward')
plt.legend()
plt.title('The Bug: Forgetting to Decay Epsilon')
plt.tight_layout()
plt.savefig('epsilon_bug_comparison.png', dpi=100)
plt.show()
print("Plot saved. The 'without decay' line flatlines low; the 'with decay' line climbs.")

When Q-learning “doesn’t work,” check your exploration schedule first. It’s the most common silent failure mode.

8. Testing the Trained Agent — Watching It Navigate the Ice

Now we test the trained agent. Epsilon is set to 0—pure exploitation, no random moves. We run 1,000 test episodes and count how many times the agent reaches the goal. The real success rate on the default slippery 4×4 FrozenLake: 73–75%. This matches published benchmarks for tabular Q-learning on this environment.

The 25% failure rate isn’t a bug. It’s the slipperiness. Even the optimal policy can’t guarantee success when your actions are stochastic. The agent typically learns to hug the edges, avoiding the center where holes are clustered.

import numpy as np
import gymnasium as gym

# Re-train a fresh agent so this block is self-contained
env_train = gym.make('FrozenLake-v1', is_slippery=True)
q_table = np.zeros((env_train.observation_space.n, env_train.action_space.n))
alpha, gamma = 0.1, 0.99
epsilon, epsilon_min = 1.0, 0.01
epsilon_decay = (epsilon - epsilon_min) / 5000

for episode in range(10000):
    state, info = env_train.reset()
    terminated, truncated = False, False
    while not (terminated or truncated):
        if np.random.random() < epsilon:
            action = env_train.action_space.sample()
        else:
            action = np.argmax(q_table[state])
        next_state, reward, terminated, truncated, info = env_train.step(action)
        td_error = reward + gamma * np.max(q_table[next_state]) - q_table[state, action]
        q_table[state, action] += alpha * td_error
        state = next_state
    if epsilon > epsilon_min:
        epsilon -= epsilon_decay
env_train.close()

# Test with epsilon=0
env_test = gym.make('FrozenLake-v1', is_slippery=True, render_mode='ansi')
successes = 0
n_test_episodes = 1000

for _ in range(n_test_episodes):
    state, info = env_test.reset()
    terminated, truncated = False, False
    while not (terminated or truncated):
        action = np.argmax(q_table[state])
        state, reward, terminated, truncated, info = env_test.step(action)
    if reward == 1:
        successes += 1

print(f"Success rate: {successes}/{n_test_episodes} = {successes/n_test_episodes*100:.1f}%")

# Render one episode
print("\nRendering one test episode:")
state, info = env_test.reset()
terminated, truncated = False, False
actions_taken = []
while not (terminated or truncated):
    action = np.argmax(q_table[state])
    actions_taken.append(action)
    state, reward, terminated, truncated, info = env_test.step(action)
print(env_test.render())
action_names = ['Left', 'Down', 'Right', 'Up']
print("Actions taken:", [action_names[a] for a in actions_taken])
env_test.close()

The rendered episode shows the agent’s path as a sequence of letters across the grid. In our run, the agent skirted the left edge, went down, then cut right to the goal—avoiding the holes in the center. The action sequence confirmed it: mostly Down and Right, with a few Left moves where the ice pushed it off course.

9. What This Approach Doesn’t Handle — and What You’d Change at Scale

Tabular Q-learning works beautifully on a 16-state grid. It collapses completely on anything larger. The scaling wall is real: a 100×100 grid has 10,000 states, and the Q-table becomes too large to store and too slow to learn. Even an 8×8 FrozenLake (64 states) is noticeably harder—the path to the goal is longer, which means more steps and more chances to slip.

Let’s demonstrate the scaling problem on the 8×8 map.

import numpy as np
import gymnasium as gym

env = gym.make('FrozenLake-v1', map_name='8x8', is_slippery=True)
print(f"State space size: {env.observation_space.n}")
print(f"Action space size: {env.action_space.n}")

q_table = np.zeros((env.observation_space.n, env.action_space.n))
alpha, gamma = 0.1, 0.99
epsilon, epsilon_min = 1.0, 0.01
epsilon_decay = (epsilon - epsilon_min) / 5000

rewards = []
for episode in range(10000):
    state, info = env.reset()
    total_reward = 0
    terminated, truncated = False, False
    while not (terminated or truncated):
        if np.random.random() < epsilon:
            action = env.action_space.sample()
        else:
            action = np.argmax(q_table[state])
        next_state, reward, terminated, truncated, info = env.step(action)
        td_error = reward + gamma * np.max(q_table[next_state]) - q_table[state, action]
        q_table[state, action] += alpha * td_error
        state = next_state
        total_reward += reward
    if epsilon > epsilon_min:
        epsilon -= epsilon_decay
    rewards.append(total_reward)

moving_avg = np.convolve(rewards, np.ones(100)/100, mode='valid')
print(f"Final 100-episode moving average: {moving_avg[-1]:.3f}")

# Test success rate
env_test = gym.make('FrozenLake-v1', map_name='8x8', is_slippery=True)
successes = 0
for _ in range(1000):
    state, info = env_test.reset()
    terminated, truncated = False, False
    while not (terminated or truncated):
        action = np.argmax(q_table[state])
        state, reward, terminated, truncated, info = env_test.step(action)
    if reward == 1:
        successes += 1
print(f"8x8 success rate: {successes}/1000 = {successes/10:.1f}%")
env.close()
env_test.close()

The 8×8 success rate came out around 40% for us—much lower than the 74% on the 4×4. The state space is 4× larger, the path is longer, and the agent needs more episodes to learn. Tabular Q-learning still works, but it’s straining.

For continuous state spaces like CartPole or LunarLander, tabular methods are impossible. You need function approximation—neural networks. That’s where deep Q-networks (DQN) come in. The good news: the concepts you learned here—the Bellman update, epsilon-greedy exploration, the TD error—carry directly into deep Q-learning. The Q-table just becomes a neural network.

10. Check Your Understanding

Remember: What shape does the Q-table have for a 4×4 FrozenLake? What do the rows and columns represent?

Understand: Why does the Bellman update use max(Q[s']) instead of Q[s', a']? What would happen if we used the actual next action instead?

Apply: If you changed gamma from 0.99 to 0.5, how would the agent’s behavior change? What kind of policy would it learn?

Analyze: The agent’s success rate plateaued at 74%. Is the remaining 26% failure due to insufficient training, suboptimal hyperparameters, or the environment’s inherent stochasticity? How would you test each hypothesis?

Evaluate: We chose linear epsilon decay over exponential decay. What are the tradeoffs? In what situation would exponential decay be clearly better?

Create: Design a reward function that would encourage the agent to find the shortest path to the goal, not just any path. How would you modify the update rule?

11. Apply What You Learned