Can torchRL handle the case that each agent has different size of observation and action in a MARL env?

[nkjcqvcpi]

I am trying to use torchRL to do a MARL project, and the agents in my env have observation and action in different shape.
I found it quite difficult to use tensordict to store these values.
I wonder if torchRL can handle this？

By the way, can torchRL works with fedreated learning frameworks like flower?

[vmoens]

You could do that in two different ways:
IMO things that have different shape should not be put at the same level unless they go through the same model (unless you need to somehow stack them, either via a nested tensor, padding or else).
If not then your data is semantically different and should be organized in separate groups.
I never saw any good justification for storing stuff that need to go through different models together and I find it personally more confusing than anything (it always looks like N different data origins put together in a single weird object to be then dismantled again and sent in the relevant model - to me that should stay in different bins all the way from the env to the model).

If you want to store things of different shapes alongside each other (again, not my preferred option) you can stack any spec using torch.stack. That will give you a spec of heterogeneous shape.
For instance

spec = torch.stack([
  CompositeSpec(obs=UnboundedContinuousTensorSpec(shape=(3,))),
  CompositeSpec(obs=UnboundedContinuousTensorSpec(shape=(4,))),
], 0)

That spec will give you a LazyStackedTensorDict, which can handle observations of different shape, dtype and even missing observations.

Otherwise you should store each group / agent separately:

spec = CompositeSpec(agent1 = UnboundedContinuousTensorSpec(shape=(3,))), agent2 = UnboundedContinuousTensorSpec(shape=(4,))))

What's the use case you have in mind by the way?

[nkjcqvcpi]

Thanks so much for your reply.

I want to learn torchRL for my future projects, and I am putting one of my previous RL project in a federated learning framework. So I try to do these things together.

I follow the instruction of https://pytorch.org/rl/tutorials/coding_ppo.html to write the env. And I already use LazyStackedTensorDict to manage these items in different shape, my obs like this:

TensorDict({'agents': 
    TensorDict({'observation': 
        LazyStackedTensorDict(*[
            TensorDict(source={"observation": obs,}, batch_size=1) 
        for obs in agents_obs], stack_dim=1)}, 
    batch_size=1)}, 
batch_size=1)

But when I call check_env_specs(), it raises error like:

Traceback (most recent call last):
File "train.py", line 77, in
check_env_specs(env)
File "lib/python3.10/site-packages/torchrl/envs/utils.py", line 470, in check_env_specs
fake_tensordict = env.fake_tensordict()
File "lib/python3.10/site-packages/torchrl/envs/common.py", line 2741, in fake_tensordict
fake_obs = observation_spec.zero()
File "lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 3752, in
key: self[key].zero(shape)
File "lib/python3.10/site-packages/torchrl/data/tensor_specs.py", line 929, in zero
return torch.stack([spec.zero(shape) for spec in self._specs], dim)
File "lib/python3.10/site-packages/torch/utils/_device.py", line 77, in torch_function
return func(*args, **kwargs)
RuntimeError: stack expects each tensor to be equal size, but got [1, 1, 328] at entry 0 and [1, 4, 6] at entry 1

It seems that observation_spec.zero() try to stack tensors in LazyStackedTensorDict.

As you said, it is no a good idea to put these tensor at the same level, then I modify my code like this

spec = CompositeSpec(agent1 = UnboundedContinuousTensorSpec(shape=(3,))), agent2 = UnboundedContinuousTensorSpec(shape=(4,))))

And I notice what I really need to care is if torchRL can work with federated learning framework like Flower, which I didn't find any docs or Q&A on the Internet.😂

Thanks again!