Skip to content

lit_sb3.buffers.replay

source module lit_sb3.buffers.replay

Classes

  • ReplayBuffer — Replay buffer used in off-policy algorithms like SAC/TD3.

source class ReplayBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, device: torch.device | str = 'auto', n_envs: int = 1, optimize_memory_usage: bool = False, handle_timeout_termination: bool = True)

Bases : BaseBuffer

Replay buffer used in off-policy algorithms like SAC/TD3.

Parameters

  • buffer_size : int — Max number of element in the buffer

  • observation_space : spaces.Space — Observation space

  • action_space : spaces.Space — Action space

  • device : torch.device | str — PyTorch device

  • n_envs : int — Number of parallel environments

  • optimize_memory_usage : bool — Enable a memory efficient variant of the replay buffer which reduces by almost a factor two the memory used, at a cost of more complexity. See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195 and https://github.com/DLR-RM/stable-baselines3/pull/28#issuecomment-637559274 Cannot be used in combination with handle_timeout_termination.

  • handle_timeout_termination : bool — Handle timeout termination (due to timelimit) separately and treat the task as infinite horizon task. https://github.com/DLR-RM/stable-baselines3/issues/284

Methods

  • add

  • sample — Sample elements from the replay buffer.

source method ReplayBuffer.add(obs: np.ndarray, next_obs: np.ndarray, action: np.ndarray, reward: np.ndarray, done: np.ndarray, infos: list[dict[str, Any]]) → None

source method ReplayBuffer.sample(batch_size: int, env: VecNormalize | None = None) → ReplayBufferSamples

Sample elements from the replay buffer.

Custom sampling when using memory efficient variant, as we should not sample the element with index self.pos See https://github.com/DLR-RM/stable-baselines3/pull/28#issuecomment-637559274.

Parameters

  • batch_size : int — Number of element to sample

  • env : VecNormalize | None — associated gym VecEnv to normalize the observations/rewards when sampling