Objective: Implement a 4x4 grid-world environment, train a Q-learning agent to navigate it while avoiding a penalty cell, and visualize both how learning progresses and what policy the agent ends up with.
Setup
Every code block below runs live in your browser using Pyodide. Click Run to execute it. The whole environment is implemented from scratch in plain numpy. No external reinforcement learning library is needed, or available, in this in-browser environment.
The grid is 4x4 (16 states, numbered 0 to 15, row-major: state 0 is top-left, state 15 is bottom-right). The agent starts at state 0. State 15 is the goal (reward +10, ends the episode). State 6 is a penalty cell (reward −10, ends the episode). Every other move costs −1, to encourage short paths.
Task 1: Build the Grid-World Environment
Define the transition function: given a state and an action (0=up, 1=down, 2=left, 3=right), compute the next state, respecting the grid's edges, and the reward for landing there.
import numpy as np
GRID_SIZE = 4
N_STATES = GRID_SIZE * GRID_SIZE
N_ACTIONS = 4 # 0=up, 1=down, 2=left, 3=right
GOAL_STATE = 15
PIT_STATE = 6
def step(state, action):
row, col = divmod(state, GRID_SIZE)
if action == 0: row = max(0, row - 1)
elif action == 1: row = min(GRID_SIZE - 1, row + 1)
elif action == 2: col = max(0, col - 1)
elif action == 3: col = min(GRID_SIZE - 1, col + 1)
next_state = row * GRID_SIZE + col
if next_state == GOAL_STATE:
return next_state, 10.0, True
if next_state == PIT_STATE:
return next_state, -10.0, True
return next_state, -1.0, False
# Quick sanity check: from state 0, moving "down" then "right" a few times
s = 0
for action in [1, 1, 3, 3]:
s, r, done = step(s, action)
print(f"action={action} -> state={s}, reward={r}, done={done}")
Task 2: Implement Epsilon-Greedy Q-Learning
Train a Q-table over many episodes using the same update rule from the lesson, with epsilon decaying over time so the agent explores heavily at first and exploits more as it learns.
import numpy as np
GRID_SIZE = 4
N_STATES = GRID_SIZE * GRID_SIZE
N_ACTIONS = 4
GOAL_STATE = 15
PIT_STATE = 6
def step(state, action):
row, col = divmod(state, GRID_SIZE)
if action == 0: row = max(0, row - 1)
elif action == 1: row = min(GRID_SIZE - 1, row + 1)
elif action == 2: col = max(0, col - 1)
elif action == 3: col = min(GRID_SIZE - 1, col + 1)
next_state = row * GRID_SIZE + col
if next_state == GOAL_STATE:
return next_state, 10.0, True
if next_state == PIT_STATE:
return next_state, -10.0, True
return next_state, -1.0, False
np.random.seed(0)
Q = np.zeros((N_STATES, N_ACTIONS))
alpha = 0.5
gamma = 0.95
n_episodes = 500
epsilon_start, epsilon_end = 1.0, 0.05
episode_returns = []
for ep in range(n_episodes):
epsilon = epsilon_start - (epsilon_start - epsilon_end) * (ep / n_episodes)
state = 0
total_reward = 0.0
done = False
steps = 0
while not done and steps < 100:
if np.random.rand() < epsilon:
action = np.random.randint(N_ACTIONS)
else:
action = np.argmax(Q[state])
next_state, reward, done = step(state, action)
best_next_q = np.max(Q[next_state])
Q[state, action] += alpha * (reward + gamma * best_next_q * (not done) - Q[state, action])
total_reward += reward
state = next_state
steps += 1
episode_returns.append(total_reward)
print("Average return, first 20 episodes:", np.mean(episode_returns[:20]))
print("Average return, last 20 episodes: ", np.mean(episode_returns[-20:]))
The (not done) multiplier on the discounted future term matters: once an episode ends, there is no
meaningful "next state" to bootstrap from, so the update should rely on the immediate reward alone at that
point.
Task 3: Plot the Learning Curve
Re-run training while recording every episode's return, then plot it to see the agent improve over time. A rolling average makes the noisy episode-to-episode swings easier to read than the raw values.
import numpy as np
import matplotlib.pyplot as plt
GRID_SIZE = 4
N_STATES = GRID_SIZE * GRID_SIZE
N_ACTIONS = 4
GOAL_STATE = 15
PIT_STATE = 6
def step(state, action):
row, col = divmod(state, GRID_SIZE)
if action == 0: row = max(0, row - 1)
elif action == 1: row = min(GRID_SIZE - 1, row + 1)
elif action == 2: col = max(0, col - 1)
elif action == 3: col = min(GRID_SIZE - 1, col + 1)
next_state = row * GRID_SIZE + col
if next_state == GOAL_STATE:
return next_state, 10.0, True
if next_state == PIT_STATE:
return next_state, -10.0, True
return next_state, -1.0, False
np.random.seed(0)
Q = np.zeros((N_STATES, N_ACTIONS))
alpha, gamma = 0.5, 0.95
n_episodes = 500
epsilon_start, epsilon_end = 1.0, 0.05
episode_returns = []
for ep in range(n_episodes):
epsilon = epsilon_start - (epsilon_start - epsilon_end) * (ep / n_episodes)
state, total_reward, done, steps = 0, 0.0, False, 0
while not done and steps < 100:
action = np.random.randint(N_ACTIONS) if np.random.rand() < epsilon else np.argmax(Q[state])
next_state, reward, done = step(state, action)
Q[state, action] += alpha * (reward + gamma * np.max(Q[next_state]) * (not done) - Q[state, action])
total_reward += reward
state = next_state
steps += 1
episode_returns.append(total_reward)
# Rolling average over a 20-episode window
window = 20
rolling = np.convolve(episode_returns, np.ones(window) / window, mode="valid")
plt.plot(rolling)
plt.xlabel("Episode")
plt.ylabel(f"Return ({window}-episode rolling average)")
plt.title("Q-learning return over training")
show_plot()
Task 4: Visualize the Learned Policy
Read off, for every state, the action with the highest Q-value, and print it as a grid of arrows. The goal and pit cells are marked specially since the agent never acts from a terminal state.
import numpy as np
GRID_SIZE = 4
N_STATES = GRID_SIZE * GRID_SIZE
N_ACTIONS = 4
GOAL_STATE = 15
PIT_STATE = 6
def step(state, action):
row, col = divmod(state, GRID_SIZE)
if action == 0: row = max(0, row - 1)
elif action == 1: row = min(GRID_SIZE - 1, row + 1)
elif action == 2: col = max(0, col - 1)
elif action == 3: col = min(GRID_SIZE - 1, col + 1)
next_state = row * GRID_SIZE + col
if next_state == GOAL_STATE:
return next_state, 10.0, True
if next_state == PIT_STATE:
return next_state, -10.0, True
return next_state, -1.0, False
np.random.seed(0)
Q = np.zeros((N_STATES, N_ACTIONS))
alpha, gamma = 0.5, 0.95
n_episodes = 500
epsilon_start, epsilon_end = 1.0, 0.05
for ep in range(n_episodes):
epsilon = epsilon_start - (epsilon_start - epsilon_end) * (ep / n_episodes)
state, done, steps = 0, False, 0
while not done and steps < 100:
action = np.random.randint(N_ACTIONS) if np.random.rand() < epsilon else np.argmax(Q[state])
next_state, reward, done = step(state, action)
Q[state, action] += alpha * (reward + gamma * np.max(Q[next_state]) * (not done) - Q[state, action])
state = next_state
steps += 1
arrows = {0: "^", 1: "v", 2: "<", 3: ">"}
print("Learned policy (^ v < > = up/down/left/right, G = goal, X = pit):\n")
for row in range(GRID_SIZE):
line = ""
for col in range(GRID_SIZE):
s = row * GRID_SIZE + col
if s == GOAL_STATE:
line += " G "
elif s == PIT_STATE:
line += " X "
else:
line += " " + arrows[np.argmax(Q[s])] + " "
print(line)
Task 5: Reflection
Go back to Task 2 or 3 and try setting epsilon_start = 0.05 instead of 1.0 (almost no
exploration from the very start), then re-run. You should see the agent converge to a noticeably worse, or
inconsistent, policy more often, since it never explores enough to discover that the path past the pit cell
leads to a much bigger reward.
This is the exploration-exploitation tradeoff from the lesson made concrete: an agent that exploits its (initially terrible, since the Q-table starts at all zeros) current beliefs too early can lock into a mediocre strategy before it has gathered enough experience to know better. The decaying epsilon schedule used in Task 2 is a simple, common way to balance the two phases without having to choose a single fixed exploration rate.