unienv 0.0.1b5__py3-none-any.whl → 0.0.1b6__py3-none-any.whl

unienv_data/storages/backend_compat.py

@@ -0,0 +1,268 @@
+from typing import Dict, Any, Optional, Tuple, Union, Generic, SupportsFloat, Type, Sequence, Mapping, TypeVar
+import numpy as np
+import copy
+
+from unienv_interface.space import Space, DictSpace
+from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils as sfu
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.backends.serialization import serialize_backend, deserialize_backend
+from unienv_interface.utils.symbol_util import *
+
+from unienv_data.base import SpaceStorage, BatchT
+
+import os
+import json
+
+WrapperBatchT = TypeVar("WrapperBatchT")
+WrapperBArrayT = TypeVar("WrapperBArrayT")
+WrapperBDeviceT = TypeVar("WrapperBDeviceT")
+WrapperBDtypeT = TypeVar("WrapperBDtypeT")
+WrapperBRngT = TypeVar("WrapperBRngT")
+
+def data_to(
+    data : Any,
+    source_backend : Optional[ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]] = None,
+    target_backend : Optional[ComputeBackend[WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]] = None,
+    target_device : Optional[WrapperBDeviceT] = None,
+):
+    if source_backend.is_backendarray(data):
+        if source_backend is not None and target_backend is not None and target_backend != source_backend:
+            data = target_backend.from_other_backend(
+                source_backend,
+                data
+            )
+        if target_device is not None:
+            data = (source_backend or target_backend).to_device(
+                data,
+                target_device
+            )
+    elif isinstance(data, Mapping):
+        data = {
+            key: data_to(value, source_backend, target_backend, target_device)
+            for key, value in data.items()
+        }
+    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
+        data = [
+            data_to(value, source_backend, target_backend, target_device)
+            for value in data
+        ]
+        try:
+            data = type(data)(data)  # Preserve the type of the original sequence
+        except:
+            pass
+    return data
+
+class ToBackendOrDeviceStorage(
+    SpaceStorage[
+        WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT
+    ],
+    Generic[
+        WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT,
+        BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType
+    ]
+):
+    # ========== Class Implementations ==========
+    @classmethod
+    def create(
+        cls,
+        single_instance_space: Space[WrapperBatchT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT],
+        inner_storage_cls : Type[SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]],
+        *args,
+        capacity : Optional[int] = None,
+        cache_path : Optional[str] = None,
+        multiprocessing : bool = False,
+        backend : Optional[ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]] = None,
+        device : Optional[BDeviceType] = None,
+        inner_storage_kwargs : Dict[str, Any] = {},
+        **kwargs
+    ) -> "ToBackendOrDeviceStorage[WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        inner_storage_path = "inner_storage" + (inner_storage_cls.single_file_ext or "")
+
+        if cache_path is not None:
+            os.makedirs(cache_path, exist_ok=True)
+
+        _inner_storage_kwargs = kwargs.copy()
+        _inner_storage_kwargs.update(inner_storage_kwargs)
+        inner_storage = inner_storage_cls.create(
+            single_instance_space.to(
+                backend=backend,
+                device=device
+            ),
+            *args,
+            cache_path=None if cache_path is None else os.path.join(cache_path, inner_storage_path),
+            capacity=capacity,
+            multiprocessing=multiprocessing,
+            **_inner_storage_kwargs
+        )
+        if (backend is None or backend == single_instance_space.backend) and (device is None or device == single_instance_space.device):
+            return inner_storage
+
+        return ToBackendOrDeviceStorage(
+            single_instance_space,
+            inner_storage,
+            inner_storage_path,
+            cache_filename=cache_path,
+        )
+
+    @classmethod
+    def load_from(
+        cls,
+        path : Union[str, os.PathLike],
+        single_instance_space : Space[Any, BDeviceType, BDtypeType, BRNGType],
+        *,
+        capacity : Optional[int] = None,
+        read_only : bool = True,
+        multiprocessing : bool = False,
+        **kwargs
+    ) -> Union[
+        "ToBackendOrDeviceStorage[WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT, BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]",
+        SpaceStorage[WrapperBatchT, WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]
+    ]:
+        metadata_path = os.path.join(path, "backend_metadata.json")
+        assert os.path.exists(metadata_path), f"Metadata file {metadata_path} does not exist"
+        with open(metadata_path, "r") as f:
+            metadata = json.load(f)
+        assert metadata["storage_type"] == cls.__name__, \
+            f"Expected storage type {cls.__name__}, but found {metadata['storage_type']}"
+        inner_storage_cls : Type[SpaceStorage] = get_class_from_full_name(metadata["inner_storage_type"])
+        inner_storage_path = metadata["inner_storage_path"]
+        inner_backend = deserialize_backend(metadata["inner_backend"])
+        inner_device = inner_backend.deserialize_device(metadata["inner_device"])
+
+        if inner_backend != single_instance_space.backend or inner_device != single_instance_space.device:
+            inner_space = single_instance_space.to(
+                backend=inner_backend if inner_backend != single_instance_space.backend else None,
+                device=inner_device if inner_device != single_instance_space.device else None
+            )
+        else:
+            inner_space = single_instance_space
+        
+        inner_storage = inner_storage_cls.load_from(
+            os.path.join(path, inner_storage_path),
+            inner_space,
+            capacity=capacity,
+            read_only=read_only,
+            multiprocessing=multiprocessing,
+            **kwargs
+        )
+
+        if inner_backend == single_instance_space.backend and inner_device == single_instance_space.device:
+            return inner_storage
+        else:
+            return ToBackendOrDeviceStorage(
+                single_instance_space,
+                inner_storage,
+                inner_storage_path,
+                cache_filename=path,
+            )
+
+    # ======== Instance Implementations ==========
+    single_file_ext = None
+
+    def __init__(
+        self,
+        single_instance_space : Space[WrapperBatchT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT],
+        inner_storage : SpaceStorage[
+            BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType
+        ],
+        inner_storage_path : Union[str, os.PathLike],
+        cache_filename : Optional[Union[str, os.PathLike]] = None,
+    ):
+        super().__init__(single_instance_space)
+        inner_backend = None if inner_storage.backend == single_instance_space.backend else inner_storage.backend
+        inner_device = None if inner_storage.device == single_instance_space.device else inner_storage.device
+        current_backend = None if inner_storage.backend == single_instance_space.backend else single_instance_space.backend
+        current_device = None if inner_storage.device == single_instance_space.device else single_instance_space.device
+
+        self._batched_instance_space = sbu.batch_space(
+            single_instance_space,
+            1
+        )
+        self._batched_inner_space = self._batched_instance_space.to(
+            backend=inner_backend,
+            device=inner_device
+        )
+
+        self.inner_storage = inner_storage
+        self.inner_storage_path = inner_storage_path
+        self._cache_filename = cache_filename
+
+        self.inner_backend = inner_backend
+        self.inner_device = inner_device
+        self.current_backend = current_backend
+        self.current_device = current_device
+    
+    @property
+    def cache_filename(self) -> Optional[Union[str, os.PathLike]]:
+        return self._cache_filename
+    
+    @property
+    def is_mutable(self) -> bool:
+        return self.inner_storage.is_mutable
+    
+    @property
+    def is_multiprocessing_safe(self) -> bool:
+        return self.inner_storage.is_multiprocessing_safe
+    
+    @property
+    def capacity(self) -> Optional[int]:
+        return self.inner_storage.capacity
+    
+    def extend_length(self, length):
+        self.inner_storage.extend_length(length)
+
+    def shrink_length(self, length):
+        self.inner_storage.shrink_length(length)
+    
+    def __len__(self):
+        return len(self.inner_storage)
+    
+    def get(self, index):
+        if self.backend.is_backendarray(index):
+            index = data_to(
+                index,
+                source_backend=self.backend,
+                target_backend=self.inner_backend,
+                target_device=self.inner_device
+            )
+        target_data = self.inner_storage.get(index)
+        return data_to(
+            target_data,
+            source_backend=self.inner_storage.backend,
+            target_backend=self.current_backend,
+            target_device=self.current_device
+        )
+    
+    def set(self, index, value):
+        target_value = data_to(
+            value,
+            source_backend=self.backend,
+            target_backend=self.inner_backend,
+            target_device=self.inner_device
+        )
+        if self.backend.is_backendarray(index):
+            index = data_to(
+                index,
+                source_backend=self.backend,
+                target_backend=self.inner_backend,
+                target_device=self.inner_device
+            )
+        self.inner_storage.set(index, target_value)
+
+    def clear(self):
+        self.inner_storage.clear()
+
+    def dumps(self, path):
+        metadata = {
+            "storage_type": __class__.__name__,
+            "inner_storage_type": get_full_class_name(type(self.inner_storage)),
+            "inner_storage_path": self.inner_storage_path,
+            "inner_backend": serialize_backend(self.inner_storage.backend),
+            "inner_device": self.inner_storage.backend.serialize_device(self.inner_storage.device),
+        }
+        self.inner_storage.dumps(os.path.join(path, self.inner_storage_path))
+        with open(os.path.join(path, "backend_metadata.json"), "w") as f:
+            json.dump(metadata, f)
+    
+    def close(self):
+        self.inner_storage.close()

unienv_data/storages/flattened.py

@@ -23,7 +23,7 @@ class FlattenedStorage(SpaceStorage[
     @classmethod
     def create(
         cls,
-        single_instance_space: Space[Any, BDeviceType, BDtypeType, BRNGType],
+        single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType],
         inner_storage_cls : Type[SpaceStorage[BArrayType, BArrayType, BDeviceType, BDtypeType, BRNGType]],
         *args,
         capacity : Optional[int] = None,
@@ -59,7 +59,7 @@ class FlattenedStorage(SpaceStorage[
     def load_from(
         cls,
         path : Union[str, os.PathLike],
-        single_instance_space : Space[Any, BDeviceType, BDtypeType, BRNGType],
+        single_instance_space : Space[BatchT, BDeviceType, BDtypeType, BRNGType],
         *,
         capacity : Optional[int] = None,
         read_only : bool = True,
@@ -95,7 +95,7 @@ class FlattenedStorage(SpaceStorage[
 
     def __init__(
         self,
-        single_instance_space: Space[Any, BDeviceType, BDtypeType, BRNGType],
+        single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType],
         inner_storage : SpaceStorage[
             BArrayType,
             BArrayType,

unienv_data/storages/hdf5.py

@@ -431,6 +431,11 @@ class HDF5Storage(SpaceStorage[
             # Convert to numpy array if it's a scalar
             if isinstance(result, (int, float)):
                 result = np.array(result)
+            if isinstance(single_instance_space, TextSpace):
+                if isinstance(result, bytes):
+                    result = result.decode('utf-8')
+                elif isinstance(result, np.ndarray):
+                    result = np.array([r.decode('utf-8') if isinstance(r, bytes) else r for r in result], dtype=object)
         return result
 
     @staticmethod
@@ -513,8 +518,8 @@ class HDF5Storage(SpaceStorage[
         reduce_io : bool = True,
         **kwargs
     ) -> "HDF5Storage":
-        assert not multiprocessing, \
-            "HDF5Storage does not support multiprocessing safe loading. Please load the storage in the main process and then share it with child processes."
+        assert read_only or not multiprocessing, \
+            "HDF5Storage does not support multiprocessing safe loading when not `read_only`. Please load the storage in the main process and then share it with child processes."
         assert os.path.exists(path), \
             f"Path {path} does not exist"
         

unienv_data/storages/image_storage.py

@@ -0,0 +1,144 @@
+from importlib import metadata
+from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequence, Union, List, Iterable, Type
+
+from unienv_interface.space import Space, BoxSpace
+from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils as sfu
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.utils.symbol_util import *
+
+from unienv_data.base import SpaceStorage
+from ._list_storage import ListStorageBase
+
+import numpy as np
+import os
+import json
+import shutil
+
+from PIL import Image
+
+class ImageStorage(ListStorageBase[
+    BArrayType,
+    BArrayType,
+    BDeviceType,
+    BDtypeType,
+    BRNGType,
+]):
+    # ========== Class Attributes ==========
+    @classmethod
+    def create(
+        cls,
+        single_instance_space: BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        *args,
+        format : str = "JPEG",
+        quality : int = 75,
+        capacity : Optional[int] = None,
+        cache_path : Optional[str] = None,
+        multiprocessing : bool = False,
+        **kwargs
+    ) -> "ImageStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        if cache_path is None:
+            raise ValueError("cache_path must be provided for ImageStorage.create")
+        assert not os.path.exists(cache_path), f"Cache path {cache_path} already exists"
+        os.makedirs(cache_path, exist_ok=True)
+        return ImageStorage(
+            single_instance_space,
+            format=format,
+            quality=quality,
+            cache_filename=cache_path,
+            capacity=capacity,
+        )
+
+    @classmethod
+    def load_from(
+        cls,
+        path : Union[str, os.PathLike],
+        single_instance_space : BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        *,
+        capacity : Optional[int] = None,
+        read_only : bool = True,
+        multiprocessing : bool = False,
+        **kwargs
+    ) -> "ImageStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        metadata_path = os.path.join(path, "image_metadata.json")
+        assert os.path.exists(metadata_path), f"Metadata file {metadata_path} does not exist"
+        with open(metadata_path, "r") as f:
+            metadata = json.load(f)
+        assert metadata["storage_type"] == cls.__name__, \
+            f"Expected storage type {cls.__name__}, but found {metadata['storage_type']}"
+        
+        format = metadata["format"]
+        quality = int(metadata["quality"])
+        if "capacity" in metadata:
+            capacity = None if metadata['capacity'] is None else int(metadata["capacity"])
+        length = None if capacity is None else metadata["length"]
+
+        return ImageStorage(
+            single_instance_space,
+            format=format,
+            quality=quality,
+            cache_filename=path,
+            mutable=not read_only,
+            capacity=capacity,
+            length=length,
+        )
+
+    # ========== Instance Implementations ==========
+    single_file_ext = None
+
+    def __init__(
+        self,
+        single_instance_space: BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        format : str,
+        quality : int,
+        cache_filename : Union[str, os.PathLike],
+        mutable : bool = True,
+        capacity : Optional[int] = None,
+        length : int = 0,
+    ):
+        super().__init__(
+            single_instance_space,
+            file_ext=format.lower(),
+            cache_filename=cache_filename,
+            mutable=mutable,
+            capacity=capacity,
+            length=length,
+        )
+        self.format = format
+        self.quality = quality
+
+    def get_from_file(self, filename : str) -> BArrayType:
+        if not os.path.exists(filename):
+            return self.backend.zeros(
+                self.single_instance_space.shape,
+                dtype=self.single_instance_space.dtype,
+                device=self.single_instance_space.device,
+            )
+        with Image.open(filename) as img:
+            rgb_img = img.convert("RGB")
+            # Ensure a compact uint8 array from PIL.
+            np_image = np.asarray(rgb_img, dtype=np.uint8)
+            return self.backend.from_numpy(np_image, dtype=self.single_instance_space.dtype, device=self.single_instance_space.device)
+    
+    def set_to_file(self, filename : str, value : BArrayType):
+        np_value = self.backend.to_numpy(value)
+        img = Image.fromarray(np_value, mode="RGB")
+        try:
+            img.save(filename, format=self.format, quality=self.quality)
+        finally:
+            img.close()
+
+    def dumps(self, path):
+        assert os.path.samefile(path, self.cache_filename), \
+            f"Dump path {path} does not match cache filename {self.cache_filename}"
+        metadata = {
+            "storage_type": __class__.__name__,
+            "format": self.format,
+            "quality": self.quality,
+            "capacity": self.capacity,
+            "length": self.length,
+        }
+        with open(os.path.join(path, "image_metadata.json"), "w") as f:
+            json.dump(metadata, f)
+
+    def close(self):
+        pass

unienv_data/storages/npz_storage.py

@@ -0,0 +1,135 @@
+from importlib import metadata
+from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequence, Union, List, Iterable, Type
+
+from unienv_interface.space import Space, BoxSpace, BinarySpace
+from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils as sfu
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.utils.symbol_util import *
+
+from unienv_data.base import SpaceStorage
+from ._list_storage import ListStorageBase
+
+import numpy as np
+import os
+import json
+import shutil
+
+from PIL import Image
+
+class NPZStorage(ListStorageBase[
+    BArrayType,
+    BArrayType,
+    BDeviceType,
+    BDtypeType,
+    BRNGType,
+]):
+    # ========== Class Attributes ==========
+    @classmethod
+    def create(
+        cls,
+        single_instance_space: BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        *args,
+        compressed : bool = True,
+        capacity : Optional[int] = None,
+        cache_path : Optional[str] = None,
+        multiprocessing : bool = False,
+        **kwargs
+    ) -> "NPZStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        assert not os.path.exists(cache_path), f"Cache path {cache_path} already exists"
+        os.makedirs(cache_path, exist_ok=True)
+        return NPZStorage(
+            single_instance_space,
+            compressed=compressed,
+            cache_filename=cache_path,
+            capacity=capacity,
+        )
+
+    @classmethod
+    def load_from(
+        cls,
+        path : Union[str, os.PathLike],
+        single_instance_space : BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        *,
+        capacity : Optional[int] = None,
+        read_only : bool = True,
+        multiprocessing : bool = False,
+        **kwargs
+    ) -> "NPZStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        metadata_path = os.path.join(path, "npz_metadata.json")
+        assert os.path.exists(metadata_path), f"Metadata file {metadata_path} does not exist"
+        with open(metadata_path, "r") as f:
+            metadata = json.load(f)
+        assert metadata["storage_type"] == cls.__name__, \
+            f"Expected storage type {cls.__name__}, but found {metadata['storage_type']}"
+        
+        compressed = metadata.get("compressed", True)
+        if "capacity" in metadata:
+            capacity = None if metadata['capacity'] is None else int(metadata["capacity"])
+        length = None if capacity is None else metadata["length"]
+
+        return NPZStorage(
+            single_instance_space,
+            compressed=compressed,
+            cache_filename=path,
+            mutable=not read_only,
+            capacity=capacity,
+            length=length,
+        )
+
+    # ========== Instance Implementations ==========
+    def __init__(
+        self,
+        single_instance_space: BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        compressed : bool,
+        cache_filename : Union[str, os.PathLike],
+        mutable : bool = True,
+        capacity : Optional[int] = None,
+        length : int = 0,
+    ):
+        assert isinstance(single_instance_space, BoxSpace) or isinstance(single_instance_space, BinarySpace), "single_instance_space must be a BoxSpace or BinarySpace"
+        super().__init__(
+            single_instance_space,
+            file_ext="npz",
+            cache_filename=cache_filename,
+            mutable=mutable,
+            capacity=capacity,
+            length=length,
+        )
+        self.compressed = compressed
+
+    def get_from_file(self, filename : str) -> BArrayType:
+        if not os.path.exists(filename):
+            return self.backend.zeros(
+                self.single_instance_space.shape,
+                dtype=self.single_instance_space.dtype,
+                device=self.single_instance_space.device,
+            )
+        
+        dat = np.load(filename, allow_pickle=False)
+        return self.backend.from_numpy(
+            dat['data'], 
+            dtype=self.single_instance_space.dtype, 
+            device=self.single_instance_space.device
+        )
+    
+    def set_to_file(self, filename : str, value : BArrayType):
+        np_value = self.backend.to_numpy(value)
+        if self.compressed:
+            np.savez_compressed(filename, data=np_value)
+        else:
+            np.savez(filename, data=np_value)
+
+    def dumps(self, path):
+        assert os.path.samefile(path, self.cache_filename), \
+            f"Dump path {path} does not match cache filename {self.cache_filename}"
+        metadata = {
+            "storage_type": __class__.__name__,
+            "compressed": self.compressed,
+            "capacity": self.capacity,
+            "length": self.length,
+        }
+        with open(os.path.join(path, "npz_metadata.json"), "w") as f:
+            json.dump(metadata, f)
+
+    def close(self):
+        pass

unienv_data/storages/pytorch.py

@@ -1,10 +1,10 @@
 import os
 import torch
-from unienv_interface.space import Space, BoxSpace
+from unienv_interface.space import Space, BoxSpace, BinarySpace
 from unienv_interface.backends import ComputeBackend
 from unienv_interface.backends.pytorch import PyTorchComputeBackend, PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType
 from unienv_data.base import SpaceStorage
-from tensordict.memmap import MemoryMappedTensor
+from unienv_data.third_party.tensordict.memmap_tensor import MemoryMappedTensor
 from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequence, Union, List, Type
 
 class PytorchTensorStorage(SpaceStorage[
@@ -20,15 +20,15 @@ class PytorchTensorStorage(SpaceStorage[
         single_instance_space : BoxSpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
         *,
         capacity : Optional[int],
-        is_memmap : bool = False,
+        is_memmap : bool = True,
         cache_path : Optional[str] = None,
         multiprocessing : bool = False,
         memmap_existok : bool = True,
     ) -> "PytorchTensorStorage":
         assert single_instance_space.backend is PyTorchComputeBackend, \
             f"Single instance space must be of type PyTorchComputeBackend, got {single_instance_space.backend}"
-        assert isinstance(single_instance_space, BoxSpace), \
-            f"Single instance space must be a BoxSpace, got {type(single_instance_space)}"
+        assert isinstance(single_instance_space, BoxSpace) or isinstance(single_instance_space, BinarySpace), \
+            f"Single instance space must be a BoxSpace or BinarySpace, got {type(single_instance_space)}"
         assert capacity is not None, "Capacity must be specified when creating a new tensor"
 
         target_shape = (capacity, *single_instance_space.shape)
@@ -63,9 +63,12 @@ class PytorchTensorStorage(SpaceStorage[
     def load_from(
         cls,
         path: Union[str, os.PathLike],
-        single_instance_space: BoxSpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+        single_instance_space: Union[
+            BoxSpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+            BinarySpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+        ],
         *,
-        is_memmap : bool = False,
+        is_memmap : bool = True,
         capacity : Optional[int] = None,
         read_only : bool = True,
         multiprocessing : bool = False,
@@ -81,7 +84,8 @@ class PytorchTensorStorage(SpaceStorage[
         target_data = MemoryMappedTensor.from_filename(
             path,
             dtype=single_instance_space.dtype,
-            shape=target_shape
+            shape=target_shape,
+            readonly=read_only,
         )
 
         if is_memmap:
@@ -111,7 +115,10 @@ class PytorchTensorStorage(SpaceStorage[
 
     def __init__(
         self,
-        single_instance_space : BoxSpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+        single_instance_space : Union[
+            BoxSpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+            BinarySpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+        ],
         data : Union[torch.Tensor, MemoryMappedTensor],
         mutable : bool = True,
     ):