Segment Trees¶

Segment trees are efficient data structures used for range queries and updates, particularly useful in prioritized experience replay. The implementation is based on OpenAI’s baselines repository and provides efficient operations for priority-based sampling in reinforcement learning.

A segment tree is a binary tree where each leaf represents an element in an array, and each internal node represents some operation (like sum or minimum) over a range of elements. This allows for O(log n) query and update operations, making it ideal for efficiently sampling experiences based on priorities in prioritized replay buffers.

The base SegmentTree class provides the foundation, while SumSegmentTree and MinSegmentTree provide specialized implementations for sum and minimum operations respectively.

from agilerl.components.segment_tree import SumSegmentTree, MinSegmentTree

# Create a sum segment tree for priority-based sampling
sum_tree = SumSegmentTree(capacity=1024)

# Create a min segment tree for finding minimum priorities
min_tree = MinSegmentTree(capacity=1024)

Classes¶

class agilerl.components.segment_tree.SegmentTree(capacity: int, operation: Callable, init_value: float)¶

Create SegmentTree.

Taken from OpenAI baselines github repository: https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py

Parameters:

capacity (int) – Capacity of segment tree
operation (Callable) – Operation to apply
init_value (float) – Initial value

get_batch(indices: ndarray) → ndarray¶

Vectorised leaf read for many indices at once.

Parameters:: indices (np.ndarray) – Leaf indices in [0, capacity).
Returns:: Leaf values, one per index.
Return type:: np.ndarray

operate(start: int = 0, end: int = 0) → float¶

Return result of applying self.operation.

Parameters:

start (int) – Start index of segment
end (int) – End index of segment

Returns:

Result of applying self.operation

Return type:

float

update_batch(indices: ndarray, values: ndarray) → None¶

Set many leaf values at once and update their ancestors.

Parameters:

indices (np.ndarray) – Leaf indices in [0, capacity).
values (np.ndarray) – New leaf values, one per index.

class agilerl.components.segment_tree.SumSegmentTree(capacity: int)¶

Create SumSegmentTree.

Taken from OpenAI baselines github repository: https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py

Parameters:: capacity (int) – Capacity of segment tree

get_batch(indices: ndarray) → ndarray¶

Vectorised leaf read for many indices at once.

Parameters:: indices (np.ndarray) – Leaf indices in [0, capacity).
Returns:: Leaf values, one per index.
Return type:: np.ndarray

operate(start: int = 0, end: int = 0) → float¶

Return result of applying self.operation.

Parameters:

start (int) – Start index of segment
end (int) – End index of segment

Returns:

Result of applying self.operation

Return type:

float

retrieve(upperbound: float) → int¶

Find the highest index i about upperbound in the tree.

Parameters:: upperbound (float) – Upper bound for cumulative sum
Returns:: Index where cumulative sum is <= upperbound
Return type:: int

retrieve_batch(upperbounds: ndarray) → ndarray¶

Vectorised retrieve() for a whole batch of upper bounds.

Parameters:: upperbounds (np.ndarray) – Upper bounds for cumulative sum, one per sample.
Returns:: Leaf indices in [0, capacity), one per upper bound.
Return type:: np.ndarray

sum(start: int = 0, end: int = 0) → float¶

Return sum of elements from start to end index.

Parameters:

start (int, optional) – Start index of range, defaults to 0
end (int, optional) – End index of range, defaults to 0 (meaning capacity)

Returns:

Sum of elements in range [start, end)

Return type:

float

update_batch(indices: ndarray, values: ndarray) → None¶

Set many leaf values at once and update their ancestors.

Parameters:

indices (np.ndarray) – Leaf indices in [0, capacity).
values (np.ndarray) – New leaf values, one per index.

class agilerl.components.segment_tree.MinSegmentTree(capacity: int)¶

Create SegmentTree.

Taken from OpenAI baselines github repository: https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py

Parameters:: capacity (int) – Capacity of segment tree

get_batch(indices: ndarray) → ndarray¶

Vectorised leaf read for many indices at once.

Parameters:: indices (np.ndarray) – Leaf indices in [0, capacity).
Returns:: Leaf values, one per index.
Return type:: np.ndarray

min(start: int = 0, end: int = 0) → float¶

Return minimum element from start to end index.

Parameters:

start (int, optional) – Start index of range, defaults to 0
end (int, optional) – End index of range, defaults to 0 (meaning capacity)

Returns:

Minimum element in range [start, end)

Return type:

float

operate(start: int = 0, end: int = 0) → float¶

Return result of applying self.operation.

Parameters:

start (int) – Start index of segment
end (int) – End index of segment

Returns:

Result of applying self.operation

Return type:

float

update_batch(indices: ndarray, values: ndarray) → None¶

Set many leaf values at once and update their ancestors.

Parameters:

indices (np.ndarray) – Leaf indices in [0, capacity).
values (np.ndarray) – New leaf values, one per index.