[RLlib] Fix bugs and speed up SegmentTree

rllib/BUILD

@@ -852,15 +852,26 @@ py_test(
 # Tag: optimizers
 # --------------------------------------------------------------------
 
-# This has bugs: See PR https://github.com/ray-project/ray/pull/7534
+py_test(
-# which fixes these and re-adds this test.
+    name = "test_optimizers",
+    tags = ["optimizers"],
+    size = "large",
+    srcs = ["optimizers/tests/test_optimizers.py"]
+)
 
-# py_test(
+py_test(
-#     name = "test_segment_tree",
+    name = "test_segment_tree",
-#     tags = ["optimizers"],
+    tags = ["optimizers"],
-#     size = "small",
+    size = "small",
-#    srcs = ["optimizers/tests/test_segment_tree.py"]
+    srcs = ["optimizers/tests/test_segment_tree.py"]
-# )
+)
+
+py_test(
+    name = "test_prioritized_replay_buffer",
+    tags = ["optimizers"],
+    size = "small",
+    srcs = ["optimizers/tests/test_prioritized_replay_buffer.py"]
+)
 
 # --------------------------------------------------------------------
 # Policies
@@ -1044,13 +1055,6 @@ py_test(
     srcs = ["tests/test_nested_spaces.py"]
 )
 
-py_test(
-    name = "tests/test_optimizers",
-    tags = ["tests_dir", "tests_dir_O"],
-    size = "large",
-    srcs = ["tests/test_optimizers.py"]
-)
-
 py_test(
     name = "tests/test_exec_api",
     tags = ["tests_dir", "tests_dir_E"],

rllib/optimizers/replay_buffer.py

@@ -227,7 +227,7 @@ class PrioritizedReplayBuffer(ReplayBuffer):
           Array of shape (batch_size,) and dtype np.int32
           idexes in buffer of sampled experiences
         """
-        assert beta > 0
+        assert beta >= 0.0
         self._num_sampled += batch_size
 
         idxes = self._sample_proportional(batch_size)

rllib/optimizers/segment_tree.py

@@ -2,140 +2,195 @@ import operator
 
 
 class SegmentTree:
-    def __init__(self, capacity, operation, neutral_element):
+    """A Segment Tree data structure.
-        """Build a Segment Tree data structure.
 
-        https://en.wikipedia.org/wiki/Segment_tree
+    https://en.wikipedia.org/wiki/Segment_tree
 
-        Can be used as regular array, but with two
+    Can be used as regular array, but with two important differences:
-        important differences:
 
-          a) setting item's value is slightly slower.
+      a) Setting an item's value is slightly slower. It is O(lg capacity),
-             It is O(lg capacity) instead of O(1).
+         instead of O(1).
-          b) user has access to an efficient `reduce`
+      b) Offers efficient `reduce` operation which reduces the tree's values
-             operation which reduces `operation` over
+         over some specified contiguous subsequence of items in the array.
-             a contiguous subsequence of items in the
+         Operation could be e.g. min/max/sum.
-             array.
 
-        Paramters
+    The data is stored in a list, where the length is 2 * capacity.
-        ---------
+    The second half of the list stores the actual values for each index, so if
-        capacity: int
+    capacity=8, values are stored at indices 8 to 15. The first half of the
-          Total size of the array - must be a power of two.
+    array contains the reduced-values of the different (binary divided)
-        operation: lambda obj, obj -> obj
+    segments, e.g. (capacity=4):
-          and operation for combining elements (eg. sum, max)
+    0=not used
-          must for a mathematical group together with the set of
+    1=reduced-value over all elements (array indices 4 to 7).
-          possible values for array elements.
+    2=reduced-value over array indices (4 and 5).
-        neutral_element: obj
+    3=reduced-value over array indices (6 and 7).
-          neutral element for the operation above. eg. float('-inf')
+    4-7: values of the tree.
-          for max and 0 for sum.
+    NOTE that the values of the tree are accessed by indices starting at 0, so
+    `tree[0]` accesses `internal_array[4]` in the above example.
+    """
+
+    def __init__(self, capacity, operation, neutral_element=None):
+        """Initializes a Segment Tree object.
+
+        Args:
+            capacity (int): Total size of the array - must be a power of two.
+            operation (operation): Lambda obj, obj -> obj
+                The operation for combining elements (eg. sum, max).
+                Must be a mathematical group together with the set of
+                possible values for array elements.
+            neutral_element (Optional[obj]): The neutral element for
+                `operation`. Use None for automatically finding a value:
+                max: float("-inf"), min: float("inf"), sum: 0.0.
         """
 
         assert capacity > 0 and capacity & (capacity - 1) == 0, \
-            "capacity must be positive and a power of 2."
+            "Capacity must be positive and a power of 2!"
-        self._capacity = capacity
+        self.capacity = capacity
-        self._value = [neutral_element for _ in range(2 * capacity)]
+        if neutral_element is None:
-        self._operation = operation
+            neutral_element = 0.0 if operation is operator.add else \
-
+                float("-inf") if operation is max else float("inf")
-    def _reduce_helper(self, start, end, node, node_start, node_end):
+        self.neutral_element = neutral_element
-        if start == node_start and end == node_end:
+        self.value = [self.neutral_element for _ in range(2 * capacity)]
-            return self._value[node]
+        self.operation = operation
-        mid = (node_start + node_end) // 2
-        if end <= mid:
-            return self._reduce_helper(start, end, 2 * node, node_start, mid)
-        else:
-            if mid + 1 <= start:
-                return self._reduce_helper(start, end, 2 * node + 1, mid + 1,
-                                           node_end)
-            else:
-                return self._operation(
-                    self._reduce_helper(start, mid, 2 * node, node_start, mid),
-                    self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1,
-                                        node_end))
 
     def reduce(self, start=0, end=None):
-        """Returns result of applying `self.operation`
+        """Applies `self.operation` to subsequence of our values.
-        to a contiguous subsequence of the array.
+
+        Subsequence is contiguous, includes `start` and excludes `end`.
 
           self.operation(
               arr[start], operation(arr[start+1], operation(... arr[end])))
 
-        Parameters
+        Args:
-        ----------
+            start (int): Start index to apply reduction to.
-        start: int
+            end (Optional[int]): End index to apply reduction to (excluded).
-          beginning of the subsequence
-        end: int
-          end of the subsequences
 
-        Returns
+        Returns:
-        -------
+            any: The result of reducing self.operation over the specified
-        reduced: obj
+                range of `self._value` elements.
-          result of reducing self.operation over the specified range of array
-          elements.
         """
         if end is None:
-            end = self._capacity - 1
+            end = self.capacity
-        if end < 0:
+        elif end < 0:
-            end += self._capacity
+            end += self.capacity
-        return self._reduce_helper(start, end, 1, 0, self._capacity - 1)
+
+        # Init result with neutral element.
+        result = self.neutral_element
+        # Map start/end to our actual index space (second half of array).
+        start += self.capacity
+        end += self.capacity
+
+        # Example:
+        # internal-array (first half=sums, second half=actual values):
+        # 0 1 2 3 | 4 5 6 7
+        # - 6 1 5 | 1 0 2 3
+
+        # tree.sum(0, 3) = 3
+        # internally: start=4, end=7 -> sum values 1 0 2 = 3.
+
+        # Iterate over tree starting in the actual-values (second half)
+        # section.
+        # 1) start=4 is even -> do nothing.
+        # 2) end=7 is odd -> end-- -> end=6 -> add value to result: result=2
+        # 3) int-divide start and end by 2: start=2, end=3
+        # 4) start still smaller end -> iterate once more.
+        # 5) start=2 is even -> do nothing.
+        # 6) end=3 is odd -> end-- -> end=2 -> add value to result: result=1
+        #    NOTE: This adds the sum of indices 4 and 5 to the result.
+
+        # Iterate as long as start != end.
+        while start < end:
+
+            # If start is odd: Add its value to result and move start to
+            # next even value.
+            if start & 1:
+                result = self.operation(result, self.value[start])
+                start += 1
+
+            # If end is odd: Move end to previous even value, then add its
+            # value to result. NOTE: This takes care of excluding `end` in any
+            # situation.
+            if end & 1:
+                end -= 1
+                result = self.operation(result, self.value[end])
+
+            # Divide both start and end by 2 to make them "jump" into the
+            # next upper level reduce-index space.
+            start //= 2
+            end //= 2
+
+            # Then repeat till start == end.
+
+        return result
 
     def __setitem__(self, idx, val):
-        # index of the leaf
+        """
-        idx += self._capacity
+        Inserts/overwrites a value in/into the tree.
-        self._value[idx] = val
+
-        idx //= 2
+        Args:
+            idx (int): The index to insert to. Must be in [0, `self.capacity`[
+            val (float): The value to insert.
+        """
+        assert 0 <= idx < self.capacity
+
+        # Index of the leaf to insert into (always insert in "second half"
+        # of the tree, the first half is reserved for already calculated
+        # reduction-values).
+        idx += self.capacity
+        self.value[idx] = val
+
+        # Recalculate all affected reduction values (in "first half" of tree).
+        idx = idx >> 1  # Divide by 2 (faster than division).
         while idx >= 1:
-            self._value[idx] = self._operation(self._value[2 * idx],
+            update_idx = 2 * idx  # calculate only once
-                                               self._value[2 * idx + 1])
+            # Update the reduction value at the correct "first half" idx.
-            idx //= 2
+            self.value[idx] = self.operation(self.value[update_idx],
+                                             self.value[update_idx + 1])
+            idx = idx >> 1  # Divide by 2 (faster than division).
 
     def __getitem__(self, idx):
-        assert 0 <= idx < self._capacity
+        assert 0 <= idx < self.capacity
-        return self._value[self._capacity + idx]
+        return self.value[idx + self.capacity]
 
 
 class SumSegmentTree(SegmentTree):
+    """A SegmentTree with the reduction `operation`=operator.add."""
+
     def __init__(self, capacity):
         super(SumSegmentTree, self).__init__(
-            capacity=capacity, operation=operator.add, neutral_element=0.0)
+            capacity=capacity, operation=operator.add)
 
     def sum(self, start=0, end=None):
-        """Returns arr[start] + ... + arr[end]"""
+        """Returns the sum over a sub-segment of the tree."""
-        return super(SumSegmentTree, self).reduce(start, end)
+        return self.reduce(start, end)
 
     def find_prefixsum_idx(self, prefixsum):
-        """Find the highest index `i` in the array such that
+        """Finds highest i, for which: sum(arr[0]+..+arr[i - i]) <= prefixsum.
-          sum(arr[0] + arr[1] + ... + arr[i - i]) <= prefixsum
 
-        if array values are probabilities, this function
+        Args:
-        allows to sample indexes according to the discrete
+            prefixsum (float): `prefixsum` upper bound in above constraint.
-        probability efficiently.
 
-        Parameters
+        Returns:
-        ----------
+            int: Largest possible index (i) satisfying above constraint.
-        perfixsum: float
-          upperbound on the sum of array prefix
-
-        Returns
-        -------
-        idx: int
-          highest index satisfying the prefixsum constraint
         """
         assert 0 <= prefixsum <= self.sum() + 1e-5
+        # Global sum node.
         idx = 1
-        while idx < self._capacity:  # while non-leaf
+
-            if self._value[2 * idx] > prefixsum:
+        # While non-leaf (first half of tree).
-                idx = 2 * idx
+        while idx < self.capacity:
+            update_idx = 2 * idx
+            if self.value[update_idx] > prefixsum:
+                idx = update_idx
             else:
-                prefixsum -= self._value[2 * idx]
+                prefixsum -= self.value[update_idx]
-                idx = 2 * idx + 1
+                idx = update_idx + 1
-        return idx - self._capacity
+        return idx - self.capacity
 
 
 class MinSegmentTree(SegmentTree):
     def __init__(self, capacity):
-        super(MinSegmentTree, self).__init__(
+        super(MinSegmentTree, self).__init__(capacity=capacity, operation=min)
-            capacity=capacity, operation=min, neutral_element=float("inf"))
 
     def min(self, start=0, end=None):
         """Returns min(arr[start], ...,  arr[end])"""
-        return super(MinSegmentTree, self).reduce(start, end)
+        return self.reduce(start, end)

rllib/optimizers/tests/old_segment_tree.py

@@ -0,0 +1,143 @@
+import operator
+
+
+class OldSegmentTree(object):
+    def __init__(self, capacity, operation, neutral_element):
+        """Build a Segment Tree data structure.
+
+        https://en.wikipedia.org/wiki/Segment_tree
+
+        Can be used as regular array, but with two
+        important differences:
+
+          a) setting item's value is slightly slower.
+             It is O(lg capacity) instead of O(1).
+          b) user has access to an efficient `reduce`
+             operation which reduces `operation` over
+             a contiguous subsequence of items in the
+             array.
+
+        Paramters
+        ---------
+        capacity: int
+          Total size of the array - must be a power of two.
+        operation: lambda obj, obj -> obj
+          and operation for combining elements (eg. sum, max)
+          must for a mathematical group together with the set of
+          possible values for array elements.
+        neutral_element: obj
+          neutral element for the operation above. eg. float('-inf')
+          for max and 0 for sum.
+        """
+
+        assert capacity > 0 and capacity & (capacity - 1) == 0, \
+            "capacity must be positive and a power of 2."
+        self._capacity = capacity
+        self._value = [neutral_element for _ in range(2 * capacity)]
+        self._operation = operation
+
+    def _reduce_helper(self, start, end, node, node_start, node_end):
+        if start == node_start and end == node_end:
+            return self._value[node]
+        mid = (node_start + node_end) // 2
+        if end <= mid:
+            return self._reduce_helper(start, end, 2 * node, node_start, mid)
+        else:
+            if mid + 1 <= start:
+                return self._reduce_helper(start, end, 2 * node + 1, mid + 1,
+                                           node_end)
+            else:
+                return self._operation(
+                    self._reduce_helper(start, mid, 2 * node, node_start, mid),
+                    self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1,
+                                        node_end))
+
+    def reduce(self, start=0, end=None):
+        """Returns result of applying `self.operation`
+        to a contiguous subsequence of the array.
+
+          self.operation(
+              arr[start], operation(arr[start+1], operation(... arr[end])))
+
+        Parameters
+        ----------
+        start: int
+          beginning of the subsequence
+        end: int
+          end of the subsequences
+
+        Returns
+        -------
+        reduced: obj
+          result of reducing self.operation over the specified range of array
+          elements.
+        """
+        if end is None:
+            end = self._capacity
+        if end < 0:
+            end += self._capacity
+        end -= 1
+        return self._reduce_helper(start, end, 1, 0, self._capacity - 1)
+
+    def __setitem__(self, idx, val):
+        # index of the leaf
+        idx += self._capacity
+        self._value[idx] = val
+        idx //= 2
+        while idx >= 1:
+            self._value[idx] = self._operation(self._value[2 * idx],
+                                               self._value[2 * idx + 1])
+            idx //= 2
+
+    def __getitem__(self, idx):
+        assert 0 <= idx < self._capacity
+        return self._value[self._capacity + idx]
+
+
+class OldSumSegmentTree(OldSegmentTree):
+    def __init__(self, capacity):
+        super(OldSumSegmentTree, self).__init__(
+            capacity=capacity, operation=operator.add, neutral_element=0.0)
+
+    def sum(self, start=0, end=None):
+        """Returns arr[start] + ... + arr[end]"""
+        return super(OldSumSegmentTree, self).reduce(start, end)
+
+    def find_prefixsum_idx(self, prefixsum):
+        """Find the highest index `i` in the array such that
+          sum(arr[0] + arr[1] + ... + arr[i - i]) <= prefixsum
+
+        if array values are probabilities, this function
+        allows to sample indexes according to the discrete
+        probability efficiently.
+
+        Parameters
+        ----------
+        perfixsum: float
+          upperbound on the sum of array prefix
+
+        Returns
+        -------
+        idx: int
+          highest index satisfying the prefixsum constraint
+        """
+        assert 0 <= prefixsum <= self.sum() + 1e-5
+        idx = 1
+        while idx < self._capacity:  # while non-leaf
+            if self._value[2 * idx] > prefixsum:
+                idx = 2 * idx
+            else:
+                prefixsum -= self._value[2 * idx]
+                idx = 2 * idx + 1
+        return idx - self._capacity
+
+
+class OldMinSegmentTree(OldSegmentTree):
+    def __init__(self, capacity):
+        super(OldMinSegmentTree, self).__init__(
+            capacity=capacity, operation=min, neutral_element=float("inf"))
+
+    def min(self, start=0, end=None):
+        """Returns min(arr[start], ...,  arr[end])"""
+
+        return super(OldMinSegmentTree, self).reduce(start, end)

rllib/optimizers/tests/test_prioritized_replay_buffer.py

@@ -0,0 +1,180 @@
+from collections import Counter
+import numpy as np
+import unittest
+
+from ray.rllib.optimizers.replay_buffer import PrioritizedReplayBuffer
+from ray.rllib.utils.test_utils import check
+
+
+class TestPrioritizedReplayBuffer(unittest.TestCase):
+    """
+    Tests insertion and (weighted) sampling of the PrioritizedReplayBuffer.
+    """
+
+    capacity = 10
+    alpha = 1.0
+    beta = 1.0
+    max_priority = 1.0
+
+    def _generate_data(self):
+        return (
+            np.random.random((4, )),  # obs_t
+            np.random.choice([0, 1]),  # action
+            np.random.rand(),  # reward
+            np.random.random((4, )),  # obs_tp1
+            np.random.choice([False, True]),  # done
+        )
+
+    def test_add(self):
+        memory = PrioritizedReplayBuffer(
+            size=2,
+            alpha=self.alpha,
+        )
+
+        # Assert indices 0 before insert.
+        self.assertEqual(len(memory), 0)
+        self.assertEqual(memory._next_idx, 0)
+
+        # Insert single record.
+        data = self._generate_data()
+        memory.add(*data, weight=0.5)
+        self.assertTrue(len(memory) == 1)
+        self.assertTrue(memory._next_idx == 1)
+
+        # Insert single record.
+        data = self._generate_data()
+        memory.add(*data, weight=0.1)
+        self.assertTrue(len(memory) == 2)
+        self.assertTrue(memory._next_idx == 0)
+
+        # Insert over capacity.
+        data = self._generate_data()
+        memory.add(*data, weight=1.0)
+        self.assertTrue(len(memory) == 2)
+        self.assertTrue(memory._next_idx == 1)
+
+    def test_update_priorities(self):
+        memory = PrioritizedReplayBuffer(size=self.capacity, alpha=self.alpha)
+
+        # Insert n samples.
+        num_records = 5
+        for i in range(num_records):
+            data = self._generate_data()
+            memory.add(*data, weight=1.0)
+            self.assertTrue(len(memory) == i + 1)
+            self.assertTrue(memory._next_idx == i + 1)
+
+        # Fetch records, their indices and weights.
+        _, _, _, _, _, weights, indices = \
+            memory.sample(3, beta=self.beta)
+        check(weights, np.ones(shape=(3, )))
+        self.assertEqual(3, len(indices))
+        self.assertTrue(len(memory) == num_records)
+        self.assertTrue(memory._next_idx == num_records)
+
+        # Update weight of indices 0, 2, 3, 4 to very small.
+        memory.update_priorities(
+            np.array([0, 2, 3, 4]), np.array([0.01, 0.01, 0.01, 0.01]))
+        # Expect to sample almost only index 1
+        # (which still has a weight of 1.0).
+        for _ in range(10):
+            _, _, _, _, _, weights, indices = memory.sample(
+                1000, beta=self.beta)
+            self.assertTrue(970 < np.sum(indices) < 1100)
+
+        # Update weight of indices 0 and 1 to >> 0.01.
+        # Expect to sample 0 and 1 equally (and some 2s, 3s, and 4s).
+        for _ in range(10):
+            rand = np.random.random() + 0.2
+            memory.update_priorities(np.array([0, 1]), np.array([rand, rand]))
+            _, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
+            # Expect biased to higher values due to some 2s, 3s, and 4s.
+            # print(np.sum(indices))
+            self.assertTrue(400 < np.sum(indices) < 800)
+
+        # Update weights to be 1:2.
+        # Expect to sample double as often index 1 over index 0
+        # plus very few times indices 2, 3, or 4.
+        for _ in range(10):
+            rand = np.random.random() + 0.2
+            memory.update_priorities(
+                np.array([0, 1]), np.array([rand, rand * 2]))
+            _, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
+            # print(np.sum(indices))
+            self.assertTrue(600 < np.sum(indices) < 850)
+
+        # Update weights to be 1:4.
+        # Expect to sample quadruple as often index 1 over index 0
+        # plus very few times indices 2, 3, or 4.
+        for _ in range(10):
+            rand = np.random.random() + 0.2
+            memory.update_priorities(
+                np.array([0, 1]), np.array([rand, rand * 4]))
+            _, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
+            # print(np.sum(indices))
+            self.assertTrue(750 < np.sum(indices) < 950)
+
+        # Update weights to be 1:9.
+        # Expect to sample 9 times as often index 1 over index 0.
+        # plus very few times indices 2, 3, or 4.
+        for _ in range(10):
+            rand = np.random.random() + 0.2
+            memory.update_priorities(
+                np.array([0, 1]), np.array([rand, rand * 9]))
+            _, _, _, _, _, _, indices = memory.sample(1000, beta=self.beta)
+            # print(np.sum(indices))
+            self.assertTrue(850 < np.sum(indices) < 1100)
+
+        # Insert n more samples.
+        num_records = 5
+        for i in range(num_records):
+            data = self._generate_data()
+            memory.add(*data, weight=1.0)
+            self.assertTrue(len(memory) == i + 6)
+            self.assertTrue(memory._next_idx == (i + 6) % self.capacity)
+
+        # Update all weights to be 1.0 to 10.0 and sample a >100 batch.
+        memory.update_priorities(
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+            np.array([0.001, 0.1, 2., 8., 16., 32., 64., 128., 256., 512.]))
+        counts = Counter()
+        for _ in range(10):
+            _, _, _, _, _, _, indices = memory.sample(
+                np.random.randint(100, 600), beta=self.beta)
+            for i in indices:
+                counts[i] += 1
+        print(counts)
+        # Expect an approximately correct distribution of indices.
+        self.assertTrue(
+            counts[9] >= counts[8] >= counts[7] >= counts[6] >= counts[5] >=
+            counts[4] >= counts[3] >= counts[2] >= counts[1] >= counts[0])
+
+    def test_alpha_parameter(self):
+        # Test sampling from a PR with a very small alpha (should behave just
+        # like a regular ReplayBuffer).
+        memory = PrioritizedReplayBuffer(size=self.capacity, alpha=0.01)
+
+        # Insert n samples.
+        num_records = 5
+        for i in range(num_records):
+            data = self._generate_data()
+            memory.add(*data, weight=np.random.rand())
+            self.assertTrue(len(memory) == i + 1)
+            self.assertTrue(memory._next_idx == i + 1)
+
+        # Fetch records, their indices and weights.
+        _, _, _, _, _, weights, indices = \
+            memory.sample(1000, beta=self.beta)
+        counts = Counter()
+        for i in indices:
+            counts[i] += 1
+        print(counts)
+        # Expect an approximately uniform distribution of indices.
+        for i in counts.values():
+            self.assertTrue(100 < i < 300)
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/optimizers/tests/test_segment_tree.py

@@ -1,4 +1,5 @@
 import numpy as np
+import timeit
 import unittest
 
 from ray.rllib.optimizers.segment_tree import SumSegmentTree, MinSegmentTree
@@ -17,6 +18,7 @@ class TestSegmentTree(unittest.TestCase):
         assert np.isclose(tree.sum(2, 3), 1.0)
         assert np.isclose(tree.sum(2, -1), 1.0)
         assert np.isclose(tree.sum(2, 4), 4.0)
+        assert np.isclose(tree.sum(2), 4.0)
 
     def test_tree_set_overlap(self):
         tree = SumSegmentTree(4)
@@ -28,6 +30,7 @@ class TestSegmentTree(unittest.TestCase):
         assert np.isclose(tree.sum(2, 3), 3.0)
         assert np.isclose(tree.sum(2, -1), 3.0)
         assert np.isclose(tree.sum(2, 4), 3.0)
+        assert np.isclose(tree.sum(2), 3.0)
         assert np.isclose(tree.sum(1, 2), 0.0)
 
     def test_prefixsum_idx(self):
@@ -92,6 +95,37 @@ class TestSegmentTree(unittest.TestCase):
         assert np.isclose(tree.min(2, -1), 4.0)
         assert np.isclose(tree.min(3, 4), 3.0)
 
+    def test_microbenchmark_vs_old_version(self):
+        """
+        Results from March 2020 (capacity=1048576):
+
+        New tree:
+        0.049599366000000256s
+        results = timeit.timeit("tree.sum(5, 60000)",
+            setup="from ray.rllib.optimizers.segment_tree import
+            SumSegmentTree; tree = SumSegmentTree({})".format(capacity),
+            number=10000)
+
+        Old tree:
+        0.13390400999999974s
+        results = timeit.timeit("tree.sum(5, 60000)",
+            setup="from ray.rllib.optimizers.tests.old_segment_tree import
+            OldSumSegmentTree; tree = OldSumSegmentTree({})".format(capacity),
+            number=10000)
+        """
+        capacity = 2**20
+        new = timeit.timeit(
+            "tree.sum(5, 60000)",
+            setup="from ray.rllib.optimizers.segment_tree import "
+            "SumSegmentTree; tree = SumSegmentTree({})".format(capacity),
+            number=10000)
+        old = timeit.timeit(
+            "tree.sum(5, 60000)",
+            setup="from ray.rllib.optimizers.tests.old_segment_tree import "
+            "OldSumSegmentTree; tree = OldSumSegmentTree({})".format(capacity),
+            number=10000)
+        self.assertGreater(old, new)
+
 
 if __name__ == "__main__":
     import pytest