MA2E: Addressing Partial Observability in Multi-Agent Reinforcement Learning with Masked Auto-Encoder

Centralized Training and Decentralized Execution (CTDE) is a widely adopted paradigm to solve cooperative multi-agent reinforcement learning (MARL) problems. Despite the successes achieved with CTDE, partial observability still limits cooperation among agents. While previous studies have attempted to overcome this challenge through communication, direct information exchanges could be restricted and introduce additional constraints. Alternatively, if an agent can infer the global information solely from local observations, it can obtain a global view without the need for communication. To this end, we propose the Multi-Agent Masked Auto-Encoder (MA 2 E), which utilizes the masked auto-encoder architecture to infer the information of other agents from partial observations. By employing masking to learn to reconstruct global information, MA 2 E serves as an inference module for individual agents within the CTDE framework. MA 2 E can be easily integrated into existing MARL algorithms and has been experimentally proven to be effective across a wide range of environments and algorithms. The code is available at https://github.com/cheesebro329/MA2E