NetWorld: Communication-Based Diffusion World Model for Multi-Agent Reinforcement Learning in Wireless Networks

TL;DR

NetWorld tackles the challenge of scalable and generalizable MARL for wireless networks by introducing a diffusion-based world model trained offline under DTDE. It uses classifier guidance to plan trajectories within the learned model and recovers executable actions via an inverse dynamics module, while mitigating cross-task heterogeneity through scenario encoders and MF-based local communication. The framework demonstrates improved data efficiency and cross-task generalization across three wireless control tasks, achieving strong few-shot adaptation with limited target-task data. This approach offers a practical path toward deployable, distributed intelligent control in large-scale, heterogeneous wireless networks, with open directions for online adaptation and cross-layer control.

Abstract

As wireless communication networks grow in scale and complexity, diverse resource allocation tasks become increasingly critical. Multi-Agent Reinforcement Learning (MARL) provides a promising solution for distributed control, yet it often requires costly real-world interactions and lacks generalization across diverse tasks. Meanwhile, recent advances in Diffusion Models (DMs) have demonstrated strong capabilities in modeling complex dynamics and supporting high-fidelity simulation. Motivated by these challenges and opportunities, we propose a Communication-based Diffusion World Model (NetWorld) to enable few-shot generalization across heterogeneous MARL tasks in wireless networks. To improve applicability to large-scale distributed networks, NetWorld adopts the Distributed Training with Decentralized Execution (DTDE) paradigm and is organized into a two-stage framework: (i) pre-training a classifier-guided conditional diffusion world model on multi-task offline datasets, and (ii) performing trajectory planning entirely within this world model to avoid additional online interaction. Cross-task heterogeneity is handled via shared latent processing for observations, two-hot discretization for task-specific actions and rewards, and an inverse dynamics model for action recovery. We further introduce a lightweight Mean Field (MF) communication mechanism to reduce non-stationarity and promote coordinated behaviors with low overhead. Experiments on three representative tasks demonstrate improved performance and sample efficiency over MARL baselines, indicating strong scalability and practical potential for wireless network optimization.

Figures (6)

• Figure 1: Comparison of task-specific MARL and NetWorld paradigms in diverse wireless resource allocation tasks.
• Figure 2: Overview of the scalable and adaptive NetWorld framework.
• Figure 3: Structural design and few-shot generalization procedure of NetWorld.
• Figure 4: Implementation procedure of NetWorld.
• Figure 5: Comparison of NetWorld with different methods across the three typical resource allocation tasks in terms of average reward.