Sample batch configuration to contain multi agent data

RLlib Configure Algorithm, Training, Evaluation, Scaling
1

Hi all,

how can I configure a sample batch to contain tuples of observations and actions for multiple agents?

Specifically I want to use the multi-agent-transformer model with a wrapper class in ray rllib (GitHub - PKU-MARL/Multi-Agent-Transformer). In the forward function a list of all agents states and actions is expected.
How can I design a sample batch to contain multi agent data? (my experience is it only contains data for a single agent)

Thanks everyone