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

unienv_data/storages/transformation.py

@@ -0,0 +1,191 @@
+from importlib import metadata
+from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequence, Union, List, Iterable, Type
+
+from unienv_interface.space import Space, BoxSpace
+from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils as sfu
+from unienv_interface.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
+
+from unienv_data.base import SpaceStorage, BatchT
+
+import numpy as np
+import os
+import json
+import pickle
+
+class TransformedStorage(SpaceStorage[
+    BatchT,
+    BArrayType,
+    BDeviceType,
+    BDtypeType,
+    BRNGType,
+]):
+    # ========== Class Attributes ==========
+    @classmethod
+    def create(
+        cls,
+        single_instance_space: Space[Any, BDeviceType, BDtypeType, BRNGType],
+        inner_storage_cls : Type[SpaceStorage[BArrayType, BArrayType, BDeviceType, BDtypeType, BRNGType]],
+        *args,
+        data_transformation : DataTransformation,
+        capacity : Optional[int] = None,
+        cache_path : Optional[str] = None,
+        **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"
+        transformed_space = data_transformation.get_target_space_from_source(single_instance_space)
+        inner_storage_path = "transformed_inner_storage" + (inner_storage_cls.single_file_ext or "")
+
+        if cache_path is not None:
+            os.makedirs(cache_path, exist_ok=True)
+
+        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
+        )
+        return TransformedStorage(
+            single_instance_space,
+            data_transformation,
+            inner_storage,
+            inner_storage_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,
+        **kwargs
+    ) -> "TransformedStorage[BArrayType, BDeviceType, BDtypeType, BRNGType]":
+        metadata_path = os.path.join(path, "transformed_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']}"
+        
+        data_transform_path = os.path.join(path, "data_transformation.pkl")
+        with open(data_transform_path, "rb") as f:
+            data_transform = pickle.load(f)
+        
+        assert isinstance(data_transform, DataTransformation)
+        transformed_space = data_transform.get_target_space_from_source(single_instance_space)
+
+        inner_storage_cls : Type[SpaceStorage] = get_class_from_full_name(metadata["inner_storage_type"])
+        inner_storage_path = metadata["inner_storage_path"]
+        inner_storage = inner_storage_cls.load_from(
+            os.path.join(path, inner_storage_path),
+            transformed_space,
+            capacity=capacity,
+            read_only=read_only,
+            **kwargs
+        )
+        return TransformedStorage(
+            single_instance_space,
+            data_transform,
+            inner_storage,
+            inner_storage_path,
+        )
+
+    # ========== Instance Implementations ==========
+    single_file_ext = None
+
+    def __init__(
+        self,
+        single_instance_space: Space[Any, BDeviceType, BDtypeType, BRNGType],
+        data_transformation : DataTransformation,
+        inner_storage : SpaceStorage[
+            BArrayType,
+            BArrayType,
+            BDeviceType,
+            BDtypeType,
+            BRNGType,
+        ],
+        inner_storage_path : Union[str, os.PathLike],
+    ):
+        super().__init__(single_instance_space)
+        transformed_space = data_transformation.get_target_space_from_source(single_instance_space)
+
+        assert inner_storage.backend == transformed_space.backend, \
+            f"Inner storage backend {inner_storage.backend} does not match single instance space backend {single_instance_space.backend}"
+        assert inner_storage.device == transformed_space.device, \
+            f"Inner storage device {inner_storage.device} does not match single instance space device {single_instance_space.device}"
+        assert inner_storage.single_instance_space == transformed_space
+
+        self._transformed_space = transformed_space
+        self._batched_transformed_space = sbu.batch_space(transformed_space, 1)
+        self._batched_instance_space = sbu.batch_space(single_instance_space, 1)
+        self.inner_storage = inner_storage
+        self.inner_storage_path = inner_storage_path
+        self.data_transformation = data_transformation
+        self.inv_data_transformation = data_transformation.direction_inverse(single_instance_space)
+        
+    @property
+    def capacity(self) -> Optional[int]:
+        return self.inner_storage.capacity
+    
+    def extend_length(self, length):
+        self.inner_storage.extend_length(length)
+
+    def shrink_length(self, length):
+        self.inner_storage.shrink_length(length)
+    
+    def __len__(self):
+        return len(self.inner_storage)
+    
+    def get_flattened(self, index):
+        dat = self.get(index)
+        if isinstance(index, int):
+            return sfu.flatten_data(self.single_instance_space, dat)
+        else:
+            return sfu.flatten_data(self.single_instance_space, dat, start_dim=1)
+
+    def get(self, index):
+        result = self.inner_storage.get(index)
+        result = self.inv_data_transformation.transform(
+            self._transformed_space if isinstance(index, int) else self._batched_transformed_space,
+            result
+        )
+        return result
+    
+    def set_flattened(self, index, value):
+        if isinstance(index, int):
+            set_value = sfu.unflatten_data(self.single_instance_space, value)
+        else:
+            set_value = sfu.unflatten_data(self._batched_instance_space, value)
+        self.set(index, set_value)
+
+    def set(self, index, value):
+        transformed_value = self.data_transformation.transform(
+            self.single_instance_space if isinstance(index, int) else self._batched_instance_space,
+            value
+        )
+        self.inner_storage.set(index, transformed_value)
+
+    def clear(self):
+        self.inner_storage.clear()
+
+    def dumps(self, path):
+        metadata = {
+            "storage_type": __class__.__name__,
+            "inner_storage_type": get_full_class_name(type(self.inner_storage)),
+            "inner_storage_path": self.inner_storage_path,
+            "transformation": get_full_class_name(type(self.data_transformation))
+        }
+        self.inner_storage.dumps(os.path.join(path, self.inner_storage_path))
+        with open(os.path.join(path, "transformed_metadata.json"), "w") as f:
+            json.dump(metadata, f)
+        with open(os.path.join(path, "data_transformation.pkl"), "wb") as f:
+            pickle.dump(self.data_transformation, f)
+
+    def close(self):
+        self.inner_storage.close()

unienv_data/transformations/image_compress.py

@@ -0,0 +1,213 @@
+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
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from PIL import Image
+import numpy as np
+import io
+
+class ImageCompressTransformation(DataTransformation):
+    has_inverse = True
+
+    def __init__(
+        self,
+        init_quality : int = 75,
+        max_size_bytes : int = 65536,
+        mode : Optional[str] = None,
+        format : str = "JPEG",
+    ) -> None:
+        """
+        Initialize JPEG compression transformation.
+        Args:
+            init_quality: Initial JPEG quality setting (1-100).
+            max_size_bytes: Maximum allowed size of compressed JPEG in bytes.
+            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.
+        """
+
+        self.init_quality = init_quality
+        self.max_size_bytes = max_size_bytes
+        self.mode = mode
+        self.format = format
+
+    @staticmethod
+    def validate_source_space(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)."
+
+    @staticmethod
+    def get_uint8_dtype(
+        backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
+    ) -> BDtypeType:
+        return backend.__array_namespace_info__().dtypes()['uint8']
+
+    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,)
+        
+        return BoxSpace(
+            source_space.backend,
+            shape=new_shape,
+            low=-source_space.backend.inf,
+            high=source_space.backend.inf,
+            dtype=self.get_uint8_dtype(source_space.backend),
+            device=source_space.device,
+        )
+    
+    def encode_to_size(self, img_array, max_bytes, min_quality=20, mode=None):
+        """
+        Encode an image (H, W, 3) or (H, W, 1) as JPEG bytes,
+        reducing quality until <= max_bytes.
+
+        Args:
+            img_array: np.ndarray, uint8, shape (H, W, 3) RGB or (H, W, 1) grayscale
+            max_bytes: maximum allowed size of JPEG file
+            min_quality: minimum JPEG quality before giving up
+
+        Returns:
+            jpeg_bytes (bytes), final_quality (int)
+        """
+        # Handle grayscale (H, W, 1) → (H, W)
+        if img_array.ndim == 3 and img_array.shape[-1] == 1:
+            img_array = np.squeeze(img_array, axis=-1)
+
+        # Create PIL Image (mode inferred automatically)
+        img = Image.fromarray(img_array, mode=mode)
+
+        quality = 95
+        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
+
+        img.close()
+        # Return lowest quality attempt if still too large
+        return image_bytes, quality
+    
+    def transform(self, source_space, data):
+        self.validate_source_space(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)
+        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,
+                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_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]
+        return ImageDecompressTransformation(
+            target_height=height,
+            target_width=width,
+            target_channels=channels,
+            mode=self.mode,
+            format=self.format,
+        )
+
+class ImageDecompressTransformation(DataTransformation):
+    has_inverse = True
+    
+    def __init__(
+        self,
+        target_height : int,
+        target_width : int,
+        target_channels : int = 3,
+        mode : Optional[str] = None,
+        format : Optional[str] = None,
+    ) -> None:
+        """
+        Initialize JPEG decompression transformation.
+        Args:
+            target_height: Height of the decompressed image.
+            target_width: Width of the decompressed image.
+            mode: Optional mode for PIL Image (e.g., "RGB", "L"). If None, inferred from input.
+            format: Image format to use for decompression (default None, which will try everything). See https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html for options.
+        """
+        self.target_height = target_height
+        self.target_width = target_width
+        self.target_channels = target_channels
+        self.mode = mode
+        self.format = format
+
+    @staticmethod
+    def validate_source_space(source_space: Space[Any, BDeviceType, BDtypeType, BRNGType]) -> None:
+        assert isinstance(source_space, BoxSpace), "JPEGDecompressTransformation only supports BoxSpace source spaces."
+        assert len(source_space.shape) >= 1, "JPEGDecompressTransformation requires source space with at least 1 dimension."
+
+    @staticmethod
+    def get_uint8_dtype(backend):
+        return ImageCompressTransformation.get_uint8_dtype(backend)
+
+    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)
+        return BoxSpace(
+            source_space.backend,
+            shape=new_shape,
+            low=0,
+            high=255,
+            dtype=self.get_uint8_dtype(source_space.backend),
+            device=source_space.device,
+        )
+    
+    def decode_bytes(self, jpeg_bytes : bytes, mode=None):
+        """
+        Decode JPEG bytes to an image array (H, W, 3).
+
+        Args:
+            jpeg_bytes: bytes of JPEG image
+
+        Returns:
+            img_array: np.ndarray, uint8, shape (H, W, 3)
+        """
+        buf = io.BytesIO(jpeg_bytes)
+        img = Image.open(buf, formats=[self.format] if self.format is not None else None)
+        if mode is not None:
+            img = img.convert(mode)
+        img_array = np.array(img)
+        img.close()
+        return img_array
+    
+    def transform(self, source_space, data):
+        self.validate_source_space(source_space)
+        data_numpy = source_space.backend.to_numpy(data)
+        flat_data_numpy = data_numpy.reshape(-1, data_numpy.shape[-1])
+        flat_decompressed_image = np.zeros((flat_data_numpy.shape[0], self.target_height, self.target_width, self.target_channels), dtype=np.uint8)
+        for i in range(flat_data_numpy.shape[0]):
+            byte_array : np.ndarray = flat_data_numpy[i]
+            flat_decompressed_image[i] = self.decode_bytes(
+                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))
+        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
+    
+    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)
+        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/backends/jax.py

@@ -1,4 +1,7 @@
-from xbarray.jax import JaxComputeBackend as XBJaxBackend
+try:
+    from xbarray.backends.jax import JaxComputeBackend as XBJaxBackend
+except ImportError:
+    from xbarray.jax import JaxComputeBackend as XBJaxBackend
 from xbarray import ComputeBackend
 from typing import Union
 import jax

unienv_interface/backends/numpy.py

@@ -1,4 +1,7 @@
-from xbarray.numpy import NumpyComputeBackend as XBNumpyBackend
+try:
+    from xbarray.backends.numpy import NumpyComputeBackend as XBNumpyBackend
+except ImportError:
+    from xbarray.numpy import NumpyComputeBackend as XBNumpyBackend
 from xbarray import ComputeBackend
 
 import numpy as np

unienv_interface/backends/pytorch.py

@@ -1,4 +1,7 @@
-from xbarray.pytorch import PytorchComputeBackend as XBPytorchBackend
+try:
+    from xbarray.backends.pytorch import PytorchComputeBackend as XBPytorchBackend
+except ImportError:
+    from xbarray.pytorch import PytorchComputeBackend as XBPytorchBackend
 from xbarray import ComputeBackend
 
 from typing import Union

unienv_interface/env_base/__init__.py

@@ -1,4 +1,5 @@
 from .env import Env
+from .vec_env import SyncVecEnv, AsyncVecEnv
 from .wrapper import Wrapper, ActionWrapper, ContextObservationWrapper
 from .funcenv import FuncEnv, FuncEnvBasedEnv
 from .funcenv_wrapper import FuncEnvWrapper

unienv_interface/env_base/env.py

@@ -90,6 +90,11 @@ class Env(abc.ABC, Generic[BArrayType, ContextType, ObsType, ActType, RenderFram
     def sample_observation(self) -> ObsType:
         return self.sample_space(self.observation_space)    
     
+    def sample_context(self) -> Optional[ContextType]:
+        if self.context_space is None:
+            return None
+        return self.sample_space(self.context_space)
+
     def update_observation_post_reset(
         self,
         old_obs: ObsType,

unienv_interface/env_base/funcenv.py

@@ -2,7 +2,7 @@ from typing import Any, Callable, Generic, TypeVar, Tuple, Dict, Optional, Suppo
 import abc
 import numpy as np
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
-from unienv_interface.space import Space
+from unienv_interface.space import Space, batch_utils as sbu
 from dataclasses import dataclass, replace as dataclass_replace
 from .env import Env, ContextType, ObsType, ActType, RenderFrame
 
@@ -121,6 +121,37 @@ class FuncEnv(
         """Close the render state."""
         raise NotImplementedError
 
+    # ========== Convenience methods ==========
+    def update_observation_post_reset(
+        self,
+        old_obs: ObsType,
+        newobs_masked: ObsType,
+        mask: BArrayType
+    ) -> ObsType:
+        assert self.batch_size is not None, "This method is used by batched environment after reset"
+        return sbu.set_at(
+            self.observation_space,
+            old_obs,
+            mask,
+            newobs_masked
+        )
+    
+    def update_context_post_reset(
+        self,
+        old_context: ContextType,
+        new_context: ContextType,
+        mask: BArrayType
+    ) -> ContextType:
+        assert self.batch_size is not None, "This method is used by batched environment after reset"
+        if self.context_space is None:
+            return None
+        return sbu.set_at(
+            self.context_space,
+            old_context,
+            mask,
+            new_context
+        )
+
     # ========== Wrapper methods ==========
     @property
     def unwrapped(self) -> "FuncEnv":

unienv_interface/env_base/funcenv_wrapper.py

@@ -110,7 +110,7 @@ class FuncEnvWrapper(
     def reset(
         self, 
         state : WrapperStateT, 
-        *,
+        *args,
         seed : Optional[int] = None,
         mask : Optional[WrapperBArrayT] = None,
         **kwargs
@@ -120,7 +120,7 @@ class FuncEnvWrapper(
         WrapperObsT,
         Dict[str, Any]
     ]:
-        return self.func_env.reset(state, seed=seed, mask=mask, **kwargs)
+        return self.func_env.reset(state, *args, seed=seed, mask=mask, **kwargs)
 
     def step(
         self,