Topic 18: Reinforcement Learning¶

Goal: Understand and implement foundational tabular Reinforcement Learning algorithms (Value Iteration and Q-Learning) from scratch to solve a GridWorld environment.

Prerequisites: Basic probability, dynamic programming concepts.

Theory Link: theory.md

In [1]:

Copied!





import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90
import random

import matplotlib.pyplot as plt
import numpy as np

%load_ext autoreload
%autoreload 2

SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)

plt.rcParams["figure.dpi"] = 90

1. Problem Setup — WHY¶

Reinforcement Learning involves an agent interacting with an environment to maximize cumulative reward. The environment is formalized as a Markov Decision Process (MDP). The agent does not know the optimal actions a priori; it must explore the environment, receive delayed rewards, and learn a policy $\pi(a|s)$ that dictates the best action to take in each state.

2. Mathematical Core — WHAT¶

The core equations that govern tabular RL methods are based on the Bellman equations.

Value Iteration (Dynamic Programming, requires known MDP dynamics $P, R$): $$ V_{k+1}(s) = \max_a \sum_{s'} P(s'|s,a) [ R(s,a,s') + \gamma V_k(s') ] $$

Q-Learning (Model-Free, learns from experience samples): $$ Q(S_t, A_t) \leftarrow Q(S_t, A_t) + \alpha \left[ R_{t+1} + \gamma \max_a Q(S_{t+1}, a) - Q(S_t, A_t) \right] $$

3. Solution Method — HOW¶

We will implement a standard GridWorld environment.

Value Iteration: We iteratively apply the Bellman optimality operator across all states until the state values converge. Then we extract the optimal policy.
Q-Learning: We allow an agent to wander the grid using an $\epsilon$-greedy policy. At each step, it observes the reward and next state, and applies the Temporal Difference (TD) update to its Q-table.

4. Implementation — BUILD¶

First, we build the GridWorld environment.

In [2]:

Copied!





class GridWorld:
    """
    A simple 2D GridWorld MDP.
    States are (row, col). Actions: 0=Up, 1=Right, 2=Down, 3=Left.
    """
    def __init__(self, rows=4, cols=4, start=(0,0), goal=(3,3), obstacles=[(1,1), (2,1)]):
        self.rows = rows
        self.cols = cols
        self.start = start
        self.goal = goal
        self.obstacles = obstacles
        self.actions = [(-1, 0), (0, 1), (1, 0), (0, -1)]  # Up, Right, Down, Left
        self.reset()
        
    def reset(self):
        self.current_state = self.start
        return self.current_state
        
    def step(self, action_idx):
        if self.current_state == self.goal:
            return self.current_state, 0.0, True
            
        dr, dc = self.actions[action_idx]
        r, c = self.current_state
        next_r, next_c = r + dr, c + dc
        
        # Bound check & obstacle check
        if (0 <= next_r < self.rows) and (0 <= next_c < self.cols) and ((next_r, next_c) not in self.obstacles):
            self.current_state = (next_r, next_c)
            
        reward = -1.0 # Step penalty
        done = (self.current_state == self.goal)
        if done:
            reward = 10.0 # Goal reward
            
        return self.current_state, reward, done
        
    def get_all_states(self):
        states = []
        for r in range(self.rows):
            for c in range(self.cols):
                if (r, c) not in self.obstacles:
                    states.append((r, c))
        return states
        
    def get_transitions(self, state, action_idx):
        # For Value Iteration: returns (prob, next_state, reward, done)
        if state == self.goal:
            return [(1.0, state, 0.0, True)]
            
        dr, dc = self.actions[action_idx]
        r, c = state
        next_r, next_c = r + dr, c + dc
        
        if (0 <= next_r < self.rows) and (0 <= next_c < self.cols) and ((next_r, next_c) not in self.obstacles):
            next_state = (next_r, next_c)
        else:
            next_state = state
            
        done = (next_state == self.goal)
        reward = 10.0 if done else -1.0
        return [(1.0, next_state, reward, done)] # Deterministic environment
class GridWorld:
    """
    A simple 2D GridWorld MDP.
    States are (row, col). Actions: 0=Up, 1=Right, 2=Down, 3=Left.
    """
    def __init__(self, rows=4, cols=4, start=(0,0), goal=(3,3), obstacles=[(1,1), (2,1)]):
        self.rows = rows
        self.cols = cols
        self.start = start
        self.goal = goal
        self.obstacles = obstacles
        self.actions = [(-1, 0), (0, 1), (1, 0), (0, -1)]  # Up, Right, Down, Left
        self.reset()
        
    def reset(self):
        self.current_state = self.start
        return self.current_state
        
    def step(self, action_idx):
        if self.current_state == self.goal:
            return self.current_state, 0.0, True
            
        dr, dc = self.actions[action_idx]
        r, c = self.current_state
        next_r, next_c = r + dr, c + dc
        
        # Bound check & obstacle check
        if (0 <= next_r < self.rows) and (0 <= next_c < self.cols) and ((next_r, next_c) not in self.obstacles):
            self.current_state = (next_r, next_c)
            
        reward = -1.0 # Step penalty
        done = (self.current_state == self.goal)
        if done:
            reward = 10.0 # Goal reward
            
        return self.current_state, reward, done
        
    def get_all_states(self):
        states = []
        for r in range(self.rows):
            for c in range(self.cols):
                if (r, c) not in self.obstacles:
                    states.append((r, c))
        return states
        
    def get_transitions(self, state, action_idx):
        # For Value Iteration: returns (prob, next_state, reward, done)
        if state == self.goal:
            return [(1.0, state, 0.0, True)]
            
        dr, dc = self.actions[action_idx]
        r, c = state
        next_r, next_c = r + dr, c + dc
        
        if (0 <= next_r < self.rows) and (0 <= next_c < self.cols) and ((next_r, next_c) not in self.obstacles):
            next_state = (next_r, next_c)
        else:
            next_state = state
            
        done = (next_state == self.goal)
        reward = 10.0 if done else -1.0
        return [(1.0, next_state, reward, done)] # Deterministic environment

Value Iteration (Dynamic Programming)¶

In [3]:

Copied!





def value_iteration(env, gamma=0.99, theta=1e-6):
    states = env.get_all_states()
    V = {s: 0.0 for s in states}
    
    while True:
        delta = 0
        for s in states:
            if s == env.goal:
                continue
            
            v_old = V[s]
            action_values = []
            for a in range(4):
                q = 0
                for prob, next_s, reward, done in env.get_transitions(s, a):
                    q += prob * (reward + gamma * V[next_s])
                action_values.append(q)
                
            V[s] = max(action_values)
            delta = max(delta, abs(v_old - V[s]))
            
        if delta < theta:
            break
            
    # Extract Q-values from converged V
    Q = {s: np.zeros(4) for s in states}
    for s in states:
        if s == env.goal: continue
        for a in range(4):
            for prob, next_s, reward, done in env.get_transitions(s, a):
                Q[s][a] += prob * (reward + gamma * V[next_s])
                
    return V, Q
def value_iteration(env, gamma=0.99, theta=1e-6):
    states = env.get_all_states()
    V = {s: 0.0 for s in states}
    
    while True:
        delta = 0
        for s in states:
            if s == env.goal:
                continue
            
            v_old = V[s]
            action_values = []
            for a in range(4):
                q = 0
                for prob, next_s, reward, done in env.get_transitions(s, a):
                    q += prob * (reward + gamma * V[next_s])
                action_values.append(q)
                
            V[s] = max(action_values)
            delta = max(delta, abs(v_old - V[s]))
            
        if delta < theta:
            break
            
    # Extract Q-values from converged V
    Q = {s: np.zeros(4) for s in states}
    for s in states:
        if s == env.goal: continue
        for a in range(4):
            for prob, next_s, reward, done in env.get_transitions(s, a):
                Q[s][a] += prob * (reward + gamma * V[next_s])
                
    return V, Q

Q-Learning (Model-Free)¶

In [4]:

Copied!





class QLearningAgent:
    def __init__(self, env, alpha=0.1, gamma=0.99, epsilon=0.1):
        self.env = env
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.Q = {s: np.zeros(4) for s in env.get_all_states()}
        
    def get_action(self, state):
        if rng.random() < self.epsilon:
            return rng.integers(0, 4)
        else:
            # Break ties randomly for max
            q_values = self.Q[state]
            max_q = np.max(q_values)
            return rng.choice(np.where(q_values == max_q)[0])
            
    def update(self, state, action, reward, next_state, done):
        best_next_q = 0.0 if done else np.max(self.Q[next_state])
        td_target = reward + self.gamma * best_next_q
        td_error = td_target - self.Q[state][action]
        self.Q[state][action] += self.alpha * td_error
class QLearningAgent:
    def __init__(self, env, alpha=0.1, gamma=0.99, epsilon=0.1):
        self.env = env
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.Q = {s: np.zeros(4) for s in env.get_all_states()}
        
    def get_action(self, state):
        if rng.random() < self.epsilon:
            return rng.integers(0, 4)
        else:
            # Break ties randomly for max
            q_values = self.Q[state]
            max_q = np.max(q_values)
            return rng.choice(np.where(q_values == max_q)[0])
            
    def update(self, state, action, reward, next_state, done):
        best_next_q = 0.0 if done else np.max(self.Q[next_state])
        td_target = reward + self.gamma * best_next_q
        td_error = td_target - self.Q[state][action]
        self.Q[state][action] += self.alpha * td_error

Let's train the Q-Learning agent.

In [5]:

Copied!





env = GridWorld()
agent = QLearningAgent(env, alpha=0.1, gamma=0.99, epsilon=0.1)

n_episodes = 500
rewards_history = []

for ep in range(n_episodes):
    state = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        action = agent.get_action(state)
        next_state, reward, done = env.step(action)
        agent.update(state, action, reward, next_state, done)
        
        state = next_state
        total_reward += reward
        
    rewards_history.append(total_reward)
env = GridWorld()
agent = QLearningAgent(env, alpha=0.1, gamma=0.99, epsilon=0.1)

n_episodes = 500
rewards_history = []

for ep in range(n_episodes):
    state = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        action = agent.get_action(state)
        next_state, reward, done = env.step(action)
        agent.update(state, action, reward, next_state, done)
        
        state = next_state
        total_reward += reward
        
    rewards_history.append(total_reward)

4.3 Deep Q-Network (DQN) from Scratch (Pure NumPy)¶

When state spaces become large or continuous, tabular representations fail due to the curse of dimensionality. A Deep Q-Network (DQN) addresses this by approximating $Q(s, a; \theta) \approx Q^\ast(s, a)$ using a neural network parameterized by $\theta$.

Key DQN Components:

Neural Network Function Approximator ($Q_\theta$): A 2-layer MLP ($d \to H \to A$) mapping state representation $s \in \mathbb{R}^{d}$ to Q-values for all actions.
Experience Replay Buffer ($\mathcal{D}$): Stores transitions $(s, a, r, s', \text{done})$ and samples random mini-batches during training. This breaks the strong temporal correlation between consecutive steps.
Target Network ($Q_{\theta^-}$): A separate set of parameters $\theta^-$ used to calculate the TD target: $$ y = r + \gamma (1 - \text{done}) \max_{a'} Q_{\theta^-}(s', a') $$ The target weights $\theta^-$ are updated periodically ($\theta^- \leftarrow \theta$), stabilizing the moving TD target.

In [6]:

Copied!





# Pure NumPy Implementation of Deep Q-Network (DQN)

def state_to_onehot(state, rows=4, cols=4):
    idx = state[0] * cols + state[1]
    onehot = np.zeros(rows * cols, dtype=np.float32)
    onehot[idx] = 1.0
    return onehot

class ReplayBuffer:
    def __init__(self, capacity=2000):
        self.capacity = capacity
        self.buffer = []
        self.ptr = 0

    def push(self, state, action, reward, next_state, done):
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)
        self.buffer[self.ptr] = (state, action, reward, next_state, done)
        self.ptr = (self.ptr + 1) % self.capacity

    def sample(self, batch_size):
        indices = rng.choice(len(self.buffer), batch_size, replace=False)
        states, actions, rewards, next_states, dones = zip(*[self.buffer[i] for i in indices])
        return (np.array(states), np.array(actions), np.array(rewards, dtype=np.float32),
                np.array(next_states), np.array(dones, dtype=np.float32))

class NeuralQNetwork:
    """2-Layer MLP Q-Function Approximator with ReLU activation."""
    def __init__(self, state_dim=16, hidden_dim=32, action_dim=4, lr=0.05):
        self.lr = lr
        self.W1 = rng.standard_normal((state_dim, hidden_dim)) * np.sqrt(2.0 / state_dim)
        self.b1 = np.zeros((1, hidden_dim))
        self.W2 = rng.standard_normal((hidden_dim, action_dim)) * np.sqrt(2.0 / hidden_dim)
        self.b2 = np.zeros((1, action_dim))

    def forward(self, x):
        self.z1 = np.dot(x, self.W1) + self.b1
        self.a1 = np.maximum(0, self.z1)
        self.q = np.dot(self.a1, self.W2) + self.b2
        return self.q

    def copy_weights_from(self, other):
        self.W1 = np.copy(other.W1)
        self.b1 = np.copy(other.b1)
        self.W2 = np.copy(other.W2)
        self.b2 = np.copy(other.b2)

    def train_step(self, x, actions, targets):
        N = x.shape[0]
        q_pred = self.forward(x)
        dq = np.zeros_like(q_pred)
        for i in range(N):
            dq[i, actions[i]] = (q_pred[i, actions[i]] - targets[i]) / N
        dW2 = np.dot(self.a1.T, dq)
        db2 = np.sum(dq, axis=0, keepdims=True)
        da1 = np.dot(dq, self.W2.T)
        dz1 = da1 * (self.z1 > 0)
        dW1 = np.dot(x.T, dz1)
        db1 = np.sum(dz1, axis=0, keepdims=True)
        self.W1 -= self.lr * dW1
        self.b1 -= self.lr * db1
        self.W2 -= self.lr * dW2
        self.b2 -= self.lr * db2
        return 0.5 * np.mean((q_pred[np.arange(N), actions] - targets)**2)

# Train DQN Agent on GridWorld
dqn_env = GridWorld()
q_net = NeuralQNetwork(state_dim=16, hidden_dim=32, action_dim=4, lr=0.05)
target_net = NeuralQNetwork(state_dim=16, hidden_dim=32, action_dim=4, lr=0.05)
target_net.copy_weights_from(q_net)

buffer = ReplayBuffer(capacity=2000)
batch_size = 32
gamma = 0.99
epsilon = 1.0
epsilon_decay = 0.995
epsilon_min = 0.05
step_count = 0
target_update_freq = 50

dqn_rewards = []
dqn_losses = []

for ep in range(300):
    s_raw = dqn_env.reset()
    s = state_to_onehot(s_raw)
    total_reward = 0
    done = False
    ep_steps = 0
    ep_loss = []
    while not done and ep_steps < 100:
        ep_steps += 1
        step_count += 1
        if rng.random() < epsilon:
            a = int(rng.integers(4))
        else:
            q_vals = q_net.forward(s.reshape(1, -1))
            a = np.argmax(q_vals[0])

        s_next_raw, r, done = dqn_env.step(a)
        s_next = state_to_onehot(s_next_raw)
        buffer.push(s, a, r, s_next, done)
        s = s_next
        total_reward += r

        if len(buffer.buffer) >= batch_size:
            b_s, b_a, b_r, b_s_next, b_d = buffer.sample(batch_size)
            target_q = target_net.forward(b_s_next)
            max_next_q = np.max(target_q, axis=1)
            y = b_r + gamma * max_next_q * (1.0 - b_d)
            loss = q_net.train_step(b_s, b_a, y)
            ep_loss.append(loss)

        if step_count % target_update_freq == 0:
            target_net.copy_weights_from(q_net)

    epsilon = max(epsilon_min, epsilon * epsilon_decay)
    dqn_rewards.append(total_reward)
    if ep_loss:
        dqn_losses.append(np.mean(ep_loss))

print(f"DQN Training Complete!")
print(f"Average Reward (last 50 episodes): {np.mean(dqn_rewards[-50:]):.2f}")
# Pure NumPy Implementation of Deep Q-Network (DQN)

def state_to_onehot(state, rows=4, cols=4):
    idx = state[0] * cols + state[1]
    onehot = np.zeros(rows * cols, dtype=np.float32)
    onehot[idx] = 1.0
    return onehot

class ReplayBuffer:
    def __init__(self, capacity=2000):
        self.capacity = capacity
        self.buffer = []
        self.ptr = 0

    def push(self, state, action, reward, next_state, done):
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)
        self.buffer[self.ptr] = (state, action, reward, next_state, done)
        self.ptr = (self.ptr + 1) % self.capacity

    def sample(self, batch_size):
        indices = rng.choice(len(self.buffer), batch_size, replace=False)
        states, actions, rewards, next_states, dones = zip(*[self.buffer[i] for i in indices])
        return (np.array(states), np.array(actions), np.array(rewards, dtype=np.float32),
                np.array(next_states), np.array(dones, dtype=np.float32))

class NeuralQNetwork:
    """2-Layer MLP Q-Function Approximator with ReLU activation."""
    def __init__(self, state_dim=16, hidden_dim=32, action_dim=4, lr=0.05):
        self.lr = lr
        self.W1 = rng.standard_normal((state_dim, hidden_dim)) * np.sqrt(2.0 / state_dim)
        self.b1 = np.zeros((1, hidden_dim))
        self.W2 = rng.standard_normal((hidden_dim, action_dim)) * np.sqrt(2.0 / hidden_dim)
        self.b2 = np.zeros((1, action_dim))

    def forward(self, x):
        self.z1 = np.dot(x, self.W1) + self.b1
        self.a1 = np.maximum(0, self.z1)
        self.q = np.dot(self.a1, self.W2) + self.b2
        return self.q

    def copy_weights_from(self, other):
        self.W1 = np.copy(other.W1)
        self.b1 = np.copy(other.b1)
        self.W2 = np.copy(other.W2)
        self.b2 = np.copy(other.b2)

    def train_step(self, x, actions, targets):
        N = x.shape[0]
        q_pred = self.forward(x)
        dq = np.zeros_like(q_pred)
        for i in range(N):
            dq[i, actions[i]] = (q_pred[i, actions[i]] - targets[i]) / N
        dW2 = np.dot(self.a1.T, dq)
        db2 = np.sum(dq, axis=0, keepdims=True)
        da1 = np.dot(dq, self.W2.T)
        dz1 = da1 * (self.z1 > 0)
        dW1 = np.dot(x.T, dz1)
        db1 = np.sum(dz1, axis=0, keepdims=True)
        self.W1 -= self.lr * dW1
        self.b1 -= self.lr * db1
        self.W2 -= self.lr * dW2
        self.b2 -= self.lr * db2
        return 0.5 * np.mean((q_pred[np.arange(N), actions] - targets)**2)

# Train DQN Agent on GridWorld
dqn_env = GridWorld()
q_net = NeuralQNetwork(state_dim=16, hidden_dim=32, action_dim=4, lr=0.05)
target_net = NeuralQNetwork(state_dim=16, hidden_dim=32, action_dim=4, lr=0.05)
target_net.copy_weights_from(q_net)

buffer = ReplayBuffer(capacity=2000)
batch_size = 32
gamma = 0.99
epsilon = 1.0
epsilon_decay = 0.995
epsilon_min = 0.05
step_count = 0
target_update_freq = 50

dqn_rewards = []
dqn_losses = []

for ep in range(300):
    s_raw = dqn_env.reset()
    s = state_to_onehot(s_raw)
    total_reward = 0
    done = False
    ep_steps = 0
    ep_loss = []
    while not done and ep_steps < 100:
        ep_steps += 1
        step_count += 1
        if rng.random() < epsilon:
            a = int(rng.integers(4))
        else:
            q_vals = q_net.forward(s.reshape(1, -1))
            a = np.argmax(q_vals[0])

        s_next_raw, r, done = dqn_env.step(a)
        s_next = state_to_onehot(s_next_raw)
        buffer.push(s, a, r, s_next, done)
        s = s_next
        total_reward += r

        if len(buffer.buffer) >= batch_size:
            b_s, b_a, b_r, b_s_next, b_d = buffer.sample(batch_size)
            target_q = target_net.forward(b_s_next)
            max_next_q = np.max(target_q, axis=1)
            y = b_r + gamma * max_next_q * (1.0 - b_d)
            loss = q_net.train_step(b_s, b_a, y)
            ep_loss.append(loss)

        if step_count % target_update_freq == 0:
            target_net.copy_weights_from(q_net)

    epsilon = max(epsilon_min, epsilon * epsilon_decay)
    dqn_rewards.append(total_reward)
    if ep_loss:
        dqn_losses.append(np.mean(ep_loss))

print(f"DQN Training Complete!")
print(f"Average Reward (last 50 episodes): {np.mean(dqn_rewards[-50:]):.2f}")

DQN Training Complete!
Average Reward (last 50 episodes): 2.92

5. Library Comparison¶

Since there is no standard tabular RL library in the Python ecosystem (unlike scikit-learn for supervised learning), we will validate our Q-learning implementation by comparing its converged Q-values against the precise, model-based Q-values computed via Value Iteration.

In [7]:

Copied!





# Run Value Iteration to get the ground-truth optimal Q-values
_, true_Q = value_iteration(env, gamma=0.99)

# 500 episodes leave rarely-visited state-action pairs far from Q*. Q-learning is
# off-policy, so a highly exploratory behavior policy still learns Q*: boost epsilon
# and keep training so every state-action pair keeps receiving updates. In this
# deterministic MDP the TD targets become exact, so Q converges tightly to Q*.
agent.epsilon = 0.5
for ep in range(4000):
    state = env.reset()
    done = False
    steps = 0
    while not done and steps < 200:
        steps += 1
        action = agent.get_action(state)
        next_state, reward, done = env.step(action)
        agent.update(state, action, reward, next_state, done)
        state = next_state
agent.epsilon = 0.1

# (a) Q-values must match Value Iteration within a tight tolerance.
# Empirically max_diff falls below 1e-4 with this schedule; 0.5 leaves a safety margin
# while still ruling out any unconverged state-action pair (initial gap is ~10).
differences = []
for s in env.get_all_states():
    if s == env.goal:
        continue
    differences.append(np.max(np.abs(true_Q[s] - agent.Q[s])))

max_diff = np.max(differences)
print(f"Maximum Q-value difference between Q-Learning and Value Iteration: {max_diff:.6f}")
assert max_diff < 0.5, f"Q-Learning has not converged to Value Iteration targets (max diff {max_diff:.4f})"

# (b) The greedy policy must agree with Value Iteration in every non-terminal reachable
# state. Compare against the *set* of optimal actions so exact ties cannot cause false failures.
for s in env.get_all_states():
    if s == env.goal:
        continue
    greedy_action = int(np.argmax(agent.Q[s]))
    assert np.isclose(true_Q[s][greedy_action], np.max(true_Q[s]), atol=1e-6), (
        f"Greedy action {greedy_action} at state {s} is suboptimal under Value Iteration"
    )
print("Greedy policy from Q-Learning matches the Value Iteration policy on all non-terminal states.")
# Run Value Iteration to get the ground-truth optimal Q-values
_, true_Q = value_iteration(env, gamma=0.99)

# 500 episodes leave rarely-visited state-action pairs far from Q*. Q-learning is
# off-policy, so a highly exploratory behavior policy still learns Q*: boost epsilon
# and keep training so every state-action pair keeps receiving updates. In this
# deterministic MDP the TD targets become exact, so Q converges tightly to Q*.
agent.epsilon = 0.5
for ep in range(4000):
    state = env.reset()
    done = False
    steps = 0
    while not done and steps < 200:
        steps += 1
        action = agent.get_action(state)
        next_state, reward, done = env.step(action)
        agent.update(state, action, reward, next_state, done)
        state = next_state
agent.epsilon = 0.1

# (a) Q-values must match Value Iteration within a tight tolerance.
# Empirically max_diff falls below 1e-4 with this schedule; 0.5 leaves a safety margin
# while still ruling out any unconverged state-action pair (initial gap is ~10).
differences = []
for s in env.get_all_states():
    if s == env.goal:
        continue
    differences.append(np.max(np.abs(true_Q[s] - agent.Q[s])))

max_diff = np.max(differences)
print(f"Maximum Q-value difference between Q-Learning and Value Iteration: {max_diff:.6f}")
assert max_diff < 0.5, f"Q-Learning has not converged to Value Iteration targets (max diff {max_diff:.4f})"

# (b) The greedy policy must agree with Value Iteration in every non-terminal reachable
# state. Compare against the *set* of optimal actions so exact ties cannot cause false failures.
for s in env.get_all_states():
    if s == env.goal:
        continue
    greedy_action = int(np.argmax(agent.Q[s]))
    assert np.isclose(true_Q[s][greedy_action], np.max(true_Q[s]), atol=1e-6), (
        f"Greedy action {greedy_action} at state {s} is suboptimal under Value Iteration"
    )
print("Greedy policy from Q-Learning matches the Value Iteration policy on all non-terminal states.")

Maximum Q-value difference between Q-Learning and Value Iteration: 0.000000
Greedy policy from Q-Learning matches the Value Iteration policy on all non-terminal states.

Reference comparison: no standard Python library ships tabular Q-learning, so the pinned, unit-tested implementation in ml_first_principles.rl_models is the reference we compare against, alongside the Value Iteration ground truth above.

In [8]:

Copied!





# Reference comparison: train the unit-tested package implementation on its GridWorldEnv
from ml_first_principles.rl_models import GridWorldEnv, QLearningAgent as RefQLearningAgent

ref_env = GridWorldEnv(grid_size=4, goal=(3, 3), trap=(1, 1))
ref_agent = RefQLearningAgent(num_states=16, num_actions=4, alpha=0.1, gamma=0.99,
                              epsilon=0.2, random_state=SEED)

for _ in range(2000):
    ref_state = ref_env.reset()  # step() raises RuntimeError after done, so reset every episode
    ref_done = False
    ref_steps = 0
    while not ref_done and ref_steps < 100:
        ref_steps += 1
        ref_action = ref_agent.select_action(ref_state)
        ref_next_state, ref_reward, ref_done, _ = ref_env.step(ref_action)
        ref_agent.update(ref_state, ref_action, ref_reward, ref_next_state, ref_done)
        ref_state = ref_next_state

# Greedy rollout: the learned policy must reach the goal (terminal reward +10; the trap gives -5)
ref_state = ref_env.reset()
ref_done, ref_reward = False, 0.0
for rollout_step in range(20):
    ref_action = int(np.argmax(ref_agent.q_table[ref_state]))
    ref_state, ref_reward, ref_done, _ = ref_env.step(ref_action)
    if ref_done:
        break

assert ref_done and ref_reward == 10.0, "Reference agent's greedy rollout failed to reach the goal"
print(f"Package QLearningAgent greedy rollout reached the goal in {rollout_step + 1} steps "
      f"(final reward {ref_reward:+.1f}).")
# Reference comparison: train the unit-tested package implementation on its GridWorldEnv
from ml_first_principles.rl_models import GridWorldEnv, QLearningAgent as RefQLearningAgent

ref_env = GridWorldEnv(grid_size=4, goal=(3, 3), trap=(1, 1))
ref_agent = RefQLearningAgent(num_states=16, num_actions=4, alpha=0.1, gamma=0.99,
                              epsilon=0.2, random_state=SEED)

for _ in range(2000):
    ref_state = ref_env.reset()  # step() raises RuntimeError after done, so reset every episode
    ref_done = False
    ref_steps = 0
    while not ref_done and ref_steps < 100:
        ref_steps += 1
        ref_action = ref_agent.select_action(ref_state)
        ref_next_state, ref_reward, ref_done, _ = ref_env.step(ref_action)
        ref_agent.update(ref_state, ref_action, ref_reward, ref_next_state, ref_done)
        ref_state = ref_next_state

# Greedy rollout: the learned policy must reach the goal (terminal reward +10; the trap gives -5)
ref_state = ref_env.reset()
ref_done, ref_reward = False, 0.0
for rollout_step in range(20):
    ref_action = int(np.argmax(ref_agent.q_table[ref_state]))
    ref_state, ref_reward, ref_done, _ = ref_env.step(ref_action)
    if ref_done:
        break

assert ref_done and ref_reward == 10.0, "Reference agent's greedy rollout failed to reach the goal"
print(f"Package QLearningAgent greedy rollout reached the goal in {rollout_step + 1} steps "
      f"(final reward {ref_reward:+.1f}).")

Package QLearningAgent greedy rollout reached the goal in 6 steps (final reward +10.0).

6. Experiments and Failures — VERIFY¶

6.1 Training Curve¶

In [9]:

Copied!





# Smooth the rewards history for better visualization
def moving_average(a, n=20):
    ret = np.cumsum(a, dtype=float)
    ret[n:] = ret[n:] - ret[:-n]
    return ret[n - 1:] / n

plt.figure(figsize=(8, 4))
plt.plot(moving_average(rewards_history))
plt.title("Q-Learning Training Curve (Smoothed)")
plt.xlabel("Episode")
plt.ylabel("Sum of Rewards")
plt.grid(True)
plt.show()
# Smooth the rewards history for better visualization
def moving_average(a, n=20):
    ret = np.cumsum(a, dtype=float)
    ret[n:] = ret[n:] - ret[:-n]
    return ret[n - 1:] / n

plt.figure(figsize=(8, 4))
plt.plot(moving_average(rewards_history))
plt.title("Q-Learning Training Curve (Smoothed)")
plt.xlabel("Episode")
plt.ylabel("Sum of Rewards")
plt.grid(True)
plt.show()

No description has been provided for this image

6.2 Visualize Policy and Q-values¶

In [10]:

Copied!





def plot_grid(Q, env, title):
    grid_v = np.zeros((env.rows, env.cols))
    policy = np.zeros((env.rows, env.cols), dtype=int)
    
    for r in range(env.rows):
        for c in range(env.cols):
            if (r, c) == env.goal:
                grid_v[r, c] = 10.0
                continue
            if (r, c) in env.obstacles:
                grid_v[r, c] = np.nan
                continue
            
            grid_v[r, c] = np.max(Q[(r, c)])
            policy[r, c] = np.argmax(Q[(r, c)])
            
    fig, ax = plt.subplots(figsize=(6,6))
    im = ax.imshow(grid_v, cmap='YlGn', vmin=-10, vmax=10)
    
    # Draw arrows
    arrow_map = {0: '↑', 1: '→', 2: '↓', 3: '←'}
    for r in range(env.rows):
        for c in range(env.cols):
            if (r, c) == env.goal:
                ax.text(c, r, 'G', ha='center', va='center', color='red', fontweight='bold', fontsize=16)
            elif (r, c) in env.obstacles:
                ax.text(c, r, 'X', ha='center', va='center', color='black', fontweight='bold', fontsize=16)
            else:
                ax.text(c, r, arrow_map[policy[r, c]], ha='center', va='center', fontsize=20)
                
    ax.set_xlabel("Grid column")
    ax.set_ylabel("Grid row")
    plt.title(title)
    plt.colorbar(im, label="State value $\\max_a Q(s, a)$")
    plt.show()

plot_grid(true_Q, env, "Value Iteration Optimal Policy & Values")
plot_grid(agent.Q, env, "Q-Learning Learned Policy & Values")
def plot_grid(Q, env, title):
    grid_v = np.zeros((env.rows, env.cols))
    policy = np.zeros((env.rows, env.cols), dtype=int)
    
    for r in range(env.rows):
        for c in range(env.cols):
            if (r, c) == env.goal:
                grid_v[r, c] = 10.0
                continue
            if (r, c) in env.obstacles:
                grid_v[r, c] = np.nan
                continue
            
            grid_v[r, c] = np.max(Q[(r, c)])
            policy[r, c] = np.argmax(Q[(r, c)])
            
    fig, ax = plt.subplots(figsize=(6,6))
    im = ax.imshow(grid_v, cmap='YlGn', vmin=-10, vmax=10)
    
    # Draw arrows
    arrow_map = {0: '↑', 1: '→', 2: '↓', 3: '←'}
    for r in range(env.rows):
        for c in range(env.cols):
            if (r, c) == env.goal:
                ax.text(c, r, 'G', ha='center', va='center', color='red', fontweight='bold', fontsize=16)
            elif (r, c) in env.obstacles:
                ax.text(c, r, 'X', ha='center', va='center', color='black', fontweight='bold', fontsize=16)
            else:
                ax.text(c, r, arrow_map[policy[r, c]], ha='center', va='center', fontsize=20)
                
    ax.set_xlabel("Grid column")
    ax.set_ylabel("Grid row")
    plt.title(title)
    plt.colorbar(im, label="State value $\\max_a Q(s, a)$")
    plt.show()

plot_grid(true_Q, env, "Value Iteration Optimal Policy & Values")
plot_grid(agent.Q, env, "Q-Learning Learned Policy & Values")

6.3 Exploration-Exploitation Tradeoff (Varying Epsilon)¶

In [11]:

Copied!





epsilons = [0.0, 0.1, 0.5, 1.0]
plt.figure(figsize=(10, 5))

for eps in epsilons:
    test_env = GridWorld()
    test_agent = QLearningAgent(test_env, alpha=0.1, gamma=0.99, epsilon=eps)
    rewards = []
    for ep in range(300):
        s = test_env.reset()
        tot = 0
        while True:
            a = test_agent.get_action(s)
            ns, r, d = test_env.step(a)
            test_agent.update(s, a, r, ns, d)
            s = ns
            tot += r
            if d or tot < -100: # cap episode length
                break
        rewards.append(tot)
    plt.plot(moving_average(rewards, n=30), label=f"eps={eps}")

plt.title("Impact of Epsilon on Learning Curve")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.legend()
plt.grid(True)
plt.show()
epsilons = [0.0, 0.1, 0.5, 1.0]
plt.figure(figsize=(10, 5))

for eps in epsilons:
    test_env = GridWorld()
    test_agent = QLearningAgent(test_env, alpha=0.1, gamma=0.99, epsilon=eps)
    rewards = []
    for ep in range(300):
        s = test_env.reset()
        tot = 0
        while True:
            a = test_agent.get_action(s)
            ns, r, d = test_env.step(a)
            test_agent.update(s, a, r, ns, d)
            s = ns
            tot += r
            if d or tot < -100: # cap episode length
                break
        rewards.append(tot)
    plt.plot(moving_average(rewards, n=30), label=f"eps={eps}")

plt.title("Impact of Epsilon on Learning Curve")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.legend()
plt.grid(True)
plt.show()

Observation:

eps=0.0: The agent exploits immediately. In a gridworld, it often gets stuck in a loop and never finds the goal (hence highly negative rewards early on, and it may never learn).
eps=1.0: The agent acts completely randomly. It finds the goal eventually but takes highly sub-optimal paths, resulting in poor average return.
eps=0.1: The agent balances exploration (to find the goal consistently) and exploitation (to use the shortest path), converging to the optimal return.

6.4 Failure Case: Divergence with Too-High Learning Rate¶

If the learning rate $\alpha$ is too high, TD updates overshoot and the Q-values become unstable.

In [12]:

Copied!





fail_env = GridWorld()
# An artificially high learning rate > 1 breaks the stability of the moving average update.
fail_agent = QLearningAgent(fail_env, alpha=1.5, gamma=0.99, epsilon=0.1)
fail_rewards = []

for ep in range(100):
    s = fail_env.reset()
    tot = 0
    for _ in range(50): # limit steps to prevent infinite divergence hanging
        a = fail_agent.get_action(s)
        ns, r, d = fail_env.step(a)
        fail_agent.update(s, a, r, ns, d)
        s = ns
        tot += r
        if d:
            break
    fail_rewards.append(tot)

plt.figure(figsize=(6, 4))
plt.plot(fail_rewards)
plt.title("Failure Case: α=1.5 causes instability")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.show()

# Check a Q-value to see if it exploded or became wildly negative
sample_q = fail_agent.Q[(0,0)]
print(f"Q-values at start state with high LR: {sample_q}")
fail_env = GridWorld()
# An artificially high learning rate > 1 breaks the stability of the moving average update.
fail_agent = QLearningAgent(fail_env, alpha=1.5, gamma=0.99, epsilon=0.1)
fail_rewards = []

for ep in range(100):
    s = fail_env.reset()
    tot = 0
    for _ in range(50): # limit steps to prevent infinite divergence hanging
        a = fail_agent.get_action(s)
        ns, r, d = fail_env.step(a)
        fail_agent.update(s, a, r, ns, d)
        s = ns
        tot += r
        if d:
            break
    fail_rewards.append(tot)

plt.figure(figsize=(6, 4))
plt.plot(fail_rewards)
plt.title("Failure Case: α=1.5 causes instability")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.show()

# Check a Q-value to see if it exploded or became wildly negative
sample_q = fail_agent.Q[(0,0)]
print(f"Q-values at start state with high LR: {sample_q}")

Q-values at start state with high LR: [1.41973029e+08 1.42071044e+08 1.42779059e+08 1.42468072e+08]

6.5 Failure Case: The Deadly Triad (Instability without Replay Buffer & Target Network)¶

Sutton & Barto define the Deadly Triad as the combination of three elements:

Function Approximation (e.g. Neural Networks)
Bootstrapping (Temporal Difference updates using $\max_{a'} Q(s', a')$)
Off-policy Training (Learning optimal policy while exploring)

When these three elements are present without stabilization mechanisms (Replay Buffer and Target Network), Q-learning is prone to severe instability, oscillations, and Q-value explosion.

Below, we train a Neural Q-Network without Replay Buffer (updating on sequential correlated steps) and without Target Network (bootstrapping directly on the fast-moving network).

In [13]:

Copied!





# Experimental Demonstration of the Deadly Triad Instability

deadly_env = GridWorld()
unstable_net = NeuralQNetwork(state_dim=16, hidden_dim=32, action_dim=4, lr=0.1)

online_q_history = []
step_cnt = 0
s_raw = deadly_env.reset()
s = state_to_onehot(s_raw)

for step in range(500):
    step_cnt += 1
    if rng.random() < 0.1:
        a = int(rng.integers(4))
    else:
        q_vals = unstable_net.forward(s.reshape(1, -1))
        a = np.argmax(q_vals[0])

    s_next_raw, r, done = deadly_env.step(a)
    s_next = state_to_onehot(s_next_raw)
    
    # Direct TD bootstrap update on current network (No Target Net, No Buffer)
    next_q = unstable_net.forward(s_next.reshape(1, -1))
    max_next_q = np.max(next_q[0])
    target = r + 0.99 * max_next_q * (1.0 - float(done))
    
    unstable_net.train_step(s.reshape(1, -1), np.array([a]), np.array([target]))
    
    # Record max Q value of start state
    q_start = unstable_net.forward(state_to_onehot((0, 0)).reshape(1, -1))
    online_q_history.append(np.max(q_start[0]))
    
    if done:
        s_raw = deadly_env.reset()
        s = state_to_onehot(s_raw)
    else:
        s = s_next

plt.figure(figsize=(10, 4))
plt.plot(online_q_history, label="Max Q(s_start) without Replay/Target Net (Deadly Triad)", color="crimson")
plt.axhline(y=10.0, color="gray", linestyle="--", label="Ground Truth Optimal Q Value (~10)")
plt.xlabel("Step")
plt.ylabel("Estimated Q-value")
plt.title("Failure Case: Q-Value Divergence / Instability under Deadly Triad")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()
# Experimental Demonstration of the Deadly Triad Instability

deadly_env = GridWorld()
unstable_net = NeuralQNetwork(state_dim=16, hidden_dim=32, action_dim=4, lr=0.1)

online_q_history = []
step_cnt = 0
s_raw = deadly_env.reset()
s = state_to_onehot(s_raw)

for step in range(500):
    step_cnt += 1
    if rng.random() < 0.1:
        a = int(rng.integers(4))
    else:
        q_vals = unstable_net.forward(s.reshape(1, -1))
        a = np.argmax(q_vals[0])

    s_next_raw, r, done = deadly_env.step(a)
    s_next = state_to_onehot(s_next_raw)
    
    # Direct TD bootstrap update on current network (No Target Net, No Buffer)
    next_q = unstable_net.forward(s_next.reshape(1, -1))
    max_next_q = np.max(next_q[0])
    target = r + 0.99 * max_next_q * (1.0 - float(done))
    
    unstable_net.train_step(s.reshape(1, -1), np.array([a]), np.array([target]))
    
    # Record max Q value of start state
    q_start = unstable_net.forward(state_to_onehot((0, 0)).reshape(1, -1))
    online_q_history.append(np.max(q_start[0]))
    
    if done:
        s_raw = deadly_env.reset()
        s = state_to_onehot(s_raw)
    else:
        s = s_next

plt.figure(figsize=(10, 4))
plt.plot(online_q_history, label="Max Q(s_start) without Replay/Target Net (Deadly Triad)", color="crimson")
plt.axhline(y=10.0, color="gray", linestyle="--", label="Ground Truth Optimal Q Value (~10)")
plt.xlabel("Step")
plt.ylabel("Estimated Q-value")
plt.title("Failure Case: Q-Value Divergence / Instability under Deadly Triad")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

7. Connections & Takeaways¶

Tabular vs Deep RL: Tabular methods (Value Iteration, tabular Q-learning) work exceptionally well for small discrete MDPs with known/finite states, but scale poorly. Deep Q-Networks (DQN) leverage function approximation to scale to large state spaces.
Stabilization Mechanisms: Deep RL requires an Experience Replay Buffer (to break sample correlation) and a Target Network (to freeze the moving TD target $Q_{\theta^-}$) to prevent the Deadly Triad divergence.
Model-Based vs Model-Free: Dynamic Programming requires full transition dynamics $P(s'|s,a)$ and $R(s,a,s')$, while Q-learning and DQN learn purely from interaction samples.
Next Steps: Proceed to continuous control algorithms (Actor-Critic, PPO) and policy-search paradigms.