`evorl.utils.ma_utils`¶

Module Contents¶

`batchify`	Batchify the data for multi-agent training.
`multi_agent_episode_done`	Check whether the multi-agent episode is done.
`unbatchify`	Unbatchify the data for multi-agent training.

evorl.utils.ma_utils.batchify(x: dict[evorl.types.AgentID, jax.Array], agent_list, padding=False) → jax.Array[source]¶

Batchify the data for multi-agent training.

Output batched data.

Parameters:

x – data from each agent, [batch_dims…, val] Note: Currently, only the last dimension is viewed as value, and the rest are batch dimensions.
agent_list – list, list of agent names
num_actors – int, number of actors
padding – bool, whether to pad the data to the same length over the last dimension. set to False if the data already has same length.

Returns:

[num_actors, batch_dims…, val]

Return type:

Concatenated data from multiple agents with shape

evorl.utils.ma_utils.multi_agent_episode_done(done: evorl.types.Done) → chex.Array[source]¶: Check whether the multi-agent episode is done.

evorl.utils.ma_utils.unbatchify(x: jax.Array, agent_list) → dict[evorl.types.AgentID, jax.Array][source]¶

Unbatchify the data for multi-agent training.

Here we assume data like actions has the same shape for each agent. (True for MaBrax)

Parameters:

x – batched data, [num_actors, batch_dims…, val] Note: Currently, only the last dimension is viewed as value, and the rest are batch dimensions.
agent_list – list, list of agent names

Returns:

data}

Return type:

Dict {agent_name