unienv 0.0.1b2__tar.gz → 0.0.1b4__tar.gz

{unienv-0.0.1b2 → unienv-0.0.1b4}/LICENSE

@@ -1,8 +1,6 @@
 MIT License
 
-Copyright (c) 2016 OpenAI
-Copyright (c) 2022 Farama Foundation
-Copyright (c) 2024 Yunhao Cao
+Copyright (c) 2025 Yunhao Cao and UniEnv Contributors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

unienv-0.0.1b4/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: unienv
+Version: 0.0.1b4
+Summary: Unified robot environment framework supporting multiple tensor and simulation backends
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv
+Project-URL: Documentation, https://github.com/UniEnvOrg/UniEnv
+Project-URL: Repository, https://github.com/UniEnvOrg/UniEnv
+Project-URL: Issues, https://github.com/UniEnvOrg/UniEnv/issues
+Project-URL: Changelog, https://github.com/UniEnvOrg/UniEnv/blob/main/CHANGELOG.md
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: xbarray>=0.0.1a8
+Requires-Dist: pillow
+Requires-Dist: cloudpickle
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Provides-Extra: gymnasium
+Requires-Dist: gymnasium>=0.29.0; extra == "gymnasium"
+Provides-Extra: video
+Requires-Dist: moviepy>=2.1; extra == "video"
+Dynamic: license-file
+
+# UniEnv
+
+Framework unifying robot environments and data APIs. UniEnv provides an universal interface for robot actors, sensors, environments, and data. 
+
+## Tensor library cross-backend Support
+
+UniEnv supports multiple tensor backends with zero-copy translation layers through the DLPack protocol, and allows you to use the same abstract compute backend interface to write custom data transformation layers, environment wrappers and other utilities. This is powered by the [XBArray](https://github.com/UniEnvOrg/XBArray) package.
+
+## Universal Robot Environment Interface
+
+UniEnv supports diverse simulation environments and real robots, built on top of the abstract environment / world interface. This allows you to reuse code across different sim and real robots.
+
+## Universal Robot Data Interface
+
+UniEnv provides a universal data interface for accessing robot data through the abstract `BatchBase` interface. We also provide a utility `ReplayBuffer` for saving data from various environments with diverse data format support, including `hdf5`, memory-mapped torch tensors, and others.
+
+## Installation
+
+Install the package with pip
+
+```bash
+pip install unienv
+```
+
+You can install optional dependencies such as `gymnasium` (for Gymnasium-compatible environments), `dev`, or `video` by running
+
+```bash
+pip install unienv[gymnasium,video]
+```
+
+## Cite
+
+If you use UniEnv in your research, please cite it as follows:
+
+```bibtex
+@software{cao_unienv,
+  author = {Cao, Yunhao AND Fang, Kuan},
+  title = {{UniEnv: Unifying Robot Environments and Data APIs}},
+  year = {2025},
+  month = oct,
+  url = {https://github.com/UniEnvOrg/UniEnv},
+  license = {MIT}
+}
+```
+
+## Acknowledgements
+
+The idea of this project is inspired by [Gymnasium](https://github.com/Farama-Foundation/Gymnasium) and its predecessor [OpenAI Gym](https://github.com/openai/gym). 
+This library is impossible without the great work of DataAPIs Consortium and their work on the [Array API Standard](https://data-apis.org/array-api/latest/). The zero-copy translation layers are powered by the [DLPack](https://github.com/dmlc/dlpack) project.

unienv-0.0.1b4/README.md

@@ -0,0 +1,49 @@
+# UniEnv
+
+Framework unifying robot environments and data APIs. UniEnv provides an universal interface for robot actors, sensors, environments, and data. 
+
+## Tensor library cross-backend Support
+
+UniEnv supports multiple tensor backends with zero-copy translation layers through the DLPack protocol, and allows you to use the same abstract compute backend interface to write custom data transformation layers, environment wrappers and other utilities. This is powered by the [XBArray](https://github.com/UniEnvOrg/XBArray) package.
+
+## Universal Robot Environment Interface
+
+UniEnv supports diverse simulation environments and real robots, built on top of the abstract environment / world interface. This allows you to reuse code across different sim and real robots.
+
+## Universal Robot Data Interface
+
+UniEnv provides a universal data interface for accessing robot data through the abstract `BatchBase` interface. We also provide a utility `ReplayBuffer` for saving data from various environments with diverse data format support, including `hdf5`, memory-mapped torch tensors, and others.
+
+## Installation
+
+Install the package with pip
+
+```bash
+pip install unienv
+```
+
+You can install optional dependencies such as `gymnasium` (for Gymnasium-compatible environments), `dev`, or `video` by running
+
+```bash
+pip install unienv[gymnasium,video]
+```
+
+## Cite
+
+If you use UniEnv in your research, please cite it as follows:
+
+```bibtex
+@software{cao_unienv,
+  author = {Cao, Yunhao AND Fang, Kuan},
+  title = {{UniEnv: Unifying Robot Environments and Data APIs}},
+  year = {2025},
+  month = oct,
+  url = {https://github.com/UniEnvOrg/UniEnv},
+  license = {MIT}
+}
+```
+
+## Acknowledgements
+
+The idea of this project is inspired by [Gymnasium](https://github.com/Farama-Foundation/Gymnasium) and its predecessor [OpenAI Gym](https://github.com/openai/gym). 
+This library is impossible without the great work of DataAPIs Consortium and their work on the [Array API Standard](https://data-apis.org/array-api/latest/). The zero-copy translation layers are powered by the [DLPack](https://github.com/dmlc/dlpack) project.

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

@@ -3,22 +3,15 @@ name = "unienv"
 description = "Unified robot environment framework supporting multiple tensor and simulation backends"
 readme = "README.md"
 license = "MIT"
-version = "0.0.1b2"
+version = "0.0.1b4"
 requires-python = ">= 3.10"
 dependencies = [
     "numpy",
     "xbarray>=0.0.1a8",
     "pillow",
-    "h5py",
+    "cloudpickle",
 ]
 
-[project.urls]
-Homepage = "https://www.quantumcookie.xyz/project/unienv/"
-Documentation = "https://readthedocs.org"
-Repository = "https://github.com/UniEnvOrg/UniEnvPy"
-Issues = "https://github.com/UniEnvOrg/UniEnvPy/issues"
-Changelog = "https://github.com/UniEnvOrg/UniEnvPy/blob/main/CHANGELOG.md"
-
 [project.optional-dependencies]
 dev = [
     "pytest",
@@ -29,12 +22,13 @@ gymnasium = [
 video = [
     "moviepy>=2.1"
 ]
-mjx = [
-    "playground",
-]
-maniskill = [
-    "mani_skill>=3.0.0b12"
-]
+
+[project.urls]
+Homepage = "https://github.com/UniEnvOrg/UniEnv"
+Documentation = "https://github.com/UniEnvOrg/UniEnv"
+Repository = "https://github.com/UniEnvOrg/UniEnv"
+Issues = "https://github.com/UniEnvOrg/UniEnv/issues"
+Changelog = "https://github.com/UniEnvOrg/UniEnv/blob/main/CHANGELOG.md"
 
 [tool.setuptools.packages.find]
 include = ["*"]

unienv-0.0.1b4/unienv.egg-info/PKG-INFO

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: unienv
+Version: 0.0.1b4
+Summary: Unified robot environment framework supporting multiple tensor and simulation backends
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv
+Project-URL: Documentation, https://github.com/UniEnvOrg/UniEnv
+Project-URL: Repository, https://github.com/UniEnvOrg/UniEnv
+Project-URL: Issues, https://github.com/UniEnvOrg/UniEnv/issues
+Project-URL: Changelog, https://github.com/UniEnvOrg/UniEnv/blob/main/CHANGELOG.md
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: xbarray>=0.0.1a8
+Requires-Dist: pillow
+Requires-Dist: cloudpickle
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Provides-Extra: gymnasium
+Requires-Dist: gymnasium>=0.29.0; extra == "gymnasium"
+Provides-Extra: video
+Requires-Dist: moviepy>=2.1; extra == "video"
+Dynamic: license-file
+
+# UniEnv
+
+Framework unifying robot environments and data APIs. UniEnv provides an universal interface for robot actors, sensors, environments, and data. 
+
+## Tensor library cross-backend Support
+
+UniEnv supports multiple tensor backends with zero-copy translation layers through the DLPack protocol, and allows you to use the same abstract compute backend interface to write custom data transformation layers, environment wrappers and other utilities. This is powered by the [XBArray](https://github.com/UniEnvOrg/XBArray) package.
+
+## Universal Robot Environment Interface
+
+UniEnv supports diverse simulation environments and real robots, built on top of the abstract environment / world interface. This allows you to reuse code across different sim and real robots.
+
+## Universal Robot Data Interface
+
+UniEnv provides a universal data interface for accessing robot data through the abstract `BatchBase` interface. We also provide a utility `ReplayBuffer` for saving data from various environments with diverse data format support, including `hdf5`, memory-mapped torch tensors, and others.
+
+## Installation
+
+Install the package with pip
+
+```bash
+pip install unienv
+```
+
+You can install optional dependencies such as `gymnasium` (for Gymnasium-compatible environments), `dev`, or `video` by running
+
+```bash
+pip install unienv[gymnasium,video]
+```
+
+## Cite
+
+If you use UniEnv in your research, please cite it as follows:
+
+```bibtex
+@software{cao_unienv,
+  author = {Cao, Yunhao AND Fang, Kuan},
+  title = {{UniEnv: Unifying Robot Environments and Data APIs}},
+  year = {2025},
+  month = oct,
+  url = {https://github.com/UniEnvOrg/UniEnv},
+  license = {MIT}
+}
+```
+
+## Acknowledgements
+
+The idea of this project is inspired by [Gymnasium](https://github.com/Farama-Foundation/Gymnasium) and its predecessor [OpenAI Gym](https://github.com/openai/gym). 
+This library is impossible without the great work of DataAPIs Consortium and their work on the [Array API Standard](https://data-apis.org/array-api/latest/). The zero-copy translation layers are powered by the [DLPack](https://github.com/dmlc/dlpack) project.

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

@@ -10,23 +10,25 @@ unienv_data/__init__.py
 unienv_data/base/__init__.py
 unienv_data/base/common.py
 unienv_data/base/storage.py
-unienv_data/base/transformations.py
 unienv_data/batches/__init__.py
 unienv_data/batches/backend_compat.py
 unienv_data/batches/combined_batch.py
 unienv_data/batches/framestack_batch.py
 unienv_data/batches/slicestack_batch.py
+unienv_data/batches/transformations.py
 unienv_data/integrations/pytorch.py
 unienv_data/replay_buffer/__init__.py
 unienv_data/replay_buffer/replay_buffer.py
 unienv_data/replay_buffer/trajectory_replay_buffer.py
 unienv_data/samplers/__init__.py
 unienv_data/samplers/multiprocessing_sampler.py
-unienv_data/samplers/slice_sampler.py
 unienv_data/samplers/step_sampler.py
-unienv_data/storages/common.py
+unienv_data/storages/dict_storage.py
+unienv_data/storages/flattened.py
 unienv_data/storages/hdf5.py
 unienv_data/storages/pytorch.py
+unienv_data/storages/transformation.py
+unienv_data/transformations/image_compress.py
 unienv_interface/__init__.py
 unienv_interface/backends/__init__.py
 unienv_interface/backends/base.py
@@ -38,17 +40,21 @@ unienv_interface/env_base/__init__.py
 unienv_interface/env_base/env.py
 unienv_interface/env_base/funcenv.py
 unienv_interface/env_base/funcenv_wrapper.py
+unienv_interface/env_base/vec_env.py
 unienv_interface/env_base/wrapper.py
 unienv_interface/func_wrapper/__init__.py
+unienv_interface/func_wrapper/frame_stack.py
 unienv_interface/func_wrapper/transformation.py
 unienv_interface/space/__init__.py
 unienv_interface/space/space.py
 unienv_interface/space/space_utils/__init__.py
 unienv_interface/space/space_utils/batch_utils.py
+unienv_interface/space/space_utils/construct_utils.py
 unienv_interface/space/space_utils/flatten_utils.py
 unienv_interface/space/space_utils/gym_utils.py
 unienv_interface/space/space_utils/serialization_utils.py
 unienv_interface/space/spaces/__init__.py
+unienv_interface/space/spaces/batched.py
 unienv_interface/space/spaces/binary.py
 unienv_interface/space/spaces/box.py
 unienv_interface/space/spaces/dict.py
@@ -64,9 +70,15 @@ unienv_interface/transformations/dict_transform.py
 unienv_interface/transformations/filter_dict.py
 unienv_interface/transformations/rescale.py
 unienv_interface/transformations/transformation.py
+unienv_interface/utils/control_util.py
+unienv_interface/utils/data_queue.py
 unienv_interface/utils/seed_util.py
+unienv_interface/utils/stateclass.py
 unienv_interface/utils/symbol_util.py
+unienv_interface/utils/vec_util.py
 unienv_interface/world/__init__.py
+unienv_interface/world/combined_funcnode.py
+unienv_interface/world/combined_node.py
 unienv_interface/world/funcnode.py
 unienv_interface/world/funcworld.py
 unienv_interface/world/node.py
@@ -81,8 +93,4 @@ unienv_interface/wrapper/frame_stack.py
 unienv_interface/wrapper/gym_compat.py
 unienv_interface/wrapper/time_limit.py
 unienv_interface/wrapper/transformation.py
-unienv_interface/wrapper/video_record.py
-unienv_maniskill/__init__.py
-unienv_maniskill/wrapper/maniskill_compat.py
-unienv_mjxplayground/__init__.py
-unienv_mjxplayground/wrapper/playground_compat.py
+unienv_interface/wrapper/video_record.py

{unienv-0.0.1b2 → unienv-0.0.1b4}/unienv.egg-info/requires.txt

@@ -1,7 +1,7 @@
 numpy
 xbarray>=0.0.1a8
 pillow
-h5py
+cloudpickle
 
 [dev]
 pytest
@@ -9,11 +9,5 @@ pytest
 [gymnasium]
 gymnasium>=0.29.0
 
-[maniskill]
-mani_skill>=3.0.0b12
-
-[mjx]
-playground
-
 [video]
 moviepy>=2.1

unienv-0.0.1b4/unienv.egg-info/top_level.txt

@@ -0,0 +1,3 @@
+dist
+unienv_data
+unienv_interface

{unienv-0.0.1b2 → unienv-0.0.1b4}/unienv_data/base/__init__.py

@@ -1,3 +1,2 @@
 from .common import BatchT, SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType, BatchBase, BatchSampler, IndexableType
-from .transformations import TransformedBatch
 from .storage import SpaceStorage

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

@@ -9,26 +9,44 @@ import dataclasses
 
 from unienv_interface.space.space_utils import batch_utils as space_batch_utils, flatten_utils as space_flatten_utils
 
+__all__ = [
+    "BatchT",
+    "BatchBase",
+    "BatchSampler",
+    "IndexableType",
+    "convert_index_to_backendarray",
+]
+
 IndexableType = Union[int, slice, EllipsisType]
 
+def convert_index_to_backendarray(
+    backend : ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
+    index : IndexableType,
+    length : int,
+    device : Optional[BDeviceType] = None,
+) -> BArrayType:
+    if isinstance(index, int):
+        return backend.asarray([index], dtype=backend.default_integer_dtype, device=device)
+    elif isinstance(index, slice):
+        return backend.arange(*index.indices(length), dtype=backend.default_integer_dtype, device=device)
+    elif index is Ellipsis:
+        return backend.arange(length, dtype=backend.default_integer_dtype, device=device)
+    else:
+        raise ValueError("Index must be an integer, slice, or Ellipsis.")
+
 BatchT = TypeVar('BatchT')
 class BatchBase(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]):
-    backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]
-    device: Optional[BDeviceType] = None
-
     # If the batch is mutable, then the data can be changed (extend_*, set_*, remove_*, etc.)
     is_mutable: bool = True
 
     def __init__(
         self,
-        single_space : Space[Any, BDeviceType, BDtypeType, BRNGType],
+        single_space : Space[BatchT, BDeviceType, BDtypeType, BRNGType],
         single_metadata_space : Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]] = None,
     ):
         self.single_space = single_space
         self.single_metadata_space = single_metadata_space
-        self._batched_space : Space[
-            BatchT, Any, BDeviceType, BDtypeType, BRNGType
-        ] = space_batch_utils.batch_space(single_space, 1)
+        self._batched_space : Space[BatchT, BDeviceType, BDtypeType, BRNGType] = space_batch_utils.batch_space(single_space, 1)
         if single_metadata_space is not None:
             self._batched_metadata_space : DictSpace[
                 BDeviceType, BDtypeType, BRNGType
@@ -36,24 +54,43 @@ class BatchBase(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BR
         else:
             self._batched_metadata_space = None
 
+    @property
+    def backend(self) -> ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]:
+        return self.single_space.backend
+    
+    @property
+    def device(self) -> Optional[BDeviceType]:
+        return self.single_space.device
+
     @abc.abstractmethod
     def __len__(self) -> int:
         raise NotImplementedError
     
     def get_flattened_at(self, idx : Union[IndexableType, BArrayType]) -> BArrayType:
-        return self.get_flattened_at_with_metadata(idx)[0]
+        unflattened_data = self.get_at(idx)
+        if isinstance(idx, int):
+            return space_flatten_utils.flatten_data(self.single_space, unflattened_data)
+        else:
+            return space_flatten_utils.flatten_data(self._batched_space, unflattened_data, start_dim=1)
 
-    @abc.abstractmethod
     def get_flattened_at_with_metadata(
         self, idx : Union[IndexableType, BArrayType]
     ) -> Tuple[BArrayType, Optional[Dict[str, Any]]]:
-        raise NotImplementedError
+        unflattened_data, metadata = self.get_at_with_metadata(idx)
+        if isinstance(idx, int):
+            return space_flatten_utils.flatten_data(self.single_space, unflattened_data), metadata
+        else:
+            return space_flatten_utils.flatten_data(self._batched_space, unflattened_data, start_dim=1), metadata
 
     def set_flattened_at(self, idx : Union[IndexableType, BArrayType], value : BArrayType) -> None:
         raise NotImplementedError
 
+    def append_flattened(self, value : BArrayType) -> None:
+        return self.extend_flattened(value[None])
+
     def extend_flattened(self, value : BArrayType) -> None:
-        raise NotImplementedError
+        unflat_data = space_flatten_utils.unflatten_data(self._batched_space, value, start_dim=1)
+        self.extend(unflat_data)
     
     def get_at(self, idx : Union[IndexableType, BArrayType]) -> BatchT:
         flattened_data = self.get_flattened_at(idx)
@@ -90,55 +127,81 @@ class BatchBase(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType, BR
     def __delitem__(self, idx : Union[IndexableType, BArrayType]) -> None:
         self.remove_at(idx)
 
+    def append(self, value : BatchT) -> None:
+        batched_data = space_batch_utils.concatenate(self._batched_space, [value])
+        self.extend(batched_data)
+
     def extend(self, value : BatchT) -> None:
         flattened_data = space_flatten_utils.flatten_data(self._batched_space, value, start_dim=1)
         self.extend_flattened(flattened_data)
 
+    def extend_from(
+        self, 
+        other : 'BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]',
+        chunk_size : int = 8,
+        tqdm : bool = False,
+    ) -> None:
+        n_total = len(other)
+        iterable_start = range(0, n_total, chunk_size)
+        if tqdm:
+            from tqdm import tqdm
+            iterable_start = tqdm(iterable_start, desc="Extending Batch")
+        for start_idx in range(0, n_total, chunk_size):
+            end_idx = min(start_idx + chunk_size, n_total)
+            data_chunk = other.get_at(slice(start_idx, end_idx))
+            self.extend(data_chunk)
+
     def close(self) -> None:
         pass
 
-    def __del__(self) -> None:
-        self.close()
-
 SamplerBatchT = TypeVar('SamplerBatchT')
 SamplerArrayType = TypeVar('SamplerArrayType')
 SamplerDeviceType = TypeVar('SamplerDeviceType')
 SamplerDtypeType = TypeVar('SamplerDtypeType')
 SamplerRNGType = TypeVar('SamplerRNGType')
-class BatchSampler(abc.ABC, Generic[
-    SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType,
-    BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType,
-]):
-    batch_size : int
-    sampled_space : Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
-    sampled_space_flat : BoxSpace[SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
-    sampled_metadata_space : Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]] = None
-
-    backend : ComputeBackend[SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]
-    device : Optional[SamplerDeviceType] = None
-
+class BatchSampler(
+    Generic[
+        SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType,
+        BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType,
+    ],
+    BatchBase[
+        SamplerBatchT, SamplerArrayType, SamplerDeviceType, SamplerDtypeType, SamplerRNGType
+    ]
+):
     data : BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]
 
     rng : Optional[SamplerRNGType] = None
     data_rng : Optional[BRNGType] = None
-
-    def get_flat_at(self, idx : SamplerArrayType) -> SamplerArrayType:
-        return self.get_flat_at_with_metadata(idx)[0]
-    
-    @abc.abstractmethod
-    def get_flat_at_with_metadata(
-        self, idx : SamplerArrayType
-    ) -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]:
-        raise NotImplementedError
     
-    def get_at(self, idx : SamplerArrayType) -> SamplerBatchT:
-        return space_flatten_utils.unflatten_data(self.sampled_space, self.get_flat_at(idx), start_dim=1)
+    is_mutable : bool = False
+
+    def __init__(
+        self,
+        single_space : Space[BatchT, BDeviceType, BDtypeType, BRNGType],
+        single_metadata_space : Optional[DictSpace[BDeviceType, BDtypeType, BRNGType]] = None,
+        batch_size : int = 1,
+    ) -> None:
+        super().__init__(single_space=single_space, single_metadata_space=single_metadata_space)
+        self.batch_size = batch_size
+        self._batched_space : Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType] = space_batch_utils.batch_space(self.single_space, batch_size)
+        self._batched_metadata_space : Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]] = space_batch_utils.batch_space(self.single_metadata_space, batch_size) if self.single_metadata_space is not None else None
+
+    def manual_seed(self, seed : int) -> None:
+        if self.rng is not None:
+            self.rng = self.backend.random.random_number_generator(seed, device=self.device)
+        if self.data_rng is not None:
+            self.data_rng = self.backend.random.random_number_generator(seed, device=self.data.device)
+
+    @property
+    def sampled_space(self) -> Space[SamplerBatchT, SamplerDeviceType, SamplerDtypeType, SamplerRNGType]:
+        return self._batched_space
     
-    def get_at_with_metadata(
-        self, idx : SamplerArrayType
-    ) -> Tuple[SamplerBatchT, Optional[Dict[str, Any]]]:
-        flat_data, metadata = self.get_flat_at_with_metadata(idx)
-        return space_flatten_utils.unflatten_data(self.sampled_space, flat_data, start_dim=1), metadata
+    @property
+    def sampled_metadata_space(self) -> Optional[DictSpace[SamplerDeviceType, SamplerDtypeType, SamplerRNGType]]:
+        return self._batched_metadata_space
+
+    def __len__(self):
+        return len(self.data)
 
     def sample_index(self) -> SamplerArrayType:
         new_rng, indices = self.backend.random.random_discrete_uniform( # (B, )
@@ -156,11 +219,11 @@ class BatchSampler(abc.ABC, Generic[
 
     def sample_flat(self) -> SamplerArrayType:
         idx = self.sample_index()
-        return self.get_flat_at(idx)
+        return self.get_flattened_at(idx)
     
     def sample_flat_with_metadata(self) -> Tuple[SamplerArrayType, Optional[Dict[str, Any]]]:
         idx = self.sample_index()
-        return self.get_flat_at_with_metadata(idx)
+        return self.get_flattened_at_with_metadata(idx)
 
     def sample(self) -> SamplerBatchT:
         idx = self.sample_index()
@@ -205,9 +268,9 @@ class BatchSampler(abc.ABC, Generic[
         n_batches = len(self.data) // self.batch_size
         num_left = len(self.data) % self.batch_size
         for i in range(n_batches):
-            yield self.get_flat_at(idx[i*self.batch_size:(i+1)*self.batch_size])
+            yield self.get_flattened_at(idx[i*self.batch_size:(i+1)*self.batch_size])
         if num_left > 0:
-            yield self.get_flat_at(idx[-num_left:])
+            yield self.get_flattened_at(idx[-num_left:])
     
     def epoch_flat_iter_with_metadata(self) -> Iterator[Tuple[SamplerArrayType, Optional[Dict[str, Any]]]]:
         if self.data_rng is not None:
@@ -217,12 +280,9 @@ class BatchSampler(abc.ABC, Generic[
         n_batches = len(self.data) // self.batch_size
         num_left = len(self.data) % self.batch_size
         for i in range(n_batches):
-            yield self.get_flat_at_with_metadata(idx[i*self.batch_size:(i+1)*self.batch_size])
+            yield self.get_flattened_at_with_metadata(idx[i*self.batch_size:(i+1)*self.batch_size])
         if num_left > 0:
-            yield self.get_flat_at_with_metadata(idx[-num_left:])
+            yield self.get_flattened_at_with_metadata(idx[-num_left:])
 
     def close(self) -> None:
         pass
-
-    def __del__(self) -> None:
-        self.close()

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

@@ -31,6 +31,7 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
         single_instance_space: Space[BatchT, BDeviceType, BDtypeType, BRNGType],
         *,
         capacity : Optional[int] = None,
+        read_only : bool = True,
         **kwargs
     ) -> "SpaceStorage[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         raise NotImplementedError
@@ -56,6 +57,17 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
     """
     cache_filename : Optional[Union[str, os.PathLike]] = None
 
+    """
+    Can the storage instance be safely used in multiprocessing environments after creation?
+    If True, the storage can be used in multiprocessing environments.
+    """
+    is_multiprocessing_safe : bool = False
+
+    """
+    Is the storage mutable? If False, the storage is read-only.
+    """
+    is_mutable : bool = True
+
     @property
     def backend(self) -> ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType]:
         return self.single_instance_space.backend
@@ -127,6 +139,3 @@ class SpaceStorage(abc.ABC, Generic[BatchT, BArrayType, BDeviceType, BDtypeType,
 
     def close(self) -> None:
         pass
-
-    def __del__(self) -> None:
-        self.close()

{unienv-0.0.1b2 → unienv-0.0.1b4}/unienv_data/batches/__init__.py

@@ -1,4 +1,5 @@
 from .backend_compat import ToBackendOrDeviceBatch
 from .combined_batch import CombinedBatch
 from .slicestack_batch import SliceStackedBatch
-from .framestack_batch import FrameStackedBatch
+from .framestack_batch import FrameStackedBatch
+from .transformations import TransformedBatch