Week 5 Dueling DQN

English translation

[Key objectives]

Understand Dueling DQN and improve learning efficiency by modifying the Q-function architecture.
Understand the roles of value and advantage streams and implement Dueling DQN

[Dueling DQN idea]

Separate the Q-function into value and advantage streams
- Value stream ( $V(s)$ ): Predict state value, indicating how good or bad state $s$ is.
- Advantage stream ( $A(s,a)$ ): Predict each action's advantage, indicating how good choosing action $a$ in state $s$ is relative to the average action.
Q-function: Combine the value and advantage streams to compute Q-values.
Equation: $Q(s,a) = V(s) + A(s,a)$

[Dueling DQN advantages]

Learn state value and action advantage separately to improve Q-value estimation accuracy.
Value Stream: learn general information about the state to efficiently estimate Q-values for various actions
Advantage Stream: learn how much better a particular action is than others, making small differences between actions easier to distinguish

[Dueling DQN network]

Network architecture

$\begin{array}{c} \boxed{\text{State }(s)} \\ \downarrow \\ \boxed{\text{Feature Extraction }} \\ \swarrow \quad \searrow \quad \quad \quad \\ \boxed{\begin{array}{c} \text{Value Stream} \\ \downarrow \\ V(s) \end{array}} \quad \boxed{\begin{array}{c} \text{Advantage Stream} \\ \downarrow \\ A(s, a_1), \dots, A(s, a_N) \end{array}} \\ \searrow \quad \swarrow \quad \quad \quad \\ \boxed{\text{Aggregation: } Q(s, a) = V(s) + \left(A(s, a) - \frac{1}{|\mathbb{A}|}\sum_{a'} A(s, a')\right)} \\ \downarrow \\ \boxed{Q(s, a_1), Q(s, a_2), \dots, Q(s, a_N)} \end{array}$
Components
1. State (s) (input):
  - Environment state information (e.g., game screen or sensor data)
2. Feature Extraction (Shared Layers):
  - Types: Convolutional layers (Conv, images), fully connected layers (FC, vectors)
  - Role: Extract useful features from input state $s$ $\rightarrow$ .
3. Value Stream (Separate Layers):
  - Types: Mainly fully connected layers (FC)
  - Role: Estimate state value ( $V(s)$ ).
  - Output: A single value (scalar)
4. Advantage Stream (Separate Layers):
  - Types: Mainly fully connected layers (FC)
  - Role: Estimate each action's advantage ( $A(s, a)$ ).
  - Output: One value per action
5. Aggregation:
  - Formula: $Q(s, a) = V(s) + \left( A(s, a) - \frac{1}{|\mathbb{A}|}\sum_{a'} A(s, a') \right)$
  - Subtract the mean advantage for stability.
  - $|\mathbb{A}|$ : Number of possible actions
  - Role: $V(s) + A(s, a) =$ Final Q-value

[Dueling DQN training process]

1. Compute Q-values

Current Q-value
- Use the main network ( $\theta$ ) to calculate Q-values for the current state
- $Q(s, a; \theta)$
Target Q-value
- Use the target network ( $\theta^-$ ) to calculate the maximum Q-value at the next state
- $y = r + \gamma \underset{a'}{max} \text{ } Q(s', a'; \theta^-)$
  - $\gamma$ : discount factor

2. Calculate the loss function

MSE loss
- Compute the MSE loss between the current Q-value and the target Q-value.
- $L(\theta) = \mathbb{E}[(Q(s, a; \theta) - y)^2]$

3. Update the networks

Backpropagation
- Use the computed loss $L(\theta)$ to update the main network parameters $\theta$ .
- $\theta \leftarrow \theta - \alpha \nabla_{\theta} L(\theta)$
  - $\alpha$ : Learning rate

4. Update the target network

Update the target network by copying parameters from the main network at fixed step intervals.
$\theta^- \leftarrow \theta$

[Exercise: CartPole with Dueling DQN]

1. Install and import libraries

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 torch
import torch.nn as nn
import torch.optim as optim
import random
from collections import deque
import numpy as np
import imageio

2. Configure the GPU and define hyperparameters

BATCH_SIZE: Number of samples per training update
GAMMA: discounting of future rewards (0~1)
EPS_START: Initial probability of a random action
EPS_END: Final probability of a random action
EPS_DECAY: decay rate of random action probability
TARGET_UPDATE: target network update interval
MEMORY_SIZE: Replay buffer capacity
LR: weight update size (learning rate)
NUM_EPISODES: total number of training episodes
RENDER_EPISODE: Statistics output interval

# Check GPU availability and set the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Hyperparameters
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 1.0
EPS_END = 0.01
EPS_DECAY = 0.995
TARGET_UPDATE = 5
MEMORY_SIZE = 30000
LR = 0.0003
NUM_EPISODES = 2000
PROGRESS_INTERVAL = 50

3. Define the Dueling DQN network

DuelingDQN class: Build a network by inheriting from PyTorch's nn.Module.
feature_layer: shared feature extraction layer that processes the input state.
value_stream: layer that estimates the state value (V).
advantage_stream: Layer estimating the relative advantage (A) of each action.
forward method: combine value and advantage to calculate the final Q-value.
(Q-value calculation): $Q(s, a) = V(s) + \left(A(s, a) - \frac{1}{|\mathbb{A}|}\sum_{a'} A(s, a')\right)$
- $Q(s, a)$ : Q-value in state $s$ for choosing action $a$
- $V(s)$ : value of state $s$
- $A(s, a)$ : advantage in state $s$ of action $a$
- $|\mathbb{A}|$ : action space size (number of possible actions)

class DuelingDQN(nn.Module):
    def __init__(self, state_size, action_size, hidden_size):
        super(DuelingDQN, self).__init__()
        self.feature_layer = nn.Sequential(
            nn.Linear(state_size, hidden_size),
            nn.ReLU()
        )
        self.value_stream = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )
        self.advantage_stream = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, action_size)
        )

    def forward(self, state):
        x = self.feature_layer(state)
        value = self.value_stream(x)
        advantage = self.advantage_stream(x)
        q_values = value + (advantage - advantage.mean(dim=1, keepdim=True))
        return q_values

4. Define the replay buffer

ReplayBuffer class: Create a fixed-size buffer using deque.
push method: Add a new experience (state, action, reward, next_state, done) to the buffer.
sample method: Randomly sample batch_size experiences from the buffer and convert them to PyTorch tensors (accounting for GPU use).
__len__ method: return the number of experiences stored in the buffer

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        state = np.array(state)
        next_state = np.array(next_state)
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        state, action, reward, next_state, done = zip(*random.sample(self.buffer, batch_size))
        return torch.tensor(np.array(state), dtype=torch.float32, device=device), \
           torch.tensor(np.array(action), dtype=torch.int64, device=device), \
           torch.tensor(np.array(reward), dtype=torch.float32, device=device), \
           torch.tensor(np.array(next_state), dtype=torch.float32, device=device), \
           torch.tensor(np.array(done), dtype=torch.float32, device=device)

    def __len__(self):
        return len(self.buffer)

5. Define the training function

Epsilon-Greedy: used to balance exploration and exploitation
- epsilon starts at EPS_START and decays exponentially toward EPS_END (using EPS_DECAY)
- If a random value is greater than epsilon, choose an action according to the learned policy (exploitation)
- If a random value is less than or equal to epsilon, choose a random action (exploration)
Replay Buffer:
- push: store each step's experience (transition) in the replay buffer
- sample: randomly sample a minibatch of experiences from the replay buffer for training
Target network:
- Used to improve training stability
- Copy the weights of policy_net to target_net every TARGET_UPDATE episodes.
Loss function (MSE loss):
- Train to minimize TD error (temporal difference error).
- TD target: $r + \gamma \max_{a'} Q_{\text{target}}(s', a')$
- Loss: $\text{MSE}(Q(s, a), r + \gamma \max_{a'} Q_{\text{target}}(s', a'))$
  - $Q(s, a)$ : Q-value predicted using policy_net
  - $r$ : Reward obtained at the current step
  - $\gamma$ : Discount factor
  - $\max_{a'} Q_{\text{target}}(s', a')$ : Maximum Q-value computed using target_net for the next state ( $s'$ )

def train(env, policy_net, target_net, optimizer, replay_buffer):
    all_rewards = []

    for episode in range(NUM_EPISODES):
        state, _ = env.reset()
        total_reward = 0
        done = False

        while not done:
            # Epsilon-Greedy Exploration
            epsilon = max(EPS_END, EPS_START * (EPS_DECAY ** episode))
            if random.random() > epsilon:
                with torch.no_grad():
                    state_tensor = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
                    q_values = policy_net(state_tensor)
                    action = q_values.argmax(dim=1).item()
            else:
                action = env.action_space.sample()

            next_state, reward, terminated, truncated, _ = env.step(action)
            total_reward += reward
            done = terminated or truncated

            # Store experiences in the Replay Buffer
            replay_buffer.push(state, action, reward, next_state, done)
            state = next_state

            # Sample from the Replay Buffer and train
            if len(replay_buffer) >= BATCH_SIZE:
                state_batch, action_batch, reward_batch, next_state_batch, done_batch = replay_buffer.sample(BATCH_SIZE)

                # Compute Double DQN target Q-values
                with torch.no_grad():
                    next_actions = policy_net(next_state_batch).argmax(dim=1, keepdim=True)
                    next_q_values = target_net(next_state_batch).gather(1, next_actions).squeeze(1)
                    target_q_values = reward_batch + (1 - done_batch) * GAMMA * next_q_values

                # Compute current Q-values
                q_values = policy_net(state_batch)
                q_values_for_actions = q_values.gather(1, action_batch.unsqueeze(1)).squeeze(1)

                # Compute the loss and optimize
                loss = nn.functional.mse_loss(q_values_for_actions, target_q_values)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

        all_rewards.append(total_reward)

        # Update the Target Network
        if (episode+1) % TARGET_UPDATE == 0:
            target_net.load_state_dict(policy_net.state_dict())

        # Print the mean reward
        if (episode+1) % PROGRESS_INTERVAL == 0:
            print(f"Episode {episode+1}: Average Reward over last {PROGRESS_INTERVAL} episodes = {np.mean(all_rewards[-PROGRESS_INTERVAL:])}")

        # Stop training if rewards remain sufficiently high for a sustained period
        if np.min(all_rewards[-TARGET_UPDATE*2:]) >= 450:
            print(f"Solved in {episode+1} episodes with reward {total_reward}!")
            break

6. Run training and inspect results

imageio: Library for generating GIF images
Initialize the environment: create CartPole with gym.make("CartPole-v1")
Create the networks policy_net (for training) and target_net (for stability), and move them to GPU/CPU.
Optimizer and replay buffer: create a Adam optimizer and a ReplayBuffer object
Call train: Start training.
Generate a GIF: save the trained agent's play as a GIF using imageio
env.close(): Close the environment.

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

policy_net = DuelingDQN(state_size, action_size, hidden_size).to(device)
target_net = DuelingDQN(state_size, action_size, hidden_size).to(device)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()

optimizer = optim.Adam(policy_net.parameters(), lr=LR)
replay_buffer = ReplayBuffer(MEMORY_SIZE)

train(env, policy_net, target_net, optimizer, replay_buffer)

Episode 50: Average Reward over last 50 episodes = 27.6
Episode 100: Average Reward over last 50 episodes = 42.24
Episode 150: Average Reward over last 50 episodes = 93.7
Episode 200: Average Reward over last 50 episodes = 79.7
Episode 250: Average Reward over last 50 episodes = 116.52
Episode 300: Average Reward over last 50 episodes = 107.16
Episode 350: Average Reward over last 50 episodes = 128.58
Episode 400: Average Reward over last 50 episodes = 122.18
Episode 450: Average Reward over last 50 episodes = 185.16
Episode 500: Average Reward over last 50 episodes = 168.18
Episode 550: Average Reward over last 50 episodes = 291.82
Solved in 575 episodes with reward 500.0!

# Render and save a GIF
frames = []
state, _ = env.reset()
done = False
total_reward = 0
while not done:
    with torch.no_grad():
        state_tensor = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
        q_values = policy_net(state_tensor)
        action = q_values.argmax(dim=1).item()
    next_state, reward, terminated, truncated, _ = env.step(action)
    total_reward += reward
    state = next_state
    done = terminated or truncated
    frames.append(env.render())

imageio.mimsave('cartpole_simulation_dueling_double_DQN.gif', frames, duration=33)
print(f"Final episode reward: {total_reward}")
env.close()

Final episode reward: 399.0