Weighted Sampling

Topics:

Definition

Weighted Sampling is the process of randomly selecting an item from a discrete set where each item has its own probability/weight rather than all items being equally likely.

Common Implementations

Cumulative Sum

The idea behind implementing weighted sampling is to make a cumulative sum of the weights and then sample from a uniform distribution over the obtained interval.
For example if 10 has a weight of 1 and 20 has a weight of 4, this means that we want to sample 20 four times as much as we sample 10. We create the cumulative sum of the weights: 1+ 4 = 5, and sample uniformly over the interval [0,5]. If we obtain a value in the interval [0,1], then the output is 10, and if we obtain a value in the interval [1,5], then the output is 20.
The downside of this implementation is that its complexity is O(n), meaning that it becomes much slower for large datasets, like the replay buffer in RL. This is because, after sampling from the uniform distribution, the algorithm would need to iterate over each interval of the cumulative sum to reach the right value. This motivates the implementation via a SumTree.

SumTree

A SumTree is a type of data structure. It is a binary tree where the value of each parent node is equal to the sum of the leaf nodes. The weighted sampling via a SumTree uses the same concept of cumulative sum and uniform sampling; it's just the implementation technique that changes to speed up the algorithm.
In the SumTree implementation, the leaf nodes correspond to the weights of the items of the distribution. This way, the top parent node is equal to the sum of all the weights, which is the upper limit of the uniform interval. The sampling and extraction process is as follows:

• take a uniform random sample between 0 an the value of the top-parent node. Call it value
• traverse the tree using value:
  • if value is less than the left child node, keep value as is and move to the left side of the tree
  • if not, subtract the value of the left child node from value and move to the right child node
  • repeat the process until a leaf node is reached, which is the weight of the sampled value

Here's the intuition with an example. In the below image, the cumulative sum of the weights is $10$. In other words, the left part of the tree has a probability of $0.5$; same for the right part. Say, for example, the value $6.5$ is obtained by sampling uniformly from$0$ to $10$. Is we think of it in terms of the cumulative sum, this means that it falls in the range [5,7]. In the SumTree, $6.5 > 5$, so the corresponding value is in the right side of the tree. We subtract the contribution of the left side of the tree and end up with $1.5$, which is less than $2$, meaning that the final weight is $2$.

References

https://www.sefidian.com/2022/11/09/sumtree-data-structure-for-prioritized-experience-replay-per-explained-with-python-code/