[Bug] Indexing issue in memory sampling function

[ADebor]

Hi @Toni-SM ,

First of all, thank you for this excellent, well-documented library!

I might be off here, but when playing with the sample method of the RandomMemory class, I've encountered the following PyTorch error: IndexError: too many indices for tensor of dimension 2 raised in the sample_by_index method of the Memory class.

This error comes from trying to index self.tensors_view[name] with a list of tensors, i.e. batch when indexes is of type torch.Tensor. When indexes is a list, it works fine.

A quick fix is to return [[self.tensors_view[name][[batch]] for name in names] for batch in batches] instead of [[self.tensors_view[name][batch] for name in names] for batch in batches] when indexes is of type torch.Tensor.

I hope this issue doesn't come from my end (I apologize in advance if that's the case), as I may be using these methods wrong.

[Toni-SM]

Hi @ADebor

Thank you very much for your words and for giving skrl a chance.

Indexes, in sample_by_index method, are expected to have only 1 dimension. I have tested the memory successfully using the following code for torch tensor, numpy array and python list

import torch
from skrl.memories.torch import RandomMemory

memory_size = 100
num_envs = 10

memory = RandomMemory(memory_size=memory_size, num_envs=num_envs)

# create tensors
memory.create_tensor(name="tensor_a", size=2)
memory.create_tensor(name="tensor_b", size=3)
memory.create_tensor(name="tensor_c", size=4)

# tests
memory.filled = True  # change status for testing

assert memory.get_tensor_names() == sorted(["tensor_a", "tensor_b", "tensor_c"])
assert len(memory) == int(memory_size * num_envs)

indexes = torch.randint(high=len(memory), size=(10,))

# indexes: torch tensor
samples0 = memory.sample_by_index(names=["tensor_a", "tensor_b", "tensor_c"], indexes=indexes)
# indexes: numpy array
samples1 = memory.sample_by_index(names=["tensor_a", "tensor_b", "tensor_c"], indexes=indexes.cpu().numpy())
# indexes: python list
samples2 = memory.sample_by_index(names=["tensor_a", "tensor_b", "tensor_c"], indexes=indexes.tolist())

assert len(samples0[0]) == len(samples1[0]) == len(samples2[0])
print(samples0[0])

---

For indexes with higher dimensions (e.g. indexes = torch.randint(high=len(memory), size=(10,1))) it will raise the following error:

Traceback (most recent call last):
  File "test_sampling.py", line 25, in <module>
    samples1 = memory.sample_by_index(names=["tensor_a", "tensor_b", "tensor_c"], indexes=indexes.cpu().numpy())
  File "/home/toni/Documents/SKRL/skrl/skrl/memories/torch/base.py", line 349, in sample_by_index
    return [[self.tensors_view[name][indexes] for name in names]]
  File "/home/toni/Documents/SKRL/skrl/skrl/memories/torch/base.py", line 349, in <listcomp>
    return [[self.tensors_view[name][indexes] for name in names]]
IndexError: too many indices for tensor of dimension 2

---

I am updating the docstring of the library components gradually and adding more useful examples and descriptions that help to disambiguate the invocations.

Please let me know if using one-dimensional indexes works for you. Or if you still need to use n-dimensional indexes in your setup.

[Toni-SM]

Close the topic due to no activity. Please reopen if still unresolved. If you have any additional comments you can always open a new discussion.