Repository: pytorch/rl

Issue:
[BUG] SliceSampler proposes indices outside storage content
## Describe the bug

With lazy map storage, slice sampler proposes indices outside the populated part of the storage.

## To Reproduce

```python
from torchrl.data import TensorDictReplayBuffer, LazyMemmapStorage
from tensordict import TensorDict
from torchrl.data.replay_buffers.samplers import SliceSampler
import torch

sampler = SliceSampler(
    slice_len=16,
    end_key="done",
    truncated_key=None,
)

rb = TensorDictReplayBuffer(
    storage=LazyMemmapStorage(20_000),
    batch_size=4 * 16,
    prefetch=10,
    sampler=sampler,
)

for _i in range(3):
    batch = TensorDict(
        {
            "obs": torch.randn(50, 3, 64, 64),
            "act": torch.randn(50, 4),
            "rew": torch.randn(50),
            "done": torch.zeros(50),
        },
        batch_size=[50],
    )
    batch["done"][-1] = 1
    rb.extend(batch)

for _i in range(10):
    rb.sample()
```

```bash
Traceback (most recent call last):
  File "/private/home/vladsobal/work/HJEPA/minimal_example.py", line 33, in <module>
    rb.sample()
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/replay_buffers.py", line 878, in sample
    data, info = super().sample(batch_size, return_info=True)
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/replay_buffers.py", line 446, in sample
    ret = self._prefetch_queue.popleft().result()
  File "/private/home/vladsobal/.conda/envs/d4rl/lib/python3.9/concurrent/futures/_base.py", line 446, in result
    return self.__get_result()
  File "/private/home/vladsobal/.conda/envs/d4rl/lib/python3.9/concurrent/futures/_base.py", line 391, in __get_result
    raise self._exception
  File "/private/home/vladsobal/.conda/envs/d4rl/lib/python3.9/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/utils.py", line 46, in decorated_fun
    output = fun(self, *args, **kwargs)
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/replay_buffers.py", line 375, in _sample
    data = self._storage[index]
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/storages.py", line 82, in __getitem__
    return self.get(item)
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/storages.py", line 836, in get
    result = super().get(index)
  File "/private/home/vladsobal/work_non_sync/rl/torchrl/data/replay_buffers/storages.py", line 548, in get
    out = storage[index]
  File "/private/home/vladsobal/work_non_sync/tensordict/tensordict/base.py", line 253, in __getitem__
    return self._index_tensordict(index)
  File "/private/home/vladsobal/work_non_sync/tensordict/tensordict/_td.py", line 760, in _index_tensordict
    source[key] = item._index_tensordict(
  File "/private/home/vladsobal/work_non_sync/tensordict/tensordict/_td.py", line 764, in _index_tensordict
    source[key] = _get_item(item, index)
  File "/private/home/vladsobal/work_non_sync/tensordict/tensordict/utils.py", line 632, in _get_item
    raise err
  File "/private/home/vladsobal/work_non_sync/tensordict/tensordict/utils.py", line 618, in _get_item
    return tensor[index]
  File "/private/home/vladsobal/work_non_sync/tensordict/tensordict/memmap.py", line 563, in __getitem__
    out = super().__getitem__(item)
IndexError: index 7880 is out of bounds for dimension 0 with size 150
```

## Expected behavior

Sampling is done correctly, indices do not point to non-existent area of storage.

## Screenshots

N/A

## System info

python 3.9.16
main branch version of torchrl

## Additional context

Add any other context about the problem here.

## Reason and Possible fixes

https://github.com/pytorch/rl/blob/528faa19fe51a65c2ab50f3eb7cf89bd7b701fd0/torchrl/data/replay_buffers/samplers.py#L793
The parentheses placement is incorrect.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/data/replay_buffers/samplers.py b/torchrl/data/replay_buffers/samplers.py
--- a/torchrl/data/replay_buffers/samplers.py
+++ b/torchrl/data/replay_buffers/samplers.py
@@ -790,7 +790,7 @@ def _get_stop_and_length(self, storage, fallback=True):
                         raise RuntimeError(
                             "Could not get a tensordict out of the storage, which is required for SliceSampler to compute the trajectories."
                         )
-                vals = self._find_start_stop_traj(end=done.squeeze())[: len(storage)]
+                vals = self._find_start_stop_traj(end=done.squeeze()[: len(storage)])
                 if self.cache_values:
                     self._cache["stop-and-length"] = vals
                 return vals


Test patch: