Puzzle it Out: Local-to-Global World Model for Offline Multi-Agent Reinforcement Learning

TL;DR

This work tackles offline multi-agent reinforcement learning by addressing out-of-distribution generalization and conservatism through a novel Local-to-Global (LOGO) world model. LOGO learns per-agent local dynamics to predict local observations and then deduces global state and rewards, augmented with an encoder-based uncertainty estimator to weight synthetic data generated for policy learning. The method yields a bound-aware, uncertainty-weighted data augmentation approach that reduces error propagation without requiring ensembles, and it integrates with existing offline MARL algorithms (e.g., MACQL). Empirical results on SMAC and MaMuJoCo show state-of-the-art performance and robustness, with notable efficiency gains and effective generalization when combined with MPC. These findings suggest LOGO as a strong generalizable baseline for offline MARL, with practical impact for scalable multi-agent decision-making under data constraints.

Abstract

Offline multi-agent reinforcement learning (MARL) aims to solve cooperative decision-making problems in multi-agent systems using pre-collected datasets. Existing offline MARL methods primarily constrain training within the dataset distribution, resulting in overly conservative policies that struggle to generalize beyond the support of the data. While model-based approaches offer a promising solution by expanding the original dataset with synthetic data generated from a learned world model, the high dimensionality, non-stationarity, and complexity of multi-agent systems make it challenging to accurately estimate the transitions and reward functions in offline MARL. Given the difficulty of directly modeling joint dynamics, we propose a local-to-global (LOGO) world model, a novel framework that leverages local predictions-which are easier to estimate-to infer global state dynamics, thus improving prediction accuracy while implicitly capturing agent-wise dependencies. Using the trained world model, we generate synthetic data to augment the original dataset, expanding the effective state-action space. To ensure reliable policy learning, we further introduce an uncertainty-aware sampling mechanism that adaptively weights synthetic data by prediction uncertainty, reducing approximation error propagation to policies. In contrast to conventional ensemble-based methods, our approach requires only an additional encoder for uncertainty estimation, significantly reducing computational overhead while maintaining accuracy. Extensive experiments across 8 scenarios against 8 baselines demonstrate that our method surpasses state-of-the-art baselines on standard offline MARL benchmarks, establishing a new model-based baseline for generalizable offline multi-agent learning.

Figures (4)

• Figure 1: The motivation of LOGO. The results in (a) and (b) demonstrate the next state prediction accuracy of baseline MAMBA and our proposed LOGO method on the SMAC 5m_vs_6m scenario, while (c) presents the performance of MACQL (Multi-Agent Conservative Q-Learning) when integrated with different world models. These results indicate that employing an inaccurate world model can adversely affect overall performance.
• Figure 2: Overall framework of LOGO. (a) We first train the LOGO world model with both prediction loss and reconstruction loss. Subsequently, during the dynamic data generation phase, the predictive model (yellow part) predicts $\boldsymbol{\hat{o}}_{t+1}$ with ($\boldsymbol{o}_t$, $\boldsymbol{a}_t$, $s_t$). These predictions are then propagated through the deductive model's (blue part) decoder (indicated by the green dashed pathway) to compute the subsequent state $s'_{t+1}$ and reward $r'_t$. Simultaneously, the state outputs from the uncertainty layer $\hat{s}_{t+1}$ along with the deduced state $s'_{t+1}$ are utilized to estimate the prediction uncertainty (indicated by the orange dashed pathway). (b) For policy training, we sample from the generated data using uncertainty-aware weighting to prioritize more reliable transitions.
• Figure 3: The prediction comparison of model-based methods. We employ Principal Component Analysis (PCA) groth2013principal to visualize and compare the ground truth next state with the predicted next state generated by world models across various model-based approaches in SMAC 6h_vs_8z map.
• Figure 4: The state and reward distribution visualization results. We employ Principal Component Analysis (PCA) groth2013principal to visualize the state and reward distributions during online testing for both the baseline model-free method (a) and LOGO (b). Each pair of plots comprises a state distribution heatmap (left), where the lighter colors correspond to higher state visitation frequency, and a corresponding reward distribution map (right), where the lighter colors correspond to higher reward value.

Theorems & Definitions (2)

• Theorem 1
• Theorem 2