[BUG] outdated PrioritizedSliceSampler implementation

[wertyuilife2]

Hi @vmoens, I encountered the bug described in issue #1969 again. I am using the same TDMPC code and the bug-fixed version (2024.3.26) of torchrl. The only difference is that I am using PrioritizedSliceSampler instead of SliceSampler.

To Reproduce

Please refer to the same reproduction method as described in issue #1969, but replace SliceSampler with PrioritizedSliceSampler.

Expected Behavior

PrioritizedSliceSampler should not fail when at capacity, just like SliceSampler.

Reason and Possible Fixes

The original bug was effectively addressed by commits e835770 and 6fb16a2. However, the changes made to fix the bug in SliceSampler have not been applied to PrioritizedSliceSampler. Specifically, SliceSampler uses the updated _sample_slices() method for sampling, but PrioritizedSliceSampler does not.

Checklist

• [x] I have checked that there is no similar issue in the repo (required)
• [x] I have read the documentation (required)
• [x] I have provided a minimal working example to reproduce the bug (required)

[vmoens]

Thanks for reporting! Should be solved by #2194 if you want to give a look!

[wertyuilife2]

Thank you for the quick response and the fix! I will test the commit and provide feedback soon.

[wertyuilife2]

My code works well with commit #2194, thanks for the fix!