[Kayvan-Zahiri]

Description

With strict_length=True, SliceSamplerWithoutReplacement permutes over every trajectory in the storage, short ones included, and _sample_slices drops the short ones only after they've been drawn. Two things go wrong:

• A batch that drew a short trajectory comes back short.
• When every trajectory in a draw is short, sample() raises RuntimeError: None of the provided indices pointed to a trajectory of sufficient length in the middle of an epoch.

With trajectory lengths [7, 2, 9, 1, 5, 12, 4, 3], slice_len=3 and batch_size=6, on main:

flat:        of 400 sample() calls, 17 raised RuntimeError; 166 of the 383 returned batches were short
batch_time:  of 400 sample() calls, 19 raised RuntimeError; 162 of the 381 returned batches were short

The strict_length docstring says short trajectories are filtered out, and the with-replacement SliceSampler already filters before it draws. Now the without-replacement sampler permutes only over trajectories with length >= seq_length and maps the drawn positions back to trajectory indices, so an epoch covers exactly the eligible trajectories. If none is long enough, it falls back to the old draw, so the existing "Did not find a single trajectory with sufficient length" error is unchanged. With strict_length=False, nothing changes.

Tests

test_slice_sampler_without_replacement_strict_skips_short is parametrized over flat and batch_time. It checks that 30 draws are all full batches, that no short trajectory is drawn, and that each epoch covers every eligible trajectory. It fails on main and passes here. The 22 without-replacement tests in test/rb/test_samplers.py pass. Across the whole file, the failure set is the same 118 on main and on this branch: all are SumSegmentTreeFp32 is not available, because the C++ extension isn't built in my environment. ufmt (black 22.3.0, usort 1.0.3) and flake8 with the repo config are clean.

I found this with a property check over SliceSampler configurations. No open issue or PR covers it. It is independent of #4583, which fixes the ends=/trajectories= boundaries. AI assistance (Claude Code) was used for the investigation and the patch.

• I have read the CONTRIBUTION guide (required)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4584

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 16 Awaiting Approval

As of commit 5b6e976 with merge base e7a05a7 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Examples Tests on Linux (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.