Variable workspace tensor sizes

[smorad]

Currently, we are unable to do the following:

from salina import Workspace
ws = Workspace()
batch_size = 5
ws.set("obs", 0, torch.zeros(batch_size, 3))
ws.set("obs", 1, torch.zeros(batch_size, 5))

due to https://github.com/facebookresearch/salina/blob/10d09bb80f78e05ddd7de58e9e24ff0f302877fb/salina/workspace.py#L45

Since tensors from sequential timesteps are stored as lists, I suspect variable-sized features should be possible. This would be immensely useful in multiagent systems as well as recurrent models (e.g. building/expanding a map as the agent explores).

[ludc]

A workspace is storing tensors of size TxBx[...] i.e T timesteps, B.= batch_size and [...] is the size of the information you want to store (for instance [...]=3x28x28 for atari games). The Workspace.set function is just a way to fill this TxBx[...] tensor over a temporal slice (and possibly to automatically adapt the size of the tensor if this corresponds to a new temporal index). What is the use-case where you want to write tensors of different sizes at different timesteps in the same variable ?

What is the use case that would need to have variable observation sizes ?

[ludc]

Ok, I think I better understand your point in the multiagent case. If the 'assert' is removed, the problem is that the workspace["obs"] function will not work since it concatenates the list of tensors to return a TxBx[...] tensor. It will have consequences on the RemoteAgent for instances (that are using TxBx[...] tensors in shared memory) and certainly other consequences...

I see two options:

• the first one would be to have 'larger' tensors (for instance workspace["obs"} and workspace["obs_size"]) but it will consume more memory
• the second one would be to remove the assert, but to take care not to use any method that needs 'full' tensors

[smorad]

I am specifically interested with using graph representations. Unfortunately, an adjacency matrix of say 50,000 by 50,000 takes 10GB of memory, even if mostly empty. Using a representation like torch.sparse_coo_tensor is significantly more efficient, but under the hood it holds variable-length tensors of indices and values.

Perhaps I could create a separate workspace to place these "non-full" tensors, so I could still use full tensor methods elsewhere.

[smorad]

I think other potential reasons for variable-sized tensors could be additional agents coming online or going offline throughout an episode. We could mask them using a fixed-size tensor, but if agents are short-lived or we do not know how many agents we will have ahead of time, then variable-sized tensors could be beneficial.

[ludc]

Ok, the Slice tensor in workspace can handle variable-sized tensors. So the adaptation just needs to prevent the user to merge sliced tensor if there are not composed of tensors of the same size. I put it on the todo list and will try to come back with a solution in the next days. In the meantime, if you remove the assert, does it works for your use case ?

[smorad]

It currently crashes when removing the assert, but I think with your fix it will be perfect!

[ludc]

Ok, let me take a look by the end of the weed-end

[smorad]

Although variable-sized tensors do not work, it seems everything "just works" with the torch.sparse_coo class! Any variable-length tensor should be realisable as a sparse tensor! For example, we can store the variable length tensor var_len_tens as follows using the dummy idx tensor:

var_len_tens = torch.tensor([1, 4, 9, 16])
idx = torch.tensor([0]).expand(1, var_len_tens.shape[0]) #Expand is zero-copy, so this doesn't take extra memory
torch.sparse_coo_tensor(indices=idx, values=var_len_tens, size=(int(1e10),))

>>> tensor(indices=tensor([[0, 0, 0, 0]]),
       values=tensor([ 1,  4,  9, 16]),
       size=(10000000000,), nnz=4, layout=torch.sparse_coo)

If we modify my original example, the following works without removing the assert:

import torch
from salina import Workspace
ws = Workspace()
batch_size = 5

s0 = torch.sparse_coo_tensor(indices=[[0, 1, 2], [0, 1, 2]], values=[1, 1, 1], size=(batch_size, 1000))
s1 = torch.sparse_coo_tensor(indices=[[3, 4], [3, 4]], values=[2, 2], size=(batch_size, 1000))

ws.set("obs", 0, s0)
ws.set("obs", 1, s1)
ws["obs"]

>>> tensor(indices=tensor([[0, 0, 0, 1, 1],
                       [0, 1, 2, 3, 4],
                       [0, 1, 2, 3, 4]]),
       values=tensor([1, 1, 1, 2, 2]),
       size=(2, 5, 1000), nnz=5, layout=torch.sparse_coo)