unienv 0.0.1b6__py3-none-any.whl → 0.0.1b8__py3-none-any.whl

{unienv-0.0.1b6.dist-info → unienv-0.0.1b8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b6
+Version: 0.0.1b8
 Summary: Unified robot environment framework supporting multiple tensor and simulation backends
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/UniEnvOrg/UniEnv

{unienv-0.0.1b6.dist-info → unienv-0.0.1b8.dist-info}/RECORD

@@ -1,22 +1,23 @@
-unienv-0.0.1b6.dist-info/licenses/LICENSE,sha256=nkklvEaJUR4QDBygz7tkEe1FMVKV1JSjnGzJNLhdIWM,1091
+unienv-0.0.1b8.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=A3RtD3Omqk0Qplsc-44ukAEzbQEU22_MkwUlC7l-HHM,13083
 unienv_data/base/storage.py,sha256=afICsO_7Zbm9azV0Jxho_z9F7JM30TUDjJM1NHETDHM,5495
 unienv_data/batches/__init__.py,sha256=Vi92f8ddgFYCqwv7xO2Pi3oJePnioJ4XrJbQVV7eIvk,234
 unienv_data/batches/backend_compat.py,sha256=tzFG8gTq0yW-J6PLvu--lCGS0lFc0QfelicJ50p_HYc,8207
-unienv_data/batches/combined_batch.py,sha256=pNrbLvU565BUDWO0pZLCnSMygmoGVCLxjC9OkLRKtLA,15330
+unienv_data/batches/combined_batch.py,sha256=iMCY7pEXf_YrIvq-G92ttaANl-SUcpzDLvH7fI-GImA,16313
 unienv_data/batches/framestack_batch.py,sha256=pdURqZeksOlbf21Nhx8kkm0gtFt6rjt2OiNWgZPdFCM,2312
 unienv_data/batches/slicestack_batch.py,sha256=Q3-gsJTvMjKTeZAHWNBTGRsws0HctsfMMTw0vylNxvA,16785
 unienv_data/batches/transformations.py,sha256=b4HqX3wZ6TuRgQ2q81Jv43PmeHGmP8cwURK_ULjGNgs,5647
-unienv_data/integrations/pytorch.py,sha256=pW5rXBXagfzwJjM_VGgg8CPXEs3e2fKgg4nY7M3dpOc,2350
+unienv_data/integrations/huggingface.py,sha256=cqaNS1Xpv5udAYgytd0zsVI8k9lSKMKz1jdxWnDseK4,1535
+unienv_data/integrations/pytorch.py,sha256=cyUypd4kZVH8WQFrrv4yFgYZJGT4MeuhvxzZMc4-1dM,4553
 unienv_data/replay_buffer/__init__.py,sha256=uVebYruIYlj8OjTYVi8UYI4gWp3S3XIdgFlHbwO260o,100
 unienv_data/replay_buffer/replay_buffer.py,sha256=8vPma5dL6jDGhI3Oo6IEvNcDYJG9Lb0Xlvxp45tQMEs,14498
 unienv_data/replay_buffer/trajectory_replay_buffer.py,sha256=cqRmzdewFS8IvJcMwxxQgwZf7TvvrViym87OaCOes3Y,24009
 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/_episode_storage.py,sha256=OpZt4P-P6LHrBR4F-tNcCFROLskWaOKWCDfoPV7qz1I,21970
+unienv_data/storages/_episode_storage.py,sha256=gdrj_bIRquyT0CIm9-0du5cCNhF-4cxwvf9u9POcZVg,22603
 unienv_data/storages/_list_storage.py,sha256=pH9xZOqXCx65NBRRD-INcP8OP-NWsI-JvdzVsPj9MSg,5225
 unienv_data/storages/backend_compat.py,sha256=BxeMJlC3FI60KLJ7QB5kF-mrGlJ6xi584Dcu4IN4Zrc,10714
 unienv_data/storages/dict_storage.py,sha256=DSqRIgo3m1XtUcLtyjYSqqpi01mr_nJOLg5BCddwPcg,13862
@@ -24,9 +25,9 @@ unienv_data/storages/flattened.py,sha256=Yf1G4D6KE36sESyDMGWKXqhFjz6Idx7N1aEhihm
 unienv_data/storages/hdf5.py,sha256=Jnls1rs7nlOOp9msmAfhuZp80OZd8S2Llls176EOUc4,27096
 unienv_data/storages/image_storage.py,sha256=4J1ZiGFHbGLHmReMztImJoDcRmiB_llD2wbMB3rdvOQ,5137
 unienv_data/storages/npz_storage.py,sha256=IP2DXbUs_ySzILne3s3hq3gwHiy9tfpWz6HcNciA8DU,4868
-unienv_data/storages/pytorch.py,sha256=bf3ys6eBlMvjyPK4XE-itENjEWq5Vm60qNwBNqJIZqg,7345
+unienv_data/storages/pytorch.py,sha256=MltiTcBuvNWQqE5-RAfjBz3gwdfATkVMTmSBR2escIE,7363
 unienv_data/storages/transformation.py,sha256=-9_jPZNpx6RXY_ojv_1UCSTa4Z9apI9V9jit8nG93oM,8133
-unienv_data/storages/video_storage.py,sha256=2vcNlghhDZWWzAdf9t0VeCMZrv-x_rYkYaCw8XV8AJA,13331
+unienv_data/storages/video_storage.py,sha256=0d3FP-BwXiq7q_s9UJE0G7v76Jfgz7Q07vX159UGVnE,22572
 unienv_data/third_party/tensordict/memmap_tensor.py,sha256=J6SkFf-FDy43XuaHLgbvDsHt6v2vYfuhRyeoV02P8vw,42589
 unienv_data/transformations/image_compress.py,sha256=f8JTY4DJEXaiu5lO77T4ROV950rh_bOZBchOF-O0tx8,13130
 unienv_interface/__init__.py,sha256=pAWqfm4l7NAssuyXCugIjekSIh05aBbOjNhwsNXcJbE,100
@@ -64,7 +65,7 @@ unienv_interface/space/spaces/text.py,sha256=ePGGJdiD3q-BAX6IHLO7HMe0OH4VrzF043K
 unienv_interface/space/spaces/tuple.py,sha256=mmJab6kl5VtQStyn754pmk0RLPSQW06Mu15Hp3Qad80,4287
 unienv_interface/space/spaces/union.py,sha256=Qisd-DdmPcGRmdhZFGiQw8_AOjYWqkuQ4Hwd-I8tdSI,4375
 unienv_interface/transformations/__init__.py,sha256=zf8NbY-HW4EgHri9PxpuelEvBpFwUtDEcJiXXhFSDNQ,435
-unienv_interface/transformations/batch_and_unbatch.py,sha256=DK-0VNFvoB9FD1qPtq9SNphQnuzRYJrm_L7VuPiLaMY,2145
+unienv_interface/transformations/batch_and_unbatch.py,sha256=LIEQ_rtAdccdw38VdmWJT_DuqdOyb7aMFcMWlyQBz2U,2164
 unienv_interface/transformations/chained_transform.py,sha256=_6E1g_8u-WAxKd-f2sHJwKQk9HTIRnulyXwHUwJP12I,2203
 unienv_interface/transformations/crop.py,sha256=sigcQcLklp3P6b6KQfP-Ja3OV1CWeusCLNKMvNNdACQ,3107
 unienv_interface/transformations/dict_transform.py,sha256=GhFSN9t3mL3gvoD_GH-np68Fo4m78YnSyHbUHeyzKcw,5540
@@ -98,7 +99,7 @@ unienv_interface/wrapper/gym_compat.py,sha256=JhLxDsO1NsJnKzKhO0MqMw9i5_1FLxoxKi
 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.1b6.dist-info/METADATA,sha256=i9zJN3WnVG3gmSSZnKI8SROtMN3U5HGUGRD4p1fu3Ho,3056
-unienv-0.0.1b6.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-unienv-0.0.1b6.dist-info/top_level.txt,sha256=wfcJ5_DruUtOEUZjEyfadaKn7B90hWqz2aw-eM3wX5g,29
-unienv-0.0.1b6.dist-info/RECORD,,
+unienv-0.0.1b8.dist-info/METADATA,sha256=eCut1OMadsIJC5XAqCZHVr9UsfdXqqrumSkH4uuNcTI,3056
+unienv-0.0.1b8.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+unienv-0.0.1b8.dist-info/top_level.txt,sha256=wfcJ5_DruUtOEUZjEyfadaKn7B90hWqz2aw-eM3wX5g,29
+unienv-0.0.1b8.dist-info/RECORD,,

{unienv-0.0.1b6.dist-info → unienv-0.0.1b8.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.1)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

unienv_data/batches/combined_batch.py

@@ -12,6 +12,102 @@ from unienv_interface.space.space_utils import batch_utils as sbu, flatten_utils
 
 from ..base.common import BatchBase, IndexableType, BatchT
 
+__all__ = [
+    "convert_single_index_to_batch",
+    "convert_index_to_batch",
+    "CombinedBatch",
+]
+
+def convert_single_index_to_batch(
+    backend : ComputeBackend,
+    index_starts : BArrayType, # (num_batches, )
+    total_length : int, # Total length of all batches
+    idx : int,
+    device : Optional[BDeviceType] = None,
+) -> Tuple[int, int]:
+    """
+    Convert a single index for this batch to a tuple of:
+        - The index of the batch
+        - The index to index into the batch
+    """
+    assert -total_length <= idx < total_length, f"Index {idx} out of bounds for batch of size {total_length}"
+    if idx < 0:
+        idx += total_length
+    batch_index = int(backend.sum(
+        idx >= index_starts
+    ) - 1)
+    return batch_index, idx - int(index_starts[batch_index])
+
+def convert_index_to_batch(
+    backend : ComputeBackend,
+    index_starts : BArrayType, # (num_batches, )
+    total_length : int, # Total length of all batches
+    idx : Union[IndexableType, BArrayType],
+    device : Optional[BDeviceType] = None,
+) -> Tuple[
+    int, 
+    Union[List[
+        Tuple[int, BArrayType, BArrayType]
+    ], int]
+]:
+    """
+    Convert an index for this batch to a tuple of:
+        - The length of the resulting array
+        - List of tuples, each containing:
+            - The index of the batch
+            - The index to index into the batch
+            - The bool mask to index into the resulting array
+    """
+    if isinstance(idx, slice):
+        idx_array = backend.arange(
+            *idx.indices(total_length),
+            dtype=backend.default_index_dtype,
+            device=device
+        )
+    elif idx is Ellipsis:
+        idx_array = backend.arange(
+            total_length,
+            dtype=backend.default_index_dtype,
+            device=device
+        )
+    elif backend.is_backendarray(idx):
+        assert len(idx.shape) == 1, "Index must be 1D"
+        assert backend.dtype_is_real_integer(idx.dtype) or backend.dtype_is_boolean(idx.dtype), \
+            f"Index must be of integer or boolean type, got {idx.dtype}"
+        if backend.dtype_is_boolean(idx.dtype):
+            assert idx.shape[0] == total_length, f"Boolean index must have the same length as the batch, got {idx.shape[0]} vs {total_length}"
+            idx_array = backend.nonzero(idx)[0]
+        else:
+            assert backend.all(backend.logical_and(-total_length <= idx, idx < total_length)), \
+                f"Index array contains out of bounds indices for batch of size {total_length}"
+            idx_array = (idx + total_length) % total_length
+    else:
+        raise ValueError(f"Invalid index type: {type(idx)}")
+    
+    assert bool(backend.all(
+        backend.logical_and(
+            -total_length <= idx_array,
+            idx_array < total_length
+        )
+    )), f"Index {idx} converted to {idx_array} is out of bounds for batch of size {total_length}"
+    
+    # Convert negative indices to positive indices
+    idx_array = backend.at(idx_array)[idx_array < 0].add(total_length)
+    idx_array_bigger = idx_array[:, None] >= index_starts[None, :] # (idx_array_shape, len(self.batches))
+    idx_array_batch_idx = backend.sum(
+        idx_array_bigger,
+        axis=-1
+    ) - 1 # (idx_array_shape, )
+    
+    result_batch_list = []
+    batch_indexes = backend.unique_values(idx_array_batch_idx)
+    for i in range(batch_indexes.shape[0]):
+        batch_index = int(batch_indexes[i])
+        result_mask = idx_array_batch_idx == batch_index
+        index_into_batch = idx_array[result_mask] - index_starts[batch_index]
+        result_batch_list.append((batch_index, index_into_batch, result_mask))
+    return idx_array.shape[0], result_batch_list
+
 class CombinedBatch(BatchBase[
     BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType
 ]):
@@ -83,18 +179,13 @@ class CombinedBatch(BatchBase[
         return self.index_caches[-1, 1]
 
     def _convert_single_index(self, idx : int) -> Tuple[int, int]:
-        """
-        Convert a single index to a tuple containing
-         - the batch index
-         - the index within the batch
-        """
-        assert -len(self) <= idx < len(self), f"Index {idx} out of bounds for batch of size {len(self)}"
-        if idx < 0:
-            idx += len(self)
-        batch_index = int(self.backend.sum(
-            idx >= self.index_caches[:, 0]
-        ) - 1)
-        return batch_index, idx - int(self.index_caches[batch_index, 0])
+        return convert_single_index_to_batch(
+            self.backend,
+            self.index_caches[:, 0],
+            len(self),
+            idx,
+            device=self.device,
+        )
 
     def _convert_index(self, idx : Union[IndexableType, BArrayType]) -> Tuple[
         int, 
@@ -102,61 +193,13 @@ class CombinedBatch(BatchBase[
             Tuple[int, BArrayType, BArrayType]
         ]
     ]:
-        """
-        Convert an index for this batch to a tuple of:
-         - The length of the resulting array
-         - List of tuples, each containing:
-             - The index of the batch
-             - The index to index into the batch
-             - The bool mask to index into the resulting array
-        """
-        if isinstance(idx, slice):
-            idx_array = self.backend.arange(
-                *idx.indices(len(self)),
-                dtype=self.backend.default_integer_dtype,
-                device=self.device
-            )
-        elif idx is Ellipsis:
-            idx_array = self.backend.arange(
-                len(self),
-                dtype=self.backend.default_integer_dtype,
-                device=self.device
-            )
-        elif self.backend.is_backendarray(idx):
-            assert len(idx.shape) == 1, "Index must be 1D"
-            assert self.backend.dtype_is_real_integer(idx.dtype) or self.backend.dtype_is_boolean(idx.dtype), \
-                f"Index must be of integer or boolean type, got {idx.dtype}"
-            if self.backend.dtype_is_boolean(idx.dtype):
-                assert idx.shape[0] == len(self), f"Boolean index must have the same length as the batch, got {idx.shape[0]} vs {len(self)}"
-                idx_array = self.backend.nonzero(idx)[0]
-            else:
-                idx_array = idx
-        else:
-            raise ValueError(f"Invalid index type: {type(idx)}")
-        
-        assert bool(self.backend.all(
-            self.backend.logical_and(
-                -len(self) <= idx_array,
-                idx_array < len(self)
-            )
-        )), f"Index {idx} converted to {idx_array} is out of bounds for batch of size {len(self)}"
-        
-        # Convert negative indices to positive indices
-        idx_array = self.backend.at(idx_array)[idx_array < 0].add(len(self))
-        idx_array_bigger = idx_array[:, None] >= self.index_caches[None, :, 0] # (idx_array_shape, len(self.batches))
-        idx_array_batch_idx = self.backend.sum(
-            idx_array_bigger,
-            axis=-1
-        ) - 1 # (idx_array_shape, )
-        
-        result_batch_list = []
-        batch_indexes = self.backend.unique_values(idx_array_batch_idx)
-        for i in range(batch_indexes.shape[0]):
-            batch_index = int(batch_indexes[i])
-            result_mask = idx_array_batch_idx == batch_index
-            index_into_batch = idx_array[result_mask] - self.index_caches[batch_index, 0]
-            result_batch_list.append((batch_index, index_into_batch, result_mask))
-        return idx_array.shape[0], result_batch_list
+        return convert_index_to_batch(
+            self.backend,
+            self.index_caches[:, 0],
+            len(self),
+            idx,
+            device=self.device,
+        )
 
     def get_flattened_at(self, idx : Union[IndexableType, BaseException]) -> BArrayType:
         if isinstance(idx, int):

unienv_data/integrations/huggingface.py

@@ -0,0 +1,47 @@
+from typing import Optional
+from datasets import Dataset as HFDataset
+from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.space.space_utils import construct_utils as scu, batch_utils as sbu
+from unienv_data.base import BatchBase, BatchT
+
+__all__ = [
+    'HFAsUniEnvDataset'
+]
+
+class HFAsUniEnvDataset(BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType]):
+    BACKEND_TO_FORMAT_MAP = {
+        "numpy": "numpy",
+        "pytorch": "torch",
+        "jax": "jax",
+    }
+
+    is_mutable = False
+
+    def __init__(
+        self,
+        hf_dataset: HFDataset,
+        backend: ComputeBackend[BArrayType, BDeviceType, BDtypeType, BRNGType],
+        device : Optional[BDeviceType] = None,
+    ) -> None:
+        kwargs = {}
+        if backend.simplified_name != 'numpy':
+            kwargs['device'] = device
+        self.hf_dataset = hf_dataset.with_format(
+            self.BACKEND_TO_FORMAT_MAP[backend.simplified_name],
+            **kwargs
+        )
+        first_data = self.hf_dataset[0]
+        super().__init__(
+            scu.construct_space_from_data(first_data, backend)
+        )
+    
+    def __len__(self) -> int:
+        return len(self.hf_dataset)
+
+    def get_at_with_metadata(self, idx):
+        if self.backend.is_backendarray(idx) and self.backend.dtype_is_boolean(idx):
+            idx = self.backend.nonzero(idx)[0]
+        return self.hf_dataset[idx], {}
+    
+    def get_at(self, idx):
+        return self.get_at_with_metadata(idx)[0]

unienv_data/integrations/pytorch.py

@@ -3,12 +3,19 @@ from typing import Optional, Union, Tuple, Dict, Any
 from torch.utils.data import Dataset
 
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.backends.pytorch import PyTorchComputeBackend, PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType
+from unienv_interface.space.space_utils import construct_utils as scu, batch_utils as sbu
 from unienv_data.base import BatchBase, BatchT
 
-class UniEnvPyTorchDataset(Dataset):
+__all__ = [
+    "UniEnvAsPyTorchDataset",
+    "PyTorchAsUniEnvDataset",
+]
+
+class UniEnvAsPyTorchDataset(Dataset):
     def __init__(
         self, 
-        batch : BatchBase[BatchT, BArrayType, BDeviceType, BDtypeType, BRNGType],
+        batch : BatchBase[BatchT, PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
         include_metadata : bool = False,
     ):
         """
@@ -60,4 +67,56 @@ class UniEnvPyTorchDataset(Dataset):
         indices = self.batch.backend.asarray(indices, dtype=self.batch.backend.default_integer_dtype, device=self.batch.device)
         if self.include_metadata:
             return self.batch.get_at_with_metadata(indices)
-        return self.batch.get_at(indices)
+        return self.batch.get_at(indices)
+
+class PyTorchAsUniEnvDataset(BatchBase[BatchT, PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType]):
+    is_mutable = False
+    def __init__(
+        self,
+        dataset: Dataset,
+    ):
+        """
+        A UniEnvPy BatchBase wrapper for PyTorch Datasets.
+
+        Args:
+            dataset (Dataset): The PyTorch Dataset to wrap.
+        """
+
+        assert len(dataset) >= 0, "The provided PyTorch Dataset must have a defined length."
+        self.dataset = dataset
+        tmp_data = dataset[0]
+        single_space = scu.construct_space_from_data(tmp_data, PyTorchComputeBackend)
+        super().__init__(
+            single_space,
+            None
+        )
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def get_at_with_metadata(self, idx):
+        if isinstance(idx, int):
+            data = self.dataset[idx]
+            return data, {}
+        elif idx is Ellipsis:
+            idx = self.backend.arange(0, len(self))
+        elif isinstance(idx, slice):
+            idx = self.backend.arange(*idx.indices(len(self)))
+        elif self.backend.is_backendarray(idx):
+            if self.backend.dtype_is_boolean(idx.dtype):
+                idx = self.backend.nonzero(idx)[0]
+            else:
+                assert self.backend.dtype_is_real_integer(idx.dtype), "Index array must be of integer or boolean type."
+                idx = (idx + len(self)) % len(self)
+        
+        idx_list = self.backend.to_numpy(idx).tolist()
+        if hasattr(self.dataset, "__getitems__"):
+            data_list = self.dataset.__getitems__(idx_list)
+        else:
+            data_list = [self.dataset[i] for i in idx_list]
+        aggregated_data = sbu.concatenate(self._batched_space, data_list)
+        return aggregated_data, {}
+
+    def get_at(self, idx):
+        data, _ = self.get_at_with_metadata(idx)
+        return data

unienv_data/storages/_episode_storage.py

@@ -188,8 +188,12 @@ class EpisodeStorageBase(SpaceStorage[
                 index = self.backend.arange(0, self.capacity, device=self.device)
             else:
                 index = self.backend.arange(0, self.length, device=self.device)
-        elif self.backend.is_backendarray(index) and self.backend.dtype_is_boolean(index.dtype):
-            index = self.backend.nonzero(index)[0]
+        elif self.backend.is_backendarray(index):
+            if self.backend.dtype_is_boolean(index.dtype):
+                index = self.backend.nonzero(index)[0]
+            else:
+                assert self.backend.dtype_is_real_integer(index.dtype), "Index array must be of integer or boolean type."
+                index = (index + (self.capacity if self.capacity is not None else self.length)) % (self.capacity if self.capacity is not None else self.length)
 
         batch_size = index.shape[0] if self.backend.is_backendarray(index) else 1
 
@@ -368,8 +372,13 @@ class EpisodeStorageBase(SpaceStorage[
                 index = self.backend.arange(0, self.capacity, device=self.device)
             else:
                 index = self.backend.arange(0, self.length, device=self.device)
-        elif self.backend.is_backendarray(index) and self.backend.dtype_is_boolean(index.dtype):
-            index = self.backend.nonzero(index)[0]
+        elif self.backend.is_backendarray(index):
+            if self.backend.dtype_is_boolean(index.dtype):
+                index = self.backend.nonzero(index)[0]
+            else:
+                assert self.backend.dtype_is_real_integer(index.dtype), "Index array must be of integer or boolean type."
+                index = (index + (self.capacity if self.capacity is not None else self.length)) % (self.capacity if self.capacity is not None else self.length)
+
         assert self.backend.is_backendarray(index) and self.backend.dtype_is_real_integer(index.dtype) and len(index.shape) == 1, "Index must be a 1D array of integers"
         sorted_indexes_arg = self.backend.argsort(index)
         sorted_indexes = index[sorted_indexes_arg]

unienv_data/storages/pytorch.py

@@ -72,6 +72,7 @@ class PytorchTensorStorage(SpaceStorage[
         capacity : Optional[int] = None,
         read_only : bool = True,
         multiprocessing : bool = False,
+        **kwargs,
     ) -> "PytorchTensorStorage":
         assert single_instance_space.backend is PyTorchComputeBackend, "PytorchTensorStorage only supports PyTorch backend"
         assert capacity is not None, "Capacity must be specified when creating a new tensor"

unienv_data/storages/video_storage.py

@@ -1,23 +1,239 @@
 from typing import Generic, TypeVar, Generic, Optional, Any, Dict, Tuple, Sequence, Union, List, Iterable, Type, Literal, cast
 from fractions import Fraction
-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.space import BoxSpace
 from unienv_interface.backends import ComputeBackend, BArrayType, BDeviceType, BDtypeType, BRNGType
+from unienv_interface.backends.pytorch import PyTorchComputeBackend
 from unienv_interface.utils.symbol_util import *
 
-from unienv_data.base import SpaceStorage
 from ._episode_storage import IndexableType, EpisodeStorageBase
 
 import numpy as np
 import os
 import json
-import shutil
+import logging
+import importlib
+import importlib.util
 
 import av
-import imageio.v3 as iio
-from imageio.plugins.pyav import PyAVPlugin
-from av.codec.hwaccel import HWAccel, hwdevices_available
+from av.codec.hwaccel import HWAccel as PyAvHWAccel
 from av.codec import codecs_available
+import av.error
+from av.video.reformatter import VideoReformatter
+
+# TorchCodec backend for video encoding / decoding
+# import av
+# from av.codec import codecs_available
+# from torchcodec.decoders import VideoDecoder, set_cuda_backend
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+LOGGER = logging.getLogger(__name__)
+
+class PyAvVideoReader:
+    HWAccel = PyAvHWAccel
+
+    @staticmethod
+    def available_hwdevices() -> List[str]:
+        from av.codec.hwaccel import hwdevices_available
+        return hwdevices_available()
+    
+    @staticmethod
+    def get_auto_hwaccel() -> Optional[HWAccel]:
+        available_hwdevices = __class__.available_hwdevices()
+        target_hwaccel = None
+        if "d3d11va" in available_hwdevices:
+            target_hwaccel = PyAvHWAccel(device_type="d3d11va", allow_software_fallback=True)
+        elif "cuda" in available_hwdevices:
+            target_hwaccel = PyAvHWAccel(device_type="cuda", allow_software_fallback=True)
+        elif "vaapi" in available_hwdevices:
+            target_hwaccel = PyAvHWAccel(device_type="vaapi", allow_software_fallback=True)
+        elif "videotoolbox" in available_hwdevices:
+            target_hwaccel = PyAvHWAccel(device_type="videotoolbox", allow_software_fallback=True)
+        return target_hwaccel
+
+    def __init__(
+        self, 
+        backend : ComputeBackend,
+        filename: str, 
+        buffer_pixel_format : Optional[str] = None,
+        hwaccel: Optional[Union[HWAccel, Literal['auto']]] = None, 
+        seek_mode: Literal['exact', 'approximate'] = 'exact',
+        device : Optional[BDeviceType] = None,
+    ):
+        if seek_mode != 'exact':
+            LOGGER.warning("PyAvVideoReader only supports 'exact' seek mode. Falling back to 'exact'.")
+        
+        if hwaccel == 'auto':
+            hwaccel = __class__.get_auto_hwaccel()
+        self.container = av.open(filename, mode='r', hwaccel=hwaccel)
+        self.video_stream = self.container.streams.video[0]
+        if buffer_pixel_format is not None:
+            self.video_reformatter = VideoReformatter()
+        else:
+            self.video_reformatter = None
+        self.buffer_pixel_format = buffer_pixel_format
+        self.total_frames = self.video_stream.frames # Sometimes this reads 0 for some containers (as they don't explicitly store it)
+        self.frame_iterator = self.container.decode(self.video_stream)
+        self.backend = backend
+        self.device = device
+    
+    def seek(self, frame_index: int):
+        self.container.seek(frame_index, any_frame=True, backward=True, stream=self.video_stream)
+        self.frame_iterator = self.container.decode(self.video_stream)
+    
+    def __next__(self):
+        try:
+            frame = next(self.frame_iterator)
+            if self.video_reformatter is not None:
+                frame = self.video_reformatter.reformat(
+                    frame,
+                    format=self.buffer_pixel_format,
+                )
+            frame = frame.to_ndarray(channel_last=True)
+        except av.error.EOFError:
+            raise StopIteration
+        if np.prod(frame.shape) == 0:
+            raise StopIteration
+    
+        return frame
+
+    def __iter__(self):
+        return self
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, *args, **kwargs):
+        self.container.close()
+    
+    def read(self, index : Union[IndexableType, BArrayType], total_length : int) -> BArrayType:
+        if isinstance(index, int):
+            self.seek(index)
+            frame_np = next(self)
+            frame = self.backend.from_numpy(frame_np)
+            if self.device is not None:
+                frame = self.backend.to_device(frame, self.device)
+            return frame
+        else:
+            if index is Ellipsis:
+                index = np.arange(total_length)
+            elif isinstance(index, slice):
+                index = np.arange(*index.indices(total_length))
+            elif self.backend.is_backendarray(index) and self.backend.dtype_is_boolean(index.dtype):
+                index = self.backend.nonzero(index)[0]
+            if self.backend.is_backendarray(index):
+                index = self.backend.to_numpy(index)
+
+            argsorted_indices = np.argsort(index)
+            sorted_index = index[argsorted_indices]
+            reserve_index = np.argsort(argsorted_indices)
+            
+            if len(index) < total_length // 2:
+                all_frames_np = []
+                past_frame_np = None
+                for frame_i in sorted_index:
+                    self.seek(frame_i)
+                    
+                    try:
+                        frame_np = next(self)
+                    except StopIteration:
+                        frame_np = past_frame_np
+                    past_frame_np = frame_np
+                    all_frames_np.append(frame_np)
+                    
+                all_frames_np = np.stack(all_frames_np, axis=0)
+                # Reorder from sorted order back to original order
+                all_frames_np = all_frames_np[reserve_index]
+            else:
+                # Create a set for O(1) lookup and a mapping from frame index to position in sorted_index
+                sorted_index_set = set(sorted_index)
+                frame_to_sorted_pos = {int(frame_idx): pos for pos, frame_idx in enumerate(sorted_index)}
+                
+                # Pre-allocate array to store frames in sorted order
+                all_frames_list = [None] * len(sorted_index)
+                past_frame_np = None
+                self.seek(0)
+                for frame_i in range(total_length):
+                    try:
+                        frame_np = next(self)
+                    except StopIteration:
+                        frame_np = past_frame_np
+                    if frame_i in sorted_index_set:
+                        # Store at the position corresponding to sorted_index order
+                        all_frames_list[frame_to_sorted_pos[frame_i]] = frame_np
+                    past_frame_np = frame_np
+                all_frames_np = np.stack(all_frames_list, axis=0)
+                # Reorder from sorted order back to original order
+                all_frames_np = all_frames_np[reserve_index]
+            all_frames = self.backend.from_numpy(all_frames_np)
+            if self.device is not None:
+                all_frames = self.backend.to_device(all_frames, self.device)
+            
+            return all_frames
+
+class TorchCodecVideoReader:
+    HWAccel = Literal['beta', 'ffmpeg']
+    def __init__(
+        self, 
+        backend : ComputeBackend,
+        filename: str, 
+        buffer_pixel_format : Optional[str] = None,
+        hwaccel: Optional[Union[HWAccel, Literal['auto']]] = None, 
+        seek_mode: Literal['exact', 'approximate'] = 'exact',
+        device : Optional[BDeviceType] = None,
+    ):
+        from torchcodec.decoders import VideoDecoder, set_cuda_backend
+        assert torch is not None, "TorchCodecVideoReader requires PyTorch and TorchCodec to be installed."
+        assert buffer_pixel_format == 'rgb24', "TorchCodecVideoReader currently only supports 'rgb24' buffer pixel format."
+        if hwaccel == 'auto':
+            hwaccel = __class__.get_auto_hwaccel()
+        if hwaccel is not None:
+            with set_cuda_backend(hwaccel):
+                self.decoder = VideoDecoder(filename, device='cuda', seek_mode=seek_mode)
+        else:
+            self.decoder = VideoDecoder(filename, seek_mode=seek_mode)
+        self.backend = backend
+        self.device = device
+        self.buffer_pixel_format = buffer_pixel_format
+
+    @staticmethod
+    def get_auto_hwaccel() -> Optional[HWAccel]:
+        assert torch is not None, "TorchCodecVideoReader requires PyTorch to be installed."
+        if torch.cuda.is_available():
+            return 'beta'
+        else:
+            return None
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, *args, **kwargs):
+        pass
+    
+    def read(self, index : Union[IndexableType, np.ndarray], total_length : int) -> np.ndarray:
+        if isinstance(index, int):
+            ret = self.decoder.get_frame_at(index).data.permute(1, 2, 0) # (C, H, W) -> (H, W, C)
+        elif index is Ellipsis:
+            ret = self.decoder.get_frames_in_range(0, len(self.decoder)).data.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        elif isinstance(index, slice):
+            start, stop, step = index.indices(len(self))
+            ret = self.decoder.get_frames_in_range(start, stop, step=step).data.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        else:
+            if self.backend.is_backendarray(index) and self.backend.dtype_is_boolean(index.dtype):
+                index = self.backend.nonzero(index)[0]
+            if self.backend.simplified_name != 'pytorch':
+                index = PyTorchComputeBackend.from_other_backend(self.backend, index).to('cpu', torch.int64)
+            ret = self.decoder.get_frames_at(index).data.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        
+        if self.backend.simplified_name != 'pytorch':
+            ret = self.backend.from_other_backend(PyTorchComputeBackend, ret)
+        if self.device is not None:
+            ret = self.backend.to_device(ret, self.device)
+        return ret
 
 class VideoStorage(EpisodeStorageBase[
     BArrayType,
@@ -30,7 +246,7 @@ class VideoStorage(EpisodeStorageBase[
     A storage for RGB or depth video data using video files
     If encoding RGB video
     - Set `buffer_pixel_format` to `rgb24`
-        - Set `file_pixel_format` to `None`
+        - Set `file_pixel_format` to `None` (especially when running with nvenc codec)
         - Set `file_ext` to anything you like (e.g., "mp4", "avi", "mkv", etc.)
     If encoding depth video
         - Set `buffer_pixel_format` to `gray16le` (You can use rescale transform inside a `TransformedStorage` to convert depth values to this format, where `dtype` should be `np.uint16`) - if in meters, set min to 0 and max to 65.535 as the multiplication factor is 1000 (i.e., depth in mm)
@@ -40,12 +256,15 @@ class VideoStorage(EpisodeStorageBase[
     """
 
     # ========== Class Attributes ==========
+    PyAV_LOG_LEVEL = av.logging.WARNING
+
     @classmethod
     def create(
         cls,
         single_instance_space: BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
         *args,
-        hardware_acceleration : Optional[Union[HWAccel, Literal['auto']]] = 'auto',
+        seek_mode : Literal['exact', 'approximate'] = 'exact',
+        hardware_acceleration : Optional[Union[Any, Literal['auto']]] = 'auto',
         codec : Union[str, Literal['auto']] = 'auto',
         file_ext : str = "mp4",
         file_pixel_format : Optional[str] = None,
@@ -64,6 +283,7 @@ class VideoStorage(EpisodeStorageBase[
         return VideoStorage(
             single_instance_space,
             cache_filename=cache_path,
+            seek_mode=seek_mode,
             hardware_acceleration=hardware_acceleration,
             codec=codec,
             file_ext=file_ext,
@@ -79,7 +299,8 @@ class VideoStorage(EpisodeStorageBase[
         path : Union[str, os.PathLike],
         single_instance_space : BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
         *,
-        hardware_acceleration : Optional[Union[HWAccel, Literal['auto']]] = 'auto',
+        seek_mode : Literal['exact', 'approximate'] = 'exact',
+        hardware_acceleration : Optional[Union[Any, Literal['auto']]] = 'auto',
         codec : Union[str, Literal['auto']] = 'auto',
         capacity : Optional[int] = None,
         read_only : bool = True,
@@ -113,6 +334,7 @@ class VideoStorage(EpisodeStorageBase[
         return VideoStorage(
             single_instance_space,
             cache_filename=path,
+            seek_mode=seek_mode,
             hardware_acceleration=hardware_acceleration,
             codec=codec,
             file_ext=file_ext,
@@ -127,31 +349,20 @@ class VideoStorage(EpisodeStorageBase[
     # ========== Instance Implementations ==========
     single_file_ext = None
 
-    @staticmethod
-    def get_auto_hwaccel() -> Optional[HWAccel]:
-        if hasattr(__class__, "_auto_hwaccel"):
-            return __class__._auto_hwaccel
-        available_hwdevices = hwdevices_available()
-        target_hwaccel = None
-        if "d3d11va" in available_hwdevices:
-            target_hwaccel = HWAccel(device_type="d3d11va", allow_software_fallback=True)
-        elif "cuda" in available_hwdevices:
-            target_hwaccel = HWAccel(device_type="cuda", allow_software_fallback=True)
-        elif "vaapi" in available_hwdevices:
-            target_hwaccel = HWAccel(device_type="vaapi", allow_software_fallback=True)
-        elif "videotoolbox" in available_hwdevices:
-            target_hwaccel = HWAccel(device_type="videotoolbox", allow_software_fallback=True)
-        __class__._auto_hwaccel = target_hwaccel
-        return target_hwaccel
-
     @staticmethod
     def get_auto_codec(
         base : Optional[str] = None
     ) -> str:
         if hasattr(__class__, "_auto_codec"):
             return __class__._auto_codec
+        
+        # --- PyAV Implementation ---
         preferred_codecs = ["av1", "hevc", "h264", "mpeg4", "vp9", "vp8"] if base is None else [base]
-        preferred_suffixes = ["_nvenc", "_amf", "_qsv"]
+        preferred_suffixes = ["_nvenc", "_vaapi", "_amf", "_qsv", "_videotoolbox"]
+
+        # --- TorchCodec Implementation ---
+        # preferred_codecs = ["hevc", "av1", "h264", "mpeg4", "vp9", "vp8"] if base is None else [base]
+        # preferred_suffixes = ["_nvenc"]
         
         target_codec = None
         for codec in preferred_codecs:
@@ -177,8 +388,10 @@ class VideoStorage(EpisodeStorageBase[
         self,
         single_instance_space: BoxSpace[BArrayType, BDeviceType, BDtypeType, BRNGType],
         cache_filename : Union[str, os.PathLike],
-        hardware_acceleration : Optional[Union[HWAccel, Literal['auto']]] = 'auto',
+        seek_mode : Literal['exact', 'approximate'] = 'exact',
+        hardware_acceleration : Optional[Union[Any, Literal['auto']]] = 'auto',
         codec : Union[str, Literal['auto']] = 'auto',
+        decode_backend : Literal['torchcodec', 'pyav', 'auto'] = 'auto',
         file_ext : str = "mp4",
         file_pixel_format : Optional[str] = None,
         buffer_pixel_format : str = "rgb24",
@@ -195,87 +408,112 @@ class VideoStorage(EpisodeStorageBase[
             capacity=capacity,
             length=length,
         )
-        self.hwaccel = None if hardware_acceleration is None else (
-            self.get_auto_hwaccel() if hardware_acceleration == 'auto' else hardware_acceleration
-        )
+        self.seek_mode = seek_mode
+        self.hwaccel = hardware_acceleration
         self.codec = self.get_auto_codec() if codec == 'auto' else codec
+        
+        if decode_backend == 'auto':
+            if importlib.util.find_spec("torchcodec"):
+                decode_backend = 'torchcodec'
+            else:
+                decode_backend = 'pyav'
+        self.decode_backend = decode_backend
+
         self.fps = fps
         self.file_pixel_format = file_pixel_format
         self.buffer_pixel_format = buffer_pixel_format
-
+        
     def get_from_file(self, filename : str, index : Union[IndexableType, BArrayType], total_length : int) -> BArrayType:
-        with iio.imopen(filename, 'r', plugin='pyav', hwaccel=self.hwaccel) as video:
-            video = cast(PyAVPlugin, video)
-            if isinstance(index, int):
-                frame_np = video.read(index=index, format=self.buffer_pixel_format)
-                frame = self.backend.from_numpy(frame_np)
-                if self.device is not None:
-                    frame = self.backend.to_device(frame, self.device)
-                return frame
-            else:
-                if index is Ellipsis:
-                    index = np.arange(total_length)
-                elif isinstance(index, slice):
-                    index = np.arange(*index.indices(total_length))
-                elif self.backend.is_backendarray(index) and self.backend.dtype_is_boolean(index.dtype):
-                    index = self.backend.nonzero(index)[0]
-                if self.backend.is_backendarray(index):
-                    index = self.backend.to_numpy(index)
-
-                argsorted_indices = np.argsort(index)
-                sorted_index = index[argsorted_indices]
-                reserve_index = np.argsort(argsorted_indices)
-                
-                if len(index) < total_length // 2:
-                    all_frames_np = []
-                    for frame_i in sorted_index:
-                        frame_np = video.read(index=frame_i, format=self.buffer_pixel_format)
-                        all_frames_np.append(frame_np)
-                    all_frames_np = np.stack(all_frames_np, axis=0)
-                    # Reorder from sorted order back to original order
-                    all_frames_np = all_frames_np[reserve_index]
-                else:
-                    # Create a set for O(1) lookup and a mapping from frame index to position in sorted_index
-                    sorted_index_set = set(sorted_index)
-                    frame_to_sorted_pos = {int(frame_idx): pos for pos, frame_idx in enumerate(sorted_index)}
-                    
-                    # Pre-allocate array to store frames in sorted order
-                    all_frames_list = [None] * len(sorted_index)
-                    past_frame_np = None
-                    video_iter = video.iter(format=self.buffer_pixel_format)
-                    for frame_i in range(total_length):
-                        try:
-                            frame_np = next(video_iter)
-                        except StopIteration:
-                            frame_np = past_frame_np
-                        if frame_i in sorted_index_set:
-                            # Store at the position corresponding to sorted_index order
-                            all_frames_list[frame_to_sorted_pos[frame_i]] = frame_np
-                        past_frame_np = frame_np
-                    all_frames_np = np.stack(all_frames_list, axis=0)
-                    # Reorder from sorted order back to original order
-                    all_frames_np = all_frames_np[reserve_index]
-                all_frames = self.backend.from_numpy(all_frames_np)
-                if self.device is not None:
-                    all_frames = self.backend.to_device(all_frames, self.device)
-                return all_frames
+        if self.decode_backend == 'pyav':
+            reader_cls = PyAvVideoReader
+        elif self.decode_backend == 'torchcodec':
+            reader_cls = TorchCodecVideoReader
+        else:
+            raise ValueError(f"Unknown decode_backend {self.decode_backend}")
+        with reader_cls(
+            backend=self.backend,
+            filename=filename,
+            buffer_pixel_format=self.buffer_pixel_format,
+            hwaccel=self.hwaccel,
+            seek_mode=self.seek_mode,
+            device=self.device,
+        ) as video_reader:
+            return video_reader.read(index, total_length)
 
+    # PyAV Implementation (Commented out due to bugs with hevc codec)
     def set_to_file(self, filename : str, value : BArrayType):
-        value_np = self.backend.to_numpy(value)
-        with iio.imopen(
-            filename,
-            'w',
-            plugin='pyav',
-        ) as video:
-            video = cast(PyAVPlugin, video)
-            video.init_video_stream(self.codec, fps=self.fps, pixel_format=self.file_pixel_format)
+        # ImageIO Implementation
+        # with iio.imopen(
+        #     filename,
+        #     'w',
+        #     plugin='pyav',
+        # ) as video:
+        #     video = cast(PyAVPlugin, video)
+        #     video.init_video_stream(self.codec, fps=self.fps, pixel_format=self.file_pixel_format)
+            
+        #     # Fix codec time base if not set:
+        #     if video._video_stream.codec_context.time_base is None:
+        #         video._video_stream.codec_context.time_base = Fraction(1 / self.fps).limit_denominator(int(2**16 - 1))
+
+        #     for i, frame in enumerate(value_np):
+        #         video.write_frame(frame, pixel_format=self.buffer_pixel_format)
+
+        # PyAV Implementation
+        logging.getLogger("libav").setLevel(self.PyAV_LOG_LEVEL)
+        with av.open(filename, mode='w') as container:
+            output_stream = container.add_stream(self.codec, rate=self.fps)
+            if len(self.single_instance_space.shape) == 3: # (H, W, C)
+                output_stream.width = self.single_instance_space.shape[1]
+                output_stream.height = self.single_instance_space.shape[0]
+            else: # (H, W)
+                output_stream.width = self.single_instance_space.shape[1]
+                output_stream.height = self.single_instance_space.shape[0]
+            if self.file_pixel_format is not None:
+                output_stream.pix_fmt = self.file_pixel_format
+            # output_stream.time_base = Fraction(1, self.fps).limit_denominator(int(2**16 - 1))
+            value_np = self.backend.to_numpy(value)
+            for i, frame_np in enumerate(value_np):
+                frame = av.VideoFrame.from_ndarray(frame_np, format=self.buffer_pixel_format, channel_last=True)
+                packets = output_stream.encode(frame)
+                if packets:
+                    container.mux(packets)
+            # Flush stream
+            packets = output_stream.encode()
+            if packets:
+                container.mux(packets)
             
-            # Fix codec time base if not set:
-            if video._video_stream.codec_context.time_base is None:
-                video._video_stream.codec_context.time_base = Fraction(1 / self.fps).limit_denominator(int(2**16 - 1))
+            # Close container
+            if hasattr(output_stream, 'close'):
+                output_stream.close()
+            # container.close()
+
+            # Restore logging level
+            av.logging.restore_default_callback()
 
-            for i, frame in enumerate(value_np):
-                video.write_frame(frame, pixel_format=self.buffer_pixel_format)
+    # TorchCodec Implementation (extremely slow, and seems to have double-write issues)
+    # def set_to_file(self, filename : str, value : BArrayType):
+    #     if self.backend.simplified_name != 'pytorch':
+    #         value_pt = PyTorchComputeBackend.from_other_backend(self.backend, value)
+    #     else:
+    #         value_pt = value
+    #     if self.codec.endswith('_nvenc'):
+    #         value_pt = value_pt.to("cuda")
+    #     else:
+    #         value_pt = value_pt.to("cpu")
+        
+    #     if len(self.single_instance_space.shape) < 3:
+    #         # Add channel dimension for grayscale images
+    #         value_pt = value_pt[..., None]  # (N, H, W) -> (N, H, W, 1)
+        
+    #     encoder = VideoEncoder(
+    #         value_pt.permute(0, 3, 1, 2), # (N, H, W, C) -> (N, C, H, W)
+    #         frame_rate=self.fps,
+    #     )
+    #     encoder.to_file(
+    #         filename,
+    #         codec=self.codec,
+    #         pixel_format=self.file_pixel_format,
+    #     )
 
     def dumps(self, path):
         assert os.path.samefile(path, self.cache_filename), \

unienv_interface/transformations/batch_and_unbatch.py

@@ -21,6 +21,7 @@ class BatchifyTransformation(DataTransformation):
                 0,
                 self.axis
             )
+        return ret
 
     def transform(self, source_space, data):
         return sbu.concatenate(