lit_sb3.buffers.recurrent.rollout

source module lit_sb3.buffers.recurrent.rollout

Classes

• RecurrentRolloutBufferSamples

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

source class RecurrentRolloutBufferSamples()

Bases : NamedTuple

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]