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

TL;DR

This paper introduces LOGO, a local-to-global world model for offline multi-agent reinforcement learning that improves data augmentation and policy generalization by leveraging local predictions and uncertainty-aware sampling.

Contribution

The paper proposes a novel local-to-global world model framework for offline MARL, enhancing prediction accuracy and data efficiency while reducing computational costs.

Findings

Outperforms state-of-the-art offline MARL methods on 8 benchmarks.

Effectively captures agent dependencies through local-to-global modeling.

Reduces computational overhead compared to ensemble-based approaches.

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.