Repository: Genesis-Embodied-AI/Genesis
File: genesis/utils/ring_buffer.py (Python)

Code before:
        """
        Copy the tensor into the next position of the ring buffer, and advance the index pointer.

        Parameters
        ----------
        tensor : torch.Tensor
            The tensor to copy into the ring buffer.
        """
        self.buffer[self._idx_ptr.value].copy_(tensor)
        self._idx_ptr.value = (self._idx_ptr.value + 1) % self.N

    def at(self, idx: int | torch.Tensor) -> torch.Tensor:
        """
        Get a view of the tensor at the given index.

        Parameters
        ----------
        idx : int | torch.Tensor
            Index of the element to get, where 0 is the latest element, 1 is the second latest, etc.
            Can be a tensor for batched indexing.
        """
        rel_idx = (self._idx_ptr.value - idx) % self.N
        if isinstance(idx, torch.Tensor):
            batch_size = len(idx) if idx.ndim > 0 else 1
            batch_indices = torch.arange(batch_size, device=gs.device, dtype=gs.tc_int)
            return self.buffer[rel_idx, batch_indices]
        return self.buffer[rel_idx]

    def get(self, idx: int) -> torch.Tensor:
        """
        Get a clone of the tensor at the given index.

        Parameters
        ----------
        idx : int
            Index of the element to get, where 0 is the latest element, 1 is the second latest, etc.
        """
        return self.at(idx).clone()

    def clone(self) -> "TensorRingBuffer":
        return TensorRingBuffer(
            self.N,
            self.buffer.shape[1:],
            dtype=self.buffer.dtype,
            buffer=self.buffer.clone(),
            idx_ptr=self._idx_ptr,
        )

    def __getitem__(self, key: int | slice | tuple) -> "TensorRingBuffer":
        """
        Enable slicing of the tensor ring buffer.

Review comment (performance):
I need to figure this out still.. It works when i do arange.. When using :batch_size and reshapes or squeeze(1) or having ifs i could not get it to work for both batched/unbatched case. I was tired when I tried last night though. Help is appreciated 🥲

Code after:
    def append(self, tensor: torch.Tensor):
        """
        Copy the tensor into the next position of the ring buffer, and advance the index pointer.

        Parameters
        ----------
        tensor : torch.Tensor
            The tensor to copy into the ring buffer.
        """
        self.buffer[self._idx].copy_(tensor)
        self._idx[()] = (self._idx + 1) % self.N

    def at(
        self, idx: int | torch.Tensor, *others_idx: int | slice | torch.Tensor, copy: bool | None = None
    ) -> torch.Tensor:
        """
        Get the value of the tensor at the given index.

        Parameters
        ----------
        idx : int | torch.Tensor
            Index of the element to get from most recent to least recent (that has not been discarded yet).
            Passing a 1D tensor for advanced (aka fancy) indexing is supported, but this is requiring allocating fresh
            memory instead of returning a view, which is less efficient.
        others_idx : int | slice | torch.Tensor, optional
            Index of the elements to extract from the selected tensor. In case of advanced indexing, this is equivalent
            but significantly more efficient than doing this extraction in a latter stage.
        copy: bool | None, optional
            If `None`, then memory will be allocated only if necessary. If `True`, then memory will be allocated
            systematically instead of returning a view. If `False`, then allocating memory is forbidden and will raise
            an exception if returning a view is impossible.
        """
        rel_idx = (self._idx - idx) % self.N
        assert len(others_idx) < self.buffer.ndim
        tensor = self.buffer[(rel_idx, *others_idx)]
        if tensor.untyped_storage().data_ptr() == self.buffer.untyped_storage().data_ptr():
            if copy:
                tensor = tensor.clone()
        elif copy == False:
            raise gs.GenesisException("Allocating memory is necessary but 'copy=False'.")
        return tensor

    def get(self, idx: int) -> torch.Tensor:
        """
        Get a clone of the tensor at the given index.

        Parameters
        ----------
        idx : int
            Index of the element to get from most recent to least recent (that has not been discarded yet).
        """