22. Prioritized Experience Replay: Sampling Probabilities and IS Weights

Problem

Uniform experience replay treats every stored transition as equally worth training on, but most transitions carry almost no learning signal once the network has already fit them well. Prioritized Experience Replay samples transitions with probability proportional to how surprising they were (typically |TD-error| + a small epsilon), so the network spends more time on the transitions it's getting most wrong. But sampling non-uniformly biases the gradient estimate, so every sampled transition's loss has to be reweighted by an importance-sampling (IS) correction to keep training unbiased.

Given a fixed set of per-transition priorities and a batch of indices already sampled proportional to priority^alpha, compute the IS weight for each sampled transition:
P(i) = priority(i)^alpha / sum_j priority(j)^alpha
w(i) = (N * P(i))^(-beta), then normalized so the largest weight in the batch is exactly 1.0.

Source: src/22_prioritized_replay_is_weights.py

def prioritized_is_weights(
    priorities: torch.Tensor, sampled_indices: torch.Tensor,
    alpha: float, beta: float,
) -> torch.Tensor: ...

>>> priorities = torch.tensor([1.0, 4.0, 9.0, 16.0])
>>> sampled = torch.tensor([3, 0])
>>> prioritized_is_weights(priorities, sampled, alpha=0.6, beta=0.4).shape
torch.Size([2])

Step-by-Step Approach

1. Turn priorities into a sampling distribution over the whole buffer: priorities.pow(alpha), normalized by its sum — this is P for every transition, not just the sampled ones.
2. Index that distribution at the sampled positions: probs[sampled_indices], giving each sampled transition's own P(i).
3. Compute the raw IS weight (N * P(i))^(-beta), where N is the total number of transitions in the buffer (priorities.shape[0]), not the sampled batch size.
4. Normalize by the batch's own max weight so the largest weight in the returned batch is exactly 1.0 — this only ever scales the effective loss down, keeping training stable regardless of alpha/beta.
5. Sanity-check the degenerate cases: alpha = 0 makes sampling uniform so every weight collapses to 1.0, and beta = 0 makes every weight 1.0 before normalization for the same reason.

The sampling probabilities have to be computed over the entire buffer (the normalizer is a sum over all N priorities) even though the function only returns weights for the handful of sampled indices — slicing before normalizing would silently compute the wrong distribution.

Reference solution

import torch


def prioritized_is_weights(
    priorities: torch.Tensor, sampled_indices: torch.Tensor,
    alpha: float, beta: float,
) -> torch.Tensor:
    # sampling distribution over the whole buffer: priority^alpha, normalized
    probs = priorities.pow(alpha)
    probs = probs / probs.sum()

    # (N * P(i))^-beta: N*P(i) is how much more/less likely i was to be
    # picked than under uniform sampling (1/N each); raising to -beta
    # shrinks the loss contribution of over-represented (high-priority)
    # transitions to compensate for how often they get trained on
    n = priorities.shape[0]
    weights = (n * probs[sampled_indices]).pow(-beta)

    # normalize by the max so weights only ever scale the loss down -- this
    # keeps the effective learning rate stable regardless of alpha/beta
    return weights / weights.max()

Key Functions & Tricks

• Tensor.pow(alpha) — raises every priority to the exponent that controls how aggressively sampling favors high-priority transitions
• probs.sum() — normalizes priority^alpha into a proper probability distribution over the whole buffer, the denominator of P(i)
• probs[sampled_indices] — fancy indexing to pull out just the sampled transitions' probabilities, after normalizing over the full buffer
• Tensor.pow(-beta) — the importance-sampling correction exponent; beta=1 fully corrects the sampling bias, beta=0 applies no correction
• weights / weights.max() — renormalizes so the largest IS weight in the batch is 1.0, a standard PER stabilization trick

How to Recognize This Pattern