[Performance] Improve geometric window sampling

[theap06]

Summary

• sample truncated geometric offsets on the storage device in constant time with a numerically stable inverse CDF
• integrate the draw into the compiled tensor kernel and avoid scalar synchronization in the history/future loops
• add distribution coverage and an offset-focused scaling benchmark

Performance

The offset benchmark is flat at roughly 8.5 microseconds for max_future=32 and max_future=4096 on CPU. In a local comparison including the maximum-availability reduction over 100,000 slots, the previous weight-vector draw took 20.0 and 72.3 microseconds respectively; the inverse-CDF draw took 19.9 microseconds in both cases.

Steady-state MPS profiling reports no item, _local_scalar_dense, or nonzero operations, and the sampled tensors remain on the storage device. Total sampling still includes an anchor-weight scan and fixed-size window construction, so this PR only makes the geometric offset draw constant-time.

Validation

• geometric trajectory sampler tests: 16 passed, including torch.compile(..., fullgraph=True)
• geometric replay-buffer benchmarks: 6 passed
• replay-buffer example completed successfully
• MPS steady-state sample/device profiler smoke test passed

Stacked on #4130.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 15 Awaiting Approval

As of commit 42df30e with merge base b0eab87 ():

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)
• 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.

[theap06]

@torchrlbot reviewer @vmoens

[github-actions]

Requested review from @vmoens (requested by @theap06).