Multi-Agent Model-Based Reinforcement Learning with Joint State-Action Learned Embeddings

TL;DR

This paper introduces a novel multi-agent reinforcement learning framework that combines joint state-action embeddings with model-based imagination to improve coordination and data efficiency in complex environments.

Contribution

It presents a new approach integrating joint state-action learned embeddings into a model-based RL framework for multi-agent systems, enhancing long-term planning.

Findings

Consistent performance improvements over baseline algorithms.

Effective use of SALE embeddings in imagination modules.

Successful application to multiple multi-agent benchmarks.

Abstract

Learning to coordinate many agents in partially observable and highly dynamic environments requires both informative representations and data-efficient training. To address this challenge, we present a novel model-based multi-agent reinforcement learning framework that unifies joint state-action representation learning with imaginative roll-outs. We design a world model trained with variational auto-encoders and augment the model using the state-action learned embedding (SALE). SALE is injected into both the imagination module that forecasts plausible future roll-outs and the joint agent network whose individual action values are combined through a mixing network to estimate the joint action-value function. By coupling imagined trajectories with SALE-based action values, the agents acquire a richer understanding of how their choices influence collective outcomes, leading to improved long-term planning and optimization under limited real-environment interactions. Empirical studies on well-established multi-agent benchmarks, including StarCraft II Micro-Management, Multi-Agent MuJoCo, and Level-Based Foraging challenges, demonstrate consistent gains of our method over baseline algorithms and highlight the effectiveness of joint state-action learned embeddings within a multi-agent model-based paradigm.