lit_sb3.utils.sequence

source module lit_sb3.utils.sequence

Sequence padding utilities for recurrent buffers.

Functions

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

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

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

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,)

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)