Support for Agent that determines how many agents are available at each timestep

[Quinticx]

I have a simulation where I would like to be able to vary the amount of total agents available at each timestep. For example, at t=0, I might have 4 agents available. At t=1, I would only need 3 agents. At t=2, I might need 5 agents. Ideally, I would like an agent to learn based on the environment how many agents are needed and set the number of available agents.

I was thinking of having a set limit to the max number of agents available at any timestep and then just masking out responses from agents that were not selected. I was wondering if there is any support for a behavior like this? Or if there was a way to actually change the number of agents per timestep rather than masking out responses.

[vmoens]

Hi @Quinticx
Thanks for the suggestion. Right now the blessed way to do this would be via masking your agents on a per-step basis. @matteobettini can provide some examples.
In the future we would like to support environments that have dynamic specs and non-masked variable agent numbers fall into that bucket.

There are pros and cons to both approaches:

Masking agents

Pros

• Masks play well with circular buffers in distributed / multiprocessed data collectors or parallel environments. Not using preallocated buffers in these classes will hard performance significantly.
• Masks will play well with contiguous storage in buffers
• Masks are unambiguous (see my point in the cons of dynamic agents)
• Masks work nicely with tensordict: the following code masks all the tensors you need in a concise way

  data = TensorDict({
      "action": torch.tensor([1, 2, 3]),
      "obs": torch.tensor([1, 2, 3]),
      ...,
      "mask": torch.tensor([True, False, True])}, batch_size=[3])
  cur_data = data[data["mask"]]

Cons

• Masks require you to know what the maximum number of agents is a priori
• Masked tensordicts require extra storage

Dynamic number of agents

### Pros

• Cheap from a memory footprint perspective
• No upper bound on the number of agents

Cons

• Not obvious at first glance. E.g. when I pass the actions to the environment, how can I tell it which action is for which agent? Using indices I guess... When I get a tensordict out of it, if it does not have the same shape as the input, how do I keep track of which agent is which? How do I populate a buffer with it? I can imagine that each of these questions can be solved independently but the comparing indices is expensive and error-prone so I expect that the engineering stack behind that will be considerable.
  I guess the equivalent of the code above will be

  data = TensorDict({
      "action": torch.tensor([1, 3]),
      "obs": torch.tensor([1, 3]),
      ...,
      "agent_idx": torch.tensor([0, 2])}, batch_size=[3])
  cur_data = data[data["agent_idx"]]

• If we need to adapt the env specs on the fly it also comes at some cost.
• When passing these variable-sized data to a model it will incur some drop of performance when using torch.compile.

I think the first thing to solve is the question of dymanic specs. If anyone has a concrete use case to share I'd love to give it a look and try to prototype a solution!