[Question] How to handle MultiDiscrete action spaces in TorchRL

[AnastasiaPsarou]

I have created a custom Parallel API PettingZoo environment with MultiDiscrete action spaces. The env.action_spec() function succeeds.

I am using the Multi-Agent PPO tutorial of TorchRL, but I’m struggling to understand how to modify the architecture so it supports MultiDiscrete action spaces. Specifically, I’d like to know how to correctly adapt the MultiAgentMLP, TensorDictModule, and ProbabilisticActor so that the policy network outputs a MultiDiscrete (or equivalently, MultiCategorical) action distribution for each agent.

Should I create number of ProbabilisticActor modules as the length of the MultiDiscrete action space? In the case where a single ProbabilisticActor module is used, which distribution class should replace Categorical to support a MultiDiscrete action space? Is there an existing script or tutorial in TorchRL that demonstrates how to handle MultiDiscrete action spaces (or MultiCategorical distributions) in a multi-agent setup?

[Xmaster6y]

Well, ProbabilisticActor does have a spec and distribution_class arguments for that. But the distribution will still be Categorical, which can handle a spec having a shape (number of actions) and a space (number of agents).

You should also have a look at how torchrl converts specs e.g.

rl/torchrl/envs/libs/gym.py

Line 414 in 80bfa6e

def convert_multidiscrete_spec(

Can't you use the default pettingzoo constructor from torchrl for your env?

rl/torchrl/envs/libs/pettingzoo.py

Line 850 in 80bfa6e

class PettingZooEnv(PettingZooWrapper):

[AnastasiaPsarou]

Thanks for the reply!

I am already wrapping my custom environment in the torchrl's PettingZoo wrapper. My action space is defined like this inside the PettingZoo environment: self.action_spaces = { agent: MultiDiscrete(nvec) for agent in self.possible_agents}, where nvec = [a, b] for each agent and later in the torchrl wrapped PettingZoo environment like this: MultiCategorical( shape=torch.Size([2, 2]), space=BoxList(boxes=[BoxList(boxes=[CategoricalBox(n=3), CategoricalBox(n=3)]), BoxList(boxes=[CategoricalBox(n=100), CategoricalBox(n=100)])]), device=cpu, dtype=torch.int64, domain=discrete)

The problem I have is with the n_agent_outputs parameter of the MultiagentMLP().

• When I pass n_agent_outputs = 2*env.action_spec.shape[-1] as defined in this tutorial I get actions of one dimension while my action space is two dimensional.
• When I pass n_agent_outputs = [n_agents, some_value] there is no bug but each agent receives an action vector of length n_agents, and the values on this vector range from 0 to some_value.
• When I pass n_agent_outputs = [a, b] I get an error about batch dimension mismatch.
• I also tried n_agent_outputs = [n_agents, (a, b)] I get an error from Torch about receiving an invalid combination of arguments.

I would greatly appreciate some guidance or recommendation on addressing this issue.

[Xmaster6y]

Okay gotcha! Basically, MultiAgentMLP does not support multi-dimensional output (at least for now). So here you have 2 options, I guess:

• If you want to choose your actions independently, then you can use n_agent_outputs=a+b, then split your logits and choose distribution_class=CompositeDistribution in ProbabilisticActor using the map arguments.
• If you want to choose your actions together, use n_agent_outputs=a*b, use a simple Categorical with ProbabilisticActor and then translate the action into 2D actions.

MultiAgentMLP could be made to output 2D actions, but ProbabilisticActor would still need 1D actions, so it's easier to stick to 1D.