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

arXiv: https://arxiv.org/abs/2602.12520 · HTML: https://mlsanity.com/paper/2602.12520/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 72 references that are in the map.
- First posted on arXiv 13 Feb 2026.

## Builds on (most cited first, 30 of 36)

- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games](https://mlsanity.com/paper/2103.01955/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Counterfactual Multi-Agent Policy Gradients](https://mlsanity.com/paper/1705.08926/) (2017)
- [Multi-Agent Reinforcement Learning: A Selective Overview of Theories and Algorithms](https://mlsanity.com/paper/1911.10635/) (2019)
- [Value-Decomposition Networks For Cooperative Multi-Agent Learning](https://mlsanity.com/paper/1706.05296/) (2017)
- [Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2003.08839/) (2020)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [QTRAN: Learning to Factorize with Transformation for Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1905.05408/) (2019)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Actor-Attention-Critic for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1810.02912/) (2018)
- [Weighted QMIX: Expanding Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2006.10800/) (2020)
- [Multiagent Cooperation and Competition with Deep Reinforcement Learning](https://mlsanity.com/paper/1511.08779/) (2015)
- [FACMAC: Factored Multi-Agent Centralised Policy Gradients](https://mlsanity.com/paper/2003.06709/) (2020)
- [Learning by Playing - Solving Sparse Reward Tasks from Scratch](https://mlsanity.com/paper/1802.10567/) (2018)
- [Benchmarking Model-Based Reinforcement Learning](https://mlsanity.com/paper/1907.02057/) (2019)
- [Continuous Deep Q-Learning with Model-based Acceleration](https://mlsanity.com/paper/1603.00748/) (2016)
- [Deep Spatial Autoencoders for Visuomotor Learning](https://mlsanity.com/paper/1509.06113/) (2015)
- [Imagination-Augmented Agents for Deep Reinforcement Learning](https://mlsanity.com/paper/1707.06203/) (2017)
- [Shared Experience Actor-Critic for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2006.07169/) (2020)
- [Provable Self-Play Algorithms for Competitive Reinforcement Learning](https://mlsanity.com/paper/2002.04017/) (2020)
- [Symmetric Graph Convolutional Autoencoder for Unsupervised Graph Representation Learning](https://mlsanity.com/paper/1908.02441/) (2019)
- [Model-Based Multi-Agent RL in Zero-Sum Markov Games with Near-Optimal Sample Complexity](https://mlsanity.com/paper/2007.07461/) (2020)
- [Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models](https://mlsanity.com/paper/2205.13817/) (2022)
- [Efficient Model-Based Multi-Agent Mean-Field Reinforcement Learning](https://mlsanity.com/paper/2107.04050/) (2021)
- [Can Increasing Input Dimensionality Improve Deep Reinforcement Learning?](https://mlsanity.com/paper/2003.01629/) (2020)
- [MASER: Multi-Agent Reinforcement Learning with Subgoals Generated from Experience Replay Buffer](https://mlsanity.com/paper/2206.10607/) (2022)

## Abstract (arXiv)

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.