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

unienv-0.0.1b3.dist-info/METADATA

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: unienv
+Version: 0.0.1b3
+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.1b3.dist-info/RECORD

@@ -0,0 +1,92 @@
+unienv-0.0.1b3.dist-info/licenses/LICENSE,sha256=nkklvEaJUR4QDBygz7tkEe1FMVKV1JSjnGzJNLhdIWM,1091
+unienv_data/__init__.py,sha256=zFxbe7aM5JvYXIK0FGnOPwWQJMN-8l_l8prB85CkcA8,95
+unienv_data/base/__init__.py,sha256=w-I8A-z7YYArkHc2ZOVGrfzfThsaDBg7aD7qMFprNM8,186
+unienv_data/base/common.py,sha256=n7GnPH-x5fwbsbwNgulg96yiwJeD86wxy2jURFbF394,12364
+unienv_data/base/storage.py,sha256=82wEDQtv0QKlyls20_SqYEkQEICbv0EwF5fcRuqeczs,5135
+unienv_data/batches/__init__.py,sha256=Vi92f8ddgFYCqwv7xO2Pi3oJePnioJ4XrJbQVV7eIvk,234
+unienv_data/batches/backend_compat.py,sha256=7Juf7nU2jYHohRzNzmGfqMMpJtFM-3oTzzLu6EbC77E,8168
+unienv_data/batches/combined_batch.py,sha256=aua1H86sa_qWrCtAAp5JIZMGtFiiKFPFkU0y5JoyElM,15325
+unienv_data/batches/framestack_batch.py,sha256=pdURqZeksOlbf21Nhx8kkm0gtFt6rjt2OiNWgZPdFCM,2312
+unienv_data/batches/slicestack_batch.py,sha256=J2EhARcPA-zz6EBnV7OLzm4yyvnZ06vrdUoPD5RkJ-o,16672
+unienv_data/batches/transformations.py,sha256=b4HqX3wZ6TuRgQ2q81Jv43PmeHGmP8cwURK_ULjGNgs,5647
+unienv_data/integrations/pytorch.py,sha256=pW5rXBXagfzwJjM_VGgg8CPXEs3e2fKgg4nY7M3dpOc,2350
+unienv_data/replay_buffer/__init__.py,sha256=uVebYruIYlj8OjTYVi8UYI4gWp3S3XIdgFlHbwO260o,100
+unienv_data/replay_buffer/replay_buffer.py,sha256=nhbC-7aHGIYhcCdmaaDdhB2U9ODAZrbKMq8dP8ffOv0,10344
+unienv_data/replay_buffer/trajectory_replay_buffer.py,sha256=fxV6FIqAHhN8opYs2WjAJMPqNRWD3iIku-4WlaydyG4,20737
+unienv_data/samplers/__init__.py,sha256=e7uunWN3r-g_2fDaMsYMe8cZcF4N-okCxqBPweQnE0s,97
+unienv_data/samplers/multiprocessing_sampler.py,sha256=FEBK8pMTnkpA0xuMkbvlv4aIdVTTubeT8BjL60BJL5o,13254
+unienv_data/samplers/step_sampler.py,sha256=ZCcrx9WbILtaR6izhIP3DhtmFcP7KQBdaYaSZ7vWwRk,3010
+unienv_data/storages/common.py,sha256=7Z6RY6cvd0TboM6qENYT3Fxs6lFLdxo3kNppUtJD_Mc,6256
+unienv_data/storages/hdf5.py,sha256=mdZdXBd6yZ6kpI4VdByl0XFTqCklHLBC8dHPNA-1cnw,25101
+unienv_data/storages/pytorch.py,sha256=q4RvJsKr6q8t4lZk9HYVpxGQXT2iFP1eDTRNQKjEYA0,6246
+unienv_data/storages/transformation.py,sha256=huVpkaKjWeQ0XUDblS10mEUBuOYZsriCoLXawyYjwnk,7713
+unienv_data/transformations/image_compress.py,sha256=dINrvmpTWy3sbqruHk0kPZG2XNyJI90ERgErXV7GamE,9131
+unienv_interface/__init__.py,sha256=pAWqfm4l7NAssuyXCugIjekSIh05aBbOjNhwsNXcJbE,100
+unienv_interface/backends/__init__.py,sha256=L7CFwCChHVL-2Dpz34pTGC37WgodfJEeDQwXscyM7FM,198
+unienv_interface/backends/base.py,sha256=1_hji1qwNAhcEtFQdAuzaNey9g5bWYj38t1sQxjnggc,132
+unienv_interface/backends/jax.py,sha256=26Wu5OQ4EEjolyZoELhlWMPNSZ7LsVoKEGpd09L80Ck,533
+unienv_interface/backends/numpy.py,sha256=6dMB2Vq7mrWukobyyGvuccluZUgjVkxr7x0hrUc_pe8,542
+unienv_interface/backends/pytorch.py,sha256=BddHmZAngsaedFlvj1mKdXpNe6AWvNwEXq_eTEUoFWA,592
+unienv_interface/backends/serialization.py,sha256=0TZlpfbP1DRB4FkM8ysDVQmn6RlYtIPisyeHjvHr7bE,2289
+unienv_interface/env_base/__init__.py,sha256=JuaVgWlg313LZpflt4LSErY94nUrfvUp0LbIPUle0MA,226
+unienv_interface/env_base/env.py,sha256=PV-AEmKwSjnFDjZFYtBW-At9w4fpm_I5C7GhfxPPrs4,4833
+unienv_interface/env_base/funcenv.py,sha256=Qwm9BP4NrsVHOr7X0l3-mbsn5IhaO3-ZVW48dLg08-k,10609
+unienv_interface/env_base/funcenv_wrapper.py,sha256=chw1iJ1RhAFMv4JAk67cttJvI9agdSm1QxNxZq0-hgM,7757
+unienv_interface/env_base/vec_env.py,sha256=bcv6NdOxt0Xp1fRMXqzFtmVw6LQ-pDj_Jvj-qaW6otQ,16116
+unienv_interface/env_base/wrapper.py,sha256=7hf4Rr2wouS0igPoahhvb2tzYY3bCaWL0NlgwpYZwQs,9734
+unienv_interface/func_wrapper/__init__.py,sha256=6BPF8O25WkIBpODVTwnOE9HGSm3KRKX6iPwFGWESlxA,123
+unienv_interface/func_wrapper/frame_stack.py,sha256=52CqAHDqwgHtOwMwxzB3Syup9kA19zdlvXCH4mI7MNU,6819
+unienv_interface/func_wrapper/transformation.py,sha256=7mdzcpjLjqtpbtXoqbkGtTMPQxoMmMsqzDWHcZLbrhs,5939
+unienv_interface/space/__init__.py,sha256=6-wLoD9mKDAfz7IuQs_Rn9DMDfDwTZ0tEhQ924libpg,99
+unienv_interface/space/space.py,sha256=mFlCcDvMgEPTXlwo_iwBlm6Eg4Bn2rrecgsfIVstdq0,4067
+unienv_interface/space/space_utils/__init__.py,sha256=GAsPoZC8YNabx3Gw5m2o4zsnG8zmA3mcuM9_lNKhiGo,121
+unienv_interface/space/space_utils/batch_utils.py,sha256=qXK7kERPXKGIYozz7lpjzVz56S9GkH6ZASfIRzCYXHY,36993
+unienv_interface/space/space_utils/construct_utils.py,sha256=Y4RpV9obY8XQ85O3r_NC1HrBk-Nm941ffRNXNL7nHgA,8323
+unienv_interface/space/space_utils/flatten_utils.py,sha256=kkHkjrsk43NDbg3Q5VAhVoIXStuRayYFO-7knsDzx4A,12289
+unienv_interface/space/space_utils/gym_utils.py,sha256=nH8EKruOKCXNrIMPUd9F4XGKCfFkhxsTmx4I1BeSgn0,15079
+unienv_interface/space/space_utils/serialization_utils.py,sha256=LWYSFN7E6tEFe8ULWm42LkFUxP_0dfTGkCcx0yl4Y8s,9530
+unienv_interface/space/spaces/__init__.py,sha256=Jap768TlwHFDDiTzHZ0qaHEFEVC1cKA2QzLlSZVQnjI,535
+unienv_interface/space/spaces/batched.py,sha256=RA8aLUSS14zBSCTm_ud18TTa-ntbIZ074xwJ0xls1Jk,3691
+unienv_interface/space/spaces/binary.py,sha256=0iQUbO37dhkznVpjhsJdwlD-KdWgCEx2H7KrybuZ_PM,3570
+unienv_interface/space/spaces/box.py,sha256=NCmileEZOKz-L3WNzZ-uwydrRFsIMdNZBwTn1vWgeI0,13316
+unienv_interface/space/spaces/dict.py,sha256=G5_iYC1Bj5DqeJ7aFlq6eRJbnpATbIRIyRu1jF_UUvk,7022
+unienv_interface/space/spaces/dynamic_box.py,sha256=HvMNgzfYwIVc5VVgCtq-8lQbNI1V1dZMI-w60AwYss4,19591
+unienv_interface/space/spaces/graph.py,sha256=KocRFLtYP5VWYpwbP6HybXH5R4jTIYJdNePKb6vhnYE,15163
+unienv_interface/space/spaces/text.py,sha256=ePGGJdiD3q-BAX6IHLO7HMe0OH4VrzF043K02eb0zXI,4443
+unienv_interface/space/spaces/tuple.py,sha256=rgZQz3EB3CLbIk9UlHBQbM6w9gssbA1izm-Qq-_Chqs,4267
+unienv_interface/space/spaces/union.py,sha256=Qisd-DdmPcGRmdhZFGiQw8_AOjYWqkuQ4Hwd-I8tdSI,4375
+unienv_interface/transformations/__init__.py,sha256=g19uGnDHMywvDAXRaqFgoWAF1vCPrbJENEpaEgtIrOw,353
+unienv_interface/transformations/batch_and_unbatch.py,sha256=ELCnNtwmgA5wpTBJZasfNSHmtf4vzydzLPmO6IGbT9o,1164
+unienv_interface/transformations/chained_transform.py,sha256=TDnUvxUKK6bXGc_sfr6ZCvvVWw7P5KX2sA9i7i2lx14,2075
+unienv_interface/transformations/dict_transform.py,sha256=ynrJrloVUix2I27Ir1mL86crT0vY5DvpiBAVxPBJup4,5357
+unienv_interface/transformations/filter_dict.py,sha256=DzR-hgHoHJObTipxwB2UrKVlTxbfIrJohaOgqdAICLY,5871
+unienv_interface/transformations/rescale.py,sha256=fM5ukWUvNvPeDO48_PRU0KyyvGhBIDxaN9XZyQ1VaQQ,4364
+unienv_interface/transformations/transformation.py,sha256=u4_9H1tvophhgG0p0F3xfkMMsRuaKY2TQmVeGoeQsJ0,1652
+unienv_interface/utils/control_util.py,sha256=lY_1EknglY3cNekWX9rYWt0ZUglaPMtIt4M5D9y0WfE,2351
+unienv_interface/utils/data_queue.py,sha256=UZiuQDOn39DB9Heu6xinrwuzAL3X8jHlDkFoSC5Phtc,5707
+unienv_interface/utils/seed_util.py,sha256=Up3nBXj7L8w-S9W5Q1U2d9accMhMf0TmHPaN6JXDVWs,677
+unienv_interface/utils/stateclass.py,sha256=xjzicPGX1UuI7q3ZAxhBCCoouKfNtLywUzQtLaT0yS4,1390
+unienv_interface/utils/symbol_util.py,sha256=NAERK-D_2MaTg2eYW-L75tbzPQN5YJIiKtM9zuQ89Sw,383
+unienv_interface/utils/vec_util.py,sha256=EIK680ReCl_rr-qKP8co5hwz8Dx-gks8SHf-CLOZSOA,373
+unienv_interface/world/__init__.py,sha256=aGuYTz8XFzW32RGkdi2b2LJ1sa0kgFrQyOR3JXDEwLQ,230
+unienv_interface/world/combined_funcnode.py,sha256=O9qWxhtMJkDVtWuGyaeEj3nKMgIyRAPqF9-5LU6yna8,10853
+unienv_interface/world/combined_node.py,sha256=tG7I9uWVxDDN6M6KeC1D14MV7YUnXYMUK9L9KXHnViA,9090
+unienv_interface/world/funcnode.py,sha256=mUpVQ_j7dVF6V7Dc435RNtMQB2LsuWRRkp-hBLOCwzc,7829
+unienv_interface/world/funcworld.py,sha256=GLp8nS0Ym1gaj7FWvD5FPkQElCgZMbpyuLsIMU0w-sw,2020
+unienv_interface/world/node.py,sha256=Qn8rErvhkRp2U0s_m_0OqDLY723w9E5W8tGdmKcP-mY,5996
+unienv_interface/world/world.py,sha256=Kl7wbNbs2YR3CjFrCLFhDB3DQUAWM6LjBwSADQtBTII,5740
+unienv_interface/wrapper/__init__.py,sha256=ZNqr-WjVRqgvIxkLkeABxpYZ6tRgJNZOzmluDeJ6W_w,614
+unienv_interface/wrapper/action_rescale.py,sha256=rTJlEHvWSuwGVX83cjfLWvszBk7B2iExX_K37vH8Wic,1231
+unienv_interface/wrapper/backend_compat.py,sha256=T6hosgu2hrZvg3xtnyELmR6Exlz-ztqdj9vdyiz7bhI,7081
+unienv_interface/wrapper/batch_and_unbatch.py,sha256=HpmnppgOKmshNlfmJYkGQYtEU7_U7q3mEdY5n4UaqEY,3457
+unienv_interface/wrapper/control_frequency_limit.py,sha256=B0E2aUbaUr2p2yIN6wT3q4rAbPYsVroioqma2qKMoC0,2322
+unienv_interface/wrapper/flatten.py,sha256=NWA5xne5j_L34oq_wT85wGvp6iHwdCSeGsk1DMugvRw,5837
+unienv_interface/wrapper/frame_stack.py,sha256=lZZh_T_AmxsRWeYSLsTU321lVgIt12MX1eWl_yRNlWg,6002
+unienv_interface/wrapper/gym_compat.py,sha256=JhLxDsO1NsJnKzKhO0MqMw9i5_1FLxoxKilWaQQyBkw,9789
+unienv_interface/wrapper/time_limit.py,sha256=VRvB00BK7deI2QtdGatqwDWmPgjgjg1E7MTvEyaW5rg,2904
+unienv_interface/wrapper/transformation.py,sha256=pQ-_YVU8WWDqSk2sONUUgQY1iigOD092KNcp1DYxoxk,10043
+unienv_interface/wrapper/video_record.py,sha256=y_nJRYgo1SeLeO_Ymg9xbbGPKm48AbU3BxZK2wd0gzk,8679
+unienv-0.0.1b3.dist-info/METADATA,sha256=bXlKeREYTsPrVKLhNKUE-Bm7BeL4QfXYfWBXh5n0rck,3033
+unienv-0.0.1b3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+unienv-0.0.1b3.dist-info/top_level.txt,sha256=wfcJ5_DruUtOEUZjEyfadaKn7B90hWqz2aw-eM3wX5g,29
+unienv-0.0.1b3.dist-info/RECORD,,

{unienv-0.0.1b1.dist-info → unienv-0.0.1b3.dist-info}/licenses/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.1b3.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+unienv_data
+unienv_interface

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_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,6 +127,10 @@ 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)
@@ -105,40 +146,49 @@ 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 +206,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 +255,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,9 +267,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

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

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

unienv_data/batches/backend_compat.py

@@ -66,7 +66,7 @@ class ToBackendOrDeviceBatch(
         )
         self.batch = batch
         self.target_backend = backend
-        self.device = device
+        self.target_device = device
     
     def __len__(self) -> int:
         return len(self.batch)
@@ -79,7 +79,18 @@ class ToBackendOrDeviceBatch(
     def backend(self) -> ComputeBackend[WrapperBArrayT, WrapperBDeviceT, WrapperBDtypeT, WrapperBRngT]:
         return self.target_backend if self.target_backend is not None else self.batch.backend
 
+    @property
+    def device(self) -> Optional[WrapperBDeviceT]:
+        return self.target_device if self.target_device is not None else self.batch.device
+
     def get_flattened_at(self, idx):
+        if self.target_backend.is_backendarray(idx):
+            idx = data_to(
+                idx,
+                source_backend=self.target_backend,
+                target_backend=self.batch.backend,
+                target_device=self.batch.device
+            )
         o_data = self.batch.get_flattened_at(idx)
         return data_to(
             o_data,
@@ -89,6 +100,13 @@ class ToBackendOrDeviceBatch(
         )
 
     def get_flattened_at_with_metadata(self, idx):
+        if self.target_backend.is_backendarray(idx):
+            idx = data_to(
+                idx,
+                source_backend=self.target_backend,
+                target_backend=self.batch.backend,
+                target_device=self.batch.device
+            )
         o_data, o_metadata = self.batch.get_flattened_at_with_metadata(idx)
         return (
             data_to(
@@ -107,6 +125,13 @@ class ToBackendOrDeviceBatch(
 
     def set_flattened_at(self, idx, value):
         assert self.is_mutable, "Batch is not mutable"
+        if self.target_backend.is_backendarray(idx):
+            idx = data_to(
+                idx,
+                source_backend=self.target_backend,
+                target_backend=self.batch.backend,
+                target_device=self.batch.device
+            )
         value = data_to(
             value,
             source_backend=self.target_backend,
@@ -126,6 +151,13 @@ class ToBackendOrDeviceBatch(
         self.batch.extend_flattened(value)
     
     def get_at(self, idx):
+        if self.target_backend.is_backendarray(idx):
+            idx = data_to(
+                idx,
+                source_backend=self.target_backend,
+                target_backend=self.batch.backend,
+                target_device=self.batch.device
+            )
         o_data = self.batch.get_at(idx)
         return (
             data_to(
@@ -137,6 +169,13 @@ class ToBackendOrDeviceBatch(
         )
     
     def get_at_with_metadata(self, idx):
+        if self.target_backend.is_backendarray(idx):
+            idx = data_to(
+                idx,
+                source_backend=self.target_backend,
+                target_backend=self.batch.backend,
+                target_device=self.batch.device
+            )
         o_data, o_metadata = self.batch.get_at_with_metadata(idx)
         return (
             data_to(
@@ -155,6 +194,13 @@ class ToBackendOrDeviceBatch(
     
     def set_at(self, idx, value):
         assert self.is_mutable, "Batch is not mutable"
+        if self.target_backend.is_backendarray(idx):
+            idx = data_to(
+                idx,
+                source_backend=self.target_backend,
+                target_backend=self.batch.backend,
+                target_device=self.batch.device
+            )
         o_value = data_to(
             value,
             source_backend=self.target_backend,

unienv_data/batches/combined_batch.py

@@ -56,8 +56,6 @@ class CombinedBatch(BatchBase[
         )
         super().__init__(single_space, new_metadata_space)
 
-        self.backend = backend
-        self.device = device
         self.is_mutable = is_mutable
         self.batches = batches
         self._build_index_cache()
@@ -248,7 +246,7 @@ class CombinedBatch(BatchBase[
             result = result_space.create_empty()
             for batch_index, index_into_batch, mask in batch_list:
                 result = sbu.set_at(
-                    self.single_space,
+                    result_space,
                     result,
                     mask,
                     self.batches[batch_index].get_at(index_into_batch),
@@ -295,7 +293,7 @@ class CombinedBatch(BatchBase[
             for batch_index, index_into_batch, mask in batch_list:
                 batch_result, metadata_result = self.batches[batch_index].get_at_with_metadata(index_into_batch)
                 result = sbu.set_at(
-                    self.single_space,
+                    result_space,
                     result,
                     mask,
                     batch_result,

unienv_data/{base → batches}/transformations.py

@@ -1,11 +1,12 @@
-from typing import Optional, Any, Union
+from typing import Optional, Any, Union, Tuple, Dict
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
 
 from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils as sfu
-from .common import *
 from unienv_interface.transformations.transformation import DataTransformation, TargetDataT, SourceDataT, SourceBArrT, SourceBDeviceT, SourceBDTypeT, SourceBDRNGT
 from unienv_interface.space import Space
 
+from ..base.common import BatchBase, BatchT, IndexableType
+
 class TransformedBatch(
     BatchBase[
         BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType

unienv_data/replay_buffer/replay_buffer.py

@@ -63,6 +63,8 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         **kwargs
     ) -> "ReplayBuffer[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         storage_path_relative = "storage" + (storage_cls.single_file_ext or "")
+        if cache_path is not None:
+            os.makedirs(cache_path, exist_ok=True)
         storage = storage_cls.create(
             single_instance_space,
             *args,
@@ -94,6 +96,7 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
         *,
         backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
         device: Optional[BDeviceType] = None,
+        read_only : bool = True,
         **storage_kwargs
     ) -> "ReplayBuffer[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]":
         with open(os.path.join(path, "metadata.json"), "r") as f:
@@ -114,6 +117,7 @@ class ReplayBuffer(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGTy
             storage_path,
             single_instance_space,
             capacity=capacity,
+            read_only=read_only,
             **storage_kwargs
         )
         return ReplayBuffer(storage, metadata["storage_path_relative"], count, offset, cache_path=path)

unienv_data/samplers/__init__.py

@@ -1,3 +1,2 @@
 from .step_sampler import StepSampler
-from .slice_sampler import SliceSampler
 from .multiprocessing_sampler import MultiprocessingSampler