[rllib] Fix segment_tree.py edge case

rllib/execution/segment_tree.py

@@ -136,7 +136,7 @@ def __setitem__(self, idx: int, val: float) -> None:
         Inserts/overwrites a value in/into the tree.
 
         Args:
-            idx: The index to insert to. Must be in [0, `self.capacity`[
+            idx: The index to insert to. Must be in [0, `self.capacity`)
             val: The value to insert.
         """
         assert 0 <= idx < self.capacity, f"idx={idx} capacity={self.capacity}"
@@ -192,6 +192,11 @@ def find_prefixsum_idx(self, prefixsum: float) -> int:
         # Global sum node.
         idx = 1
 
+        # Edge case when prefixsum can clip into the invalid regions
+        #   https://github.com/ray-project/ray/issues/54284
+        if prefixsum >= self.value[idx]:
+            prefixsum = self.value[idx] - 1e-5
+
         # While non-leaf (first half of tree).
         while idx < self.capacity:
             update_idx = 2 * idx

rllib/utils/replay_buffers/tests/test_segment_tree_replay_buffer_api.py

@@ -1,7 +1,9 @@
 import numpy as np
 import unittest
 
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
 from ray.rllib.execution.segment_tree import SumSegmentTree, MinSegmentTree
+from ray.rllib.utils.replay_buffers import PrioritizedEpisodeReplayBuffer
 
 
 class TestSegmentTree(unittest.TestCase):
@@ -94,6 +96,52 @@ def test_max_interval_tree(self):
         assert np.isclose(tree.min(2, -1), 4.0)
         assert np.isclose(tree.min(3, 4), 3.0)
 
+    @staticmethod
+    def _get_episode(episode_len=None, id_=None, with_extra_model_outs=False):
+        eps = SingleAgentEpisode(id_=id_, observations=[0.0], infos=[{}])
+        ts = np.random.randint(1, 200) if episode_len is None else episode_len
+        for t in range(ts):
+            eps.add_env_step(
+                observation=float(t + 1),
+                action=int(t),
+                reward=0.1 * (t + 1),
+                infos={},
+                extra_model_outputs=(
+                    {k: k for k in range(2)} if with_extra_model_outs else None
+                ),
+            )
+        eps.is_terminated = np.random.random() > 0.5
+        eps.is_truncated = False if eps.is_terminated else np.random.random() > 0.8
+        return eps
+
+    def test_find_prefixsum_idx(self, buffer_size=80):
+        """Fix edge case related to https://github.com/ray-project/ray/issues/54284"""
+        replay_buffer = PrioritizedEpisodeReplayBuffer(capacity=buffer_size)
+        sum_segment = replay_buffer._sum_segment
+
+        for i in range(10):
+            replay_buffer.add(self._get_episode(id_=str(i), episode_len=10))
+
+        assert sum_segment.capacity >= buffer_size
+
+        # standard cases
+        for sample in np.linspace(0, sum_segment.sum(), 50):
+            prefixsum_idx = sum_segment.find_prefixsum_idx(sample)
+            assert (
+                prefixsum_idx in replay_buffer._tree_idx_to_sample_idx
+            ), f"{sum_segment.sum()=}, {sample=}, {prefixsum_idx=}"
+
+        # edge cases
+        for sample in [
+            sum_segment.sum() - 0.00001,
+            sum_segment.sum(),
+            sum_segment.sum() + 0.00001,
+        ]:
+            prefixsum_idx = sum_segment.find_prefixsum_idx(sample)
+            assert (
+                prefixsum_idx in replay_buffer._tree_idx_to_sample_idx
+            ), f"{sum_segment.sum()=}, {sample=}, {prefixsum_idx=}"
+
 
 if __name__ == "__main__":
     import pytest