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

unienv_data/transformations/image_compress.py

@@ -6,9 +6,31 @@ from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, B
 from PIL import Image
 import numpy as np
 import io
+import math
+
+# https://stackoverflow.com/questions/3471663/jpeg-compression-ratio
+JPEG_QUALITY_COMPRESSION_MAP = {
+    "quality": np.array([55, 60, 65, 70, 75, 80, 85, 90, 95, 100], dtype=int),
+    "compression_ratio": np.array([43.27, 36.90, 34.24, 31.50, 26.00, 25.06, 19.08, 14.30, 9.88, 5.27], dtype=float),
+    "conservative_ratio": 0.6,
+}
+def get_jpeg_compression_ratio(init_quality : int) -> int:
+    if init_quality <= JPEG_QUALITY_COMPRESSION_MAP['quality'][0]:
+        return math.floor(JPEG_QUALITY_COMPRESSION_MAP['compression_ratio'][0] * JPEG_QUALITY_COMPRESSION_MAP['conservative_ratio'])
+    if init_quality >= JPEG_QUALITY_COMPRESSION_MAP['quality'][-1]:
+        return math.floor(JPEG_QUALITY_COMPRESSION_MAP['compression_ratio'][-1] * JPEG_QUALITY_COMPRESSION_MAP['conservative_ratio'])
+
+    for i in range(1, len(JPEG_QUALITY_COMPRESSION_MAP['quality'])):
+        if init_quality <= JPEG_QUALITY_COMPRESSION_MAP['quality'][i]:
+            q_low = JPEG_QUALITY_COMPRESSION_MAP['quality'][i - 1]
+            q_high = JPEG_QUALITY_COMPRESSION_MAP['quality'][i]
+            r_low = JPEG_QUALITY_COMPRESSION_MAP['compression_ratio'][i - 1]
+            r_high = JPEG_QUALITY_COMPRESSION_MAP['compression_ratio'][i]
+            ratio = r_low + (r_high - r_low) * (init_quality - q_low) / (q_high - q_low)
+            return math.floor(ratio * JPEG_QUALITY_COMPRESSION_MAP['conservative_ratio'])
 
 CONSERVATIVE_COMPRESSION_RATIOS = {
-    "JPEG": 10, # https://stackoverflow.com/questions/3471663/jpeg-compression-ratio
+    "JPEG": get_jpeg_compression_ratio,
 }
 
 class ImageCompressTransformation(DataTransformation):
@@ -18,8 +40,10 @@ class ImageCompressTransformation(DataTransformation):
         self,
         init_quality : int = 70,
         max_size_bytes : Optional[int] = None,
+        compression_ratio : Optional[float] = None,
         mode : Optional[str] = None,
         format : str = "JPEG",
+        last_channel : bool = True,
     ) -> None:
         """
         Initialize JPEG compression transformation.
@@ -29,21 +53,25 @@ 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."
+        assert not (max_size_bytes is not None and compression_ratio is not None), "Specify either max_size_bytes or compression_ratio, not both."
+        assert max_size_bytes is not None or compression_ratio 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.compression_ratio = compression_ratio if compression_ratio is not None else (CONSERVATIVE_COMPRESSION_RATIOS.get(format, None)(init_quality) if max_size_bytes is None else None)
         self.mode = mode
         self.format = format
+        self.last_channel = last_channel
 
-    @staticmethod
-    def validate_source_space(source_space: Space[Any, BDeviceType, BDtypeType, BRNGType]) -> None:
+    def validate_source_space(self, source_space: Space[Any, BDeviceType, BDtypeType, BRNGType]) -> None:
         assert isinstance(source_space, BoxSpace), "JPEGCompressTransformation only supports BoxSpace source spaces."
-        assert len(source_space.shape) >= 3 and (
-            source_space.shape[-1] == 3 or
-            source_space.shape[-1] == 1
-        ), "JPEGCompressTransformation only supports BoxSpace source spaces with shape (..., H, W, 1 or 3)."
+        if not self.last_channel:
+            assert len(source_space.shape) >= 2
+        else:
+            assert len(source_space.shape) >= 3 and (
+                source_space.shape[-1] == 3 or
+                source_space.shape[-1] == 1
+            ), "JPEGCompressTransformation only supports BoxSpace source spaces with shape (..., H, W, 1 or 3)."
 
     @staticmethod
     def get_uint8_dtype(
@@ -52,14 +80,22 @@ class ImageCompressTransformation(DataTransformation):
         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]
+        if self.last_channel:
+            H, W, C = source_space.shape[-3], source_space.shape[-2], source_space.shape[-1]
+        else:
+            H, W = source_space.shape[-2], source_space.shape[-1]
+            C = 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)
         
         max_compressed_size = self._get_max_compressed_size(source_space)
-        new_shape = source_space.shape[:-3] + (max_compressed_size,)
+
+        if not self.last_channel:
+            new_shape = source_space.shape[:-2] + (max_compressed_size,)
+        else:
+            new_shape = source_space.shape[:-3] + (max_compressed_size,)
         
         return BoxSpace(
             source_space.backend,
@@ -108,7 +144,11 @@ class ImageCompressTransformation(DataTransformation):
         
         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:])
+        if not self.last_channel:
+            flat_data_numpy = data_numpy.reshape(-1, *data_numpy.shape[-2:])
+        else:
+            flat_data_numpy = data_numpy.reshape(-1, *data_numpy.shape[-3:])
+        
         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]
@@ -119,16 +159,26 @@ class ImageCompressTransformation(DataTransformation):
             )
             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] + (max_compressed_size, ))
+        
+        if not self.last_channel:
+            compressed_data = flat_compressed_data.reshape(data_numpy.shape[:-2] + (max_compressed_size, ))
+        else:
+            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
     
     def direction_inverse(self, source_space = None):
         assert source_space is not None, "Source space must be provided to get inverse transformation."
         self.validate_source_space(source_space)
-        height = source_space.shape[-3]
-        width = source_space.shape[-2]
-        channels = source_space.shape[-1]
+        
+        if not self.last_channel:
+            height = source_space.shape[-2]
+            width = source_space.shape[-1]
+            channels = None
+        else:
+            height = source_space.shape[-3]
+            width = source_space.shape[-2]
+            channels = source_space.shape[-1]
         return ImageDecompressTransformation(
             target_height=height,
             target_width=width,
@@ -136,6 +186,12 @@ class ImageCompressTransformation(DataTransformation):
             mode=self.mode,
             format=self.format,
         )
+    
+    def __setstate__(self, state):
+        # for backward compatibility
+        self.__dict__.update(state)
+        if 'last_channel' not in state:
+            self.last_channel = True
 
 class ImageDecompressTransformation(DataTransformation):
     has_inverse = True
@@ -144,7 +200,7 @@ class ImageDecompressTransformation(DataTransformation):
         self,
         target_height : int,
         target_width : int,
-        target_channels : int = 3,
+        target_channels : Optional[int] = 3,
         mode : Optional[str] = None,
         format : Optional[str] = None,
     ) -> None:
@@ -174,6 +230,9 @@ class ImageDecompressTransformation(DataTransformation):
     def get_target_space_from_source(self, source_space):
         self.validate_source_space(source_space)
         new_shape = source_space.shape[:-1] + (self.target_height, self.target_width, self.target_channels)
+        if self.target_channels is None:
+            new_shape = new_shape[:-1]
+        
         return BoxSpace(
             source_space.backend,
             shape=new_shape,
@@ -212,7 +271,10 @@ class ImageDecompressTransformation(DataTransformation):
                 byte_array.tobytes(),
                 mode=self.mode
             )
-        decompressed_image = flat_decompressed_image.reshape(data_numpy.shape[:-1] + (self.target_height, self.target_width, self.target_channels))
+        if self.target_channels is None:
+            decompressed_image = flat_decompressed_image.reshape(data_numpy.shape[:-1] + (self.target_height, self.target_width))
+        else:
+            decompressed_image = flat_decompressed_image.reshape(data_numpy.shape[:-1] + (self.target_height, self.target_width, self.target_channels))
         decompressed_image_backend = source_space.backend.from_numpy(decompressed_image, dtype=self.get_uint8_dtype(source_space.backend), device=source_space.device)
         return decompressed_image_backend
     
@@ -223,4 +285,5 @@ class ImageDecompressTransformation(DataTransformation):
             max_size_bytes=source_space.shape[-1],
             mode=self.mode,
             format=self.format if self.format is not None else "JPEG",
+            last_channel=self.target_channels is not None,
         )

unienv_interface/space/space_utils/batch_utils.py

@@ -334,7 +334,7 @@ def swap_batch_dims_in_data(
 ) -> Any:
     if backend.is_backendarray(data):
         return _tensor_transpose(backend, data, dim1, dim2)
-    elif isinstance(data, np.ndarray) and data.dtype != object:
+    elif isinstance(data, np.ndarray):
         return _tensor_transpose(NumpyComputeBackend, data, dim1, dim2)
     elif isinstance(data, GraphInstance):
         return GraphInstance(
@@ -734,6 +734,8 @@ def _get_at_tuple(space: TupleSpace, items: typing.Tuple[Any, ...], index : Arra
 
 @get_at.register(BatchedSpace)
 def _get_at_batched(space: BatchedSpace, items: np.ndarray, index: ArrayAPIGetIndex) -> typing.Union[np.ndarray, Any]:
+    if space.backend.is_backendarray(index):
+        index = space.backend.to_numpy(index)
     return items[index]
 
 @singledispatch
@@ -830,6 +832,8 @@ def _set_at_batched(
     value: typing.Union[np.ndarray, Any],
 ) -> np.ndarray:
     new_data = items.copy()
+    if space.backend.is_backendarray(index):
+        index = space.backend.to_numpy(index)
     new_data[index] = value
     return new_data
 

unienv_interface/space/spaces/dict.py

@@ -34,6 +34,12 @@ class DictSpace(Space[Dict[str, Any], BDeviceType, BDtypeType, BRNGType]):
         new_spaces: Dict[str, Space[Any, BDeviceType, BDtypeType, BRNGType]] = {}
 
         for key, space in spaces.items():
+            if not isinstance(space, Space) and isinstance(space, Mapping):
+                space = DictSpace(
+                    backend,
+                    space,
+                    device
+                )
             assert isinstance(
                 space, Space
             ), f"Dict space element is not an instance of Space: key='{key}', space={space}"

unienv_interface/transformations/__init__.py

@@ -1,6 +1,8 @@
 from .transformation import DataTransformation
+from .identity import IdentityTransformation
 from .rescale import RescaleTransformation
 from .filter_dict import DictIncludeKeyTransformation, DictExcludeKeyTransformation
 from .batch_and_unbatch import BatchifyTransformation, UnBatchifyTransformation
 from .dict_transform import DictTransformation
-from .chained_transform import ChainedTransformation
+from .chained_transform import ChainedTransformation
+from .crop import CropTransformation

unienv_interface/transformations/batch_and_unbatch.py

@@ -6,26 +6,65 @@ from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, B
 
 class BatchifyTransformation(DataTransformation):
     has_inverse = True
+
+    def __init__(
+        self,
+        axis : int = 0
+    ) -> None:
+        self.axis = axis
     
     def get_target_space_from_source(self, source_space):
-        return sbu.batch_space(source_space, 1)
+        ret = sbu.batch_space(source_space, 1)
+        if self.axis != 0:
+            ret = sbu.swap_batch_dims(
+                ret,
+                0,
+                self.axis
+            )
+        return ret
 
     def transform(self, source_space, data):
         return sbu.concatenate(
             source_space,
-            [data]
+            [data],
+            axis=self.axis
         )
     
     def direction_inverse(self, source_space = None):
-        return UnBatchifyTransformation()
+        return UnBatchifyTransformation(axis=self.axis)
 
 class UnBatchifyTransformation(DataTransformation):
     has_inverse = True
+
+    def __init__(
+        self,
+        axis : int = 0
+    ) -> None:
+        self.axis = axis
     
     def get_target_space_from_source(self, source_space):
+        if self.axis != 0:
+            source_space = sbu.swap_batch_dims(
+                source_space,
+                0,
+                self.axis
+            )
+        assert sbu.batch_size(source_space) == 1, "Cannot unbatch space with batch size > 1"
         return next(iter(sbu.unbatch_spaces(source_space)))
 
     def transform(self, source_space, data):
+        if self.axis != 0:
+            source_space = sbu.swap_batch_dims(
+                source_space,
+                0,
+                self.axis
+            )
+            data = sbu.swap_batch_dims_in_data(
+                source_space.backend,
+                data,
+                0,
+                self.axis
+            )
         return sbu.get_at(
             source_space,
             data,
@@ -33,4 +72,4 @@ class UnBatchifyTransformation(DataTransformation):
         )
     
     def direction_inverse(self, source_space = None):
-        return BatchifyTransformation()
+        return BatchifyTransformation(axis=self.axis)

unienv_interface/transformations/chained_transform.py

@@ -45,17 +45,18 @@ class ChainedTransformation(DataTransformation):
     ) -> Optional["ChainedTransformation"]:
         if not self.has_inverse:
             return None
+
+        source_spaces = [source_space]
+        for transformation in self.transformations:
+            next_space = transformation.get_target_space_from_source(source_spaces[-1])
+            source_spaces.append(next_space)
         
-        inverse_mapping = {
-            key: transformation.direction_inverse(source_space)
-            for key, transformation in self.mapping.items()
-        }
+        inverted_transformations = []
+        for i in reversed(range(len(self.transformations))):
+            inverted_transformations.append(self.transformations[i].direction_inverse(source_spaces[i]))
 
         return ChainedTransformation(
-            transformations=[
-                transformation.direction_inverse(source_space)
-                for transformation in reversed(self.transformations)
-            ]
+            inverted_transformations
         )
 
     def close(self):

unienv_interface/transformations/crop.py

@@ -0,0 +1,69 @@
+from unienv_interface.space.space_utils import batch_utils as sbu
+from unienv_interface.transformations import DataTransformation
+from unienv_interface.space import Space, BoxSpace, TextSpace
+from typing import Union, Any, Optional, Tuple, List
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from .identity import IdentityTransformation
+
+class CropTransformation(DataTransformation):
+    has_inverse = True
+
+    def __init__(
+        self,
+        crop_low : Union[int, float, BArrayType],
+        crop_high : Union[int, float, BArrayType],
+    ) -> None:
+        """
+        Initialize Crop Transformation.
+        Args:
+            crop_low: Lower bound for cropping the data.
+            crop_high: Upper bound for cropping the data.
+        """
+        self.crop_low = crop_low
+        self.crop_high = crop_high
+
+    def validate_source_space(self, source_space : Space[Any, BDeviceType, BDtypeType, BRNGType]) -> None:
+        assert isinstance(source_space, BoxSpace), "CropTransformation only supports Box spaces"
+
+    def get_crop_range(self, source_space : BoxSpace[Any, BDeviceType, BDtypeType, BRNGType]) -> Tuple[BArrayType, BArrayType]:
+        new_low = self.crop_low
+        if source_space.backend.is_backendarray(new_low):
+            if len(new_low.shape) < len(source_space.shape):
+                new_low = source_space.backend.reshape(
+                    new_low,
+                    (-1,)*(len(source_space.shape) - len(new_low.shape)) + new_low.shape
+                )
+            new_low = source_space.backend.astype(new_low, source_space.dtype)
+            if source_space.device is not None:
+                new_low = source_space.backend.to_device(new_low, source_space.device)
+        new_high = self.crop_high
+        if source_space.backend.is_backendarray(new_high):
+            if len(new_high.shape) < len(source_space.shape):
+                new_high = source_space.backend.reshape(
+                    new_high,
+                    (-1,)*(len(source_space.shape) - len(new_high.shape)) + new_high.shape
+                )
+            new_high = source_space.backend.astype(new_high, source_space.dtype)
+            if source_space.device is not None:
+                new_high = source_space.backend.to_device(new_high, source_space.device)
+        return new_low, new_high
+
+    def get_target_space_from_source(self, source_space):
+        self.validate_source_space(source_space)
+        crop_low, crop_high = self.get_crop_range(source_space)
+        return BoxSpace(
+            backend=source_space.backend,
+            low=crop_low,
+            high=crop_high,
+            shape=source_space.shape,
+            dtype=source_space.dtype,
+            device=source_space.device,
+        )
+
+    def transform(self, source_space, data):
+        self.validate_source_space(source_space)
+        crop_low, crop_high = self.get_crop_range(source_space)
+        return source_space.backend.clip(data, crop_low, crop_high)
+    
+    def direction_inverse(self, source_space = None):
+        return IdentityTransformation()

unienv_interface/transformations/dict_transform.py

@@ -135,12 +135,18 @@ class DictTransformation(DataTransformation):
 
         inverse_mapping = {}
         for key, transformation in self.mapping.items():
-            inverse_mapping[key] = transformation.direction_inverse(
-                None if source_space is None else get_chained_value(
+            if source_space is not None:
+                current_source = get_chained_value(
                     source_space,
                     key.split('/'),
                     ignore_missing_keys=self.ignore_missing_keys
                 )
+                if current_source is None:
+                    continue
+            else:
+                current_source = None
+            inverse_mapping[key] = transformation.direction_inverse(
+                current_source
             )
         
         return DictTransformation(

unienv_interface/transformations/identity.py

@@ -0,0 +1,16 @@
+from .transformation import DataTransformation, TargetDataT
+from unienv_interface.space import BoxSpace
+from typing import Union, Any, Optional
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+
+class IdentityTransformation(DataTransformation):
+    has_inverse = True
+
+    def get_target_space_from_source(self, source_space):
+        return source_space
+
+    def transform(self, source_space, data):
+        return data
+    
+    def direction_inverse(self, source_space = None):
+        return self

unienv_interface/transformations/rescale.py

@@ -11,8 +11,8 @@ def _get_broadcastable_value(
     if isinstance(value, (int, float)):
         return value
     else:
-        assert target_ndim <= len(value), "Value must have at least as many dimensions as target space"
-        target_shape = tuple([1] * (target_ndim - len(value)) + list(value.shape))
+        assert target_ndim >= len(value.shape), "Target space must have at least as many dimensions as the value"
+        target_shape = tuple([1] * (target_ndim - len(value.shape)) + list(value.shape))
         return backend.reshape(value, target_shape)
 
 class RescaleTransformation(DataTransformation):
@@ -22,6 +22,7 @@ class RescaleTransformation(DataTransformation):
         new_low : Union[BArrayType,float] = -1.0,
         new_high : Union[BArrayType,float] = 1.0,
         new_dtype : Optional[BDtypeType] = None,
+        nan_to : Optional[Union[float, int, BArrayType]] = None,
     ):
         # assert isinstance(source_space, BoxSpace), "RescaleTransformation only supports Box action spaces"
         # assert source_space.backend.dtype_is_real_floating(source_space.dtype), "RescaleTransformation only supports real-valued floating spaces"
@@ -31,6 +32,7 @@ class RescaleTransformation(DataTransformation):
         self.new_high = new_high
         self._new_span = new_high - new_low
         self.new_dtype = new_dtype
+        self.nan_to = nan_to
 
     def get_target_space_from_source(self, source_space):
         assert isinstance(source_space, BoxSpace), "RescaleTransformation only supports Box spaces"
@@ -67,14 +69,26 @@ class RescaleTransformation(DataTransformation):
             source_space.backend,
             self.new_low,
             target_ndim
-        ), source_space.backend.device(data))
+        ), source_space.backend.device(data)) if source_space.backend.is_backendarray(self.new_low) else self.new_low
         target_high = source_space.backend.to_device(_get_broadcastable_value(
             source_space.backend,
             self.new_high,
             target_ndim
-        ), source_space.backend.device(data))
+        ), source_space.backend.device(data)) if source_space.backend.is_backendarray(self.new_high) else self.new_high
         scaling_factor = (target_high - target_low) / (source_space._high - source_space._low)
         target_data = (data - source_space._low) * scaling_factor + target_low
+        
+        if self.nan_to is not None:
+            target_data = source_space.backend.where(
+                source_space.backend.isnan(target_data),
+                self.nan_to,
+                target_data
+            )
+
+        if self.new_dtype is not None:
+            if source_space.backend.dtype_is_real_integer(self.new_dtype) and source_space.backend.dtype_is_real_floating(target_data.dtype):
+                target_data = source_space.backend.round(target_data)
+            target_data = source_space.backend.astype(target_data, self.new_dtype)
         return target_data
     
     def direction_inverse(self, source_space = None):
@@ -95,4 +109,9 @@ class RescaleTransformation(DataTransformation):
             new_low=new_low,
             new_high=new_high,
             new_dtype=source_space.dtype
-        )
+        )
+    
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        if not hasattr(self, "nan_to"):
+            self.nan_to = None

unienv_interface/wrapper/backend_compat.py

@@ -31,7 +31,7 @@ def data_to(
             key: data_to(value, source_backend, target_backend, target_device)
             for key, value in data.items()
         }
-    elif isinstance(data, Sequence):
+    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
         data = [
             data_to(value, source_backend, target_backend, target_device)
             for value in data