Skip to content

lit_sb3.data.rollout

source module lit_sb3.data.rollout

PyTorch IterableDataset wrappers that bridge RL data collection.

with Lightning's DataLoader-based training loop.

Classes

  • RolloutDataset — IterableDataset wrapping a RolloutBuffer for on-policy algorithms.

source class RolloutDataset(rollout_buffer: RolloutBuffer, batch_size: int, n_epochs: int)

Bases : IterableDataset

IterableDataset wrapping a RolloutBuffer for on-policy algorithms.

The dataset yields mini-batches from the rollout buffer. It is intended to be iterated once per RL training iteration: - on_train_epoch_start() collects a fresh rollout and fills the buffer - Each iter call yields n_epochs * ceil(buffer_size / batch_size) batches (i.e., n_epochs full passes over the buffer data)

Because this holds a reference to the buffer (not a copy), it automatically serves freshly collected data each epoch without needing to be recreated.

Parameters

  • rollout_buffer : RolloutBuffer — RolloutBuffer to iterate over.

  • batch_size : int — Mini-batch size for each gradient step.

  • n_epochs : int — Number of passes over the rollout buffer per RL iteration.