lit_sb3.buffers

source package lit_sb3.buffers

Classes

• BaseBuffer — Base class that represent a buffer (rollout or replay).

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

• DictRolloutBuffer — Dict Rollout buffer used in on-policy algorithms like A2C/PPO.

• MaskableDictRolloutBuffer — Dict Rollout buffer used in on-policy algorithms like A2C/PPO.

• MaskableDictRolloutBufferSamples

• MaskableRolloutBuffer — Rollout buffer that also stores the invalid action masks associated with each observation.

• MaskableRolloutBufferSamples

• NStepReplayBuffer — Replay buffer used for computing n-step returns in off-policy algorithms like SAC/DQN.

• RecurrentDictRolloutBuffer — Dict Rollout buffer used in on-policy algorithms like A2C/PPO.

• RecurrentRolloutBuffer — Rollout buffer that also stores the LSTM cell and hidden states.

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

• RolloutBuffer — Rollout buffer used in on-policy algorithms like A2C/PPO.

Functions

• create_sequencers — Create the utility function to chunk data into.

• pad — Chunk sequences and pad them to have constant dimensions.

• pad_and_flatten — Pad and flatten the sequences of scalar values,.

source class BaseBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, device: torch.device | str = 'auto', n_envs: int = 1)

Bases : ABC

Base class that represent a buffer (rollout or replay).

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 to which the values will be converted

• n_envs : int — Number of parallel environments

Methods

• swap_and_flatten — Swap and then flatten axes 0 (buffer_size) and 1 (n_envs).

• size — Returns.

• add — Add elements to the buffer.

• extend — Add a new batch of transitions to the buffer.

• reset — Reset the buffer.

• sample — Args.

• to_torch — Convert a numpy array to a PyTorch tensor.

source staticmethod BaseBuffer.swap_and_flatten(arr: np.ndarray) → np.ndarray

Swap and then flatten axes 0 (buffer_size) and 1 (n_envs).

to convert shape from [n_steps, n_envs, ...] (when ... is the shape of the features) to [n_steps * n_envs, ...] (which maintain the order).

Parameters

• arr : np.ndarray

• arr : np.ndarray — The array to process.

source method BaseBuffer.size() → int

Returns.

The current size of the buffer.

source method BaseBuffer.add(*args, **kwargs) → None

Add elements to the buffer.

Raises

• NotImplementedError

source method BaseBuffer.extend(*args, **kwargs) → None

Add a new batch of transitions to the buffer.

source method BaseBuffer.reset() → None

Reset the buffer.

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

Args.

batch_size: Number of element to sample env: associated gym VecEnv to normalize the observations/rewards when sampling.

source method BaseBuffer.to_torch(array: np.ndarray, copy: bool = True) → torch.Tensor

Convert a numpy array to a PyTorch tensor.

Note: it copies the data by default.

Parameters

• array : np.ndarray

• copy : bool — Whether to copy or not the data (may be useful to avoid changing things by reference). This argument is inoperative if the device is not the CPU.

• array : np.ndarray — The array to convert to a tensor.

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

Bases : ReplayBuffer

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

Extends the ReplayBuffer to use dictionary observations.

Parameters

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

• observation_space : spaces.Dict — 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 Disabled for now (see https://github.com/DLR-RM/stable-baselines3/pull/243#discussion_r531535702)

• 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 DictReplayBuffer.add(obs: dict[str, np.ndarray], next_obs: dict[str, np.ndarray], action: np.ndarray, reward: np.ndarray, done: np.ndarray, infos: list[dict[str, Any]]) → None

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

Sample elements from the replay buffer.

Parameters

• batch_size : int — Number of element to sample

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

source class DictRolloutBuffer(buffer_size: int, observation_space: spaces.Dict, action_space: spaces.Space, device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)

Bases : RolloutBuffer

Dict Rollout buffer used in on-policy algorithms like A2C/PPO.

Extends the RolloutBuffer to use dictionary observations.

It corresponds to buffer_size transitions collected using the current policy. This experience will be discarded after the policy update. In order to use PPO objective, we also store the current value of each state and the log probability of each taken action.

The term rollout here refers to the model-free notion and should not be used with the concept of rollout used in model-based RL or planning. Hence, it is only involved in policy and value function training but not action selection.

Parameters

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

• observation_space : spaces.Dict — Observation space

• action_space : spaces.Space — Action space

• device : torch.device | str — PyTorch device

• gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to Monte-Carlo advantage estimate when set to 1.

• gamma : float — Discount factor

• n_envs : int — Number of parallel environments

Methods

• reset

• add — Args.

• get

source method DictRolloutBuffer.reset() → None

source method DictRolloutBuffer.add(obs: dict[str, np.ndarray], action: np.ndarray, reward: np.ndarray, episode_start: np.ndarray, value: torch.Tensor, log_prob: torch.Tensor) → None

Args.

obs: Observation action: Action reward: episode_start: Start of episode signal. value: estimated value of the current state following the current policy. log_prob: log probability of the action following the current policy.

source method DictRolloutBuffer.get(batch_size: int | None = None) → Generator[DictRolloutBufferSamples, None, None]

source class MaskableDictRolloutBuffer(buffer_size: int, observation_space: spaces.Dict, action_space: spaces.Space, device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)

Bases : DictRolloutBuffer

Dict Rollout buffer used in on-policy algorithms like A2C/PPO.

Extends the RolloutBuffer to use dictionary observations.

It corresponds to buffer_size transitions collected using the current policy. This experience will be discarded after the policy update. In order to use PPO objective, we also store the current value of each state and the log probability of each taken action.

The term rollout here refers to the model-free notion and should not be used with the concept of rollout used in model-based RL or planning. Hence, it is only involved in policy and value function training but not action selection.

Parameters

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

• observation_space : spaces.Dict — Observation space

• action_space : spaces.Space — Action space

• device : torch.device | str — PyTorch device

• gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to classic advantage when set to 1.

• gamma : float — Discount factor

• n_envs : int — Number of parallel environments

Methods

• reset

• add — Add transition to buffer.

• get

source method MaskableDictRolloutBuffer.reset() → None

Raises

• TypeError

source method MaskableDictRolloutBuffer.add(*args, action_masks: np.ndarray | None = None, **kwargs) → None

Add transition to buffer.

Parameters

• *args — Positional arguments forwarded to parent add.

• action_masks : np.ndarray | None — Masks applied to constrain the choice of possible actions.

• **kwargs — Keyword arguments forwarded to parent add.

source method MaskableDictRolloutBuffer.get(batch_size: int | None = None) → Generator[MaskableDictRolloutBufferSamples, None, None]

source class MaskableDictRolloutBufferSamples()

Bases : NamedTuple

source class MaskableRolloutBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)

Bases : RolloutBuffer

Rollout buffer that also stores the invalid action masks associated with each observation.

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

• gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to classic advantage when set to 1.

• gamma : float — Discount factor

• n_envs : int — Number of parallel environments

Methods

• reset

• add — Add transition to buffer.

• get

source method MaskableRolloutBuffer.reset() → None

Raises

• TypeError

source method MaskableRolloutBuffer.add(*args, action_masks: np.ndarray | None = None, **kwargs) → None

Add transition to buffer.

Parameters

• *args — Positional arguments forwarded to parent add.

• action_masks : np.ndarray | None — Masks applied to constrain the choice of possible actions.

• **kwargs — Keyword arguments forwarded to parent add.

source method MaskableRolloutBuffer.get(batch_size: int | None = None) → Generator[MaskableRolloutBufferSamples, None, None]

source class MaskableRolloutBufferSamples()

Bases : NamedTuple

source class NStepReplayBuffer(*args, n_steps: int = 3, gamma: float = 0.99, **kwargs)

Bases : ReplayBuffer

Replay buffer used for computing n-step returns in off-policy algorithms like SAC/DQN.

The n-step return combines multiple steps of future rewards, discounted by the discount factor gamma. This can help improve sample efficiency and credit assignment.

This implementation uses the same storage space as a normal replay buffer, and NumPy vectorized operations at sampling time to efficiently compute the n-step return, without requiring extra memory.

This implementation is inspired by

• https://github.com/younggyoseo/FastTD3
• https://github.com/DLR-RM/stable-baselines3/pull/81

It avoids potential issues such as

• https://github.com/younggyoseo/FastTD3/issues/6

Parameters

• buffer_size — Max number of element in the buffer

• observation_space — Observation space

• action_space — Action space

• device — PyTorch device

• n_envs — Number of parallel environments

• optimize_memory_usage — Not supported

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

• n_steps : int — Number of steps to accumulate rewards for n-step returns

• gamma : float — Discount factor for future rewards

source class RecurrentDictRolloutBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, hidden_state_shape: tuple[int, int, int, int], device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)

Bases : DictRolloutBuffer

Dict Rollout buffer used in on-policy algorithms like A2C/PPO.

Extends the RecurrentRolloutBuffer to use dictionary observations.

Parameters

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

• observation_space : spaces.Space — Observation space

• action_space : spaces.Space — Action space

• hidden_state_shape : tuple[int, int, int, int] — Shape of the buffer that will collect lstm states

• device : torch.device | str — PyTorch device

• gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to classic advantage when set to 1.

• gamma : float — Discount factor

• n_envs : int — Number of parallel environments

Methods

• reset

• add — Args.

• get

source method RecurrentDictRolloutBuffer.reset() → None

source method RecurrentDictRolloutBuffer.add(*args, lstm_states: RNNStates, **kwargs) → None

Args.

hidden_states: LSTM cell and hidden state.

source method RecurrentDictRolloutBuffer.get(batch_size: int | None = None) → Generator[RecurrentDictRolloutBufferSamples, None, None]

source class RecurrentRolloutBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, hidden_state_shape: tuple[int, int, int, int], device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)

Bases : RolloutBuffer

Rollout buffer that also stores the LSTM cell and hidden states.

Parameters

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

• observation_space : spaces.Space — Observation space

• action_space : spaces.Space — Action space

• hidden_state_shape : tuple[int, int, int, int] — Shape of the buffer that will collect lstm states (n_steps, lstm.num_layers, n_envs, lstm.hidden_size)

• device : torch.device | str — PyTorch device

• gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to classic advantage when set to 1.

• gamma : float — Discount factor

• n_envs : int — Number of parallel environments

Methods

• reset

• add — Args.

• get

source method RecurrentRolloutBuffer.reset() → None

source method RecurrentRolloutBuffer.add(*args, lstm_states: RNNStates, **kwargs) → None

Args.

hidden_states: LSTM cell and hidden state.

source method RecurrentRolloutBuffer.get(batch_size: int | None = None) → Generator[RecurrentRolloutBufferSamples, None, None]

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

source class RolloutBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)

Bases : BaseBuffer

Rollout buffer used in on-policy algorithms like A2C/PPO.

It corresponds to buffer_size transitions collected using the current policy. This experience will be discarded after the policy update. In order to use PPO objective, we also store the current value of each state and the log probability of each taken action.

The term rollout here refers to the model-free notion and should not be used with the concept of rollout used in model-based RL or planning. Hence, it is only involved in policy and value function training but not action selection.

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

• gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to classic advantage when set to 1.

• gamma : float — Discount factor

• n_envs : int — Number of parallel environments

Methods

• reset

• compute_returns_and_advantage — Post-processing step: compute the lambda-return (TD(lambda) estimate).

• add — Args.

• get

source method RolloutBuffer.reset() → None

source method RolloutBuffer.compute_returns_and_advantage(last_values: torch.Tensor, dones: np.ndarray) → None

Post-processing step: compute the lambda-return (TD(lambda) estimate).

and GAE(lambda) advantage.

Uses Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438) to compute the advantage. To obtain Monte-Carlo advantage estimate (A(s) = R - V(S)) where R is the sum of discounted reward with value bootstrap (because we don't always have full episode), set gae_lambda=1.0 during initialization.

The TD(lambda) estimator has also two special cases

• TD(1) is Monte-Carlo estimate (sum of discounted rewards)
• TD(0) is one-step estimate with bootstrapping (r_t + gamma * v(s_{t+1}))

For more information, see discussion in https://github.com/DLR-RM/stable-baselines3/pull/375.

Parameters

• last_values : torch.Tensor — state value estimation for the last step (one for each env)

• dones : np.ndarray — if the last step was a terminal step (one bool for each env).

source method RolloutBuffer.add(obs: np.ndarray, action: np.ndarray, reward: np.ndarray, episode_start: np.ndarray, value: torch.Tensor, log_prob: torch.Tensor) → None

Args.

obs: Observation action: Action reward: episode_start: Start of episode signal. value: estimated value of the current state following the current policy. log_prob: log probability of the action following the current policy.

source method RolloutBuffer.get(batch_size: int | None = None) → Generator[RolloutBufferSamples, None, None]

source create_sequencers(episode_starts: np.ndarray, env_change: np.ndarray, device: torch.device) → tuple[np.ndarray, Callable, Callable]

Create the utility function to chunk data into.

sequences and pad them to create fixed size tensors.

Parameters

• episode_starts : np.ndarray — Indices where an episode starts

• env_change : np.ndarray — Indices where the data collected come from a different env (when using multiple env for data collection)

• device : torch.device — PyTorch device

Returns

• tuple[np.ndarray, Callable, Callable] — Indices of the transitions that start a sequence, pad and pad_and_flatten utilities tailored for this batch (sequence starts and ends indices are fixed)

source pad(seq_start_indices: np.ndarray, seq_end_indices: np.ndarray, device: torch.device, tensor: np.ndarray, padding_value: float = 0.0) → torch.Tensor

Chunk sequences and pad them to have constant dimensions.

Parameters

• seq_start_indices : np.ndarray — Indices of the transitions that start a sequence

• seq_end_indices : np.ndarray — Indices of the transitions that end a sequence

• device : torch.device — PyTorch device

• tensor : np.ndarray — Tensor of shape (batch_size, *tensor_shape)

• padding_value : float — Value used to pad sequence to the same length (zero padding by default)

Returns

• torch.Tensor — (n_seq, max_length, *tensor_shape)

source pad_and_flatten(seq_start_indices: np.ndarray, seq_end_indices: np.ndarray, device: torch.device, tensor: np.ndarray, padding_value: float = 0.0) → torch.Tensor

Pad and flatten the sequences of scalar values,.

while keeping the sequence order. From (batch_size, 1) to (n_seq, max_length, 1) -> (n_seq * max_length,).

Parameters

• seq_start_indices : np.ndarray — Indices of the transitions that start a sequence

• seq_end_indices : np.ndarray — Indices of the transitions that end a sequence

• device : torch.device — PyTorch device (cpu, gpu, ...)

• tensor : np.ndarray — Tensor of shape (max_length, n_seq, 1)

• padding_value : float — Value used to pad sequence to the same length (zero padding by default)

Returns

• torch.Tensor — (n_seq * max_length,) aka (padded_batch_size,)