unienv 0.0.1b7__tar.gz → 0.0.1b8__tar.gz

{unienv-0.0.1b7/unienv.egg-info → unienv-0.0.1b8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b7
+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.1b7 → unienv-0.0.1b8}/pyproject.toml

@@ -3,7 +3,7 @@ name = "unienv"
 description = "Unified robot environment framework supporting multiple tensor and simulation backends"
 readme = "README.md"
 license = "MIT"
-version = "0.0.1b7"
+version = "0.0.1b8"
 requires-python = ">= 3.10"
 dependencies = [
     "numpy",

{unienv-0.0.1b7 → unienv-0.0.1b8/unienv.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unienv
-Version: 0.0.1b7
+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.1b7 → unienv-0.0.1b8}/unienv.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ 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/huggingface.py
 unienv_data/integrations/pytorch.py
 unienv_data/replay_buffer/__init__.py
 unienv_data/replay_buffer/replay_buffer.py

{unienv-0.0.1b7 → unienv-0.0.1b8}/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-0.0.1b8/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-0.0.1b8/unienv_data/integrations/pytorch.py

@@ -0,0 +1,122 @@
+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
+
+__all__ = [
+    "UniEnvAsPyTorchDataset",
+    "PyTorchAsUniEnvDataset",
+]
+
+class UniEnvAsPyTorchDataset(Dataset):
+    def __init__(
+        self, 
+        batch : BatchBase[BatchT, PyTorchArrayType, PyTorchDeviceType, PyTorchDtypeType, PyTorchRNGType],
+        include_metadata : bool = False,
+    ):
+        """
+        A PyTorch Dataset wrapper for UniEnvPy batches.
+        Note that UniEnv's `BatchBase` will automatically collate data when indexed with batches, and therefore in the dataloader you can set `collate_fn=None`.
+        
+        Args:
+            batch (BatchBase): The UniEnvPy batch to wrap.
+        """
+        self.batch = batch
+        self.include_metadata = include_metadata
+
+    def __len__(self) -> int:
+        """
+        Get the length of the dataset.
+
+        Returns:
+            int: The number of items in the dataset.
+        """
+        return len(self.batch)
+
+    def __getitem__(self, index) -> Union[BatchT, Tuple[BatchT, Dict[str, Any]]]:
+        """
+        Get an item from the dataset.
+
+        Args:
+            index (int): The index of the item to retrieve.
+
+        Returns:
+            Union[BatchT, Tuple[BatchT, Dict[str, Any]]]: The batch data at the specified index,
+            optionally including metadata if `include_metadata` is True.
+        """
+        assert isinstance(index, int), "Index must be an integer."
+        if self.include_metadata:
+            return self.batch.get_at_with_metadata(index)
+        return self.batch.get_at(index)
+
+    def __getitems__(self, indices: list[int]) -> list[Union[BatchT, Tuple[BatchT, Dict[str, Any]]]]:
+        """
+        Get multiple items from the dataset.
+
+        Args:
+            indices (list[int]): The indices of the items to retrieve.
+
+        Returns:
+            Union[BatchT, Tuple[BatchT, Dict[str, Any]]]: Batch data at the specified indices, 
+            optionally including metadata if `include_metadata` is True.
+        """
+        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)
+
+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-0.0.1b7 → unienv-0.0.1b8}/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-0.0.1b7 → unienv-0.0.1b8}/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"