Skip to content

[Feature] Dispatch RandomSampler/SliceSampler to without-replacement variants via replacement=False#3749

Open
vmoens wants to merge 1 commit into
pytorch:mainfrom
vmoens:feature/sampler-replacement-arg
Open

[Feature] Dispatch RandomSampler/SliceSampler to without-replacement variants via replacement=False#3749
vmoens wants to merge 1 commit into
pytorch:mainfrom
vmoens:feature/sampler-replacement-arg

Conversation

@vmoens
Copy link
Copy Markdown
Collaborator

@vmoens vmoens commented May 13, 2026

Summary

Adds a replacement keyword to the with-replacement sampler classes so users can opt into the without-replacement variant from the simpler constructor:

RandomSampler(replacement=False)               # -> SamplerWithoutReplacement
RandomSampler(replacement=False, drop_last=True, shuffle=False)
SliceSampler(replacement=False, slice_len=5)   # -> SliceSamplerWithoutReplacement
SliceSampler(replacement=False, slice_len=5, drop_last=True)

Implementation: Sampler gains an ABCMeta-derived metaclass _SamplerMeta whose __call__ pops replacement from kwargs and dispatches via a small registry _REPLACEMENT_DISPATCH ({RandomSampler: SamplerWithoutReplacement, SliceSampler: SliceSamplerWithoutReplacement}). When replacement=True (default) the metaclass is a transparent no-op. Samplers that have no without-replacement variant (e.g. PrioritizedSampler) raise TypeError on replacement=False.

The existing SamplerWithoutReplacement / SliceSamplerWithoutReplacement / PrioritizedSliceSampler classes are untouched and remain the canonical entry points; the new keyword is purely additive.

Test plan

  • test/test_rb.py::TestSamplers::test_replacement_kwarg_* (6 new tests covering type identity, default forwarding, kwarg forwarding, isinstance invariants, subclass non-interference, error on unsupported variants, and end-to-end use in a ReplayBuffer for both RandomSampler and SliceSampler dispatch).
  • Full test/test_rb.py suite passes (3929 passed, no regressions).
  • Pickle round-trip confirmed for both dispatch and non-dispatch cases.

🤖 Generated with Claude Code

…variant via replacement=False

A metaclass on Sampler intercepts the ``replacement`` keyword: when False,
``RandomSampler(...)`` is dispatched to ``SamplerWithoutReplacement(...)`` and
``SliceSampler(...)`` to ``SliceSamplerWithoutReplacement(...)``. Additional
kwargs (``drop_last``, ``shuffle``) are forwarded to the variant constructor.
Subclasses without a registered variant raise ``TypeError`` when invoked with
``replacement=False``.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented May 13, 2026

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

❌ 2 New Failures, 1 Pending

As of commit 5c12071 with merge base cc31dc3 (image):

NEW FAILURES - The following jobs have failed:

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Feature New feature ReplayBuffers

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant