Prioritized Experience Replay (PER)

English translation

[Problems with a conventional replay buffer]

Conventional Replay Buffer:
- Randomly sample all experiences with equal probability (uniform random sampling)
- Problems
  - Inefficiency: not all experiences are equally important for learning
    - Some experiences have very large TD errors (differences between predictions and actual values) -> important for learning!
    - Some experiences have almost no TD error -> offer little help for learning
  - Slow learning: important experiences are not sufficiently used for learning

[Core Prioritized Experience Replay (PER) concept: use important experiences more often]

Goal: Learn more often from experiences with large TD errors -> Maximize sampling efficiency and accelerate learning.
Key idea
- Use TD error: Use the difference between the current Q-network prediction and the actual target (TD error).
  - TD error = (actual target) - (current Q-network prediction)

$\delta_i = r_i + \gamma \max_{a'} Q(s'_i, a'; \theta^{-}) - Q(s_i, a_i; \theta)$

Useful experience: Larger TD error -> Greater potential to improve the Q-network substantially -> Learn from it more often!
Assign priorities
- Criterion: Based on the absolute TD error ( $|\delta_i|$ )
- Principle: Larger TD errors increase the sampling probability (priority $\propto$ ).
- Formula
  - $p_i = (|\delta_i| + \epsilon)^\alpha$
    - $p_i$ : Priority of experience $i$ (affects sampling probability)
    - $\delta_i$ : TD error of experience $i$
    - $\epsilon$ : Small positive constant (prevents priority from becoming 0 and guarantees a minimum sampling probability for every experience)
    - $\alpha$ : Controls prioritization strength (0: uniform sampling; 1: proportional to TD error)
- Example:
  - $\alpha = 0.6, \epsilon = 0.001$
  - Experience with TD error ( $\delta$ ) of 1 >> experience with TD error of 0.1 (sampled with much higher probability)

[Importance sampling weights: correcting PER bias]

Problem with PER: Bias.
- Learning too often from experiences with large TD errors -> Risk of overfitting
- Rarely learning from experiences with small TD errors -> Missing information
Solution: Use importance sampling weights (IS weights)
- Purpose: Correct bias caused by PER
- Method: Weight each experience's update during learning.
  - Frequently sampled experience with large TD error -> Lower weight
  - Rarely sampled experience with small TD error -> Higher weight
- Formula: $w_i = \left( \frac{1}{N} \cdot \frac{1}{P(i)} \right)^\beta$
  - $w_i$ : Weight for transition $i$
  - $N$ : Replay buffer size
  - $P(i)$ : Probability of sampling transition $i$ ( $P(i) = \frac{p_i}{\sum_k p_k}$ )
  - Parameter: $\beta$ (0–1)
    - $\beta = 0$ : No correction (same as ordinary replay)
    - $\beta = 1$ : Full correction
    - Early training: Start with a small $\beta$ -> Gradually increase to 1 (annealing)
    - Reason: The Q-network is unstable early in training -> Bias correction may instead interfere with learning.

[Effects of PER]

Faster learning
- Focus learning on important experiences -> more efficient Q-network updates
- Reduce unnecessary computation -> save time
Better performance
- Reduce overestimation (quickly correct experiences with large TD errors).
- Improve training stability.
Improved sample efficiency
- Better performance can be achieved with fewer samples.

[Activity: Dueling Double DQN with PER (Prioritized Experience Replay)]

1. Import libraries

gym: reinforcement learning environment library
torch: PyTorch (neural network construction and training)
numpy: numerical operations and NumPy arrays
imageio: generate GIF images
tqdm: Display training progress.
matplotlib: visualize training results with plots

import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
import random
from collections import deque
import imageio
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

2. Configure the GPU and define hyperparameters

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

# Hyperparameters (adjusted for the CartPole environment)
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

# PER hyperparameters
ALPHA = 0.6
BETA_START = 0.4
BETA_DECAY = 0.999

3. Define the Dueling DQN network

Define the Dueling DQN network: separate value and advantage streams to calculate Q-values

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. Implement Prioritized Experience Replay (PER)

Define the PrioritizedReplayBuffer class for prioritized experience replay
- push: Store a new experience (transition) in the buffer and initialize its priority to the maximum.
  - states, actions, rewards, next_states, dones: Store each component as a NumPy array for faster sampling.
- sample: Sample experiences according to priority and return importance sampling weights (IS weights).
  - priorities: Use np.array for efficient priority calculation and sampling.
- update_priorities: Update the priorities of sampled experiences based on TD errors.

class PrioritizedReplayBuffer:
    def __init__(self, capacity, alpha=ALPHA):
        self.capacity = capacity
        self.alpha = alpha
        self.buffer = deque(maxlen=capacity)
        self.priorities = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        state = np.array(state)
        next_state = np.array(next_state)
        max_priority = max(self.priorities, default=1) # Initially 1, then max
        self.buffer.append((state, action, reward, next_state, done))
        self.priorities.append(max_priority)

    def sample(self, batch_size, beta=BETA_START):
        priorities = np.array(self.priorities)
        probs = priorities ** self.alpha
        probs /= probs.sum()

        indices = np.random.choice(len(self.buffer), batch_size, p=probs, replace=False)
        samples = [self.buffer[idx] for idx in indices]

        total = len(self.buffer)
        weights = (total * probs[indices]) ** (-beta)
        weights /= weights.max()  # Normalize for stability
        weights = torch.tensor(weights, dtype=torch.float32, device=device)

        state, action, reward, next_state, done = zip(*samples)

        return torch.tensor(np.array(state), dtype=torch.float32, device=device), \
               torch.tensor(action, dtype=torch.int64, device=device), \
               torch.tensor(reward, dtype=torch.float32, device=device), \
               torch.tensor(np.array(next_state), dtype=torch.float32, device=device), \
               torch.tensor(done, dtype=torch.float32, device=device), \
               weights, indices # Return importance sampling weights and sampled indices

    def update_priorities(self, indices, errors, offset=0.1):
        for idx, error in zip(indices, errors):
            self.priorities[idx] = (abs(error) + offset) ** self.alpha

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

5. Define the training function

Description: train the agent with Dueling Double DQN with PER
- Epsilon-greedy: gradually decrease epsilon to balance exploration and exploitation
- Double DQN: use Double DQN to calculate target Q-values (reduce overestimation)
  - Select the next state's action with policy_net and evaluate its Q-value with target_net
- Prioritized Experience Replay
  - After computing TD errors, call replay_buffer.update_priorities to update priorities
  - Include the importance weights returned by replay_buffer.sample (weights) in the loss calculation
- Training progress: display visually with tqdm

def train(env, policy_net, target_net, optimizer, replay_buffer, num_episodes=NUM_EPISODES):
    all_rewards = []
    epsilon = EPS_START
    beta = BETA_START
    for episode in tqdm(range(num_episodes), desc="Training"):
        state, _ = env.reset()
        total_reward = 0
        done = False

        while not done:
            # Epsilon-Greedy Action Selection
            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, weights, indices = replay_buffer.sample(BATCH_SIZE, beta)

                # 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
                errors = (target_q_values - q_values_for_actions).detach().cpu().numpy()
                replay_buffer.update_priorities(indices, errors) # update priorities
                loss = (weights * (q_values_for_actions - target_q_values)**2).mean()

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

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

        # Update epsilon, beta
        epsilon = max(EPS_END, epsilon * EPS_DECAY)
        beta = min(1.0, beta * BETA_DECAY)

        # Print the mean reward
        all_rewards.append(total_reward)
        if (episode+1) % PROGRESS_INTERVAL == 0:
            print(f"Episode {episode+1}: Average Reward = {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!")
            break

    return all_rewards

6. Run training and inspect results

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 = PrioritizedReplayBuffer(MEMORY_SIZE)

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

# Plot training results
plt.figure(figsize=(10, 5))
plt.plot(all_rewards)
plt.title('Rewards per Episode')
plt.xlabel('Episode')
plt.ylabel('Reward')
plt.show()

Training:   3%|▎         | 51/2000 [00:08<06:17,  5.17it/s]

Episode 50: Average Reward = 27.52

Training:   5%|▌         | 100/2000 [00:23<09:17,  3.41it/s]

Episode 100: Average Reward = 50.72

Training:   8%|▊         | 150/2000 [00:42<10:35,  2.91it/s]

Episode 150: Average Reward = 55.6

Training:  10%|█         | 200/2000 [01:15<18:20,  1.64it/s]

Episode 200: Average Reward = 89.44

Training:  12%|█▎        | 250/2000 [02:05<38:10,  1.31s/it]

Episode 250: Average Reward = 114.56

Training:  15%|█▌        | 300/2000 [03:01<42:24,  1.50s/it]

Episode 300: Average Reward = 113.38

Training:  18%|█▊        | 350/2000 [04:23<55:10,  2.01s/it]

Episode 350: Average Reward = 141.38

Training:  20%|██        | 400/2000 [06:13<38:40,  1.45s/it]

Episode 400: Average Reward = 168.92

Training:  22%|██▎       | 450/2000 [08:33<1:59:55,  4.64s/it]

Episode 450: Average Reward = 218.8

Training:  25%|██▌       | 500/2000 [12:14<2:22:00,  5.68s/it]

Episode 500: Average Reward = 355.84

Training:  28%|██▊       | 550/2000 [15:57<1:14:46,  3.09s/it]

Episode 550: Average Reward = 344.56

Training:  30%|███       | 600/2000 [19:10<2:04:54,  5.35s/it]

Episode 600: Average Reward = 309.68

Training:  32%|███▎      | 650/2000 [22:30<1:47:21,  4.77s/it]

Episode 650: Average Reward = 322.32

Training:  35%|███▌      | 700/2000 [26:07<1:39:18,  4.58s/it]

Episode 700: Average Reward = 347.04

Training:  38%|███▊      | 750/2000 [29:52<1:51:12,  5.34s/it]

Episode 750: Average Reward = 357.04

Training:  38%|███▊      | 756/2000 [30:36<50:22,  2.43s/it]

Solved in 757 episodes!

<Figure size 1000x500 with 1 Axes>

# Render the final policy rollout and save it as 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_dueling_double_dqn_per.gif', frames, duration=33)
print(f"Validation Episode Reward: {total_reward}")
env.close()

Validation Episode Reward: 500.0