Joint-embedding world models improve multi-agent reinforcement learning performance
MA-JEPA: Joint-Embedding World Models for Multi-Agent Reinforcement Learning
Machine LearningArtificial IntelligenceMultiagent Systems
Summary
Training AI agents to work together can be slow because they need lots of practice to learn good strategies. The authors study a new way to teach multiple agents by having them predict each other's future observations in a shared hidden space. Their method, called MA-JEPA, uses a special model that learns from imagined experiences instead of just real ones, making training more efficient. When tested on a popular multi-agent game, their approach matched or beat strong existing methods in four of eight scenarios.
What this means in practice
- •For multiplayer game developers: Improve AI teammates that learn faster and coordinate better by training on imagined multi-agent experiences.
- •For robotics teams: Build cooperative robot controllers that efficiently learn joint tasks by predicting each robot’s future observations in shared models.
Authors
Brandon Gary Kaplowitz, Osaze James Obahor, Christian Schroeder de Witt
Abstract
World models improve sample efficiency by training policies on imagined trajectories, but their usefulness depends on learning representations that capture the information needed for future control. We study whether self-supervised joint-embedding prediction (JEPA) can provide this learning signal for multi-agent reinforcement learning. We introduce MA-JEPA, a stochastic world model that replaces observation reconstruction with prediction of target representations, enabling model-based multi-agent reinforcement learning with centralized training and decentralized execution. A categorical latent state and a causal Transformer are trained with posterior and action-conditioned dynamics prediction objectives and are then used for actor-critic learning from latent imagination. A training-only joint predictor conditions on all agents' local states and actions to predict each agent's next local observation embedding. These predictions are passed through the same local posterior used during real interaction with a centralized critic that is used only for value learning, with execution remaining decentralized. Our experiments show that this architecture performs strongly on SMAC, matching or exceeding the strongest reported comparator mean win rate on four of eight evaluated maps.