Skip to content

unienv_data.storages.metadata_routing

Build a DictStorage layout from the Space.metadata of a DictSpace.

Camera images are best stored by an image / video storage, everything else by a tensor storage. Instead of listing the keys by hand::

from unienv_data.replay_buffer import ReplayBuffer
from unienv_data.storages import DictStorage, dict_storage_kwargs_from_space

rb = ReplayBuffer.create(
    DictStorage, obs_space,
    cache_path="cache/rb", capacity=None,
    **dict_storage_kwargs_from_space(obs_space),
)

routes every kind == "image" uint8 leaf of obs_space to an ImageStorage and the remaining leaves to one HDF5Storage.

dict_storage_kwargs_from_space 

dict_storage_kwargs_from_space(space: DictSpace, *, kind_storage_cls: Optional[Mapping[str, Optional[Type[SpaceStorage]]]] = None, key_storage_cls: Optional[Mapping[str, Optional[Type[SpaceStorage]]]] = None, default_storage_cls: Optional[Type[SpaceStorage]] = None, kind_kwargs: Optional[Mapping[str, Mapping[str, Any]]] = None, key_kwargs: Optional[Mapping[str, Mapping[str, Any]]] = None, default_kwargs: Optional[Mapping[str, Any]] = None, nested_separator: str = '/') -> Dict[str, Any]

DictStorage.create keyword arguments that route each leaf of space by its metadata.

Every leaf of the (nested) DictSpace is assigned, in order of precedence:

  1. key_storage_cls[key] if its key (nested keys joined by nested_separator) is listed (None sends it to the default storage);
  2. kind_storage_cls[leaf.metadata["kind"]] if its kind is listed and the storage can take it. The default is {"image": ImageStorage}, i.e. JPEG files (ImageStorage's defaults; use kind_kwargs={"image": {"format": "PNG"}} for lossless files or map "image" to VideoStorage). Image leaves must be (H, W, 3) uint8 with an "HWC" (or undeclared) layout, others (e.g. CHW or float images) go to the default storage;
  3. otherwise the default storage, default_storage_cls (HDF5Storage if not given, which needs h5py and a NumPy-backend space), which receives all remaining leaves as one DictSpace (key "*").

Keyword arguments for the sub-storages are merged from kind_kwargs[kind] and key_kwargs[key] (default_kwargs / key_kwargs["*"] for the default storage). A VideoStorage leaf whose metadata declares fps gets it as the fps argument unless one is given.

Returns {"storage_cls_map", "key_kwargs", "nested_separator"}, to be passed as ReplayBuffer.create(DictStorage, space, **kwargs) or DictStorage.create(space, **kwargs). Image and video storages write files, so a cache_path is required when any leaf is routed to one.

create_dict_storage_from_space 

create_dict_storage_from_space(space: DictSpace, *, capacity: Optional[int] = None, cache_path: Optional[str] = None, multiprocessing: bool = False, **routing_kwargs: Any) -> DictStorage

DictStorage.create with the layout of :func:dict_storage_kwargs_from_space (which takes routing_kwargs).