Bug in Usage of SumSegmentTree

I was porting your RainbowDQN agent to DeepMind's bsuite repo, and Codex caught a subtle bug in your SumSegmentTree implementation/usage. I believe it is correct!

SumSegmentTree uses heap indexing with 2 * capacity - 1 nodes and places leaves at idx + capacity - 1. That only preserves left-to-right leaf order when the leaf count is power-of-two padded. For non-power-of-two capacities, retrieve() walks a tree whose subtrees no longer correspond to contiguous replay-index ranges, so prefix-sum sampling returns wrong indices. A minimal counterexample is capacity=3: equal priorities should map prefix masses 0.5, 1.5, 2.5 to indices 0, 1, 2, but the current layout does not preserve that ordering.