SumTree
SumTree(capacity: int) -> NoneA binary tree data structure where the parent node is the sum of its children. Used for efficient prioritized sampling.
Initialize the SumTree.
Parameters:
capacity(int) – Maximum number of leaf nodes.
Properties
total_priority
total_priority: floatGet the total priority (root value).
Returns:
float– Total priority.
Methods
get_leaf
get_leaf(values: float | np.ndarray) -> tuple[np.ndarray, np.ndarray]Find the leaf nodes for given values.
Parameters:
values(float | ndarray) – Value or values to search for.
Returns:
tuple[ndarray, ndarray]– A tuple of (data indices, priorities).
max_priority
max_priority() -> floatGet the maximum priority in leaf nodes.
Returns:
float– Maximum priority.
recompute_max
recompute_max() -> NoneRecompute the true maximum priority across the currently active leaves. This is O(N_active) and should be called rarely (only when the last leaf at the previously tracked maximum drops below it).
update
update(indices: int | np.ndarray, priorities: float | np.ndarray) -> NoneUpdate the priority of one or more nodes.
IMPORTANT: If the same index appears multiple times in a batch, only the last occurrence should define the final leaf value. We must deduplicate before computing and propagating changes to avoid double-counting.
Also keep the tracked maximum priority in sync with the leaf values without scanning all leaves on every update. We do this by maintaining a count of how many leaves currently equal the tracked maximum; when this count drops to zero, we recompute the true maximum across the active leaves.
Parameters:
indices(int | ndarray) – Index or indices of the leaf nodes (data indices).priorities(float | ndarray) – New priority value or values.
Values
capacity
capacity = capacitydata_idx
data_idx = 0size
size = 0tree
tree = np.zeros(2 * capacity - 1)