Fix Rainbow segment tree sampling for non-power-of-two capacities

[discobot]

Description

Fixes #546.

SumSegmentTree.retrieve() returns permuted indices for any non-power-of-two capacity, including the default --buffer-size 1000000.

One correction to the thread: this code did not come from stable-baselines3 — neither SB3 nor sb3-contrib contains a SegmentTree (GitHub code search returns zero hits in both, while e.g. ReplayBuffer returns plenty). The in-file comment credits openai/baselines segment_tree.py, whose constructor asserts "capacity must be positive and a power of 2."; the adaptation in #509 dropped that invariant. So there is no upstream to report to, and this PR fixes it here directly.

The fix pads the leaf array to the next power of two (0.0 in the sum tree, float("inf") in the min tree), keeping the public API unchanged. retrieve() also clamps its result to capacity - 1: float round-off can step one leaf past the last real index when the sampled value is within an ulp of total(), and unlike the old layout the padded one would otherwise turn that into an out-of-range index. sample() now reads sampled priorities through a new SumSegmentTree.get() accessor instead of indexing the tree array with the old layout's offset.

Worth noting for anyone assessing impact on past results: the marginal sampling distribution was never broken — each index still occupied one contiguous mass interval, just in permuted order — so prioritized sampling stayed exactly proportional to priorities. The new test suite includes a 200k-draw proportionality test that passes on both the old and new code, alongside prefix-order regression tests (capacities 1–33) that fail on master. The new test file is wired into the core CI job (CI lists test files explicitly), and the existing test_rainbow_atari smoke (--buffer-size 10, itself non-power-of-two) passes locally on CPU.

Types of changes

• Bug fix
• New feature
• New algorithm
• Documentation

Checklist:

• I've read the CONTRIBUTION guide (required).
• I have ensured pre-commit run --all-files passes (required).
• I have updated the tests accordingly (if applicable).
• I have updated the documentation and previewed the changes via mkdocs serve.
  • I have explained note-worthy implementation details.
  • I have explained the logged metrics.
  • I have added links to the original paper and related papers.

If you need to run benchmark experiments for a performance-impacting changes:

• I have contacted @vwxyzjn to obtain access to the openrlbenchmark W&B team.
• I have used the benchmark utility to submit the tracked experiments to the openrlbenchmark/cleanrl W&B project, optionally with --capture_video.
• I have performed RLops with python -m openrlbenchmark.rlops.
  • For new feature or bug fix:
    • I have used the RLops utility to understand the performance impact of the changes and confirmed there is no regression.
  • For new algorithm:
    • I have created a table comparing my results against those from reputable sources (i.e., the original paper or other reference implementation).
  • I have added the learning curves generated by the python -m openrlbenchmark.rlops utility to the documentation.
  • I have added links to the tracked experiments in W&B, generated by python -m openrlbenchmark.rlops ....your_args... --report, to the documentation.

[vercel]

Someone is attempting to deploy a commit to the Costa Huang's projects Team on Vercel.

A member of the Team first needs to authorize it.