Learning to ground multi-agent reinforcement learning with masked multi-agent AutoEncoding.

Abstract

Data efficiency is a persistent bottleneck in multi-agent reinforcement learning (MARL), especially under partial observability where each agent must act on limited local observations. Existing approaches often struggle to fully exploit the shared latent structure across agents, leading to slow convergence. We propose Masked Multi-Agent Autoencoding (MMAAE), a novel self-supervised representation learning framework for centralized training with decentralized execution (CTDE) MARL. MMAAE introduces a dual-space masking mechanism that simultaneously performs random masking and reconstruction in both the raw observation space and the learned representation space. By employing an asymmetric autoencoder to reconstruct masked observation tokens and aligning reconstructed features with target representations via a contrastive objective, MMAAE forces the encoder to capture robust inter-agent dependencies and global state information. Crucially, this auxiliary task is optimized jointly with the policy, requiring no additional environment interactions. We integrate MMAAE with popular MARL backbones (MAPPO and MAT) and evaluate it across diverse benchmarks: SMAC, Multi-Agent MuJoCo, and MAQC. Empirical results demonstrate that MMAAE significantly improves sample efficiency, reducing Time-to-Threshold by up to 42% and increasing early-training AUC by up to 38%, while matching or exceeding the final performance of strong baselines. Ablation studies further confirm that the dual-masking strategy is essential, as removing either observation- or representation-space masking leads to a marked drop in performance.