[BUG] KeyError for Inverse Transform in ReplayBuffer

Describe the bug

The inverse transforms on the collected data in the replay buffer are executed on the nested tensordict with keys "_data" and "index". When passing to the inverse transforms data = self._transform.inv(data) we get a KeyError:
KeyError: 'key "(\'next\', \'done\')" not found in TensorDict with keys [\'_data\', \'index\']'.

Changing the line to data.set("_data", self._transform.inv(data.get("_data"))) solves the issue. Following works as well: data = self._transform.inv(data.get("_data"))
However, the second example now does not have the nested tensordict with keys "_data" and "index". I suppose those are important even though the script runs without errors.

@vmoens What do you think? elif case hints that we indeed only want the data and not the nested tensordict with "_data" and "index" keys.

    def extend(self, data: Sequence) -> torch.Tensor:
        """Extends the replay buffer with one or more elements contained in an iterable.

        Args:
            data (iterable): collection of data to be added to the replay
                buffer.

        Returns:
            Indices of the data added to the replay buffer.
        """
        if self._transform is not None and is_tensor_collection(data):
            data = self._transform.inv(data)
        elif self._transform is not None and len(self._transform):
            # Accepts transforms that act on "data" key
            data = self._transform.inv(TensorDict({"data": data}, [])).get("data")
        with self._replay_lock:
            index = self._writer.extend(data)
            self._sampler.extend(index)
        return index

To Reproduce

Steps to reproduce the behavior.

Basically, any example script where you add an inverse transform to the replay buffer like Reward2Go.

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)