unienv 0.0.1b3__tar.gz → 0.0.1b4__tar.gz

{unienv-0.0.1b3/unienv.egg-info → unienv-0.0.1b4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b3
+Version: 0.0.1b4
 Summary: Unified robot environment framework supporting multiple tensor and simulation backends
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv

{unienv-0.0.1b3 → unienv-0.0.1b4}/pyproject.toml

@@ -3,7 +3,7 @@ name = "unienv"
 description = "Unified robot environment framework supporting multiple tensor and simulation backends"
 readme = "README.md"
 license = "MIT"
-version = "0.0.1b3"
+version = "0.0.1b4"
 requires-python = ">= 3.10"
 dependencies = [
     "numpy",

{unienv-0.0.1b3 → unienv-0.0.1b4/unienv.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b3
+Version: 0.0.1b4
 Summary: Unified robot environment framework supporting multiple tensor and simulation backends
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv.egg-info/SOURCES.txt

@@ -23,7 +23,8 @@ unienv_data/replay_buffer/trajectory_replay_buffer.py
 unienv_data/samplers/__init__.py
 unienv_data/samplers/multiprocessing_sampler.py
 unienv_data/samplers/step_sampler.py
-unienv_data/storages/common.py
+unienv_data/storages/dict_storage.py
+unienv_data/storages/flattened.py
 unienv_data/storages/hdf5.py
 unienv_data/storages/pytorch.py
 unienv_data/storages/transformation.py

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_data/base/common.py

@@ -135,12 +135,25 @@ class BatchBase(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BR
         flattened_data = space_flatten_utils.flatten_data(self._batched_space, value, start_dim=1)
         self.extend_flattened(flattened_data)
 
+    def extend_from(
+        self, 
+        other : 'BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]',
+        chunk_size : int = 8,
+        tqdm : bool = False,
+    ) -> None:
+        n_total = len(other)
+        iterable_start = range(0, n_total, chunk_size)
+        if tqdm:
+            from tqdm import tqdm
+            iterable_start = tqdm(iterable_start, desc="Extending Batch")
+        for start_idx in range(0, n_total, chunk_size):
+            end_idx = min(start_idx + chunk_size, n_total)
+            data_chunk = other.get_at(slice(start_idx, end_idx))
+            self.extend(data_chunk)
+
     def close(self) -> None:
         pass
 
-    def __del__(self) -> None:
-        self.close()
-
 SamplerBatchT = TypeVar('SamplerBatchT')
 SamplerArrayType = TypeVar('SamplerArrayType')
 SamplerDeviceType = TypeVar('SamplerDeviceType')
@@ -273,6 +286,3 @@ class BatchSampler(
 
     def close(self) -> None:
         pass
-
-    def __del__(self) -> None:
-        self.close()

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_data/base/storage.py

@@ -57,6 +57,17 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
     """
     cache_filename : Optional[Union[str, os.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 : bool = False
+
+    """
+    Is the storage mutable? If False, the storage is read-only.
+    """
+    is_mutable : bool = True
+
     @property
     def backend(self) -> ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]:
         return self.single_instance_space.backend
@@ -128,6 +139,3 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
 
     def close(self) -> None:
         pass
-
-    def __del__(self) -> None:
-        self.close()

unienv-0.0.1b4/unienv_data/storages/dict_storage.py

@@ -0,0 +1,341 @@
+from importlib import metadata
+from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequence, Union, List, Iterable, Type, Callable, Mapping
+
+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.utils.symbol_util import *
+
+from unienv_data.base import SpaceStorage, BatchT
+
+import numpy as np
+import os
+import json
+
+def map_transform(
+    data : Dict[str, Any],
+    value_map : Dict[str, Any],
+    fn : Callable[[str, Any, Any], Any], # (str, data, value_map) -> transformed data
+    prefix : str = "",
+) -> Tuple[
+    Dict[str, Any], # Transformed data
+    Dict[str, Any], # Residual data
+]:
+    transformed_data = {}
+    residual_data = {}
+    for key, value in data.items() if isinstance(data, Mapping) else data.spaces.items():
+        full_key = prefix + key
+        if full_key in value_map:
+            transformed_data[key] = fn(full_key, value, value_map[full_key])
+        elif isinstance(value, Mapping) or isinstance(value, DictSpace):
+            sub_transformed, sub_residual = map_transform(
+                value,
+                value_map,
+                fn,
+                prefix=full_key + "/",
+            )
+            if len(sub_transformed) > 0:
+                transformed_data[key] = sub_transformed
+            if len(sub_residual) > 0:
+                residual_data[key] = sub_residual
+        else:
+            residual_data[key] = value
+    if len(residual_data) > 0 and (prefix + "*") in value_map:
+        residual_transformed = fn(prefix + "*", residual_data, value_map[prefix + "*"])
+        if isinstance(residual_transformed, Mapping) or isinstance(residual_transformed, DictSpace):
+            for key, value in residual_transformed.items():
+                transformed_data[key] = value
+        residual_data = {}
+    return transformed_data, residual_data
+
+def get_chained_residual_space(
+    space : DictSpace[BDeviceType, BDtypeType, BRNGType],
+    all_keys : List[str],
+    prefix : str = "",
+) -> DictSpace[BDeviceType, BDtypeType, BRNGType]:
+    residual_spaces = {}
+
+    if len(residual_spaces) > 0 and (prefix + "*") in all_keys:
+        return DictSpace(
+            space.backend,
+            {},
+            device=space.device,
+        )
+
+    for key, subspace in space.spaces.items():
+        full_key = prefix + key
+        if full_key in all_keys:
+            continue
+        elif isinstance(subspace, DictSpace):
+            sub_residual = get_chained_residual_space(
+                subspace,
+                all_keys,
+                prefix=full_key + "/",
+            )
+            if len(sub_residual.spaces) > 0:
+                residual_spaces[key] = sub_residual
+        else:
+            residual_spaces[key] = subspace
+
+    return DictSpace(
+        space.backend,
+        residual_spaces,
+        device=space.device,
+    )
+
+def get_chained_space(
+    space : DictSpace[BDeviceType, BDtypeType, BRNGType],
+    key_chain : str,
+    all_keys : List[str],
+) -> Space[Any, BDeviceType, BDtypeType, BRNGType]:
+    if key_chain.endswith("*"):
+        prefix = key_chain[:-1]
+        subspace = get_chained_residual_space(
+            get_chained_space(
+                space,
+                prefix,
+                all_keys,
+            ) if len(prefix) > 0 else space,
+            [key for key in all_keys if key != key_chain],
+            prefix=prefix,
+        )
+        return subspace
+    key_chain = key_chain.split("/")
+    current_space : Space[Any, BDeviceType, BDtypeType, BRNGType]
+    current_space = space
+    for key in key_chain:
+        if len(key) == 0:
+            continue
+        assert isinstance(current_space, DictSpace), \
+            f"Expected DictSpace while traversing key chain, but got {type(current_space)}"
+        current_space = current_space.spaces[key]
+    return current_space
+
+class DictStorage(SpaceStorage[
+    Dict[str, Any],
+    BArrayType,
+    BDeviceType,
+    BDtypeType,
+    BRNGType,
+]):
+    # ========== Class Attributes ==========
+    @classmethod
+    def create(
+        cls,
+        single_instance_space: Space[Any, BDeviceType, BDtypeType, BRNGType],
+        storage_cls_map : Dict[
+            str,
+            Type[SpaceStorage],
+        ],
+        *args,
+        capacity : Optional[int] = None,
+        cache_path : Optional[str] = None,
+        key_kwargs : Dict[str, Any] = {},
+        type_kwargs : Dict[Type[SpaceStorage[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], Dict[str, Any]] = {},
+        **kwargs
+    ) -> "DictStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        if cache_path is not None:
+            os.makedirs(cache_path, exist_ok=True)
+
+        storage_map = {}
+        all_keys = list(storage_cls_map.keys())
+        for key, sub_storage_cls in storage_cls_map.items():
+            sub_storage_path = key.replace("/", ".").replace("*", "_default") + (sub_storage_cls.single_file_ext or "")
+            subspace = get_chained_space(single_instance_space, key, all_keys)
+            sub_kwargs = kwargs.copy()
+            if sub_storage_cls in type_kwargs:
+                sub_kwargs.update(type_kwargs[sub_storage_cls])
+            if key in key_kwargs:
+                sub_kwargs.update(key_kwargs[key])
+            storage_map[key] = sub_storage_cls.create(
+                subspace,
+                *args,
+                cache_path=None if cache_path is None else os.path.join(cache_path, sub_storage_path),
+                capacity=capacity,
+                **sub_kwargs
+            )
+
+        return DictStorage(
+            single_instance_space,
+            storage_map,
+            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,
+        key_kwargs : Dict[str, Any] = {},
+        type_kwargs : Dict[Type[SpaceStorage[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], Dict[str, Any]] = {},
+        **kwargs
+    ) -> "DictStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        metadata_path = os.path.join(path, "dict_storage_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']}"
+        
+        storage_map_metadata = metadata["storage_map"]
+        storage_map = {}
+
+        all_keys = list(storage_map_metadata.keys())
+        for key, storage_meta in storage_map_metadata.items():
+            storage_cls : Type[SpaceStorage] = get_class_from_full_name(storage_meta["type"])
+            storage_path = storage_meta["path"]
+            
+            subspace = get_chained_space(single_instance_space, key, all_keys)
+            
+            sub_kwargs = kwargs.copy()
+            if storage_cls in type_kwargs:
+                sub_kwargs.update(type_kwargs[storage_cls])
+            if key in key_kwargs:
+                sub_kwargs.update(key_kwargs[key])
+            storage_map[key] = storage_cls.load_from(
+                os.path.join(path, storage_path),
+                subspace,
+                capacity=capacity,
+                read_only=read_only,
+                **sub_kwargs
+            )
+
+        return DictStorage(
+            single_instance_space,
+            storage_map,
+            cache_filename=path,
+        )
+
+    # ========== Instance Implementations ==========
+    single_file_ext = None
+
+    def __init__(
+        self,
+        single_instance_space: DictSpace[BDeviceType, BDtypeType, BRNGType],
+        storage_map : Dict[
+            str,
+            SpaceStorage[
+                BArrayType,
+                BArrayType,
+                BDeviceType,
+                BDtypeType,
+                BRNGType,
+            ],
+        ],
+        cache_filename: Optional[Union[str, os.PathLike]] = None,
+    ):
+        assert len(storage_map) > 0, "Storage map cannot be empty"
+        first_storage = next(iter(storage_map.values()))
+        init_capacity = first_storage.capacity
+        init_len = len(first_storage)
+        for key, storage in storage_map.items():
+            assert storage.capacity == init_capacity, \
+                f"All storages must have the same capacity, but storage {key} has capacity {storage.capacity} while first storage has capacity {init_capacity}"
+            assert len(storage) == init_len, \
+                f"All storages must have the same length, but storage {key} has length {len(storage)} while first storage has length {init_len}"
+
+        super().__init__(single_instance_space)
+        self._batched_instance_space = sbu.batch_space(single_instance_space, 1)
+        self.storage_map = storage_map
+        self._cache_filename = cache_filename if all(
+            storage.cache_filename is not None for storage in storage_map.values()
+        ) else None
+
+    @property
+    def cache_filename(self) -> Optional[Union[str, os.PathLike]]:
+        return self._cache_filename
+
+    @property
+    def is_mutable(self) -> bool:
+        return all(storage.is_mutable for storage in self.storage_map.values())
+
+    @property
+    def is_multiprocessing_safe(self) -> bool:
+        return all(storage.is_multiprocessing_safe for storage in self.storage_map.values())
+
+    @property
+    def capacity(self) -> Optional[int]:
+        return next(iter(self.storage_map.values())).capacity
+    
+    def extend_length(self, length):
+        for storage in self.storage_map.values():
+            storage.extend_length(length)
+
+    def shrink_length(self, length):
+        for storage in self.storage_map.values():
+            storage.shrink_length(length)
+
+    def __len__(self):
+        return len(next(iter(self.storage_map.values())))
+
+    def get_flattened(self, index):
+        unflat_data = self.get(index)
+        if isinstance(index, int):
+            flat_data = sfu.flatten_data(self.single_instance_space, unflat_data)
+        else:
+            flat_data = sfu.flatten_data(self._batched_instance_space, unflat_data, start_dim=1)
+        return flat_data
+
+    def get(self, index):
+        result, residual = map_transform(
+            self.single_instance_space,
+            self.storage_map,
+            lambda key, space, storage: storage.get(index)
+        )
+        assert len(residual) == 0, f"Some spaces do not have corresponding storage: {residual}"
+        return result
+    
+    def set_flattened(self, index, value):
+        if isinstance(index, int):
+            unflat_data = sfu.unflatten_data(self.single_instance_space, value)
+        else:
+            unflat_data = sfu.unflatten_data(self._batched_instance_space, value, start_dim=1)
+        self.set(index, unflat_data)
+
+    def set(self, index, value):
+        _, residual = map_transform(
+            value,
+            self.storage_map,
+            lambda key, data, storage: storage.set(index, data)
+        )
+        assert len(residual) == 0, f"Some spaces do not have corresponding storage: {residual}"
+
+    def get_subspace_by_key(
+        self,
+        key: str,
+    ) -> Space[Any, BDeviceType, BDtypeType, BRNGType]:
+        return get_chained_space(
+            self.single_instance_space,
+            key,
+            list(self.storage_map.keys()),
+        )
+
+    def clear(self):
+        for storage in self.storage_map.values():
+            storage.clear()
+
+    def dumps(self, path):
+        os.makedirs(path, exist_ok=True)
+
+        storage_map_metadata = {}
+        for key, storage in self.storage_map.items():
+            sub_storage_path = key.replace("/", ".").replace("*", "_default") + (storage.single_file_ext or "")
+            storage_map_metadata[key] = {
+                "type": get_full_class_name(type(storage)),
+                "path": sub_storage_path,
+            }
+            storage.dumps(os.path.join(path, sub_storage_path))
+        
+        metadata = {
+            "storage_type": __class__.__name__,
+            "storage_map": storage_map_metadata,
+        }
+        with open(os.path.join(path, "dict_storage_metadata.json"), "w") as f:
+            json.dump(metadata, f)
+
+    def close(self):
+        for storage in self.storage_map.values():
+            storage.close()

unienv-0.0.1b3/unienv_data/storages/common.py → unienv-0.0.1b4/unienv_data/storages/flattened.py

@@ -3,9 +3,7 @@ from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequen
 
 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.env_base.env import ContextType, ObsType, ActType
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
-from unienv_interface.backends.numpy import NumpyComputeBackend
 from unienv_interface.utils.symbol_util import *
 
 from unienv_data.base import SpaceStorage, BatchT
@@ -31,7 +29,7 @@ class FlattenedStorage(SpaceStorage[
         capacity : Optional[int] = None,
         cache_path : Optional[str] = None,
         **kwargs
-    ) -> "FlattenedStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+    ) -> "FlattenedStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         flattened_space = sfu.flatten_space(single_instance_space)
         inner_storage_path = "inner_storage" + (inner_storage_cls.single_file_ext or "")
 
@@ -49,6 +47,7 @@ class FlattenedStorage(SpaceStorage[
             single_instance_space,
             inner_storage,
             inner_storage_path,
+            cache_filename=cache_path,
         )
 
     @classmethod
@@ -60,7 +59,7 @@ class FlattenedStorage(SpaceStorage[
         capacity : Optional[int] = None,
         read_only : bool = True,
         **kwargs
-    ) -> "FlattenedStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+    ) -> "FlattenedStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         metadata_path = os.path.join(path, "flattened_metadata.json")
         assert os.path.exists(metadata_path), f"Metadata file {metadata_path} does not exist"
         with open(metadata_path, "r") as f:
@@ -81,6 +80,7 @@ class FlattenedStorage(SpaceStorage[
             single_instance_space,
             inner_storage,
             inner_storage_path,
+            cache_filename=path,
         )
 
     # ========== Instance Implementations ==========
@@ -97,6 +97,7 @@ class FlattenedStorage(SpaceStorage[
             BRNGType,
         ],
         inner_storage_path : Union[str, os.PathLike],
+        cache_filename : Optional[Union[str, os.PathLike]] = None,
     ):
         super().__init__(single_instance_space)
         assert inner_storage.backend == single_instance_space.backend, \
@@ -109,7 +110,20 @@ class FlattenedStorage(SpaceStorage[
         self._batched_instance_space = sbu.batch_space(single_instance_space, 1)
         self.inner_storage = inner_storage
         self.inner_storage_path = inner_storage_path
-        
+        self._cache_filename = cache_filename
+
+    @property
+    def cache_filename(self) -> Optional[Union[str, os.PathLike]]:
+        return self._cache_filename if self.inner_storage.cache_filename is not None else None
+
+    @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

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_data/storages/hdf5.py

@@ -2,7 +2,6 @@ from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequen
 
 from unienv_interface.space import Space, BoxSpace, DictSpace, TextSpace, BinarySpace
 from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils as sfu
-from unienv_interface.env_base.env import ContextType, ObsType, ActType
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
 from unienv_interface.backends.numpy import NumpyComputeBackend, NumpyArrayType, NumpyDeviceType, NumpyDtypeType, NumpyRNGType
 from unienv_interface.utils.symbol_util import *
@@ -498,7 +497,7 @@ class HDF5Storage(SpaceStorage[
             capacity=capacity,
             reduce_io=reduce_io,
         )
-    
+
     @classmethod
     def load_from(
         cls, 
@@ -562,6 +561,20 @@ class HDF5Storage(SpaceStorage[
         assert self.capacity is None or self._len == self.capacity, \
             f"If the storage has a fixed capacity, the length must match the capacity. Expected {self.capacity}, got {self._len}"
     
+    @property
+    def is_mutable(self) -> bool:
+        return self.root.file.mode != 'r'
+
+    @property
+    def is_multiprocessing_safe(self) -> bool:
+        return not self.is_mutable
+
+    @property
+    def cache_filename(self) -> Optional[Union[str, os.PathLike]]:
+        if isinstance(self.root, h5py.File):
+            return self.root.filename
+        return None
+
     def extend_length(self, length):
         assert self.capacity is None, \
             "Cannot extend length of a storage with fixed capacity"
@@ -644,4 +657,30 @@ class HDF5Storage(SpaceStorage[
     def close(self):
         if isinstance(self.root, h5py.File):
             self.root.close()
-        self.root = None
+        self.root = None
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        if (self.root, h5py.File):
+            state['filename'] = self.root.filename
+            state['mode'] = self.root.file.mode
+        else:
+            state['filename'] = self.root.file.filename
+            state['mode'] = self.root.file.mode
+            state['full_name'] = self.root.name
+        del state['root']
+        return state
+    
+    def __setstate__(self, state):
+        if 'filename' and 'mode' in state:
+            self.root = h5py.File(
+                state['filename'],
+                mode=state['mode']
+            )
+            if 'full_name' in state:
+                self.root = self.root[state['full_name']]
+                del state['full_name']
+            
+            del state['filename']
+            del state['mode']
+        self.__dict__.update(state)

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_data/storages/pytorch.py

@@ -1,7 +1,6 @@
 import os
 import torch
 from unienv_interface.space import Space, BoxSpace
-from unienv_interface.env_base.env import ContextType, ObsType, ActType
 from unienv_interface.backends import ComputeBackend
 from unienv_interface.backends.pytorch import PyTorchComputeBackend, PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType
 from unienv_data.base import SpaceStorage
@@ -24,6 +23,7 @@ class PytorchTensorStorage(SpaceStorage[
         is_memmap : bool = False,
         cache_path : Optional[str] = None,
         memmap_existok : bool = True,
+        multiprocessing : bool = False,
     ) -> "PytorchTensorStorage":
         assert single_instance_space.backend is PyTorchComputeBackend, \
             f"Single instance space must be of type PyTorchComputeBackend, got {single_instance_space.backend}"
@@ -54,8 +54,10 @@ class PytorchTensorStorage(SpaceStorage[
                 dtype=single_instance_space.dtype,
                 device=single_instance_space.device
             )
-        
-        return PytorchTensorStorage(single_instance_space, data)
+            if multiprocessing:
+                data = data.share_memory_()
+
+        return PytorchTensorStorage(single_instance_space, data, mutable=True)
 
     @classmethod
     def load_from(
@@ -66,11 +68,15 @@ class PytorchTensorStorage(SpaceStorage[
         is_memmap : bool = False,
         capacity : Optional[int] = None,
         read_only : bool = True,
+        multiprocessing : bool = False,
     ) -> "PytorchTensorStorage":
         assert single_instance_space.backend is PyTorchComputeBackend, "PytorchTensorStorage only supports PyTorch backend"
         assert capacity is not None, "Capacity must be specified when creating a new tensor"
         assert os.path.exists(path), "File does not exist"
 
+        if is_memmap and not read_only:
+            assert os.access(path, os.W_OK), "File is not writable, cannot open in read-write mode"
+
         target_shape = (capacity, *single_instance_space.shape)
         target_data = MemoryMappedTensor.from_filename(
             path,
@@ -88,11 +94,14 @@ class PytorchTensorStorage(SpaceStorage[
                 dtype=single_instance_space.dtype,
                 device=single_instance_space.device
             )
-            data.copy_(target_data)
+            if multiprocessing:
+                data = data.share_memory_()
+            data = data.copy_(target_data)
 
         return PytorchTensorStorage(
             single_instance_space,
-            data
+            data,
+            mutable=not read_only
         )
 
     # ========== Instance Implementations ==========
@@ -104,6 +113,7 @@ class PytorchTensorStorage(SpaceStorage[
         self,
         single_instance_space : BoxSpace[PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
         data : Union[torch.Tensor, MemoryMappedTensor],
+        mutable : bool = True,
     ):
         assert single_instance_space.shape == data.shape[1:], \
             f"Single instance space shape {single_instance_space.shape} does not match data shape {data.shape[1:]}"
@@ -111,6 +121,7 @@ class PytorchTensorStorage(SpaceStorage[
             single_instance_space
         )
         self.data = data
+        self._mutable = mutable
 
     @property
     def device(self) -> Optional[PyTorchDeviceType]:
@@ -122,6 +133,14 @@ class PytorchTensorStorage(SpaceStorage[
             return self.data.filename
         return None
     
+    @property
+    def is_mutable(self) -> bool:
+        return self._mutable
+
+    @property
+    def is_multiprocessing_safe(self) -> bool:
+        return self.data.is_shared()
+    
     @property
     def capacity(self) -> int:
         return self.data.shape[0]
@@ -134,9 +153,11 @@ class PytorchTensorStorage(SpaceStorage[
         return self.data[index]
 
     def set(self, index : Union[int, slice, torch.Tensor], value : torch.Tensor) -> None:
+        assert self.is_mutable, "Storage is not mutable"
         self.data[index] = value
 
     def clear(self) -> None:
+        assert self.is_mutable, "Storage is not mutable"
         pass
     
     def dumps(self, path: Union[str, os.PathLike]) -> None:

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_data/storages/transformation.py

@@ -3,9 +3,7 @@ from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequen
 
 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.env_base.env import ContextType, ObsType, ActType
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
-from unienv_interface.backends.numpy import NumpyComputeBackend
 from unienv_interface.utils.symbol_util import *
 from unienv_interface.transformations import DataTransformation
 

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_interface/world/funcnode.py

@@ -21,7 +21,6 @@ class FuncWorldNode(ABC, Generic[
     """
 
     name : str
-    world : FuncWorld[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]
     control_timestep : Optional[float] = None
     context_space : Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
     observation_space : Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
@@ -29,6 +28,7 @@ class FuncWorldNode(ABC, Generic[
     has_reward : bool = False
     has_termination_signal : bool = False
     has_truncation_signal : bool = False
+    world : Optional[FuncWorld[WorldStateT, BArrayType, BDeviceType, BDtypeType, BRNGType]] = None
 
     @property
     def backend(self) -> ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]:

{unienv-0.0.1b3 → unienv-0.0.1b4}/unienv_interface/world/node.py

@@ -8,7 +8,7 @@ from .world import World
 
 class WorldNode(ABC, Generic[ContextType, ObsType, ActType, BArrayType, BDeviceType, BDtypeType, BRNGType]):
     """
-    Each `WorldNode` in the simulated / real world will manage some aspect of the environment.
+    Each `WorldNode` in the simulated / real world will manage some aspect of the environment. This can include sensors, robots, or other entities that interact with the world.
     How the methods in this class will be called once environment resets:
     `World.reset(...)` -> `WorldNode.reset(...)` -> `WorldNode.after_reset(...)` -> `WorldNode.get_observation(...)` -> World can start stepping normally
     How the methods in this class will be called during a environment step:
@@ -16,7 +16,6 @@ class WorldNode(ABC, Generic[ContextType, ObsType, ActType, BArrayType, BDeviceT
     """
 
     name : str
-    world : World[BArrayType, BDeviceType, BDtypeType, BRNGType]
     control_timestep : Optional[float] = None
     context_space : Optional[Space[ContextType, BDeviceType, BDtypeType, BRNGType]] = None
     observation_space : Optional[Space[ObsType, BDeviceType, BDtypeType, BRNGType]] = None
@@ -24,6 +23,7 @@ class WorldNode(ABC, Generic[ContextType, ObsType, ActType, BArrayType, BDeviceT
     has_reward : bool = False
     has_termination_signal : bool = False
     has_truncation_signal : bool = False
+    world : Optional[World[BArrayType, BDeviceType, BDtypeType, BRNGType]] = None
 
     @property
     def backend(self) -> ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]: