[Feature] Add geometric trajectory window sampler

[vmoens]

Stack

Depends on #4180, which adds the reusable fragmented-trajectory index and SliceSampler(fragmented=True). This PR is targeted at that branch so its diff contains only the geometric window policy and its public integration.

Summary

• add GeometricTrajectoryWindowSampler on top of the shared fragmented-trajectory index
• draw one shared geometric future offset per batch and sample anchors uniformly from eligible stored steps
• return fixed-shaped history/future windows plus validity_mask, future_offset, and anchor_index metadata
• keep the sampling kernel tensor-only and compatible with torch.compile(fullgraph=True)
• add an executable replay-buffer example, focused tests, API documentation, and cached plus write-and-sample benchmarks

Semantics

One geometric future offset k is drawn per batch, then anchors are sampled uniformly from stored trajectory steps with complete requested history and at least k consecutive future steps. max_future bounds the draw and fixes the returned shape to [batch_size, history + max_future + 1].

Positions before trajectory step zero and after t+k repeat a valid storage index and are marked False in validity_mask. Missing non-negative steps remain unavailable rather than being treated as padding.

Performance

Local CPU benchmark with batch_size=64, history=8, and max_future=32 after stacking on #4180:

path	1k items	100k items
cached sample, median	169 us	317 us
tracked write + sample, median	305 us	456 us

Validation

• pytest test/rb/test_samplers.py -k GeometricTrajectoryWindowSampler -q (15 passed)
• pytest benchmarks/test_replaybuffer_benchmark.py -k geometric_trajectory_window_sampler --benchmark-only --benchmark-min-rounds=3 -q (4 passed)
• executable replay-buffer example
• default torch.compile smoke test and fullgraph=True test
• CPU/device-residency smoke tests
• Ruff checks and git diff --check

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[gtnv]

a write during _full_rebuild mixed 27/32 windows. the smaller rebuild-on-revision path passed 14 behavior tests and both probes. can we replace the pending-repair path and recheck after rebuilding?

Suggested change

	def _maybe_refresh_index(self, storage: Storage) -> None:
	revision = int(storage._mutation_revision)
	if self._trajectory_positions is None or self._cache_storage_id != id(storage):
	self._full_rebuild(storage, revision)
	return
	if self._pending_indices and self._pending_storage_id == id(storage):
	changed_revisions = {
	pending_revision
	for pending_revision in self._pending_revisions
	if pending_revision > self._cache_revision
	}
	expected_revisions = set(range(self._cache_revision + 1, revision + 1))
	if changed_revisions == expected_revisions:
	self._apply_pending(storage, revision)
	else:
	self._full_rebuild(storage, revision)
	return
	if self._cache_revision != revision:
	self._full_rebuild(storage, revision)
	def _maybe_refresh_index(self, storage: Storage) -> None:
	while True:
	revision = int(storage._mutation_revision)
	if (
	self._trajectory_positions is not None
	and self._cache_storage_id == id(storage)
	and self._cache_revision == revision
	):
	return
	self._full_rebuild(storage, revision)
	if int(storage._mutation_revision) == revision:
	return