Treat a tuple/list of Tensors as a single Tensor

Question

I'm using Pytorch for some robotics Reinforcement Learning tasks. I'd like to use both images and information about the state as observations for this task. The implementation I'm using does not directly support this so I'm making some amendments. Expected observations are either state, as a 1 dimensional Tensor, or images as a 3 dimensional Tensor (channels, width, height). In my task I would like the observation to be a tuple of Tensors.

In many places in my codebase, the observation is of course expected to be a single Tensor, not a tuple of Tensors. Is there an easy way to treat a tuple of Tensors as a single Tensor?

For example, I would like:

observation.to(device)

to work as normal when observation is a single Tensor, and call .to(device) on each Tensor when observation is a tuple of Tensors.

It should be simple enough to create a data type that can support this, but I'm wondering does such a data type already exist? I haven't found anything so far.

Answer 1

If your tensors are all of the same size, you can use torch.stack to concatenate them into one tensor with one more dimension.

Example:

>>> import torch
>>> a=torch.randn(2,1)
>>> b=torch.randn(2,1)
>>> c=torch.randn(2,1)
>>> a
tensor([[ 0.7691],
        [-0.0297]])
>>> b
tensor([[ 0.4844],
        [-0.9142]])
>>> c
tensor([[ 0.0210],
        [-1.1543]])
>>> torch.stack((a,b,c))
tensor([[[ 0.7691],
         [-0.0297]],

        [[ 0.4844],
         [-0.9142]],

        [[ 0.0210],
         [-1.1543]]])

You can then use torch.unbind to go the other direction.