Stop Sampling Randomly: Mastering Prioritized Experience Replay (PER)
Stop Sampling Randomly: Mastering Prioritized Experience Replay (PER)
In the world of Deep Reinforcement Learning (DRL), the Experience Replay Buffer is a cornerstone of stability. By storing past transitions and sampling them randomly, we break the temporal correlation of data and prevent the agent from "forgetting" old experiences.
But there is a fundamental flaw in standard DQN: it treats all memories as equal.
Imagine an agent learning to navigate a complex maze. It spends 99% of its time hitting walls (low-information transitions) and 1% of its time finding the exit (high-information transitions). In a uniform replay buffer, the agent will sample the "hitting walls" experience 99% of the time, drastically slowing down convergence.
Enter Prioritized Experience Replay (PER).
The Core Intuition: Learning from Surprise
The central thesis of PER is that not all experiences are equally valuable. The most useful transitions are those where the agent's current prediction is most wrong.
In RL, this "wrongness" is quantified by the Temporal Difference (TD) Error. A high TD error indicates that the transition was "surprising" and that the agent has the most to learn from it. By prioritizing transitions with higher TD errors, we focus the agent's training on its weakest points.
The Mathematical Framework
To implement this without introducing instability, PER relies on three key formulas:
-
The Priority Score: We use the magnitude of the TD error $\delta$ as the priority $p_i$. To ensure no transition has zero probability of being sampled, we add a small constant $\epsilon$: $$\text{Proportional Prioritization: } p_i = |\delta_i| + \epsilon$$
-
Sampling Probability: To control how aggressive the prioritization is, we introduce a hyperparameter $\alpha$. When $\alpha=0$, we have uniform sampling; when $\alpha=1$, we have full prioritization: $$P(i) = \frac{p_i^\alpha}{\sum_k p_k^\alpha}$$
-
Importance Sampling (IS) Weights: Prioritizing high-error samples introduces a bias—the model thinks high-error states are more common than they actually are. To correct this, we scale the gradient update using IS weights: $$w_i = \left( \frac{1}{N} \cdot \frac{1}{P(i)} \right)^\beta$$ Where $\beta$ is annealed from an initial value to 1.0 over time.
System Architecture
Implementing PER efficiently requires more than just a list; it requires a data structure that allows for fast sampling and fast updates. The SumTree is the gold standard here, reducing the complexity of sampling and updating from $O(N)$ to $O(\log N)$.
Production-Ready Implementation
Below is the complete implementation using PyTorch and OpenAI Gym. Note the use of the SumTree to maintain the priority distribution.
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import namedtuple
import gym
Transition = namedtuple('Transition', ('state', 'action', 'next_state', 'reward', 'done'))
class SumTree:
"""Binary tree where each node is the sum of its children for O(log N) sampling."""
def __init__(self, capacity):
self.capacity = capacity
self.tree = np.zeros(2 * capacity - 1)
self.data = np.zeros(capacity, dtype=object)
self.write = 0
self.n_entries = 0
def _propagate(self, idx, change):
parent = (idx - 1) // 2
self.tree[parent] += change
if parent != 0:
self._propagate(parent, change)
def update(self, idx, p):
change = p - self.tree[idx]
self.tree[idx] = p
self._propagate(idx, change)
def add(self, p, data):
idx = self.write + self.capacity - 1
self.data[self.write] = data
self.update(idx, p)
self.write = (self.write + 1) % self.capacity
if self.n_entries < self.capacity:
self.n_entries += 1
def get(self, s):
idx = self._retrieve(0, s)
data_idx = idx - self.capacity + 1
return (idx, self.tree[idx], self.data[data_idx])
def _retrieve(self, idx, s):
left = 2 * idx + 1
right = left + 1
if left >= len(self.tree):
return idx
if s <= self.tree[left]:
return self._retrieve(left, s)
else:
return self._retrieve(right, s - self.tree[left])
@property
def total_priority(self):
return self.tree[0]
class PrioritizedReplayBuffer:
def __init__(self, capacity, alpha=0.6, beta=0.4, beta_increment=0.001):
self.tree = SumTree(capacity)
self.alpha = alpha
self.beta = beta
self.beta_increment = beta_increment
self.max_p = 1.0
def push(self, state, action, next_state, reward, done):
data = Transition(state, action, next_state, reward, done)
self.tree.add(self.max_p**self.alpha, data)
def sample(self, batch_size):
batch, idxs, priorities = [], [], []
segment = self.tree.total_priority / batch_size
self.beta = np.min([1., self.beta + self.beta_increment])
for i in range(batch_size):
s = random.uniform(segment * i, segment * (i + 1))
(idx, p, data) = self.tree.get(s)
priorities.append(p)
idxs.append(idx)
batch.append(data)
sampling_probabilities = np.array(priorities) / self.tree.total_priority
is_weights = np.power(self.tree.n_entries * sampling_probabilities, -self.beta)
is_weights /= is_weights.max()
return batch, idxs, torch.FloatTensor(is_weights)
def update_priorities(self, idxs, errors):
for idx, error in zip(idxs, errors):
p = (np.abs(error) + 1e-5)**self.alpha
self.tree.update(idx, p)
self.max_p = max(self.max_p, np.abs(error) + 1e-5)
class QNetwork(nn.Module):
def __init__(self, state_dim, action_dim):
super(QNetwork, self).__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, 64), nn.ReLU(),
nn.Linear(64, 64), nn.ReLU(),
nn.Linear(64, action_dim)
)
def forward(self, x): return self.net(x)
class PERAgent:
def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99):
self.state_dim, self.action_dim, self.gamma = state_dim, action_dim, gamma
self.model = QNetwork(state_dim, action_dim)
self.target_model = QNetwork(state_dim, action_dim)
self.target_model.load_state_dict(self.model.state_dict())
self.optimizer = optim.Adam(self.model.parameters(), lr=lr)
self.memory = PrioritizedReplayBuffer(capacity=10000)
def select_action(self, state, epsilon=0.1):
if random.random() < epsilon:
return random.randint(0, self.action_dim - 1)
state = torch.FloatTensor(state).unsqueeze(0)
with torch.no_grad():
return self.model(state).argmax().item()
def train_step(self, batch_size):
if self.memory.tree.n_entries < batch_size: return None
batch, idxs, is_weights = self.memory.sample(batch_size)
states = torch.FloatTensor(np.array([t.state for t in batch]))
actions = torch.LongTensor(np.array([t.action for t in batch])).unsqueeze(1)
next_states = torch.FloatTensor(np.array([t.next_state for t in batch]))
rewards = torch.FloatTensor(np.array([t.reward for t in batch])).unsqueeze(1)
dones = torch.FloatTensor(np.array([t.done for t in batch])).unsqueeze(1)
curr_q = self.model(states).gather(1, actions)
with torch.no_grad():
next_q = self.target_model(next_states).max(1)[0].unsqueeze(1)
target_q = rewards + (1 - dones) * self.gamma * next_q
td_errors = (target_q - curr_q).detach().squeeze().numpy()
# Apply Importance Sampling weights to the loss
loss = (is_weights.unsqueeze(1) * nn.MSELoss(reduction='none')(curr_q, target_q)).mean()
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.memory.update_priorities(idxs, td_errors)
return loss.item()
def update_target(self):
self.target_model.load_state_dict(self.model.state_dict())Key Takeaways for Implementation
1. The $\alpha$ and $\beta$ Trade-off
- $\alpha$ (Prioritization Degree): If you set $\alpha$ too high, the agent may overfit to a small set of "hard" samples, leading to instability. Start with $0.6$.
- $\beta$ (Bias Correction): Bias correction is most critical at the end of training. This is why we anneal $\beta$ from $0.4 \to 1.0$.
2. Complexity Analysis
| Operation | Uniform Replay | Prioritized Replay (SumTree) |
|---|---|---|
| Push | $O(1)$ | $O(\log N)$ |
| Sample | $O(1)$ | $O(\log N)$ |
| Update | N/A | $O(\log N)$ |
3. When to use PER?
PER is most effective in sparse reward environments. If your agent only receives a reward once every 1,000 steps, PER ensures that the single successful trajectory is replayed frequently enough to propagate the value back to the starting state.