[BugFix] Heterogeneous envs compatibility

[vmoens]

Depends on pytorch/tensordict#400

cc @matteobettini can you have a look at the env I designed in the tests to see if it makes sense?

[matteobettini]

Test env makes sense.

Would it make sense to support agents with also different number of dims?
For example, torch.stack([spec with shape [3,2], spec with shape [7,]])

[vmoens]

Is this something we should expect to happen or can we consider that if someone is doing that they're on the wrong track?
Happy to support it if there's a use case (even hypothetical)

[matteobettini]

I can imagine one agent having a camera (eg shape 64,64,3) and the other agent having normal position obs (eg shape 10)

[vmoens]

I can imagine one agent having a camera (eg shape 64,64,3) and the other agent having normal position obs (eg shape 10)

I would name one "pixels" and the other one "position"
You would not expect to pass both in the same neural net right?
I guess that implicitly we're assuming that keys point to content that is semantically similar

[matteobettini]

So in case we have

• M agents with cameras 64,64,3
• N agents with lidars 30,2
• P agents with normal obs 10

The agent dimension will be different for each key.

So at that point we cannot group these into agent_fetures

in the extreme case we could have each agent with a semantically different space and then each agent has its own key

[vmoens]

So at that point we cannot group these into agent_fetures

What do you mean?

You can do

observation_spec = CompositeSpec(agent_features=torch.stack([
     CompositeSpec(pixels=Spec(...)),
     CompositeSpec(lidar=OtherSpec(...)),
], 0), shape=(2,)

If you index that spec, it will give you 2 different sub-specs. The keys don't have to match.
You can sample from it etc

>>> from torchrl.data import CompositeSpec, UnboundedContinuousTensorSpec as Spec
>>> import torch
>>> spec1 = CompositeSpec(obs1=Spec(shape=(3,)))
>>> spec2 = CompositeSpec(obs2=Spec(shape=(4,)))
>>> s = torch.stack([spec1, spec2], 0)
>>> rs = s.rand()
>>> print(rs)
LazyStackedTensorDict(
    fields={
    },
    batch_size=torch.Size([2]),
    device=cpu,
    is_shared=False)
>>> print(rs[0])
CompositeSpec(
    obs1: UnboundedContinuousTensorSpec(
        shape=torch.Size([3]),
        space=None,
        device=cpu,
        dtype=torch.float32,
        domain=continuous), device=cpu, shape=torch.Size([]))
>>> print(rs[1])
CompositeSpec(
    obs2: UnboundedContinuousTensorSpec(
        shape=torch.Size([4]),
        space=None,
        device=cpu,
        dtype=torch.float32,
        domain=continuous), device=cpu, shape=torch.Size([]))

[matteobettini]

Ah. I wasn’t aware stacking was so advanced. Basically it can abstract also over heterogenous dicts and not only heterogeneous lists. Looks good then.

[vmoens]

We should fix prints though, put in parenthesis fields that are present in some but not all children or smth

[vmoens]

Depends on pytorch/tensordict#416

[matteobettini]

@vmoens Would it be possible to merge a version of step_mdp that works with nested key before this whole PR? the fact that this function does not currently support nested is causing me all sorts of issues with collectors and so on

[vmoens]

Sure makes sense
I believe this PR can be split in several smaller ones, let's see what we can do