[0]
[0, 1]
[0, 1, 2]
...
[0, 1, 2, 3, 4]
[1, 2, 3, 4, 5]
...
[6, 7, 8, 9]
[7, 8, 9]
[8, 9]
[9]

[0, 1, 2, 3, 4]
[1, 2, 3, 4, 5]
...
[6, 7, 8, 9]
[7, 8, 9]
[8, 9]
[9]

from torchrl.collectors.utils import split_trajectories
from torchrl.data import SliceSampler, TensorDictReplayBuffer, LazyMemmapStorage
from tensordict import TensorDict
import torch
torch.manual_seed(0)
data = TensorDict(
    {"obs": torch.arange(1, 11).repeat(10),
     "eps": torch.arange(100) // 10 + 1} ,
    [100]
)

for N in (2, 4, ):
    rb = TensorDictReplayBuffer(sampler=SliceSampler(num_slices=10, traj_key="eps", span=(N, N)), batch_size=50, storage=LazyMemmapStorage(100))
    rb.extend(data)

    for _ in range(10):
        sample = rb.sample()
        sample = split_trajectories(sample)
        assert (sample["next", "truncated"].squeeze(-1).sum(-1) == 1).all()
        assert ((sample["obs"] == 0).sum(-1) <= N).all(), sample["obs"]
        assert ((sample["eps"] == 0).sum(-1) <= N).all()
        for i in range(sample.shape[0]):
            curr_eps = sample[i]["eps"]
            curr_eps = curr_eps[curr_eps != 0]
            assert curr_eps.unique().numel() == 1
        print(sample)
        print('episode')
        print(sample["eps"])
        print('obs')
        print(sample["obs"])
        print('trunc')
        print(sample["next", "truncated"].int().squeeze())