[Feature Request] Split and truncate trajectories

Motivation

For sequence models/recurrent RL, we often want a long sequence to be split and padded into equal-sized segments of shape (batch, segment_length). split_trajectories handles the padding, but does not account for splitting long sequences (i.e. sequence length > segment length should be split into multiple segments). I propose that I add such a method to tensordict.

Solution

We can do something like

# Note this is untested and likely incorrect, but you get the idea
def truncate_trajectories(td, segment_length, mask_key=('collector', 'mask'), traj_id_key=('collector', 'traj_ids')):
    if traj_id_key is not None:
        del td[traj_id_key]
    lengths = td[mask_key].sum(dim=1)
    truncated_lengths = lengths % segment_length
    num_segments = torch.sum(lengths // segment_length)
    batch_index = torch.repeat_interleave(torch.arange(num_segments), truncated_lengths)
    time_index = torch.arange(segment_length).repeat(num_segments)
    indices = torch.stack([batch_index, time_index], dim=0)

    for k, v in list(tensordict.items()):
        td[k] = torch.zeros_like(v, shape=(num_segments, segment_length, *v.shape[2:])).scatter_(
            dim=0, index=indices, src=v,
        )

    return td

The usage would be something like

td = split_trajectories(td)
td = truncate_trajectories(td)

Alternatives

We could also add a max_segment_length arguments to split_trajectories and do this sort of logic within split_trajectories.

Additional context

Add any other context or screenshots about the feature request here.

Checklist

• I have checked that there is no similar issue in the repo (required)