Skip to content

SumTree ​

python
SumTree(capacity: int) -> None

A 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.

Source

Properties ​

total_priority ​

python
total_priority: float

Get the total priority (root value).

Returns:

  • float – Total priority.

Source

Methods ​

get_leaf ​

python
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).

Source

max_priority ​

python
max_priority() -> float

Get the maximum priority in leaf nodes.

Returns:

  • float – Maximum priority.

Source

recompute_max ​

python
recompute_max() -> None

Recompute 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).

Source

update ​

python
update(indices: int | np.ndarray, priorities: float | np.ndarray) -> None

Update 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.

Source

Values ​

capacity ​

python
capacity = capacity

Source

data_idx ​

python
data_idx = 0

Source

size ​

python
size = 0

Source

tree ​

python
tree = np.zeros(2 * capacity - 1)

Source