[Feature Design] Multi-agent API for TorchRL

[matteobettini]

Hello,

This is the official issue for discussing the multi-agent API of torchrl for environments and collectors.

There has been a lot of discussion on wether the agent dimension should be part of the batch_size or not.

In this design doc we propose an API in which the agent dimension is added to the specs, but not to the batch_size of an environment. This allows to have part of the specs which are shared by all agents, like a global state, a global reward or a global done. It also allows to have these attributes as per-agent attributes.

n_workers = 4
n_vectorized_envs = 32
batch_size = (n_workers, n_vectorized_envs)

n_agents = 2
n_obs_per_agent = 6
n_state_features = 18

state_spec = UnboundedContinuousTensorSpec(
    torch.Size((
        n_state_features,
    ))
) # shape = (n_state_features,)
obs_spec = UnboundedContinuousTensorSpec(
    torch.Size((
        n_obs_per_agent,
    ))
)  # shape = (n_obs_per_agent,)
action_spec = UnboundedContinuousTensorSpec(
    torch.Size((2,)),
)  # shape = (2,)
reward_spec = UnboundedContinuousTensorSpec(
    torch.Size((1,)),
) # shape = (1,)
done_spec = UnboundedContinuousTensorSpec(
    torch.Size((1,)),
)  # shape = (1,)

# Create multi-agent specs. In this case specs are identical for all agents but they could differ
multi_agent_obs_spec = torch.stack([obs_spec.clone() for _ in range(n_agents)], dim=0) # UnboundedContinuousTensorSpec with shape = (n_agents, n_obs_per_agent)
multi_agent_action_spec = torch.stack([action_spec.clone() for _ in range(n_agents)], dim=0)  # UnboundedContinuousTensorSpec with  shape = (n_agents, 2)
multi_agent_reward_spec = torch.stack([reward_spec.clone() for _ in range(n_agents)], dim=0)  # UnboundedContinuousTensorSpec with shape = (n_agents, 1)
multi_agent_done_spec = torch.stack([done_spec.clone() for _ in range(n_agents)],dim=0)  # UnboundedContinuousTensorSpec with shape = (n_agents, 1)

input_spec = CompositeSpec(
    action=multi_agent_action_spec
).expand(batch_size) # 
# input_spec.shape = (*batch_size),
# input_spec["action"].shape = (*batch_size, n_agents, 2)

output_spec = CompositeSpec(
    observation=multi_agent_obs_spec,
    done=multi_agent_done_spec, # or done_spec in case of a single done flag for all agents
    reward=multi_agent_reward_spec, # or reward_spec in case of a global reward (DecPOMDP)
    state=state_spec, # this is shared by all agents
).expand(batch_size) 
# output_spec.shape = (*batch_size)
# output_spec["observation"].shape = (*batch_size, n_agents, n_obs_per_agent),
# output_spec["done"].shape = (*batch_size, n_agents, 1),
# output_spec["reward"].shape = (*batch_size, n_agents, 1),
# output_spec["state"].shape = (*batch_size, n_state_features)

Collectors will then consider the batch_size dimensions as dimensions for replicated environments (e.g. parallel workers, vectorized envs, etc) and will count a multi-agent frame (with data for all agents) as 1 frame.

Tagging som folks in my lab for opinions @smorad @Acciorocketships @janblumenkamp

@vmoens if you can tag the other people interested in multi-agent torchrl

[btx0424]

I agree. This was the point I tried to make back in #883. It is common for CTDE paradigms to have some central/global stuff. This way we also have more flexibility with Transforms.

Additionally I think it would be convenient to have an optional helper which automates stacking and expanding.

Rough looking:

class EnvBase:

    @property
    def agent_spec(self):
        if not hasattr(self, "_agent_spec"):
            self._agent_spec = {}
        return _AgentSpecView(self)

@dataclass
class AgentSpec:
    """
    `TensorSpec`s from each individual agent's view, for the ease of policy construction.
    """
    n: int
    observation_spec: TensorSpec
    action_spec: TensorSpec
    reward_spec: TensorSpec
    state_spec: TensorSpec = None

class _AgentSpecView(Mapping[str, AgentSpec]):
    def __init__(self, env):
        self.env = env
    
    def __setitem__(self, key, value):
        """
        Stack, expand the individual specs and update them into env.
        Here `key` is the name assigned to the type of that agent.
        """

    def __getitem__(self, key):
        """Return a view of the specs, maybe through selecting or indexing"""

so that for a custom Env we can do

class MyEnv(EnvBase):
    def __init__(...):
        self.agent_spec["predator"] = ...
        self.agent_spec["prey"] = ...

env = MyEnv(...)
policy_module = TensorDictSequential(
    make_actor(agent_spec) for k, v in env.agent_spec.items()
)

This way we save the mental trouble of parsing the whole Env's specs when constructing a joint policy.

[matteobettini]

Yes, i am definitely am open to have helpers that will simplify keeping track of agent names in the specs. I think we can look into this once we have the basic functionality going.

[btx0424]

Sure. Also, I think the collectors currently do not need further modification since we can always have a joint policy passed in. What other main functionalities do you think are missing right now?

One thing I can think of is reset handling when some agents can reach done before the whole environment. The current implementation of _step_and_maybe_reset would try to do a reset anyway since done_or_terminated.any() is true. A solution would be checking against done_spec, e.g., see if done is of shape [*batch_size, 1]. If not, i.e., it's [*batch_size, *] we can by default fall back to take done.all(-1).

[matteobettini]

The idea, as in my post above is that
output_spec["done"].shape = (*batch_size, n_agents, 1),

So the collectors (when they find any done in the td (or also all the time)) will call reset with a reset flag

_reset = done.clone()

the env will then know what to reset (if only some agents are done in some vectorized dims).

Multiagent environments will write their own logic on how to consider the _reset tensor. (see the vmas wrapper for an example of how these vectorized/multiagent resets work)

[matteobettini]

What we are missing right now is mainly everything out of the sampling part. Losses, algos and models.

As soon as I can I'll get to work on this for a full multi-agent example.

[btx0424]

I am happy to know what algorithms and models you have in mind. If the above-mentioned batch shape change is soon to be implemented I can start to add mappo right away (following your idea above). Our team also has been needing a fast env to do sanity checks and performance tuning lately and VMAS is a good fit :)

Previously the main obstacle was the inconvenience to add a central obs due to the current batch size setting.

[matteobettini]

Yes if you have questions and stuff about vmas let me know.

I am a little busy before the end of march but the idea yes was to start with a ppo example (so all mappo, ippo, etc )

[matteobettini]

#1027 Shows some examples of multiagent training in VMAS