Prioritized Experience Replay (Brief Overview)

Experience Replay is a technique used in the standard Deep Q-Network approach to break correlations between consecutive samples and reuse past experiences. The replay buffer stores transitions $(s, a, r, s')$ and the agent samples mini-batches uniformly at random from this buffer to perform Q-network updates. While effective, this uniform sampling treats all transitions as equally important. Intuitively, however, some experiences might offer much more learning potential than others. Imagine an agent encountering a situation that leads to a surprisingly high or low reward, or where its current Q-value prediction is way off. These 'surprising' transitions seem more valuable for learning than routine ones where the predictions are already accurate.

Prioritized Experience Replay (PER) moves away from uniform sampling. Instead of picking transitions randomly, PER samples transitions proportionally to their learning potential, typically measured by the magnitude of their Temporal Difference (TD) error. The TD error represents how "surprising" a transition was to the network:

$$\delta_t = r_t + \gamma \max_{a'} Q_{\text{target}}(s_{t+1}, a') - Q(s_t, a_t)$$

A large absolute TD error, $|\delta_t|$, indicates that the network's prediction for the state-action pair $(s_t, a_t)$ was inaccurate compared to the observed reward and the estimated value of the next state. These are precisely the transitions we want the agent to focus on.

Assigning Priorities

In PER, each transition $i$ in the replay buffer is assigned a priority $p_i$. A common way to define this priority is based on the absolute TD error:

$$p_i = |\delta_i| + \epsilon$$

Here, $\epsilon$ is a small positive constant added to ensure that transitions with zero TD error still have a non-zero probability of being sampled.

The probability $P(i)$ of sampling transition $i$ is then defined based on its priority:

$$P(i) = \frac{p_i^\alpha}{\sum_k p_k^\alpha}$$

The exponent $\alpha \ge 0$ is a hyperparameter that controls the degree of prioritization. When $\alpha = 0$, we recover the original uniform sampling strategy ($P(i) = 1/N$, where $N$ is the buffer size). As $\alpha$ increases, the sampling becomes more heavily skewed towards transitions with high TD errors. Efficiently sampling according to these probabilities often involves specialized data structures like SumTrees.

Correcting for Bias with Importance Sampling

Sampling transitions based on their priority introduces a bias because the updates no longer reflect the original distribution of experiences. Transitions with high TD errors are over-represented in the training batches. To counteract this bias, PER uses Importance Sampling (IS) weights when calculating the loss for the sampled transitions.

The IS weight $w_i$ for transition $i$ corrects for its non-uniform sampling probability $P(i)$:

$$w_i = \left( \frac{1}{N \cdot P(i)} \right)^\beta$$

Here, $N$ is the size of the replay buffer. The hyperparameter $\beta \in [0, 1]$ controls how much correction is applied. $\beta = 1$ fully compensates for the non-uniform probabilities, while $\beta = 0$ applies no correction. In practice, $\beta$ is often annealed from an initial value (e.g., 0.4) towards 1 over the course of training.

These weights are then incorporated into the loss function, typically by multiplying the TD error $\delta_i$ by $w_i$ during the gradient update step. This ensures that while we sample important transitions more often, the magnitude of their update is scaled down to prevent overfitting to these specific samples and to maintain unbiased estimates in expectation. For numerical stability, the weights are often normalized by dividing by the maximum weight in the mini-batch: $w_i \leftarrow w_i / \max_j w_j$.

Putting it into Practice

• Priority Updates: When a transition is sampled and its TD error is computed for the loss calculation, that TD error is then used to update the transition's priority in the replay buffer.
• New Transitions: When new transitions are added to the buffer, they are typically assigned maximum priority to ensure they are seen by the learning algorithm at least once.

Summary

Prioritized Experience Replay is an enhancement to the standard experience replay mechanism used in DQN. By sampling transitions based on their TD error magnitude, it focuses the learning process on the most informative experiences. While it introduces some complexity through priority calculation, non-uniform sampling, and importance sampling weights, PER often leads to significant improvements in learning speed and data efficiency compared to uniform sampling. It's frequently combined with other DQN improvements like Double DQN and Dueling Networks to build highly effective RL agents.