Skip to content

lit_sb3.buffers.base

source module lit_sb3.buffers.base

Classes

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

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.