rlightning.buffer.rollout_buffer

Rollout buffer implementation for on-policy data collection.

Defines a buffer that samples transitions/episodes, validates batch sizing, and clears stored data after sampling. Integrates with the buffer registry and supports sharded/distributed sampling via the base buffer.

class rlightning.buffer.rollout_buffer.RolloutBuffer(config: BufferConfig, obs_preprocessor: Preprocessor | None = <function default_obs_preprocessor>, reward_preprocessor: Preprocessor | None = <function default_reward_preprocessor>, env_ret_preprocess_fn: Callable | None = <function default_env_ret_preprocess_fn>, policy_resp_preprocess_fn: Callable | None = <function default_policy_resp_preprocess_fn>, preprocess_fn: Callable | None = <function default_preprocess_fn>, postprocess_fn: Callable | None = <function default_postprocess_fn>)[source]

Bases: DataBuffer

Rollout Buffer for on-policy algorithms

sample(batch_size: int | None = None, shuffle: bool | None = True, drop_last: bool | None = True) List[Dict] | List[MockModule('ray.ObjectRef')][source]

Sample a batch of data (transitions/truncated_episodes) from the buffer.

Sharded storage support: - Each storage samples indices based on its own size. - DistributedSampler then splits those indices to the workers bound to that storage.

Returns:

int, “indices”: np.ndarray}.

Return type:

dict mapping worker rank -> {“storage_idx”