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

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
 
@@ -33,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"
@@ -42,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,
@@ -64,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")
@@ -87,6 +91,7 @@ class TransformedStorage(SpaceStorage[
             transformed_space,
             capacity=capacity,
             read_only=read_only,
+            multiprocessing=multiprocessing,
             **kwargs
         )
         return TransformedStorage(
@@ -142,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/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_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]:

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[