lit_sb3.data.replay

source module lit_sb3.data.replay

PyTorch IterableDataset wrappers that bridge RL data collection.

with Lightning's DataLoader-based training loop.

Classes

• ReplayBufferDataset — Finite IterableDataset for off-policy algorithms.

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

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 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.