[BUG] Operations on Spec of Stacked CompositeSpec with different keys crashes

[hyerra]

Describe the bug

I'm trying to print a CompositeSpec, but the current implementation crashes with a KeyError. This is related to #1201.
cc: @matteobettini

To Reproduce

Please try running the following code example:

import torch
from torchrl.data.tensor_specs import (
    CompositeSpec,
    UnboundedContinuousTensorSpec,
)

specs = [CompositeSpec({
    "observation_0": UnboundedContinuousTensorSpec(
        shape=torch.Size([128, 128, 3]),
        device="cpu",
        dtype=torch.float32)}),
    CompositeSpec({
        "observation_1": UnboundedContinuousTensorSpec(
        shape=torch.Size([128, 128, 3]),
        device="cpu",
        dtype=torch.float32)}), CompositeSpec({
    "observation_2": UnboundedContinuousTensorSpec(
        shape=torch.Size([128, 128, 3]),
        device="cpu",
        dtype=torch.float32)}), CompositeSpec({
    "observation_3": UnboundedContinuousTensorSpec(
        shape=torch.Size([4]),
        device="cpu",
        dtype=torch.float32)})]

spec = torch.stack(specs, dim=0)
print(spec)

You will see the crash. Here is a traceback that I've seen:

Traceback (most recent call last):
  File "/Users/hyerra/Desktop/test.py", line 27, in <module>
    print(spec)
  File "/Users/hyerra/Library/Caches/pypoetry/virtualenvs/.../lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 3165, in __repr__
    [indent(f"{k}: {repr(item)}", 4 * " ") for k, item in self.items()]
  File "/Users/hyerra/Library/Caches/pypoetry/virtualenvs/.../lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 3165, in <listcomp>
    [indent(f"{k}: {repr(item)}", 4 * " ") for k, item in self.items()]
  File "/Users/hyerra/Library/Caches/pypoetry/virtualenvs/.../lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 3139, in items
    yield key, self[key]
  File "/Users/hyerra/Library/Caches/pypoetry/virtualenvs/.../lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 760, in __getitem__
    [composite_spec[item] for composite_spec in self._specs], dim=self.dim
  File "/Users/hyerra/Library/Caches/pypoetry/virtualenvs/.../lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 760, in <listcomp>
    [composite_spec[item] for composite_spec in self._specs], dim=self.dim
  File "/Users/hyerra/Library/Caches/pypoetry/virtualenvs/.../lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 2624, in __getitem__
    return self._specs[idx]
KeyError: 'observation_0'

Expected behavior

I should be able to print the spec, especially if it was created successfully.

System info

Describe the characteristic of your environment:

import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)

None 1.24.3 3.10.11 (main, Apr 7 2023, 07:24:53) [Clang 14.0.0 (clang-1400.0.29.202)] darwin

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)

[hyerra]

I just updated this with a simpler, reproducible example

[vmoens]

Thanks for reporting!

Printing heterogeneous lazy stacks is a nightmare. In tensordict we do our best to print the common keys and every dim that match, putting a -1 for those which don't.
If a key is missing in one sub TD we simply don't print it. Another option would be to print it in parenthesis or similar syntax.
Same could apply to Composite.

Do you know what you would expect to see in this print?

[hyerra]

Yea I can imagine it's complicated! I think there are two approaches that work better depending on situation. I think the version you described above works really well in cases where most of the stacked specs have similar attributes/dimensions. However, in cases where the stacked specs are expected to contain different information, I would have expected to see like a tree like structure showing the whole tree. Maybe something like this:
Lazy...(
CompositeSpec(
observation_1: Unbounded..(),
),
CompositeSpec(
observation_2: Unbounded..(),
), ...
)
The Lazy part signifies it's a lazy stack and then the multiple entries allows you to show each field without having to select only the common keys.

I do see how this could get really verbose though if every spec has the same attributes and is similar. Did you have something different in mind?

[hyerra]

Also, @vmoens I think this issue might be more significant than I initially thought. I was playing around with this some more, and operations like unsqueeze also crash. For example, if I try to do:

spec = torch.stack(specs, dim=0).unsqueeze(dim=0)

This also crashes with the KeyError: 'observation_0'

[vmoens]

Got it, i'll find some bandwidth to fix all this!

[matteobettini]

My proposition for the print is

LazyStackedCompositeSpec(
    camera: BoundedTensorSpec(
        shape=torch.Size([3, 32, 32, 3]),
        space=ContinuousBox(
            minimum=Tensor(shape=torch.Size([3, 32, 32, 3]), device=cpu, dtype=torch.float32, contiguous=True), 
            maximum=Tensor(shape=torch.Size([3, 32, 32, 3]), device=cpu, dtype=torch.float32, contiguous=True)),
        device=cpu,
        dtype=torch.float32,
        domain=continuous),
    vector: LazyStackedTensorSpec(
         shape=torch.Size([3, -1]), space=None, device=cpu, dtype=torch.float32, domain=continuous), device=cpu, shape=torch.Size([3]),
    
   lazy_keys={
        0 -> lidar: UnboundedContinuousTensorSpec(
                  shape=torch.Size([20]),
                  device=cpu,
                  dtype=torch.float32,
                  domain=continuous),
              obs_0: UnboundedContinuousTensorSpec(
                  shape=torch.Size([1]),
                  device=cpu,
                  dtype=torch.float32,
                  domain=continuous),
        1 -> lidar: UnboundedContinuousTensorSpec(
                  shape=torch.Size([20]),
                  device=cpu,
                  dtype=torch.float32,
                  domain=continuous),
              obs_1: UnboundedContinuousTensorSpec(
                  shape=torch.Size([ 1,2]),
                  device=cpu,
                  dtype=torch.float32,
                  domain=continuous),
        2 ->  obs_2: UnboundedContinuousTensorSpec(
                  shape=torch.Size([ 1,2,3]),
                  device=cpu,
                  dtype=torch.float32,
                  domain=continuous),
    }

 device=cpu, shape=torch.Size([3]))