How to define observation spec for Multi agent, multiple groups and batch size = 32?

[zoetsekas]

In my multi agent environment that is based on EnvBase the observation spec is defined as:
batch size is 32
input features is 90
I have mutlple groups, but each group has one agent,
Group A has agent A, Group B has agent B

     self.batch_size = torch.Size([32])
      base_observation_spec = Bounded(
            low=-1.0,
            high=1.0,
            shape=torch.Size([32, self.feature_length]),
            device=device,
            dtype=torch.float32,
            domain=continuous
        )

        atomic_observation_spec = Composite(
            observation=base_observation_spec, shape=self.batch_size
        )

        self.observation_spec = Composite(
            {
                x: atomic_observation_spec for x in self.agent_names
            }, shape=self.batch_size
        )

The error I get is:

RuntimeError: TensorDictModule failed with operation
MultiAgentMLP(
MLP(
(0): Linear(in_features=90, out_features=256, bias=True)
(1): Softmax(dim=None)
(2): Linear(in_features=256, out_features=256, bias=True)
(3): Softmax(dim=None)
(4): Linear(in_features=256, out_features=1, bias=True)
),
n_agents=1,
share_params=True,
centralized=False,
agent_dim=-2)
in_keys=[('NAIL', 'observation')]
out_keys=[('NAIL', 'param')].

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/opt/project/trading/training/model/rl/multi_agent/ppo/torch_rl/trainer.py", line 322, in
train()
File "/opt/project/trading/training/model/rl/multi_agent/ppo/torch_rl/trainer.py", line 170, in train
collector = SyncDataCollector(
File "/usr/local/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 741, in init
self._make_final_rollout()
File "/usr/local/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 835, in _make_final_rollout
policy_output = self.policy(policy_input)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 314, in wrapper
return func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/utils.py", line 359, in wrapper
result = func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/sequence.py", line 481, in forward
tensordict_exec = self._run_module(module, tensordict_exec, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/sequence.py", line 456, in _run_module
tensordict = module(tensordict, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 314, in wrapper
return func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/utils.py", line 359, in wrapper
result = func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/sequence.py", line 481, in forward
tensordict_exec = self._run_module(module, tensordict_exec, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/sequence.py", line 456, in _run_module
tensordict = module(tensordict, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 314, in wrapper
return func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/utils.py", line 359, in wrapper
result = func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/probabilistic.py", line 624, in forward
tensordict_out = self.get_dist_params(tensordict, tensordict_out, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/probabilistic.py", line 576, in get_dist_params
return tds(tensordict, tensordict_out, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 314, in wrapper
return func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/utils.py", line 359, in wrapper
result = func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/sequence.py", line 481, in forward
tensordict_exec = self._run_module(module, tensordict_exec, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/sequence.py", line 456, in _run_module
tensordict = module(tensordict, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 314, in wrapper
return func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/utils.py", line 359, in wrapper
result = func(_self, tensordict, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1089, in forward
raise err from RuntimeError(
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1060, in forward
raise err
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1046, in forward
tensors = self._call_module(tensors, **kwargs)
File "/usr/local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1010, in _call_module
out = self.module(*tensors, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
return forward_call(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torchrl/modules/models/multiagent.py", line 138, in forward
inputs = self._pre_forward_check(inputs)
File "/usr/local/lib/python3.10/site-packages/torchrl/modules/models/multiagent.py", line 449, in _pre_forward_check
raise ValueError(
ValueError: Multi-agent network expected input with shape[-2]=1, but got torch.Size([32, 90])