Q-Learning
Q-Learning is a model-free, off-policy algorithm that learns the optimal action-value function $Q^*(s,a)$ directly from experience. It stores Q-values in a table (one row per state, one column per action) and updates them using the Bellman equation as a learning target.
Q-Learning on FrozenLake
import numpy as np
import gymnasium as gym
import matplotlib.pyplot as plt
# FrozenLake: navigate a 4x4 grid, avoid holes, reach the goal
# States: 16 (4x4 positions), Actions: 4 (Left, Down, Right, Up)
env = gym.make("FrozenLake-v1", is_slippery=False) # deterministic for learning
n_states = env.observation_space.n # 16
n_actions = env.action_space.n # 4
# ── Hyperparameters
LR = 0.8 # learning rate alpha
GAMMA = 0.95 # discount factor
EPS_START = 1.0 # initial exploration rate
EPS_END = 0.01
EPS_DECAY = 0.995
N_EPISODES = 2000
Q = np.zeros((n_states, n_actions)) # Q-table: (16, 4) initialised to 0
eps = EPS_START
ep_rewards, ep_lengths = [], []
for ep in range(N_EPISODES):
state, _ = env.reset()
total_r = 0
for step in range(200):
# ── Action selection (epsilon-greedy)
if np.random.random() < eps:
action = env.action_space.sample()
else:
action = np.argmax(Q[state])
# ── Step environment
next_state, reward, terminated, truncated, _ = env.step(action)
# ── Q-Learning update (Bellman target)
td_target = reward + GAMMA * np.max(Q[next_state]) * (not terminated)
td_error = td_target - Q[state, action]
Q[state, action] += LR * td_error
state = next_state
total_r += reward
if terminated or truncated:
break
eps = max(EPS_END, eps * EPS_DECAY)
ep_rewards.append(total_r)
ep_lengths.append(step + 1)
# ── Evaluate final policy
wins = sum(ep_rewards[-200:]) / 200
print(f"Win rate (last 200 eps): {wins:.1%}")
print(f"Final epsilon: {eps:.4f}")
print("
Learned Q-table (max Q per state):")
print(Q.max(axis=1).reshape(4,4).round(3))
# ── Plot learning curve (rolling average)
window = 50
rolling = np.convolve(ep_rewards, np.ones(window)/window, mode="valid")
plt.figure(figsize=(10, 4))
plt.plot(ep_rewards, alpha=0.3, color="#4a8fa8", lw=0.8)
plt.plot(rolling, color="#4a8fa8", lw=2, label=f"{window}-ep rolling avg")
plt.xlabel("Episode"); plt.ylabel("Reward (0=fail, 1=success)")
plt.title("Q-Learning on FrozenLake-v1 (is_slippery=False)")
plt.legend(); plt.grid(alpha=0.3); plt.tight_layout(); plt.show()Limits of Tabular Q-Learning
A Q-table works when the state space is small and discrete. In most real problems — Atari games (thousands of pixels = millions of states), robotics (continuous joint angles) — a table is completely infeasible. The solution: approximate $Q(s,a)$ with a neural network.
Deep Q-Network (DQN)
DQN (Mnih et al., 2015) was the first algorithm to achieve human-level performance on Atari games from raw pixels. It replaces the Q-table with a neural network $Q(s,a;\theta)$ and introduces two key tricks to stabilise training:
import numpy as np
import gymnasium as gym
import tensorflow as tf
from tensorflow import keras
from collections import deque
import random, time
# ── Replay Buffer
class ReplayBuffer:
def __init__(self, capacity=10_000):
self.buf = deque(maxlen=capacity)
def push(self, state, action, reward, next_state, done):
self.buf.append((state, action, reward, next_state, done))
def sample(self, batch_size):
batch = random.sample(self.buf, batch_size)
s, a, r, ns, d = zip(*batch)
return (np.array(s, dtype=np.float32),
np.array(a),
np.array(r, dtype=np.float32),
np.array(ns, dtype=np.float32),
np.array(d, dtype=np.float32))
def __len__(self): return len(self.buf)
# ── DQN Network
def build_dqn(n_states, n_actions, hidden=(128, 64)):
return keras.Sequential([
keras.layers.Input(shape=(n_states,)),
keras.layers.Dense(hidden[0], activation="relu"),
keras.layers.Dense(hidden[1], activation="relu"),
keras.layers.Dense(n_actions), # one Q-value per action, no activation
])
# ── DQN Agent
class DQNAgent:
def __init__(self, n_states, n_actions,
lr=1e-3, gamma=0.99, eps=1.0, eps_min=0.01, eps_decay=0.995,
batch_size=64, target_update=100, buffer_size=10_000):
self.n_actions = n_actions
self.gamma = gamma
self.eps = eps
self.eps_min = eps_min
self.eps_decay = eps_decay
self.batch_size = batch_size
self.target_update= target_update
self.step_count = 0
self.buffer = ReplayBuffer(buffer_size)
self.online = build_dqn(n_states, n_actions)
self.target = build_dqn(n_states, n_actions)
self.target.set_weights(self.online.get_weights())
self.optimizer = keras.optimizers.Adam(lr)
self.loss_fn = keras.losses.MeanSquaredError()
def act(self, state):
if np.random.random() < self.eps:
return np.random.randint(self.n_actions)
q = self.online(state[np.newaxis], training=False)[0].numpy()
return np.argmax(q)
def train_step(self):
if len(self.buffer) < self.batch_size:
return None
s, a, r, ns, done = self.buffer.sample(self.batch_size)
# Compute Bellman targets using target network
next_q = self.target(ns, training=False).numpy()
targets = r + self.gamma * np.max(next_q, axis=1) * (1 - done)
with tf.GradientTape() as tape:
q_values = self.online(s, training=True)
# Only update Q for the taken action
action_mask = tf.one_hot(a, self.n_actions)
q_taken = tf.reduce_sum(q_values * action_mask, axis=1)
loss = self.loss_fn(targets, q_taken)
grads = tape.gradient(loss, self.online.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.online.trainable_variables))
self.step_count += 1
if self.step_count % self.target_update == 0:
self.target.set_weights(self.online.get_weights()) # sync target
self.eps = max(self.eps_min, self.eps * self.eps_decay)
return loss.numpy()
# ── Training loop on CartPole
env = gym.make("CartPole-v1")
n_states = env.observation_space.shape[0] # 4
n_actions = env.action_space.n # 2
agent = DQNAgent(n_states, n_actions)
rewards_history = []
t0 = time.time()
for episode in range(500):
state, _ = env.reset()
ep_reward = 0
for _ in range(500):
action = agent.act(state)
next_state, r, terminated, truncated, _ = env.step(action)
done = terminated or truncated
agent.buffer.push(state, action, r, next_state, float(done))
agent.train_step()
state = next_state
ep_reward += r
if done: break
rewards_history.append(ep_reward)
if (episode + 1) % 50 == 0:
avg = np.mean(rewards_history[-50:])
print(f"Ep {episode+1:4d} Avg reward: {avg:7.1f} eps: {agent.eps:.3f}")
print(f"
Training took {time.time()-t0:.1f}s")
# ── Plot
window = 20
rolling = np.convolve(rewards_history, np.ones(window)/window, mode="valid")
plt.figure(figsize=(10, 4))
plt.plot(rewards_history, alpha=0.3, color="#5c8a58", lw=0.8)
plt.plot(rolling, color="#5c8a58", lw=2, label=f"{window}-ep rolling avg")
plt.axhline(195, color="red", ls="--", label="Solved threshold (195)")
plt.xlabel("Episode"); plt.ylabel("Total Reward")
plt.title("DQN on CartPole-v1"); plt.legend(); plt.grid(alpha=0.3)
plt.tight_layout(); plt.show()DQN Improvements
Using Stable-Baselines3
For production RL, use Stable-Baselines3 — battle-tested implementations of DQN, PPO, SAC, and more:
# pip install stable-baselines3
from stable_baselines3 import DQN, PPO
from stable_baselines3.common.evaluation import evaluate_policy
import gymnasium as gym
env = gym.make("CartPole-v1")
# ── DQN from Stable-Baselines3
dqn = DQN(
"MlpPolicy", env,
learning_rate=1e-3,
buffer_size=10_000,
learning_starts=1000,
batch_size=64,
tau=1.0,
gamma=0.99,
target_update_interval=100,
verbose=1,
)
dqn.learn(total_timesteps=50_000)
mean_r, std_r = evaluate_policy(dqn, env, n_eval_episodes=20)
print(f"DQN: mean reward = {mean_r:.1f} ± {std_r:.1f}")
# ── PPO (better for continuous action spaces)
env_cont = gym.make("LunarLander-v2")
ppo = PPO("MlpPolicy", env_cont, verbose=1)
ppo.learn(total_timesteps=100_000)
mean_r, _ = evaluate_policy(ppo, env_cont, n_eval_episodes=10)
print(f"PPO LunarLander: {mean_r:.1f}")Summary
- Q-Learning: tabular TD method. Update $Q(s,a)$ towards $r + \gamma\max_{a'}Q(s',a')$. Converges to $Q^*$ for small discrete spaces.
- DQN: neural net approximator for $Q^*(s,a)$. Two key tricks: experience replay + target network.
- Double DQN, Dueling DQN, Prioritised Replay all improve stability and performance.
- Use Stable-Baselines3 for production RL — don't re-implement from scratch.