unienv_data¶
BatchBase(single_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], single_metadata_space: Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]] = None)
Bases: ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Base interface for dataset-like collections of structured samples.
BatchBase stores the per-sample space definition and provides a uniform
API for retrieving either structured values or their flattened backend-array
representation. Concrete subclasses decide where data lives and which
mutation operations are supported.
Initialize the batch contract for one logical sample.
get_flattened_at(idx: Union[IndexableType, BArrayType]) -> BArrayType
Fetch samples as flattened backend arrays.
get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]
Fetch flattened samples together with optional per-sample metadata.
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
get_at(idx: Union[IndexableType, BArrayType]) -> BatchT
Fetch samples in their structured form.
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Optional[Dict[str, Any]]]
Fetch structured samples together with optional metadata.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
BatchSampler(single_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], single_metadata_space: Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]] = None, batch_size: int = 1)
Bases: Generic[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], BatchBase[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
Read-only batch wrapper that samples mini-batches from another batch.
Configure the sampling batch shape without binding the source data yet.
property
¶
sampled_space: Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
property
¶
sampled_metadata_space: Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]]
get_flattened_at(idx: Union[IndexableType, BArrayType]) -> BArrayType
Fetch samples as flattened backend arrays.
get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]
Fetch flattened samples together with optional per-sample metadata.
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
get_at(idx: Union[IndexableType, BArrayType]) -> BatchT
Fetch samples in their structured form.
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Optional[Dict[str, Any]]]
Fetch structured samples together with optional metadata.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
manual_seed(seed: int) -> None
Reset sampler RNG state, including the optional data-index RNG.
sample_flat_with_metadata() -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]
Sample a flattened batch together with optional metadata.
sample_with_metadata() -> Tuple[SamplerBatchT, Optional[Dict[str, Any]]]
Sample a structured batch together with optional metadata.
epoch_iter() -> Iterator[SamplerBatchT]
Iterate once over a random permutation of the source data.
epoch_iter_with_metadata() -> Iterator[Tuple[SamplerBatchT, Optional[Dict[str, Any]]]]
Iterate once over shuffled structured batches plus metadata.
epoch_flat_iter() -> Iterator[SamplerArrayType]
Iterate once over shuffled batches in flattened form.
epoch_flat_iter_with_metadata() -> Iterator[Tuple[SamplerArrayType, Optional[Dict[str, Any]]]]
Iterate once over shuffled flattened batches plus metadata.
SpaceStorage(single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType])
Bases: ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
SpaceStorage is an abstract base class for storages that hold instances of a specific space.
It provides a common interface for creating, loading, and managing the storage of instances of a given space.
Note that if you want your space storage to support multiprocessing, you need to check / implement __getstate__ and __setstate__ methods to ensure that the storage can be pickled and unpickled correctly.
Bind the storage to the space of one stored element.
class-attribute
instance-attribute
¶
single_file_ext: Optional[str] = None
The total capacity (number of single instances) of the storage. If None, the storage has unlimited capacity.
class-attribute
instance-attribute
¶
capacity: Optional[int] = None
The cache path for the storage. If None, the storage will not use caching.
class-attribute
instance-attribute
¶
cache_filename: Optional[Union[str, PathLike]] = None
Can the storage instance be safely used in multiprocessing environments after creation? If True, the storage can be used in multiprocessing environments.
class-attribute
instance-attribute
¶
is_multiprocessing_safe: bool = False
Is the storage mutable? If False, the storage is read-only.
classmethod
¶
create(single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], *args, capacity: Optional[int], cache_path: Optional[Union[str, PathLike]] = None, multiprocessing: bool = False, **kwargs) -> SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a new storage instance for samples from single_instance_space.
classmethod
¶
load_from(path: Union[str, PathLike], single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], *, capacity: Optional[int] = None, read_only: bool = True, multiprocessing: bool = False, **kwargs) -> SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Restore a storage instance previously written to disk.
append(value) -> None
Append one structured sample to the storage.
Appended values may be immediately visible or only visible after
mark_segment_end() depending on the storage implementation.
extend_length(length: int) -> None
This is used by capacity = None storages to extend the length of the storage If this is called on a storage with a fixed capacity, we will simply ignore the call.
shrink_length(length: int) -> None
This is used by capacity = None storages to shrink the length of the storage If this is called on a storage with a fixed capacity, we will simply ignore the call.
abstractmethod
¶
get(index: Union[IndexableType, BArrayType]) -> BatchT
Read one or more samples from storage.
abstractmethod
¶
set(index: Union[IndexableType, BArrayType], value: BatchT) -> None
Write one or more samples into storage.
open_segment_writer() -> SegmentWriter
A writer that builds one segment of this storage off to the side; see SegmentWriter.
Unlike mark_segment_start() / append(), any number of writers can be open at a time.
The default keeps the samples in memory until the commit; storages override it to stream instead.
clear() -> None
Clear all data inside the storage and set the length to 0 if the storage has unlimited capacity. For storages with fixed capacity, this should reset the storage to its initial state.
abstractmethod
¶
dumps(path: Union[str, PathLike]) -> None
Dumps the storage to the specified path.
This is used for storages that have a single file extension (e.g. .pt for PyTorch).
SegmentWriter(storage: SpaceStorage)
Bases: ABC
Builds one contiguous segment of a storage off to the side.
Values appended to a writer land at an index range that is only chosen when the writer is committed,
so several writers of one storage can be filled at the same time (e.g. one per environment of a batched rollout)
and committed in any order. Obtain one with SpaceStorage.open_segment_writer().
finish() -> None
Declare the segment complete (no more appends) and release what only appending needs (e.g. encoders).
Called by commit() if needed; finish early when the commit has to wait.
abstractmethod
¶
commit(start_index: int) -> None
Write the segment to the storage indices [start_index, start_index + len(self))
(modulo the capacity of fixed-capacity storages). Storages with unlimited capacity must already be long enough.
The writer cannot be used afterwards.
BufferedSegmentWriter(storage: SpaceStorage)
Bases: SegmentWriter
Keeps the appended samples in memory and writes them with one set() on commit.
finish() -> None
Declare the segment complete (no more appends) and release what only appending needs (e.g. encoders).
Called by commit() if needed; finish early when the commit has to wait.
MappedSegmentWriter(storage: SpaceStorage, inner_writer: SegmentWriter, transform: Callable[[Any], Any])
Bases: SegmentWriter
Forwards to a writer of an inner storage, mapping every appended sample with transform.
finish() -> None
Declare the segment complete (no more appends) and release what only appending needs (e.g. encoders).
Called by commit() if needed; finish early when the commit has to wait.
ToBackendOrDeviceBatch(batch: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], backend: Optional[ComputeBackend[WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]] = None, device: Optional[WrapperBDeviceT] = None)
Bases: Generic[WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], BatchBase[WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]
property
¶
backend: ComputeBackend[WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
close()
Release this view. The wrapped batch is not owned by the view and stays open;
close it through whoever created it (BatchBase.__del__ calls close(), so
closing the source here would tear it down whenever a view is garbage-collected).
CombinedBatch(batches: Sequence[BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]], device: Optional[BDeviceType] = None, use_thread_pool: bool = False)
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[Any]) -> BatchBase
Lazy view over one field of the samples that combines the children's column views, so
that children able to read a single column cheaply (replay buffers, stacked views, ...)
do so instead of fetching whole samples. Values, spaces, metadata and the flattened
layout are those of SubItemBatch(self, nested_keys), which is returned when the
children's column views cannot be combined. The view is read-only.
get_flattened_at_with_metadata(idx: Union[IndexableType, BaseException]) -> Tuple[BArrayType, Dict[str, Any]]
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Dict[str, Any]]
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
close() -> None
Release this view. The child batches are not owned by the view and stay open;
close them through whoever created it (BatchBase.__del__ calls close(), so
closing the children here would tear it down whenever a view is garbage-collected).
SliceStackedBatch(batch: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], fixed_offset: BArrayType, get_valid_mask_function: Optional[Callable[[SliceStackedBatch, BArrayType, BatchT, Dict[str, Any]], BArrayType]] = None, fill_invalid_data: bool = True, stack_metadata: bool = False)
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
A batch that stacks frames with given fixed offsets. This is a read-only batch, since it is a view of the original batch. If you want to change the data, you should mutate the containing batch instead.
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
expand_index(index: BArrayType) -> BArrayType
Expand indexes to slice the data Args: index (BArrayType): A 1D tensor of indices to expand. Can be boolean or integer. Returns: BArrayType: A 2D tensor of indices, where each row corresponds to the expanded indices for each batch item.
get_valid_mask_flattened(expanded_idx: BArrayType, data: BArrayType, metadata: Dict[str, Any]) -> BArrayType
Get the valid mask for the flattened data. Args: expanded_idx (B, T): A 2D tensor of indices to slice the data. data (B, T, *D): The data to slice. Returns: Valid Mask (B, T)
get_valid_mask(expanded_idx: BArrayType, data: BatchT, metadata: Dict[str, Any]) -> BArrayType
Get the valid mask for the data. Args: expanded_idx (B, T): A 2D tensor of indices to slice the data. data (BatchT): The data to slice. Returns: Valid Mask (B, T)
fill_data_with_stack_mask(space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]], data: Union[BArrayType, BatchT, Any], valid_mask: BArrayType) -> Union[BArrayType, BatchT, Any]
Fill the data with the mask as if the frames were frame-stacked.
data is updated in place where the backend allows it (backend.at(...).set;
JAX and read-only arrays are copied). This is safe because data is the result
of fetching an index array from the wrapped batch, a fresh gather: batches wrapped
by a SliceStackedBatch must not return arrays aliasing their storage from
get_at* / get_flattened_at* with index arrays.
Args:
space (Optional[Space]): The space to fill the data with. If None, the data is assumed to be a backend array.
data (Union[BArrayType, BatchT, Any]): The data to fill sized (B, T, *D)
valid_mask (BArrayType): The mask to fill the data with, sized (B, T)
Returns:
Union[BArrayType, BatchT, Any]: The filled data, sized (B, T, *D)
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Dict[str, Any]]
get_column(nested_keys: Sequence[Any]) -> BatchBase
Lazy view over one field of the stacked samples.
Stacks self.batch.get_column(nested_keys) instead of stacking whole samples and
dropping the other fields afterwards, so only that column (plus the episode column the
built-in valid-mask helpers read) is gathered. Values, spaces, metadata and the
flattened layout are those of SubItemBatch(self, nested_keys), which is returned
when that is not possible (a custom valid mask function, which may read any field, a
mask read from metadata, or a column view with different metadata).
close() -> None
Release this view. The wrapped batch is not owned by the view and stays open;
close it through whoever created it (BatchBase.__del__ calls close(), so
closing the source here would tear it down whenever a view is garbage-collected).
staticmethod
¶
get_valid_mask_function_with_episodeid_key(episode_id_key: Union[str, int] = 'episode_id', is_in_metadata: bool = False) -> Callable[[SliceStackedBatch, BArrayType, BatchT], BArrayType]
staticmethod
¶
get_valid_mask_function_with_episode_end_key(episode_end_key: Union[str, int] = 'episode_end', is_in_metadata: bool = False) -> Callable[[SliceStackedBatch, BArrayType, BatchT], BArrayType]
FrameStackedBatch(batch: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], prefetch_horizon: int = 0, postfetch_horizon: int = 0, get_valid_mask_function: Optional[Callable[[SliceStackedBatch, BArrayType, BatchT, Dict[str, Any]], BArrayType]] = None, fill_invalid_data: bool = True, stack_metadata: bool = False)
Bases: SliceStackedBatch[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
A batch that stacks frames in a sliding window manner. This batch allows for prefetching and postfetching of frames, which can be useful for training models that require temporal context (e.g. Diffusion Policy, ACT, etc.) This is a read-only batch, since it is a view of the original batch. If you want to change the data, you should mutate the containing batch instead.
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Dict[str, Any]]
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[Any]) -> BatchBase
Lazy view over one field of the stacked samples.
Stacks self.batch.get_column(nested_keys) instead of stacking whole samples and
dropping the other fields afterwards, so only that column (plus the episode column the
built-in valid-mask helpers read) is gathered. Values, spaces, metadata and the
flattened layout are those of SubItemBatch(self, nested_keys), which is returned
when that is not possible (a custom valid mask function, which may read any field, a
mask read from metadata, or a column view with different metadata).
close() -> None
Release this view. The wrapped batch is not owned by the view and stays open;
close it through whoever created it (BatchBase.__del__ calls close(), so
closing the source here would tear it down whenever a view is garbage-collected).
expand_index(index: BArrayType) -> BArrayType
Expand indexes to slice the data Args: index (BArrayType): A 1D tensor of indices to expand. Can be boolean or integer. Returns: BArrayType: A 2D tensor of indices, where each row corresponds to the expanded indices for each batch item.
get_valid_mask_flattened(expanded_idx: BArrayType, data: BArrayType, metadata: Dict[str, Any]) -> BArrayType
Get the valid mask for the flattened data. Args: expanded_idx (B, T): A 2D tensor of indices to slice the data. data (B, T, *D): The data to slice. Returns: Valid Mask (B, T)
get_valid_mask(expanded_idx: BArrayType, data: BatchT, metadata: Dict[str, Any]) -> BArrayType
Get the valid mask for the data. Args: expanded_idx (B, T): A 2D tensor of indices to slice the data. data (BatchT): The data to slice. Returns: Valid Mask (B, T)
fill_data_with_stack_mask(space: Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]], data: Union[BArrayType, BatchT, Any], valid_mask: BArrayType) -> Union[BArrayType, BatchT, Any]
Fill the data with the mask as if the frames were frame-stacked.
data is updated in place where the backend allows it (backend.at(...).set;
JAX and read-only arrays are copied). This is safe because data is the result
of fetching an index array from the wrapped batch, a fresh gather: batches wrapped
by a SliceStackedBatch must not return arrays aliasing their storage from
get_at* / get_flattened_at* with index arrays.
Args:
space (Optional[Space]): The space to fill the data with. If None, the data is assumed to be a backend array.
data (Union[BArrayType, BatchT, Any]): The data to fill sized (B, T, *D)
valid_mask (BArrayType): The mask to fill the data with, sized (B, T)
Returns:
Union[BArrayType, BatchT, Any]: The filled data, sized (B, T, *D)
staticmethod
¶
get_valid_mask_function_with_episodeid_key(episode_id_key: Union[str, int] = 'episode_id', is_in_metadata: bool = False) -> Callable[[SliceStackedBatch, BArrayType, BatchT], BArrayType]
staticmethod
¶
get_valid_mask_function_with_episode_end_key(episode_end_key: Union[str, int] = 'episode_end', is_in_metadata: bool = False) -> Callable[[SliceStackedBatch, BArrayType, BatchT], BArrayType]
SubIndexedBatch(batch: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], sub_indexes: Union[IndexableType, BArrayType])
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
A batch that offers a view of the original batch with a fixed index range. This does not support batch extension.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
SubItemBatch(batch: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], sub_indexes: Sequence[Any])
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
A batch that offers a view of the original batch with a fixed index. This is a read-only batch, since it is a view of the original batch. If you want to change the data, you should mutate the containing batch instead.
get_flattened_at(idx: Union[IndexableType, BArrayType]) -> BArrayType
Fetch samples as flattened backend arrays.
get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]
Fetch flattened samples together with optional per-sample metadata.
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
TransformedBatch(batch: BatchBase[SourceDataT, SourceBArrT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT], transformation: DataTransformation, metadata_transformation: Optional[DataTransformation] = None)
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
instance-attribute
¶
metadata_transformation = metadata_transformation if self.batch.single_metadata_space is not None else None
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BArrayType, Optional[Dict[str, Any]]]
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
get_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> Tuple[BatchT, Optional[Dict[str, Any]]]
close()
Release this view. The wrapped batch (and the transformation) are not owned by the view and stay open;
close them through whoever created it (BatchBase.__del__ calls close(), so
closing the source here would tear it down whenever a view is garbage-collected).
StepSampler(data: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], batch_size: int, seed: Optional[int] = None, device: Optional[BDeviceType] = None)
Bases: BatchSampler[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Samples mini-batches of individual steps (uniformly, with replacement) from data.
Performance notes:
- For data held in RAM (e.g. an in-memory
ReplayBufferand views over it), sampling in-process is usually faster thanMultiprocessingSampler: worker processes only pay off when producing a batch is expensive (decoding, transformations, storage reads), since every batch is copied to the main process. sample_flat()/get_flattened_at()flatten a sample into one array of a common dtype, souint8images next tofloat32fields are cast tofloat32(4x the bytes and a conversion per batch). Prefersample()for image data.- Read single fields through
data.get_column(keys)rather than sampling whole samples.
property
¶
sampled_space: Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
property
¶
sampled_metadata_space: Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]]
instance-attribute
¶
data_rng = self.backend.random.random_number_generator(seed, device=data.device)
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
manual_seed(seed: int) -> None
Reset sampler RNG state, including the optional data-index RNG.
sample_flat_with_metadata() -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]
Sample a flattened batch together with optional metadata.
sample_with_metadata() -> Tuple[SamplerBatchT, Optional[Dict[str, Any]]]
Sample a structured batch together with optional metadata.
epoch_iter() -> Iterator[SamplerBatchT]
Iterate once over a random permutation of the source data.
epoch_iter_with_metadata() -> Iterator[Tuple[SamplerBatchT, Optional[Dict[str, Any]]]]
Iterate once over shuffled structured batches plus metadata.
epoch_flat_iter() -> Iterator[SamplerArrayType]
Iterate once over shuffled batches in flattened form.
epoch_flat_iter_with_metadata() -> Iterator[Tuple[SamplerArrayType, Optional[Dict[str, Any]]]]
Iterate once over shuffled flattened batches plus metadata.
MultiprocessingSampler(sampler: BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], n_workers: int = 4, n_buffers: int = 8, ctx: Optional[BaseContext] = None, doze_time: float = 0.005, daemon: Optional[bool] = None, shared_memory: bool = False)
Bases: BatchSampler[SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Run sampler's sample* functions (and epoch iteration over its get*_at
functions) in n_workers background processes, keeping n_buffers batches queued.
Workers pay off when producing a batch is expensive (decoding, transformations,
storage reads). For data already in RAM (e.g. an in-memory ReplayBuffer),
sampling in-process is usually faster (e.g. 0.7 ms vs 3 ms for 256 samples of 48 KiB
with 4 workers): every batch is otherwise copied to the parent through shared memory.
Workers run their torch ops single-threaded (see _limit_worker_intraop_threads).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
shared_memory
|
bool
|
For |
False
|
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
Overwrite existing samples using flattened data.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
extend_flattened(value: BArrayType) -> None
Append a batched block of flattened samples.
set_at(idx: Union[IndexableType, BArrayType], value: BatchT) -> None
Overwrite existing samples using structured data.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
get_column(nested_keys: Sequence[str]) -> BatchBase[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a nested field inside each sample.
sample_flat_with_metadata() -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]
Sample a flattened batch together with optional metadata.
sample_with_metadata() -> Tuple[SamplerBatchT, Optional[Dict[str, Any]]]
Sample a structured batch together with optional metadata.
epoch_iter() -> Iterator[SamplerBatchT]
Iterate once over a random permutation of the source data.
epoch_iter_with_metadata() -> Iterator[Tuple[SamplerBatchT, Optional[Dict[str, Any]]]]
Iterate once over shuffled structured batches plus metadata.
epoch_flat_iter() -> Iterator[SamplerArrayType]
Iterate once over shuffled batches in flattened form.
epoch_flat_iter_with_metadata() -> Iterator[Tuple[SamplerArrayType, Optional[Dict[str, Any]]]]
Iterate once over shuffled flattened batches plus metadata.
ReplayBuffer(storage: SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], storage_path_relative: str, count: int = 0, offset: int = 0, cache_path: Optional[Union[str, PathLike]] = None, multiprocessing: bool = False, *, maintain_segment_metadata: bool = False, physical_segments: Optional[Sequence[Tuple[int, int]]] = None, active_segment_start_physical: Optional[int] = None, active_segment_length: int = 0)
Bases: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Ring-buffer style batch storage built on top of a SpaceStorage.
The buffer exposes the standard BatchBase interface while handling
round-robin overwrite semantics, persistence metadata, and optional
multiprocessing-safe counters.
With multiprocessing=True writers serialize on a shared lock and publish count and offset
together; readers in other processes take no lock and always map indices with a consistent pair.
A read that races with a write can still return rows that the write is overwriting (the oldest samples
of a full buffer), as with any lock-free ring buffer.
Wrap a storage object with replay-buffer indexing semantics.
append_flattened(value: BArrayType) -> None
Append one flattened sample to the batch.
remove_at(idx: Union[IndexableType, BArrayType]) -> None
Remove one or more samples from the batch.
extend_from(other: BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType], chunk_size: int = 8, tqdm: bool = False) -> None
Copy data from another batch in bounded-size chunks.
get_slice(idx: Union[IndexableType, BArrayType]) -> BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a lazy view over a subset of indices.
staticmethod
¶
create(storage_cls: Type[SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]], single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType], *args, cache_path: Optional[Union[str, PathLike]] = None, capacity: Optional[int] = None, multiprocessing: bool = False, maintain_segment_metadata: bool = False, **kwargs) -> ReplayBuffer[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Create a new replay buffer and its backing storage.
staticmethod
¶
get_length_from_path(path: Union[str, PathLike]) -> Optional[int]
staticmethod
¶
get_capacity_from_path(path: Union[str, PathLike]) -> Optional[int]
staticmethod
¶
get_space_from_path(path: Union[str, PathLike], *, backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], device: Optional[BDeviceType] = None) -> Optional[Space[BatchT, BDeviceType, BDtypeType, BRNGType]]
staticmethod
¶
load_from(path: Union[str, PathLike], *, backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType], device: Optional[BDeviceType] = None, read_only: bool = True, multiprocessing: bool = False, single_instance_space: Optional[Space[BatchT, BDeviceType, BDtypeType, BRNGType]] = None, **storage_kwargs) -> ReplayBuffer[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
Load a replay buffer plus its backing storage from disk.
If single_instance_space is provided, it is validated against the
space reconstructed from the on-disk metadata and, if compatible,
reused instead of keeping the reconstructed copy. This avoids
allocating duplicate Space objects when many buffers with
identical spaces are loaded simultaneously (e.g. for
CombinedBatch).
dumps(path: Union[str, PathLike])
Persist the replay buffer metadata and its backing storage.
get_flattened_at_with_metadata(idx: Union[IndexableType, BArrayType]) -> BArrayType
set_flattened_at(idx: Union[IndexableType, BArrayType], value: BArrayType) -> None
open_segment_writer() -> SegmentWriter
A writer that builds one segment off to the side; commit_segment() then appends it at the tail.
Unlike append() / mark_segment_end(), any number of writers can be open at a time
(e.g. one per environment of a batched rollout), and they can be committed in any order.
Storages that support it (e.g. VideoStorage) encode the samples as they are appended.
commit_segment(writer: SegmentWriter) -> None
Append the segment built by writer (from open_segment_writer()) at the tail, as extend() would.
ReplayBufferCollectionWrapper(env: Env[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType], replay_buffers: Union[ReplayBuffer, Sequence[ReplayBuffer]], replay_mode: Literal['auto', 'shared', 'per_env'] = 'auto', transition_builder: Optional[Union[Callable[..., Any], DataTransformation]] = None, validate_transition: bool = False, dump_on_episode_end: bool = False, dump_on_reset_boundary: bool = False, dump_path: Optional[PathLike] = None)
Bases: Wrapper[BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType, BArrayType, ContextType, ObsType, ActType, RenderFrame, BDeviceType, BDtypeType, BRNGType]
Automatically record every transition into one or more replay buffers.
Parameters¶
env : Env
The environment to wrap.
replay_buffers : ReplayBuffer | Sequence[ReplayBuffer]
One replay buffer (shared mode) or one per environment slot (per‑env
mode).
replay_mode : "auto", "shared" or "per_env"
* "shared" - a single flat replay buffer.
* "per_env" - one buffer per batched slot, with episode segments.
* "auto" - resolve based on batch_size and buffer count.
transition_builder : callable or DataTransformation, optional
Either a callable (cached_obs, cached_context, action, next_obs,
reward, terminated, truncated, env_index, rb_single_space) ->
transition, or a DataTransformation that maps a canonical
raw transition dict to the replay buffer's format. When None
the default DictSpace‑based builder is used.
validate_transition : bool
When True every built transition is validated against the target
replay buffer's space (single‑instance for append / batched for
shared extend). Off by default for speed.
dump_on_episode_end : bool
When True, automatically call dumps(path) on the replay
buffer(s) after an episode ends (i.e., when terminated or truncated
is true). Requires dump_path to be set.
dump_on_reset_boundary : bool
When True, automatically call dumps(path) on reset boundaries
(full reset or masked reset). Requires dump_path to be set.
dump_path : str | os.PathLike | None
Filesystem path where the replay buffer(s) should be dumped when
auto-dump is triggered. May be reassigned after construction to
redirect subsequent dumps. Required if dump_on_episode_end or
dump_on_reset_boundary is True. Repeated dumps overwrite
the same path.
classmethod
¶
shared(env: Env, replay_buffer: ReplayBuffer, **kwargs: Any) -> 'ReplayBufferCollectionWrapper'
Create a wrapper in shared mode with a single replay buffer.
classmethod
¶
per_env(env: Env, replay_buffers: Sequence[ReplayBuffer], **kwargs: Any) -> 'ReplayBufferCollectionWrapper'
Create a wrapper in per‑env mode with one buffer per slot.
reset(*args: Any, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs: Any) -> Tuple[ContextType, ObsType, Dict[str, Any]]
step(action: ActType) -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]
reset_async(*args: Any, mask: Optional[BArrayType] = None, seed: Optional[int] = None, **kwargs: Any) -> None
step_wait() -> Tuple[ObsType, Union[SupportsFloat, BArrayType], Union[bool, BArrayType], Union[bool, BArrayType], Dict[str, Any]]
set_replay_buffers(replay_buffers: Union[ReplayBuffer, Sequence[ReplayBuffer]], *, replay_mode: Optional[Literal['auto', 'shared', 'per_env']] = None, require_same_space: bool = True, allow_mid_episode: bool = False, finalize_old_segments: bool = False, dump_old: bool = False) -> List[ReplayBuffer]
Swap the replay buffer(s) without rebuilding the wrapper.
Parameters¶
replay_buffers : ReplayBuffer | Sequence[ReplayBuffer]
The new replay buffer(s) to use.
replay_mode : "auto", "shared", "per_env", or None
If provided, change the replay mode. If None, keep the
current mode.
require_same_space : bool
If True, the new buffers must have the same single_space
as the old buffers. Default is True.
allow_mid_episode : bool
If True, allow swapping even if this splits an episode
across old/new buffers. Default is False.
finalize_old_segments : bool
If True, close any open segments in the old buffers before
swapping. Default is False.
dump_old : bool
If True, dump the old buffers before swapping. Requires
dump_path to be set. Only does this safely after
segment finalization where needed. Default is False.
Returns¶
List[ReplayBuffer] The old replay buffer(s) that were replaced.
Raises¶
ValueError
If the new buffers don't match the expected layout or space.
RuntimeError
If there are open segments and allow_mid_episode and
finalize_old_segments are both False.