Skip to content

lit_sb3.data

source package lit_sb3.data

Classes

source class InfiniteReplayBufferDataset()

Bases : IterableDataset

Unbounded IterableDataset used when trainer-side RL budget control is authoritative.

Yields dummy integer tokens forever. The actual data is NOT provided by this dataset - instead, training_step handles replay buffer sampling. This dataset exists purely to keep Lightning advancing batches until the trainer or callback-driven RL budget requests a stop.

source class ReplayBufferDataset(total_steps: int)

Bases : IterableDataset

Finite IterableDataset for off-policy algorithms.

Yields a fixed number of dummy integer tokens. The actual data is NOT provided by this dataset - instead, training_step handles replay buffer sampling. This dataset simply drives the Lightning training loop for a specified number of steps.

Parameters

  • total_steps : int — Total number of training steps (= timer for entire learn()).

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.