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

unienv_data/storages/dict_storage.py

@@ -12,6 +12,24 @@ import numpy as np
 import os
 import json
 
+
+def _merge_nested_mappings(
+    primary: Mapping[str, Any],
+    secondary: Mapping[str, Any],
+) -> Mapping[str, Any]:
+    """Merge secondary into primary without clobbering explicitly matched keys."""
+    merged: Dict[str, Any] = dict(primary)
+    for merge_key, merge_value in secondary.items():
+        if (
+            merge_key in merged
+            and isinstance(merged[merge_key], Mapping)
+            and isinstance(merge_value, Mapping)
+        ):
+            merged[merge_key] = _merge_nested_mappings(merged[merge_key], merge_value)
+        elif merge_key not in merged:
+            merged[merge_key] = merge_value
+    return merged
+
 def map_transform(
     data : Dict[str, Any],
     value_map : Dict[str, Any],
@@ -44,7 +62,10 @@ def map_transform(
         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
+                if key in transformed_data and isinstance(transformed_data[key], Mapping) and isinstance(value, Mapping):
+                    transformed_data[key] = _merge_nested_mappings(transformed_data[key], value)
+                elif key not in transformed_data:
+                    transformed_data[key] = value
         residual_data = {}
     return transformed_data, residual_data
 
@@ -52,7 +73,7 @@ def get_chained_residual_space(
     space : DictSpace[BDeviceType, BDtypeType, BRNGType],
     all_keys : List[str],
     prefix : str = "",
-) -> DictSpace[BDeviceType, BDtypeType, BRNGType]:
+) -> Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]]:
     residual_spaces = {}
 
     if len(residual_spaces) > 0 and (prefix + "*") in all_keys:
@@ -72,10 +93,13 @@ def get_chained_residual_space(
                 all_keys,
                 prefix=full_key + "/",
             )
-            if len(sub_residual.spaces) > 0:
+            if sub_residual is not None and len(sub_residual.spaces) > 0:
                 residual_spaces[key] = sub_residual
         else:
             residual_spaces[key] = subspace
+    
+    if len(residual_spaces) == 0:
+        return None
 
     return DictSpace(
         space.backend,
@@ -87,7 +111,7 @@ def get_chained_space(
     space : DictSpace[BDeviceType, BDtypeType, BRNGType],
     key_chain : str,
     all_keys : List[str],
-) -> Space[Any, BDeviceType, BDtypeType, BRNGType]:
+) -> Optional[Space[Any, BDeviceType, BDtypeType, BRNGType]]:
     if key_chain.endswith("*"):
         prefix = key_chain[:-1]
         subspace = get_chained_residual_space(
@@ -106,8 +130,8 @@ def get_chained_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)}"
+        if not isinstance(current_space, DictSpace) or key not in current_space.spaces:
+            return None
         current_space = current_space.spaces[key]
     return current_space
 
@@ -130,6 +154,7 @@ class DictStorage(SpaceStorage[
         *args,
         capacity : Optional[int] = None,
         cache_path : Optional[str] = None,
+        multiprocessing : bool = False,
         key_kwargs : Dict[str, Any] = {},
         type_kwargs : Dict[Type[SpaceStorage[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], Dict[str, Any]] = {},
         **kwargs
@@ -142,6 +167,8 @@ class DictStorage(SpaceStorage[
         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)
+            if subspace is None:
+                continue
             sub_kwargs = kwargs.copy()
             if sub_storage_cls in type_kwargs:
                 sub_kwargs.update(type_kwargs[sub_storage_cls])
@@ -152,6 +179,7 @@ class DictStorage(SpaceStorage[
                 *args,
                 cache_path=None if cache_path is None else os.path.join(cache_path, sub_storage_path),
                 capacity=capacity,
+                multiprocessing=multiprocessing,
                 **sub_kwargs
             )
 
@@ -169,6 +197,7 @@ class DictStorage(SpaceStorage[
         *,
         capacity : Optional[int] = None,
         read_only : bool = True,
+        multiprocessing : bool = False,
         key_kwargs : Dict[str, Any] = {},
         type_kwargs : Dict[Type[SpaceStorage[Any, BArrayType, BDeviceType, BDtypeType, BRNGType]], Dict[str, Any]] = {},
         **kwargs
@@ -189,7 +218,9 @@ class DictStorage(SpaceStorage[
             storage_path = storage_meta["path"]
             
             subspace = get_chained_space(single_instance_space, key, all_keys)
-            
+            if subspace is None:
+                continue
+
             sub_kwargs = kwargs.copy()
             if storage_cls in type_kwargs:
                 sub_kwargs.update(type_kwargs[storage_cls])
@@ -200,6 +231,7 @@ class DictStorage(SpaceStorage[
                 subspace,
                 capacity=capacity,
                 read_only=read_only,
+                multiprocessing=multiprocessing,
                 **sub_kwargs
             )
 

unienv_data/storages/flattened.py

@@ -28,6 +28,8 @@ class FlattenedStorage(SpaceStorage[
         *args,
         capacity : Optional[int] = None,
         cache_path : Optional[str] = None,
+        multiprocessing : bool = False,
+        inner_storage_kwargs : Dict[str, Any] = {},
         **kwargs
     ) -> "FlattenedStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         flattened_space = sfu.flatten_space(single_instance_space)
@@ -36,12 +38,15 @@ class FlattenedStorage(SpaceStorage[
         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(
             flattened_space,
             *args,
             cache_path=None if cache_path is None else os.path.join(cache_path, inner_storage_path),
             capacity=capacity,
-            **kwargs
+            multiprocessing=multiprocessing,
+            **_inner_storage_kwargs
         )
         return FlattenedStorage(
             single_instance_space,
@@ -58,6 +63,7 @@ class FlattenedStorage(SpaceStorage[
         *,
         capacity : Optional[int] = None,
         read_only : bool = True,
+        multiprocessing : bool = False,
         **kwargs
     ) -> "FlattenedStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         metadata_path = os.path.join(path, "flattened_metadata.json")
@@ -74,6 +80,7 @@ class FlattenedStorage(SpaceStorage[
             flattened_space,
             capacity=capacity,
             read_only=read_only,
+            multiprocessing=multiprocessing,
             **kwargs
         )
         return FlattenedStorage(

unienv_data/storages/hdf5.py

@@ -458,6 +458,7 @@ class HDF5Storage(SpaceStorage[
         single_instance_space, 
         capacity, 
         cache_path = None, 
+        multiprocessing : bool = False,
         initial_capacity : Optional[int] = None,
         compression : Union[
             Dict[str, Any],
@@ -476,6 +477,8 @@ class HDF5Storage(SpaceStorage[
     ) -> "HDF5Storage":
         assert cache_path is not None, \
             "cache_path must be provided for HDF5Storage"
+        assert not multiprocessing, \
+            "HDF5Storage does not support multiprocessing safe creation. Please create the storage in the main process and then load it in child processes."
         root = h5py.File(
             cache_path,
             "w",
@@ -506,9 +509,12 @@ class HDF5Storage(SpaceStorage[
         *, 
         capacity = None, 
         read_only = True,
+        multiprocessing : bool = False,
         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 os.path.exists(path), \
             f"Path {path} does not exist"
         

unienv_data/storages/pytorch.py

@@ -22,8 +22,8 @@ class PytorchTensorStorage(SpaceStorage[
         capacity : Optional[int],
         is_memmap : bool = False,
         cache_path : Optional[str] = None,
-        memmap_existok : bool = True,
         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}"

unienv_data/storages/transformation.py

@@ -31,6 +31,8 @@ class TransformedStorage(SpaceStorage[
         data_transformation : DataTransformation,
         capacity : Optional[int] = None,
         cache_path : Optional[str] = None,
+        multiprocessing : bool = False,
+        inner_storage_kwargs : Dict[str, Any] = {},
         **kwargs
     ) -> "TransformedStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
         assert data_transformation.has_inverse, "To transform storages (potentially to save space), you need to use inversible data transformations"
@@ -40,12 +42,15 @@ class TransformedStorage(SpaceStorage[
         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(
             transformed_space,
             *args,
             cache_path=None if cache_path is None else os.path.join(cache_path, inner_storage_path),
             capacity=capacity,
-            **kwargs
+            multiprocessing=multiprocessing,
+            **_inner_storage_kwargs
         )
         return TransformedStorage(
             single_instance_space,
@@ -62,6 +67,7 @@ class TransformedStorage(SpaceStorage[
         *,
         capacity : Optional[int] = None,
         read_only : bool = True,
+        multiprocessing : bool = False,
         **kwargs
     ) -> "TransformedStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
         metadata_path = os.path.join(path, "transformed_metadata.json")
@@ -85,6 +91,7 @@ class TransformedStorage(SpaceStorage[
             transformed_space,
             capacity=capacity,
             read_only=read_only,
+            multiprocessing=multiprocessing,
             **kwargs
         )
         return TransformedStorage(
@@ -140,6 +147,14 @@ class TransformedStorage(SpaceStorage[
     def __len__(self):
         return len(self.inner_storage)
     
+    @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
+
     def get_flattened(self, index):
         dat = self.get(index)
         if isinstance(index, int):

unienv_data/transformations/image_compress.py

@@ -7,13 +7,17 @@ from PIL import Image
 import numpy as np
 import io
 
+CONSERVATIVE_COMPRESSION_RATIOS = {
+    "JPEG": 10, # https://stackoverflow.com/questions/3471663/jpeg-compression-ratio
+}
+
 class ImageCompressTransformation(DataTransformation):
     has_inverse = True
 
     def __init__(
         self,
-        init_quality : int = 75,
-        max_size_bytes : int = 65536,
+        init_quality : int = 70,
+        max_size_bytes : Optional[int] = None,
         mode : Optional[str] = None,
         format : str = "JPEG",
     ) -> None:
@@ -25,9 +29,11 @@ class ImageCompressTransformation(DataTransformation):
             mode: Optional mode for PIL Image (e.g., "RGB", "L"). If None, inferred from input.
             format: Image format to use for compression (default "JPEG"). See https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html for options.
         """
+        assert max_size_bytes is not None or format in CONSERVATIVE_COMPRESSION_RATIOS, "Either max_size_bytes must be specified or format must have a conservative compression ratio defined."
 
         self.init_quality = init_quality
         self.max_size_bytes = max_size_bytes
+        self.compression_ratio = CONSERVATIVE_COMPRESSION_RATIOS.get(format, None) if max_size_bytes is None else None
         self.mode = mode
         self.format = format
 
@@ -45,9 +51,15 @@ class ImageCompressTransformation(DataTransformation):
     ) -> BDtypeType:
         return backend.__array_namespace_info__().dtypes()['uint8']
 
+    def _get_max_compressed_size(self, source_space : BoxSpace):
+        H, W, C = source_space.shape[-3], source_space.shape[-2], source_space.shape[-1]
+        return self.max_size_bytes if self.max_size_bytes is not None else (H * W * C // self.compression_ratio) + 1
+
     def get_target_space_from_source(self, source_space):
         self.validate_source_space(source_space)
-        new_shape = source_space.shape[:-3] + (self.max_size_bytes,)
+        
+        max_compressed_size = self._get_max_compressed_size(source_space)
+        new_shape = source_space.shape[:-3] + (max_compressed_size,)
         
         return BoxSpace(
             source_space.backend,
@@ -78,14 +90,14 @@ class ImageCompressTransformation(DataTransformation):
         # Create PIL Image (mode inferred automatically)
         img = Image.fromarray(img_array, mode=mode)
 
-        quality = 95
+        quality = self.init_quality
         while quality >= min_quality:
             buf = io.BytesIO()
             img.save(buf, format=self.format, quality=quality)
             image_bytes = buf.getvalue()
             if len(image_bytes) <= max_bytes:
                 return image_bytes, quality
-            quality -= 5
+            quality -= 10
 
         img.close()
         # Return lowest quality attempt if still too large
@@ -93,19 +105,21 @@ class ImageCompressTransformation(DataTransformation):
     
     def transform(self, source_space, data):
         self.validate_source_space(source_space)
+        
+        max_compressed_size = self._get_max_compressed_size(source_space)
         data_numpy = source_space.backend.to_numpy(data)
         flat_data_numpy = data_numpy.reshape(-1, *data_numpy.shape[-3:])
-        flat_compressed_data = np.zeros((flat_data_numpy.shape[0], self.max_size_bytes), dtype=np.uint8)
+        flat_compressed_data = np.zeros((flat_data_numpy.shape[0], max_compressed_size), dtype=np.uint8)
         for i in range(flat_data_numpy.shape[0]):
             img_array = flat_data_numpy[i]
             image_bytes, _ = self.encode_to_size(
                 img_array,
-                self.max_size_bytes,
+                max_compressed_size,
                 mode=self.mode
             )
             byte_array = np.frombuffer(image_bytes, dtype=np.uint8)
             flat_compressed_data[i, :len(byte_array)] = byte_array
-        compressed_data = flat_compressed_data.reshape(data_numpy.shape[:-3] + (self.max_size_bytes, ))
+        compressed_data = flat_compressed_data.reshape(data_numpy.shape[:-3] + (max_compressed_size, ))
         compressed_data_backend = source_space.backend.from_numpy(compressed_data, dtype=self.get_uint8_dtype(source_space.backend), device=source_space.device)
         return compressed_data_backend
     
@@ -206,7 +220,6 @@ class ImageDecompressTransformation(DataTransformation):
         assert source_space is not None, "Source space must be provided to get inverse transformation."
         self.validate_source_space(source_space)
         return ImageCompressTransformation(
-            init_quality=75,
             max_size_bytes=source_space.shape[-1],
             mode=self.mode,
             format=self.format if self.format is not None else "JPEG",

unienv_interface/func_wrapper/frame_stack.py

@@ -8,7 +8,7 @@ from unienv_interface.utils import seed_util
 from unienv_interface.env_base.funcenv import FuncEnv, ContextType, ObsType, ActType, RenderFrame, StateType, RenderStateType
 from unienv_interface.env_base.funcenv_wrapper import *
 from unienv_interface.space import Space
-from unienv_interface.utils.data_queue import FuncSpaceDataQueue, SpaceDataQueueState
+from unienv_interface.utils.framestack_queue import FuncSpaceDataQueue, SpaceDataQueueState
 from unienv_interface.utils.stateclass import StateClass, field
 
 class FuncFrameStackWrapperState(

unienv_interface/space/space_utils/flatten_utils.py

@@ -192,14 +192,20 @@ def unflatten_data(space : Space, data : BArrayType, start_dim : int = 0) -> Any
 @flatten_data.register(BinarySpace)
 def _flatten_data_common(space: typing.Union[BoxSpace, BinarySpace], data: BArrayType, start_dim : int = 0) -> BArrayType:
     assert -len(space.shape) <= start_dim <= len(space.shape)
-    return space.backend.reshape(data, data.shape[:start_dim] + (-1,))
+    dat = space.backend.reshape(data, data.shape[:start_dim] + (-1,))
+    if isinstance(space, BinarySpace):
+        dat = space.backend.astype(dat, space.backend.default_integer_dtype)
+    return dat
 
 @unflatten_data.register(BoxSpace)
 @unflatten_data.register(BinarySpace)
 def _unflatten_data_common(space: typing.Union[BoxSpace, BinarySpace], data: Any, start_dim : int = 0) -> BArrayType:
     assert -len(space.shape) <= start_dim <= len(space.shape)
     unflat_dat = space.backend.reshape(data, data.shape[:start_dim] + space.shape[start_dim:])
-    unflat_dat = space.backend.astype(unflat_dat, space.dtype)
+    if isinstance(space, BinarySpace):
+        unflat_dat = space.backend.astype(unflat_dat, space.dtype if space.dtype is not None else space.backend.default_boolean_dtype)
+    else:
+        unflat_dat = space.backend.astype(unflat_dat, space.dtype)
     return unflat_dat
 
 @flatten_data.register(DynamicBoxSpace)

unienv_interface/space/spaces/tuple.py

@@ -40,7 +40,7 @@ class TupleSpace(Space[Tuple[Any, ...], BDeviceType, BDtypeType, BRNGType]):
             return self
         
         new_device = device if backend is not None else (device or self.device)
-        return Tuple(
+        return TupleSpace(
             backend=backend or self.backend,
             spaces=[space.to(backend, new_device) for space in self.spaces],
             device=new_device
@@ -93,11 +93,11 @@ class TupleSpace(Space[Tuple[Any, ...], BDeviceType, BDtypeType, BRNGType]):
 
     def __eq__(self, other: Any) -> bool:
         """Check whether ``other`` is equivalent to this instance."""
-        return isinstance(other, Tuple) and self.spaces == other.spaces
+        return isinstance(other, TupleSpace) and self.spaces == other.spaces
     
-    def __copy__(self) -> "Tuple[BDeviceType, BDtypeType, BRNGType]":
+    def __copy__(self) -> "TupleSpace[BDeviceType, BDtypeType, BRNGType]":
         """Create a shallow copy of the Dict space."""
-        return Tuple(
+        return TupleSpace(
             backend=self.backend,
             spaces=copy.copy(self.spaces),
             device=self.device

unienv_interface/transformations/image_resize.py

@@ -0,0 +1,106 @@
+from unienv_interface.space.space_utils import batch_utils as sbu
+from .transformation import DataTransformation, TargetDataT
+from unienv_interface.space import Space, BoxSpace
+from typing import Union, Any, Optional
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+
+class ImageResizeTransformation(DataTransformation):
+    has_inverse = True
+    
+    def __init__(
+        self,
+        new_height: int,
+        new_width: int
+    ):
+        self.new_height = new_height
+        self.new_width = new_width
+
+    def _validate_source_space(self, source_space : Space[Any, BDeviceType, BDtypeType, BRNGType]) -> BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType]:
+        assert isinstance(source_space, BoxSpace), \
+            f"ImageResizeTransformation only supports BoxSpace, got {type(source_space)}"
+        assert len(source_space.shape) >= 3, \
+            f"ImageResizeTransformation only supports spaces with at least 3 dimensions (H, W, C), got shape {source_space.shape}"
+        assert source_space.shape[-3] > 0 and source_space.shape[-2] > 0, \
+            f"ImageResizeTransformation requires positive height and width, got shape {source_space.shape}"
+        return source_space
+
+    def get_target_space_from_source(self, source_space):
+        source_space = self._validate_source_space(source_space)
+
+        backend = source_space.backend
+        new_shape = (
+            *source_space.shape[:-3],
+            self.new_height,
+            self.new_width,
+            source_space.shape[-1]
+        )
+        new_low = backend.min(source_space.low, axis=(-3, -2), keepdims=True)
+        new_high = backend.max(source_space.high, axis=(-3, -2), keepdims=True)
+
+        return BoxSpace(
+            source_space.backend,
+            new_low,
+            new_high,
+            dtype=source_space.dtype,
+            device=source_space.device,
+            shape=new_shape
+        )
+
+    def transform(self, source_space, data):
+        source_space = self._validate_source_space(source_space)
+        backend = source_space.backend
+        if backend.simplified_name == "jax":
+            target_shape = (
+                *data.shape[:-3],
+                self.new_height,
+                self.new_width,
+                source_space.shape[-1]
+            )
+            import jax.image
+            resized_data = jax.image.resize(
+                data,
+                shape=target_shape,
+                method='bilinear',
+                antialias=True
+            )
+        elif backend.simplified_name == "pytorch":
+            import torch.nn.functional as F
+            # PyTorch expects (B, C, H, W)
+            data_permuted = backend.permute_dims(data, (*range(len(data.shape[:-3])), -1, -3, -2))
+            resized_data_permuted = F.interpolate(
+                data_permuted,
+                size=(self.new_height, self.new_width),
+                mode='bilinear',
+                align_corners=False,
+                antialias=True
+            )
+            # Permute back to original shape
+            resized_data = backend.permute_dims(resized_data_permuted, (*range(len(resized_data_permuted.shape[:-3])), -2, -1, -3))
+        elif backend.simplified_name == "numpy":
+            import cv2
+            flat_data = backend.reshape(data, (-1, *source_space.shape[-3:]))
+            resized_flat_data = []
+            for i in range(flat_data.shape[0]):
+                img = flat_data[i]
+                resized_img = cv2.resize(
+                    img,
+                    (self.new_width, self.new_height),
+                    interpolation=cv2.INTER_LINEAR
+                )
+                resized_flat_data.append(resized_img)
+            resized_flat_data = backend.stack(resized_flat_data, axis=0)
+            resized_data = backend.reshape(
+                resized_flat_data,
+                (*data.shape[:-3], self.new_height, self.new_width, source_space.shape[-1])
+            )
+        else:
+            raise ValueError(f"Unsupported backend: {backend.simplified_name}")
+        return resized_data
+
+    def direction_inverse(self, source_space = None):
+        assert source_space is not None, "Inverse transformation requires source_space"
+        source_space = self._validate_source_space(source_space)
+        return ImageResizeTransformation(
+            new_height=source_space.shape[-3],
+            new_width=source_space.shape[-2]
+        )

unienv_interface/transformations/iter_transform.py

@@ -0,0 +1,92 @@
+from typing import Union, Any, Optional, Mapping, List, Callable, Dict
+
+from unienv_interface.space.space_utils import batch_utils as sbu
+from unienv_interface.space import Space, DictSpace, TupleSpace
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+
+import copy
+from .transformation import DataTransformation, TargetDataT
+
+def default_is_leaf_fn(space : Space[Any, BDeviceType, BDtypeType, BRNGType]):
+    return not isinstance(space, (DictSpace, TupleSpace))
+
+class IterativeTransformation(DataTransformation):
+    def __init__(
+        self,
+        transformation: DataTransformation,
+        is_leaf_node_fn: Callable[[Space[Any, BDeviceType, BDtypeType, BRNGType]], bool] = default_is_leaf_fn,
+        inv_is_leaf_node_fn: Callable[[Space[Any, BDeviceType, BDtypeType, BRNGType]], bool] = default_is_leaf_fn
+    ):
+        self.transformation = transformation
+        self.is_leaf_node_fn = is_leaf_node_fn
+        self.inv_is_leaf_node_fn = inv_is_leaf_node_fn
+        self.has_inverse = transformation.has_inverse
+
+    def get_target_space_from_source(
+        self, 
+        source_space : Space[Any, BDeviceType, BDtypeType, BRNGType]
+    ):
+        if self.is_leaf_node_fn(source_space):
+            return self.transformation.get_target_space_from_source(source_space)
+        elif isinstance(source_space, DictSpace):
+            rsts = {
+                key: self.get_target_space_from_source(subspace)
+                for key, subspace in source_space.spaces.items()
+            }
+            backend = source_space.backend if len(rsts) == 0 else next(iter(rsts.values())).backend
+            device = source_space.device if len(rsts) == 0 else next(iter(rsts.values())).device
+            return DictSpace(
+                backend,
+                rsts,
+                device=device
+            )
+        elif isinstance(source_space, TupleSpace):
+            rsts = tuple(
+                self.get_target_space_from_source(subspace)
+                for subspace in source_space.spaces
+            )
+            backend = source_space.backend if len(rsts) == 0 else next(iter(rsts)).backend
+            device = source_space.device if len(rsts) == 0 else next(iter(rsts)).device
+            return TupleSpace(
+                backend,
+                rsts,
+                device=device
+            )
+        else:
+            raise ValueError(f"Unsupported space type: {type(source_space)}")
+
+    def transform(
+        self, 
+        source_space: Space,
+        data: Union[Mapping[str, Any], BArrayType]
+    ) -> Union[Mapping[str, Any], BArrayType]:
+        if self.is_leaf_node_fn(source_space):
+            return self.transformation.transform(source_space, data)
+        elif isinstance(source_space, DictSpace):
+            return {
+                key: self.transform(subspace, data[key])
+                for key, subspace in source_space.spaces.items()
+            }
+        elif isinstance(source_space, TupleSpace):
+            return tuple(
+                self.transform(subspace, data[i])
+                for i, subspace in enumerate(source_space.spaces)
+            )
+        else:
+            raise ValueError(f"Unsupported space type: {type(source_space)}")
+
+    def direction_inverse(
+        self,
+        source_space = None,
+    ) -> Optional["IterativeTransformation"]:
+        if not self.has_inverse:
+            return None
+
+        return IterativeTransformation(
+            self.transformation.direction_inverse(),
+            is_leaf_node_fn=self.inv_is_leaf_node_fn,
+            inv_is_leaf_node_fn=self.is_leaf_node_fn
+        )
+
+    def close(self):
+        self.transformation.close()

unienv_interface/utils/symbol_util.py

@@ -5,9 +5,15 @@ __all__ = [
     "get_class_from_full_name",
 ]
 
+REMAP = {
+    "unienv_data.storages.common.FlattenedStorage": "unienv_data.storages.flattened.FlattenedStorage",
+}
+
 def get_full_class_name(cls : Type) -> str:
     return f"{cls.__module__}.{cls.__qualname__}"
 
 def get_class_from_full_name(full_name : str) -> Type:
+    if full_name in REMAP:
+        full_name = REMAP[full_name]
     module_name, class_name = full_name.rsplit(".", 1)
-    return getattr(__import__(module_name, fromlist=[class_name]), class_name)
+    return getattr(__import__(module_name, fromlist=[class_name]), class_name)

unienv_interface/wrapper/frame_stack.py

@@ -7,7 +7,7 @@ from unienv_interface.space.space_utils import batch_utils as sbu
 from unienv_interface.env_base.env import Env, ContextType, ObsType, ActType, RenderFrame, BArrayType, BDeviceType, BDtypeType, BRNGType
 from unienv_interface.env_base.wrapper import ContextObservationWrapper, ActionWrapper, WrapperContextT, WrapperObsT, WrapperActT
 from unienv_interface.space import Space, DictSpace
-from unienv_interface.utils.data_queue import SpaceDataQueue
+from unienv_interface.utils.framestack_queue import SpaceDataQueue
 
 class FrameStackWrapper(
     ContextObservationWrapper[