pantheonrl.common.trajsaver.transitions_collate_fn

transitions_collate_fn(batch)

Custom torch.utils.data.DataLoader collate_fn for TransitionsMinimal. Use this as the collate_fn argument to DataLoader if using an instance of TransitionsMinimal as the dataset argument.

Parameters:

batch (Sequence[Mapping[str, ndarray]]) –

Return type:

Dict[str, ndarray | Tensor]