Week 3 Deep Q-Network (DQN)

[Key objectives]

• Understand and implement the core components of DQN (Deep Q-Network)
• Understand experience replay (replay buffer), the ε-greedy policy, and the target network

[DQN algorithm overview]

• DQN combines Q-learning and neural networks to perform reinforcement learning in complex environments
• Approximate the Q-function with a Q-network

$$\begin{array}{c} \boxed{\text{Input: State } (s)} \\ \downarrow \\ \boxed{\text{Layer 1: Fully Connected (Linear), } \phi = W_1 s + b_1} \\ \downarrow \\ \boxed{\text{Activation Function: ReLU, } a = \max(0, \phi)} \\ \downarrow \\ \boxed{\text{Layer 2: Fully Connected (Linear), } v = W_2 a + b_2} \\ \downarrow \\ \boxed{\text{Output: State Value } (V(s))} \end{array}$$

• Main DQN components:
  • Experience replay (replay buffer)
  • ε-greedy (epsilon-greedy) policy
  • Target network

[Experience replay (replay buffer)]

• Store past experiences in memory and sample them randomly to reduce correlations between training data and improve training stability
• Transition: Store experience in the form $(s_t, a_t, r_t, s_{t+1}, done)$
  • $s_t$: Current state
  • $a_t$: Action
  • $r_t$: Reward
  • $s_{t+1}$: Next state
  • $done$: Whether the episode has ended
• Sampling: randomly form minibatches from stored experiences for training

[ε-greedy (epsilon-greedy) policy]

• Balance exploration and exploitation
• Choose a random action with probability ε and the action with the highest current Q-value with probability 1 - ε.
• Gradually decrease ε as training progresses to reduce exploration and increase exploitation.

[Target network]

• Improve the stability of Q-value updates
• Use separate main and target networks
  • Main network: Used to predict and update Q-values
  • Target network: Used to compute Q-values, with parameters copied from the main network at fixed intervals
• Fixing the target network reduces TD-error variability to stabilize training

[DQN training update]

• Loss function: Train to minimize TD (temporal difference) error.
• TD target: $r + \gamma \max_{a'} Q(s', a'; \theta^-)$
  • $r$: Immediate reward
  • $\gamma$: Discount factor
  • $s'$: Next state
  • $a'$: Optimal action selected using the target network in the next state
  • $\theta^-$: Target network parameters
• TD error: $TD ; Error = Q(s, a; \theta) - (r + \gamma \max_{a'} Q(s', a'; \theta^-))$
  • $\theta$: Parameters of the main network
• Loss Function: $Loss = \mathbb{E}[(TD ; Error)^2]$
• Update: adjust main network parameters to minimize the loss function

[Action selection with an ε-greedy policy]

• Choose a random action with probability $\epsilon$
• Choose the action with the highest current Q-value with probability $1 - \epsilon$

[Exercise: implement DQN (Deep Q-Network) (CartPole)]

1. Environment setup

• gym: Provides the CartPole environment
• torch: Deep learning models and computations
• numpy: Array and numerical operations
• collections.deque: implement an experience replay buffer

import gymnasium as gym
import random
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque

env = gym.make('CartPole-v1', render_mode='rgb_array')
state_size = env.observation_space.shape[0]
action_size = env.action_space.n

2. Define the DQN model

• DQNAgent class: Define the neural network model.
• Use 3 fully connected layers (nn.Linear).
• Apply the ReLU activation function (torch.relu)
• Input: state; output: Q-values for each action

class DQNAgent(nn.Module):
    def __init__(self, state_size, action_size):
        super(DQNAgent, self).__init__()
        self.fc1 = nn.Linear(state_size, 128)  # Input layer
        self.fc2 = nn.Linear(128, 128)  # Hidden layer
        self.fc3 = nn.Linear(128, action_size)  # Output layer

    def forward(self, x):
        x = torch.relu(self.fc1(x))  # Apply the activation function (ReLU)
        x = torch.relu(self.fc2(x))
        return self.fc3(x)  # Return Q-values

3. Set hyperparameters

• Set the learning rate (learning_rate), discount factor (gamma), exploration probability (epsilon), etc.
• Set the minibatch size (batch_size) and memory size (memory_size).
• Set the target network update interval (target_update_frequency)

learning_rate = 0.0003
gamma = 0.99  # Discount factor
epsilon = 1.0  # Initial exploration probability
epsilon_min = 0.01
epsilon_decay = 0.995
batch_size = 128
memory_size = 30000
target_update_frequency = 5  # Target network update interval

4. Create the DQN agent and target network

• Create DQNAgent model instances (agent, target_agent)
• Initialize the target network (copy the agent network's weights).
• Set the Adam optimizer and MSE loss function.

agent = DQNAgent(state_size, action_size)
target_agent = DQNAgent(state_size, action_size)
target_agent.load_state_dict(agent.state_dict()) # Copy Q-network parameters
optimizer = optim.Adam(agent.parameters(), lr=learning_rate)
criterion = nn.MSELoss()

5. Initialize the experience replay buffer

• Initialize the experience replay buffer (memory) using deque.

memory = deque(maxlen=memory_size)

6. Training function

• Sample a minibatch from the experience replay buffer.
• Extract state, action, reward, next state, and termination status
• Use the main network to select an action and the target network to evaluate its Q-value
• Calculate the loss function by comparing predicted and target Q-values
• Backpropagation and weight updates

def train(batch_size):
  # Sample a minibatch from the experience replay buffer
  minibatch = random.sample(memory, batch_size)

  # Extract states, actions, rewards, next states, and episode end indicators from the minibatch
  states = torch.tensor(np.array([transition[0] for transition in minibatch]), dtype=torch.float32)
  actions = torch.tensor(np.array([transition[1] for transition in minibatch]), dtype=torch.long)
  rewards = torch.tensor(np.array([transition[2] for transition in minibatch]), dtype=torch.float32)
  next_states = torch.tensor(np.array([transition[3] for transition in minibatch]), dtype=torch.float32)
  dones = torch.tensor(np.array([transition[4] for transition in minibatch]), dtype=torch.float32)

  # Compute target Q-values (using the Target Network)
  target_q_values = target_agent(next_states).max(1)[0].detach() # Obtain the maximum Q-value from the target network
  targets = rewards + (gamma * target_q_values * (1 - dones)) # If done = 1, set the next-state Q-value to 0

  # Compute predicted Q-values
  predicted_q_values = agent(states).gather(1, actions.unsqueeze(1)) # Obtain Q-values computed by the trained network

  # Compute the loss and backpropagate
  loss = criterion(predicted_q_values.squeeze(), targets)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()

6. Training loop

• Repeat episodes episodes times.
• Reset the environment and initialize the state for each episode
• Choose an action with an ε-greedy policy.
• Execute an action in the environment and obtain the next state, reward, and termination status.
• Add experiences to the replay buffer
• Call the training function when the buffer size reaches batch_size.
• Update the target network and decrease exploration probability.
• Print episode results

# Training loop
episodes = 100
for e in range(episodes):
    state, _ = env.reset()
    state = torch.tensor(state, dtype=torch.float32)
    done = False
    total_reward = 0

    while not done:
        # Select an action with an ε-greedy policy
        if random.random() < epsilon:
            action = env.action_space.sample()
        else:
            action = agent(state).argmax().item()
        # Execute an action in the environment and obtain the next state, reward, and episode end status
        next_state, reward, terminated, truncated, _ = env.step(action)
        next_state = torch.tensor(next_state, dtype=torch.float32)
        total_reward += reward
        done = terminated or truncated

        memory.append((state, action, reward, next_state, done))

        if len(memory) >= batch_size:
            train(batch_size)

        state = next_state

        # Update the target network
        if e % target_update_frequency == 0:
            target_agent.load_state_dict(agent.state_dict())

        # Decrease exploration probability
        if epsilon > epsilon_min:
            epsilon *= epsilon_decay

    # Print intermediate results
    if (e+1) % 20 == 0:
        print(f"Episode {e + 1}/{episodes}, Total Reward: {total_reward}, Epsilon: {epsilon:.3f}")

Episode 20/100, Total Reward: 11.0, Epsilon: 0.220
Episode 40/100, Total Reward: 24.0, Epsilon: 0.067
Episode 60/100, Total Reward: 118.0, Epsilon: 0.010
Episode 80/100, Total Reward: 158.0, Epsilon: 0.010
Episode 100/100, Total Reward: 199.0, Epsilon: 0.010

7. Evaluate the trained agent (rendering with GIF)

import imageio
import os

frames = []
state, _ = env.reset()
state = torch.tensor(state, dtype=torch.float32)
done = False
for i in range(3000): # 3000 steps simulation
  action = agent(state).argmax().item()
  next_state, reward, terminated, truncated, _ = env.step(action)
  next_state = torch.tensor(next_state, dtype=torch.float32)
  state = next_state
  done = terminated or truncated
  frames.append(env.render())

  if done:
    break

imageio.mimsave('cartpole_simulation.gif', frames, duration=33) # save to gif file
print("GIF saved as cartpole_simulation.gif")

GIF saved as cartpole_simulation.gif