Structured Representation Learning in Model-Based Multi-Agent Reinforcement Learning

In this paper, we propose a novel value-based multi-agent reinforcement learning (MARL) method that integrates structured representation learning into world model imagination. We use latent states from recurrent and representation models to build the agent networks, injecting permutation-invariant and permutation-equivariant inductive biases into the networks and constraining the network embeddings to a product of probability simplices. These well-structured representations of the agents transfer smoothly across different team compositions. Moreover, we present a world model as the imagination module of our framework and apply the simplicial embeddings to the model latents. The world model supplies the mixing network with imagined trajectories, guiding agent learning through well-shaped representations and improving sample efficiency. Across different MARL benchmarks, our method demonstrates consistent gains over value-based and actor-critic baseline methods in both sample efficiency and asymptotic performance. We perform ablation experiments, design-space studies, and analyses of learned representations to highlight the importance of the main framework components within a model-based multi-agent paradigm.