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().
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.
BufferedSegmentWriter
¶
BufferedSegmentWriter(storage: SpaceStorage)
Bases: SegmentWriter
Keeps the appended samples in memory and writes them with one set() on commit.
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.
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.
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.
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.
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).