Merge remote-tracking branch 'origin/main' into openx

# Conflicts:
#	test/test_rb.py

Author: vmoens

test/test_rb.py

@@ -1595,6 +1595,7 @@ def test_slice_sampler(
                 "obs": torch.randn((3, 4, 5)).expand(100, 3, 4, 5),
                 "act": torch.randn((20,)).expand(100, 20),
                 "steps": steps,
+                "count": torch.arange(100),
                 "other": torch.randn((20, 50)).expand(100, 20, 50),
                 done_key: done,
             },
@@ -1621,7 +1622,8 @@ def test_slice_sampler(
             num_slices = batch_size // slice_len
         trajs_unique_id = set()
         too_short = False
-        for _ in range(20):
+        count_unique = set()
+        for _ in range(10):
             index, info = sampler.sample(storage, batch_size=batch_size)
             if _data_prefix:
                 samples = storage._storage["_data"][index]
@@ -1640,6 +1642,14 @@ def test_slice_sampler(
             trajs_unique_id = trajs_unique_id.union(
                 samples["another_episode"].view(-1).tolist()
             )
+            count_unique = count_unique.union(samples.get("count").view(-1).tolist())
+            if len(count_unique) == 100:
+                # all items have been sampled
+                break
+        else:
+            raise AssertionError(
+                f"Not all items can be sampled: {set(range(100))-count_unique} are missing"
+            )
         if strict_length:
             assert not too_short
         else:

torchrl/data/replay_buffers/samplers.py

@@ -776,7 +776,7 @@ def _sample_slices(
         relative_starts = (
             (
                 torch.rand(num_slices, device=lengths.device)
-                * (lengths[traj_idx] - seq_length)
+                * (lengths[traj_idx] - seq_length + 1)
             )
             .floor()
             .to(start_idx.dtype)