[rllib] Fix segment_tree.py edge case #57599

pseudo-rnd-thoughts · 2025-10-09T13:26:59Z

Why are these changes needed?

SegmentTree is a component of the rllib PrioritizedEpisodeReplayBuffer however for extreme edge case prefix sum values then find_prefixsum_idx will return invalid out of bounds value.
I couldn't find a bug, rather if the prefixsum_value is equal to the SegmentTree.sum() then traversing down the tree could cause it to return invalid indexes.

I've added unittests to reproduce the original error and check against it

Related issue number

Close #54284

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run pre-commit jobs to lint the changes in this PR. (pre-commit setup)
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: Mark Towers <mark@anyscale.com>

gemini-code-assist

Code Review

This pull request addresses an edge case in SegmentTree.find_prefixsum_idx that could lead to returning an invalid index, particularly when the prefixsum value aligns with the sum of a subtree. The fix involves a small adjustment to the prefixsum value. The addition of a unit test to reproduce and verify the fix is a great step. My feedback focuses on making the logical fix more robust and less dependent on a fixed magic number, which could be brittle under certain conditions.

rllib/execution/segment_tree.py

Signed-off-by: Mark Towers <mark@anyscale.com>

simonsays1980

Awesome PR @pseudo-rnd-thoughts!! Thanks for fixing this hard-to-find one.

simonsays1980 · 2025-10-14T13:30:04Z

rllib/execution/segment_tree.py


        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`)


simonsays1980 · 2025-10-14T13:31:41Z

rllib/execution/segment_tree.py


+        # Edge case when prefixsum can clip into the invalid regions
+        #   https://github.com/ray-project/ray/issues/54284
+        if prefixsum >= self.value[idx]:


Awesome! Simple fix for a big problem :)

rllib/utils/replay_buffers/tests/test_segment_tree_replay_buffer_api.py

simonsays1980 · 2025-10-14T15:36:52Z

rllib/utils/replay_buffers/tests/test_segment_tree_replay_buffer_api.py

+            sum_segment.sum() + 0.00001,
+        ]:
+            prefixsum_idx = sum_segment.find_prefixsum_idx(sample)
+            assert (


Again, like above. Could we use the TestCase checking methods or the one from us? Or is there a specific reason why you would prefer the assert here?

I'm more familiar with pytest rather than unittest and the rest of the file uses assert rather than self.assertTrue.

I've updated to use TestCase implementation

simonsays1980 · 2025-10-14T15:36:55Z

rllib/utils/replay_buffers/tests/test_segment_tree_replay_buffer_api.py

+            ), f"{sum_segment.sum()=}, {sample=}, {prefixsum_idx=}"
+
+        # edge cases
+        for sample in [


Very nice! Could we add a comment of why this case could cause problems on the SumSegment?

Signed-off-by: Mark Towers <mark@anyscale.com>

simonsays1980

LGTM

# Conflicts: # rllib/utils/replay_buffers/tests/test_segment_tree_replay_buffer_api.py

Signed-off-by: Mark Towers <mark@anyscale.com>

cursor · 2025-10-21T11:56:28Z

rllib/execution/segment_tree.py

+        # 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


Bug: Prefix Sum Edge Case Handling Fails

The edge case handling in _sample_prefixsum incorrectly modifies the prefixsum. It clamps valid input values that are slightly above the total sum, and can also result in a negative prefixsum when the total sum is very small or zero. This leads to assertion failures and logically inconsistent states for tree traversal.

## Why are these changes needed? `SegmentTree` is a component of the rllib `PrioritizedEpisodeReplayBuffer` however for extreme edge case prefix sum values then `find_prefixsum_idx` will return invalid out of bounds value. I couldn't find a bug, rather if the `prefixsum_value` is equal to the `SegmentTree.sum()` then traversing down the tree could cause it to return invalid indexes. I've added unittests to reproduce the original error and check against it ## Related issue number Close ray-project#54284 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Mark Towers <mark@anyscale.com> Co-authored-by: Mark Towers <mark@anyscale.com> Co-authored-by: simonsays1980 <simon.zehnder@gmail.com> Signed-off-by: xgui <xgui@anyscale.com>

## Why are these changes needed? `SegmentTree` is a component of the rllib `PrioritizedEpisodeReplayBuffer` however for extreme edge case prefix sum values then `find_prefixsum_idx` will return invalid out of bounds value. I couldn't find a bug, rather if the `prefixsum_value` is equal to the `SegmentTree.sum()` then traversing down the tree could cause it to return invalid indexes. I've added unittests to reproduce the original error and check against it ## Related issue number Close #54284 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Mark Towers <mark@anyscale.com> Co-authored-by: Mark Towers <mark@anyscale.com> Co-authored-by: simonsays1980 <simon.zehnder@gmail.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

## Why are these changes needed? `SegmentTree` is a component of the rllib `PrioritizedEpisodeReplayBuffer` however for extreme edge case prefix sum values then `find_prefixsum_idx` will return invalid out of bounds value. I couldn't find a bug, rather if the `prefixsum_value` is equal to the `SegmentTree.sum()` then traversing down the tree could cause it to return invalid indexes. I've added unittests to reproduce the original error and check against it ## Related issue number Close ray-project#54284 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Mark Towers <mark@anyscale.com> Co-authored-by: Mark Towers <mark@anyscale.com> Co-authored-by: simonsays1980 <simon.zehnder@gmail.com>

## Why are these changes needed? `SegmentTree` is a component of the rllib `PrioritizedEpisodeReplayBuffer` however for extreme edge case prefix sum values then `find_prefixsum_idx` will return invalid out of bounds value. I couldn't find a bug, rather if the `prefixsum_value` is equal to the `SegmentTree.sum()` then traversing down the tree could cause it to return invalid indexes. I've added unittests to reproduce the original error and check against it ## Related issue number Close ray-project#54284 ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [x] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Mark Towers <mark@anyscale.com> Co-authored-by: Mark Towers <mark@anyscale.com> Co-authored-by: simonsays1980 <simon.zehnder@gmail.com> Signed-off-by: Aydin Abiar <aydin@anyscale.com>

[rllib] Fix segment_tree.py edge case

719e1ec

Signed-off-by: Mark Towers <mark@anyscale.com>

pseudo-rnd-thoughts assigned simonsays1980 Oct 9, 2025

pseudo-rnd-thoughts requested a review from a team as a code owner October 9, 2025 13:27

This comment was marked as outdated.

Sign in to view

gemini-code-assist bot reviewed Oct 9, 2025

View reviewed changes

rllib/execution/segment_tree.py Outdated Show resolved Hide resolved

[rllib] Address cursors and gemini's problems

b2d75de

Signed-off-by: Mark Towers <mark@anyscale.com>

This comment was marked as outdated.

Sign in to view

pseudo-rnd-thoughts added the rllib RLlib related issues label Oct 9, 2025

simonsays1980 requested changes Oct 14, 2025

View reviewed changes

simonsays1980 mentioned this pull request Oct 14, 2025

[RLlib] Unexpected KeyError while training SAC #54284

Closed

Code review by Simon

62166e9

Signed-off-by: Mark Towers <mark@anyscale.com>

This comment was marked as outdated.

Sign in to view

simonsays1980 approved these changes Oct 14, 2025

View reviewed changes

simonsays1980 added rllib-algorithms An RLlib algorithm/Trainer is not learning. go add ONLY when ready to merge, run all tests labels Oct 14, 2025

Merge branch 'master' into issue-54284

e7eb784

simonsays1980 enabled auto-merge (squash) October 14, 2025 17:08

github-actions bot disabled auto-merge October 14, 2025 17:08

pseudo-rnd-thoughts and others added 3 commits October 21, 2025 10:51

Merge branch 'master' into issue-54284

a076f58

Merge branch 'master' into issue-54284

3e14aa4

# Conflicts: # rllib/utils/replay_buffers/tests/test_segment_tree_replay_buffer_api.py

pre-commit

31f5500

Signed-off-by: Mark Towers <mark@anyscale.com>

cursor bot reviewed Oct 21, 2025

View reviewed changes

simonsays1980 merged commit 48724e2 into ray-project:master Oct 21, 2025
6 checks passed

[rllib] Fix segment_tree.py edge case #57599

[rllib] Fix segment_tree.py edge case #57599

Conversation

pseudo-rnd-thoughts commented Oct 9, 2025

Why are these changes needed?

Related issue number

Checks

Uh oh!

This comment was marked as outdated.

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

simonsays1980 left a comment

Choose a reason for hiding this comment

Uh oh!

simonsays1980 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

simonsays1980 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

simonsays1980 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

pseudo-rnd-thoughts Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simonsays1980 Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

pseudo-rnd-thoughts Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

simonsays1980 left a comment

Choose a reason for hiding this comment

Uh oh!

cursor bot Oct 21, 2025

Choose a reason for hiding this comment

Bug: Prefix Sum Edge Case Handling Fails

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pseudo-rnd-thoughts Oct 14, 2025 •

edited

Loading