Conditional Diffusion Model for Multi-Agent Dynamic Task Decomposition

TL;DR

The paper tackles the challenge of coordinating many agents under partial observability by learning dynamic task decompositions. It introduces CD^3T, a two-level hierarchical MARL framework that uses a diffusion model to learn latent action representations, clusters them into subtasks, and employs a subtask selector plus subtask policies to guide coordination. A diffusion-based latent representation and an attention-driven credit assignment mechanism underpin robust, interpretable decomposition and scalable learning, validated by strong results on LBF, SMAC, and SMACv2. The approach yields improved performance, reduced effective action spaces, and insightful visualizations of dynamic subtasks, suggesting practical impact for complex cooperative tasks in uncertain environments.

Abstract

Task decomposition has shown promise in complex cooperative multi-agent reinforcement learning (MARL) tasks, which enables efficient hierarchical learning for long-horizon tasks in dynamic and uncertain environments. However, learning dynamic task decomposition from scratch generally requires a large number of training samples, especially exploring the large joint action space under partial observability. In this paper, we present the Conditional Diffusion Model for Dynamic Task Decomposition (C$\text{D}^\text{3}$T), a novel two-level hierarchical MARL framework designed to automatically infer subtask and coordination patterns. The high-level policy learns subtask representation to generate a subtask selection strategy based on subtask effects. To capture the effects of subtasks on the environment, C$\text{D}^\text{3}$T predicts the next observation and reward using a conditional diffusion model. At the low level, agents collaboratively learn and share specialized skills within their assigned subtasks. Moreover, the learned subtask representation is also used as additional semantic information in a multi-head attention mixing network to enhance value decomposition and provide an efficient reasoning bridge between individual and joint value functions. Experimental results on various benchmarks demonstrate that C$\text{D}^\text{3}$T achieves better performance than existing baselines.

Figures (11)

• Figure 1: The overall framework of C$\text{D}^\text{3}$T. We first derive a latent action representation $z_{a_i}$ for each agent from its action space, conditioned on its local observation $o_i$ and other agents' one-hot actions $\boldsymbol{a}_{-i}$, to pretrain a diffusion model. Latent representations are then clustered to define subtask-specific action spaces. The subtask selector and subtask policy share the same architecture with different parameters. At every $\Delta T$ steps, the selector assigns a subtask to each agent and estimates the joint Q-value $Q_{tot}^\Phi$ using the global state $s_t$ and subtask representation $\boldsymbol{z}_\phi$, while the subtask policy computes $Q_{tot}$ with $s_t$ and the action representation $\boldsymbol{z}_a$.
• Figure 2: Performance comparison with baselines on LBF.
• Figure 3: Performance comparison with baselines on easy, hard, and super hard scenarios.
• Figure 4: Ablation studies of C$\text{D}^\text{3}$T on SMAC benchmark.
• Figure 5: The process of subtask generation through subtask representations. (a) illustrates that moving northward or eastward has a comparable effect on the environment by directing agents toward the enemies, whereas moving southward or westward moves agents away. (b) depicts the distribution of action representations in a two-dimensional space after PCA projection. (c) visualizes the formation of subtasks derived from action representations following clustering.

Theorems & Definitions (1)

• Theorem 1