Skip to content

unienv_data.base.storage

SegmentWriter

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().

storage instance-attribute

storage = storage

length instance-attribute

length = 0

is_finished instance-attribute

is_finished = False

append

append(value) -> None

Append one structured sample to the segment.

finish

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.

commit 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.

abort

abort() -> None

Drop the segment without writing it.

BufferedSegmentWriter

BufferedSegmentWriter(storage: SpaceStorage)

Bases: SegmentWriter

Keeps the appended samples in memory and writes them with one set() on commit.

storage instance-attribute

storage = storage

length instance-attribute

length = 0

is_finished instance-attribute

is_finished = False

commit

commit(start_index: int) -> None

abort

abort() -> None

append

append(value) -> None

Append one structured sample to the segment.

finish

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

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.

inner_writer instance-attribute

inner_writer = inner_writer

transform instance-attribute

transform = transform

storage instance-attribute

storage = storage

length instance-attribute

length = 0

is_finished instance-attribute

is_finished = False

commit

commit(start_index: int) -> None

abort

abort() -> None

append

append(value) -> None

Append one structured sample to the segment.

finish

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.

SpaceStorage

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.

single_file_ext 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.

capacity class-attribute instance-attribute

capacity: Optional[int] = None

The cache path for the storage. If None, the storage will not use caching.

cache_filename 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.

is_multiprocessing_safe class-attribute instance-attribute

is_multiprocessing_safe: bool = False

Is the storage mutable? If False, the storage is read-only.

is_mutable class-attribute instance-attribute

is_mutable: bool = True

backend property

backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]

device property

device: Optional[BDeviceType]

single_instance_space instance-attribute

single_instance_space = single_instance_space

has_open_segment property

has_open_segment: bool

pending_segment_length property

pending_segment_length: int

create 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.

load_from 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.

get_segments

get_segments()

mark_segment_start

mark_segment_start(start_index: int) -> None

append

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.

mark_segment_end

mark_segment_end() -> None

abort_segment

abort_segment() -> None

extend_length

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

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.

get abstractmethod

get(index: Union[IndexableType, BArrayType]) -> BatchT

Read one or more samples from storage.

set abstractmethod

set(index: Union[IndexableType, BArrayType], value: BatchT) -> None

Write one or more samples into storage.

open_segment_writer

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

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.

dumps 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).

close

close() -> None