NetWorld: Communication-Based Diffusion Model

• The paper introduces NetWorld, a novel multi-agent reinforcement learning framework that integrates conditional diffusion, classifier guidance, and mean-field communication.
• It employs a two-stage pipeline combining offline pre-training with in-model trajectory planning to achieve sample efficiency and robust performance in data-constrained network scenarios.
• Empirical validation shows improved average rewards and scalability while theoretical guarantees ensure bounded modeling errors and stable decentralized planning.

The Communication-based Diffusion World Model (NetWorld) is a model-based framework for multi-agent reinforcement learning (MARL) in large-scale wireless communication networks. NetWorld integrates conditional diffusion models, classifier guidance, inverse dynamics, and a lightweight mean-field communication mechanism to enable sample-efficient, scalable trajectory planning and few-shot generalization across heterogeneous tasks. The model operates under the Distributed Training with Decentralized Execution (DTDE) paradigm, specifically tailored for data- and communication-constrained network scenarios (Meng et al., 31 Jan 2026, Meng et al., 27 Oct 2025).

1. Framework Overview

NetWorld adopts a two-stage pipeline combining offline conditional diffusion model pre-training and in-model trajectory planning, with the following principal components:

• Pre-training: Given multi-task offline datasets $$\mathcal{D} = \{(o^{(i)}_{1:T}, a^{(i)}_{1:T}, r^{(i)}_{1:T})\}_{i=1}^N$$, the system jointly trains:
  • A classifier-guided conditional diffusion world model $p_\theta(\mathbf{x}_{0:T} \mid c)$ over latent trajectories, with classifier $p_\phi(c \mid \mathbf{x})$ guiding towards high-return solutions.
  • A scenario encoder $E_\omega$ mapping scaled raw observations to latent codes.
  • An inverse dynamics model $f_\psi$ for action recovery.
• Trajectory Planning: At execution, agents encode observations, sample latent trajectories fully within the learned model using classifier-guided denoising, and decode the best candidate trajectory into actions, without requiring further online environment interaction.

The DTDE paradigm is central: during training, all model parameters are shared, and each agent accesses only its local observations and communication-limited mean-field messages. In execution, agents act independently using local inference and single-hop neighbor summaries, precluding any central coordinator (Meng et al., 31 Jan 2026).

2. Conditional Diffusion Model and Classifier Guidance

The core of NetWorld is a conditional diffusion model operating over agent-specific latent trajectories. The diffusion process is constructed as follows:

• Forward (noising) process: For each agent and time step, latent states $\mathbf{x}_t$ are diffused via

$$q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\,\mathbf{x}_{t-1},\,\beta_t \mathbf{I}).$$

• Reverse (denoising) process: The model learns the reverse trajectory via

$$p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t, c) = \mathcal{N}(\mathbf{x}_{t-1}; \mu_\theta(\mathbf{x}_t, t, c), \sigma_t^2 \mathbf{I}),$$

where $c$ denotes the two-hot encoded, discounted return used for conditioning.

• Classifier guidance: Trajectory generation is steered towards high-reward regions by augmenting the mean with a gradient of the return-classifier:

$$\tilde{\mu}_\theta(\mathbf{x}_t) = \mu_\theta(\mathbf{x}_t) + \alpha_t \sigma_t^2 \nabla_{\mathbf{x}_t} \log p_\phi(c \mid \mathbf{x}_t),$$

where $\alpha_t$ is a guidance strength parameter (Meng et al., 31 Jan 2026, Meng et al., 27 Oct 2025).

This procedure enables precise control over the anticipated returns of sampled trajectories, mitigating sample inefficiency and poor generalization within heterogeneous domains.

3. Latent Processing, Inverse Dynamics, and Discretization

Latent transformation and action decoding are handled through a structured pipeline:

• Shared latent representation: Observations $o_t$ are scaled with $$\mathrm{symlog}(x) = \mathrm{sign}(x) \log(1 + |x|)$$, then embedded to $\mathbf{x}_t \in \mathbb{R}^d$ via encoder $E_\omega$; the same encoder processes mean-field messages.
• Two-hot discretization: For all actions and rewards, the model discretizes continuous values into a probability distribution over two adjacent bins in the symlog domain. For $y \in [b_i, b_{i+1}]$, the two-hot vector $\pi$ is defined by

$$\pi_i = \frac{b_{i+1}-y}{b_{i+1}-b_i},\quad \pi_{i+1} = \frac{y-b_i}{b_{i+1}-b_i}.$$

• Inverse dynamics model: $f_\psi$ maps latent state pairs $(\mathbf{x}_t, \mathbf{x}_{t+1})$ to the predicted two-hot action encoding, trained by minimizing cross-entropy with ground truth encodings.

This design unifies action and state representation across heterogeneous MARL tasks and allows for efficient action decoding during model-based planning (Meng et al., 31 Jan 2026).

4. Mean Field Communication for Distributed Coordination

To address inherent non-stationarity and scalability limitations in decentralized MARL, NetWorld employs a mean-field (MF) communication mechanism:

• Neighborhood aggregation: Each agent $i$ computes latent embedding $h^{(i)}_t$ and receives $1$-hop neighbor embeddings, aggregating them as

$$m^{(i)}_t = \frac{1}{|N(i)|} \sum_{j \in N(i)} h^{(j)}_t.$$

• Input fusion: The pair $(h^{(i)}_t, m^{(i)}_t)$ is concatenated and presented at every denoising/classification step, supplying each agent with a summary of local peer behavior.
• Communication efficiency: Only local means are shared, yielding $O(\dim(o)\,|N(i)|)$ communication overhead per agent per step.

This scheme substantially improves multi-agent coordination, as confirmed by ablations showing marked performance degradation when mean-field aggregation is omitted (Meng et al., 27 Oct 2025). Theoretical results provide an explicit bound on the KL divergence between the true and mean-field guided trajectory distributions, ensuring modeling fidelity (Meng et al., 27 Oct 2025).

5. Training Procedure and Objectives

The learning procedure consists of multi-task pre-training followed by rapid few-shot adaptation:

• Combined loss objective:

$$\mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{diff}} + \lambda_{\mathrm{cls}} \mathcal{L}_{\mathrm{cls}} + \lambda_{\mathrm{inv}} \mathcal{L}_{\mathrm{inv}} + \text{(regularization)},$$

where - $\mathcal{L}_{\mathrm{diff}}$: Diffusion denoising loss (MSE between real and predicted noise) - $\mathcal{L}_{\mathrm{cls}}$: Classifier cross-entropy loss - $\mathcal{L}_{\mathrm{inv}}$: Inverse dynamics cross-entropy - $\lambda_{\mathrm{cls}}$, $\lambda_{\mathrm{inv}}$: Hyperparameters

• Pre-training: Conducted on multi-task offline expert datasets, with each gradient step updating all model components using mini-batches containing trajectories from various tasks.
• Few-shot adaptation: For a new target task, fine-tuning is performed using only 10 expert trajectories, retraining $E_\omega$, $\theta$, $\phi$, and $f_\psi$ over a few epochs (Meng et al., 31 Jan 2026).

6. Trajectory Planning and Execution

During deployment, each agent iteratively plans in its local world model without further real-environment rollouts:

1. State encoding: Current observation $o_\mathrm{curr}$ is encoded into $\mathbf{x}_0$, and the mean-field message $m_0$ is computed.
2. Trajectory sampling: Multiple candidate latent trajectories are sampled using classifier-guided diffusion, starting from $\mathbf{x}_0$.
3. Trajectory selection: Each candidate is scored via the classifier; the highest-scoring trajectory is selected.
4. Action decoding: The first-step transition is decoded into an action using the inverse dynamics model.
5. Application and iteration: The decoded action is executed, new local observation is observed, and the process repeats (Meng et al., 31 Jan 2026).

This process is repeated independently for each agent at every time step; agents coordinate implicitly through MF summaries without a centralized controller.

7. Empirical Validation, Scalability, and Limitations

Extensive experiments demonstrate that NetWorld achieves improved average return, sample efficiency, and robustness compared to state-of-the-art offline DTDE MARL baselines (e.g., MA-CQL, MA-IQL, MA-DT, MA-DD). Table 1 summarizes core empirical results:

Experiment Scenario	NetWorld Protocol	Reward Gain over Baselines
Coordinated Beamforming (19)	30× fewer trajectories	10–15% higher avg. reward
RB Scheduling (8–9)	Few-shot adaptation	Smoother/faster convergence
Network Slicing (19)	Offline fine-tuning	Robust to heterogeneity

NetWorld retains its strengths under increasing network scale, with performance degradation observed if mean-field communication is ablated (Meng et al., 27 Oct 2025). Planning horizon and diffusion steps affect computational complexity, but fast ODE solvers and classifier guidance stabilize sample quality and sampling speed. Limitations include computational demands for long-horizon planning, reliance on offline coverage, and higher cost versus direct policy-based inference (Meng et al., 31 Jan 2026, Meng et al., 27 Oct 2025).

Applicability is not limited to wireless networking; the architecture generalizes to multi-robot path planning, traffic signal control, energy management, and any cooperative task with local communications and distributed control constraints.

8. Theoretical Guarantees and Broader Context

The mean-field approximation and classifier-guided diffusion are theoretically shown to incur bounded modeling errors. Specifically, with appropriate Lipschitz and bounded-difference assumptions on classifier and diffusion score function, the drift difference and resulting KL divergence between ideal and mean-field-guided trajectory distributions are explicitly controlled (Meng et al., 27 Oct 2025). This guarantees the reliability and convergence of distributed planning even as network size scales.

Summary insights highlight the model's sample efficiency, stable convergence, scalability with low communication overhead, and practical feasibility for MARL under strong decentralization constraints. A plausible implication is that similar architectures could underpin future MARL systems for resource-constrained, privacy-sensitive, and large-scale distributed applications in wireless networks and beyond (Meng et al., 31 Jan 2026, Meng et al., 27 Oct 2025).