Prioritized Experience Replay (PER)

Standard Experience Replay, as discussed in the context of DQN, treats all transitions stored in the replay buffer as equally important. When sampling a mini-batch for training, each transition $(s, a, r, s')$ has an equal probability of being selected. However, intuitively, not all experiences offer the same learning potential. Some transitions might be routine or confirm what the agent already knows well, while others might represent surprising outcomes or critical moments where the agent's current Q-value estimates are significantly inaccurate. Learning from these "surprising" transitions could lead to faster and more effective updates.

Prioritized Experience Replay (PER) builds on this intuition by sampling transitions based on their significance, giving preference to those from which the agent can learn the most. The core idea is to use the magnitude of the Temporal Difference (TD) error as a proxy for how "surprising" or informative a transition is. Recall the TD error for a transition $(s, a, r, s')$ using target network parameters $\theta^-$ and current network parameters $\theta$:

$$\delta = r + \gamma \max_{a'} Q(s', a'; \theta^-) - Q(s, a; \theta)$$

A large absolute TD error, $|\delta|$, indicates a large discrepancy between the predicted Q-value $Q(s, a; \theta)$ and the estimated target value $r + \gamma \max_{a'} Q(s', a'; \theta^-)$. Such transitions represent situations where the current Q-function is likely inaccurate and thus provide a strong learning signal.

Assigning Priorities

In PER, instead of uniform sampling, we assign a priority $p_i$ to each transition $i$ in the replay buffer. A common choice for this priority is directly related to 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. This prevents them from never being replayed.

Once priorities are assigned, we convert them into sampling probabilities $P(i)$ for each transition $i$:

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

The exponent $\alpha \ge 0$ controls the degree of prioritization.

• If $\alpha = 0$, then $p_i^\alpha = 1$ for all $i$ (assuming $p_i > 0$), and we recover uniform sampling: $P(i) = 1/N$, where $N$ is the number of transitions in the buffer.
• If $\alpha = 1$, we have direct proportional prioritization based on $p_i$.
• Values of $\alpha$ between 0 and 1 allow interpolation between uniform and full prioritization, providing a way to temper the focus on high-priority samples.

Efficient Sampling with SumTrees

Calculating the normalization term $\sum_k p_k^\alpha$ and sampling proportionally from potentially millions of transitions can be computationally expensive if done naively. PER typically employs a specialized data structure called a SumTree (a variant of a segment tree) for efficient implementation.

A SumTree is a binary tree where each leaf node stores the priority $p_i^\alpha$ of a single transition. Each internal node stores the sum of the priorities of its children. The root node, therefore, stores the total sum $\sum_k p_k^\alpha$.

A simplified SumTree structure. Leaf nodes hold transition priorities (e.g., $p_i^\alpha$). Parent nodes hold the sum of their children's priorities. The root holds the total sum. Sampling involves drawing a random value in $[0, \text{Root Sum}]$ and traversing the tree to find the corresponding leaf/transition.

This structure allows for:

1. Efficient Sampling: To sample a transition, generate a random number $z$ between 0 and the total sum stored at the root. Then, traverse the tree downwards. At each node, go left if $z$ is less than the sum of the left child, otherwise subtract the left child's sum from $z$ and go right. This process finds the transition corresponding to the sampled value $z$ in $O(\log N)$ time.
2. Efficient Priority Updates: When a transition's TD error (and thus its priority) is updated after being used for training, only the priorities of its ancestors in the tree need to be updated. This also takes $O(\log N)$ time.

Correcting the Bias with Importance Sampling

Prioritized sampling is not without consequences. By preferentially selecting transitions with high TD errors, we introduce bias. The network will primarily see transitions where it performs poorly, potentially distorting the expected value updates because the sampled distribution no longer matches the distribution under which the experiences were generated.

To counteract this bias, PER incorporates Importance Sampling (IS) weights into the learning update. The intuition is to down-weight the updates for transitions that were sampled more frequently than they would have been under uniform sampling. The IS weight $w_i$ for transition $i$ is calculated as:

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

Here, $N$ is the size of the replay buffer, $P(i)$ is the probability of sampling transition $i$ under the prioritized scheme, and $\beta \ge 0$ is an exponent that controls how much correction is applied.

• If $\beta = 0$, no correction is applied ($w_i = 1$).
• If $\beta = 1$, the bias introduced by prioritization is fully corrected.

In practice, $\beta$ is often annealed from an initial value (e.g., 0.4) linearly towards 1 over the course of training. Starting with a smaller $\beta$ allows the agent to benefit more strongly from the prioritized samples early on when the Q-function is still poorly estimated, while increasing $\beta$ later helps reduce bias as the estimates become more refined.

For stability, these weights are typically normalized by dividing by the maximum weight in the mini-batch:

$$w_i \leftarrow \frac{w_i}{\max_j w_j}$$

This ensures that the updates are only scaled down, not potentially scaled up significantly. The normalized weight $w_i$ is then used to scale the TD error (or the loss) for transition $i$ during the gradient update step. For instance, using mean squared error loss, the contribution of transition $i$ would be weighted:

$$\text{Loss}_i = w_i \cdot \delta_i^2$$

Summary of PER

Prioritized Experience Replay modifies the standard DQN training loop in the following ways:

1. Store Priorities: When adding a new transition to the replay buffer, assign it a maximum priority initially to ensure it gets sampled at least once. Store priorities alongside transitions, often using a SumTree.
2. Prioritized Sampling: Sample mini-batches from the buffer using probabilities $P(i) \propto p_i^\alpha$ derived from the stored priorities.
3. Compute IS Weights: For each sampled transition $i$, calculate its importance sampling weight $w_i = (N \cdot P(i))^{-\beta}$. Normalize these weights.
4. Weighted Update: Compute the TD error $\delta_i$ for each transition in the mini-batch. Use the weighted TD error $w_i \cdot \delta_i$ (or incorporate $w_i$ into the loss calculation) when performing the gradient descent step on the Q-network.
5. Update Priorities: Update the priorities $p_i$ in the SumTree for the sampled transitions using their newly computed absolute TD errors $|\delta_i|$.

By focusing learning on transitions that are most surprising or informative, PER often leads to significant improvements in data efficiency and faster convergence compared to standard Experience Replay with uniform sampling. However, it introduces additional complexity in implementation (SumTree management) and new hyperparameters ($\alpha$, $\beta$, $\epsilon$) that need tuning. It represents a sophisticated enhancement that directly addresses the efficiency of learning from stored experiences in DQN.