Multi-Agent Coordination via Support-Preserving Distillation

Sangmin Lee; Youngju Na; Chanmi Lee; Sung-eui Yoon

cs.LGOct 7, 2026

Multi-Agent Coordination via Support-Preserving Distillation

Authors: Sangmin Lee, Youngju Na, Chanmi Lee, Sung-eui Yoon

Organizations: School of Computing, KAIST

Abstract

Offline MARL increasingly relies on generative policies to model multimodal joint behavior, typically by distilling a centralized teacher into decentralized one-step actors under the CTDE. We identify a failure mode at the teacher training stage: standard flow-based teachers pair noise with replay targets independently, so nearby noise samples can be routed toward conflicting coordination modes. The teacher then produces samples between valid modes, and because the distillation loss regresses each local actor onto the conditional mean of the teacher's output given local input, this error is not absorbed but propagated to the student. To remove this teacher-side artifact, we propose Mode-Support Semi-Discrete Optimal Transport (MoSDOT), which summarizes multimodal replay into a finite mode support with prescribed capacities and uses conditional semi-discrete optimal transport to assign each noise sample to a single mode before teacher training. We additionally study a shared-randomness variant that uses a shared noise component at execution to expose the residual gap intrinsic to strict-product execution. On controlled diagnostics and offline MARL benchmarks, MoSDOT improves endpoint quality and routing consistency, particularly on datasets exhibiting multimodal joint behavior.

Figures & tables

Figure 1: Overview of source-to-mode routing. (Left) Standard Flow Policy: independent noise–target pairing can entangle nearby source samples with incompatible coordination modes, producing off-support fan mass. (Middle) MoSDOT: semi-discrete optimal transport partitions noise space into capacity-matched cells assigned to coordinated behavior modes. (Right) Route quality is measured by off-support avoidance ( fan-free endpoints ), consistency of source-to-mode routing ( clean route consistency ), endpoint accuracy to valid modes ( strict success ), and balanced mass across modes ( balanced coverage ). Full metric definitions are provided in Appendix E.6 .

Figure 2: MoSDOT-based framework for multi-agent distillation. (a) Replay joint actions are summarized into a mode support

\{\mathbf{y}_{k}(\mathbf{x})\}

with capacities

\{b_{k}(\mathbf{x})\}

, and conditional SDOT partitions the noise space into Laguerre cells

\Gamma_{\textbf{x}}^{-1}(k)

assigned to each mode. (b) The centralized teacher

\mu_{\psi}

is distilled into decentralized one-step actors

\{\mu_{w_{i}}\}

, each using only its local observation and noise. An optional shared signal

h

correlates local noise to mitigate the strict-product gap.

Figure 3: Landmark policy trajectories. Black $\times$ marks denote the six anchor mode endpoints, and colors indicate the three anchor pairs. The top row shows centralized joint teacher rollouts, and the bottom row shows one-step factorized students distilled under independent local noise. Baseline students show fan mass between modes, while MoSDOT reduces fan mass and retains the structure.

Figure 6: Source-space partitions in Agent 1’s noise plane.

Figure 5: XOR strict-product visualization. (a) Target support placing equal mass on $(-1,+1)$ and $(+1,-1)$ . (b) MAC-Flow : bridge mass spreads across all quadrants under strict-product execution. (c) Ours w/o shared signal : clean teacher, residual strict-product gap. (d) Ours with shared signal : gap closed. Blue: factored student; pink: joint teacher.

Benchmark	Scenario	Dataset	Offline / BC baselines			Diffusion policies			Flow / OT policies
Benchmark	Scenario	Dataset	BC	MABCQ	MACQL	Diffusion BC	MADiff	DoF	Flow BC	MAC-Flow	MoSDOT
SMACv1	3m	Good	$16.0$ $\pm\,1.0$	$3.7$ $\pm\,1.1$	$19.1$ $\pm\,0.1$	$19.5$ $\pm\,0.5$	$19.3$ $\pm\,0.5$	$19.8$ $\pm\,0.2$	$20.0$ $\pm\,0.0$	$19.8$ $\pm\,0.2$	$20.0$ $\pm\,0.0$
		Medium	$8.2$ $\pm\,0.8$	$4.0$ $\pm\,1.0$	$13.7$ $\pm\,0.3$	$13.3$ $\pm\,0.7$	$16.4$ $\pm\,2.6$	$18.6$ $\pm\,1.2$	$14.7$ $\pm\,1.5$	$18.0$ $\pm\,3.2$	$18.8$ $\pm\,1.5$
		Poor	$4.4$ $\pm\,0.1$	$3.4$ $\pm\,1.0$	$4.2$ $\pm\,0.1$	$4.2$ $\pm\,0.2$	$10.3$ $\pm\,6.1$	$10.9$ $\pm\,1.1$	$4.5$ $\pm\,0.1$	$10.6$ $\pm\,2.2$	$16.4$ $\pm\,4.9$
	8m	Good	$16.7$ $\pm\,0.4$	$4.8$ $\pm\,0.6$	$18.9$ $\pm\,0.9$	$19.4$ $\pm\,0.5$	$18.9$ $\pm\,1.1$	$19.6$ $\pm\,0.3$	$19.5$ $\pm\,0.2$	$19.7$ $\pm\,0.3$	$20.0$ $\pm\,0.0$
		Medium	$10.7$ $\pm\,0.5$	$5.6$ $\pm\,0.6$	$15.5$ $\pm\,1.5$	$18.6$ $\pm\,0.6$	$16.8$ $\pm\,1.6$	$18.6$ $\pm\,0.8$	$18.2$ $\pm\,0.8$	$19.4$ $\pm\,0.6$	$20.0$ $\pm\,0.0$

Table 2: Performance on discrete-action offline MARL benchmarks (SMACv1 and SMACv2). Mean over

6

seeds,

\pm 2\sigma

. Best and second-best in sage and powder blue .

Scenario	Dataset	Offline MARL					Generative policies
Scenario	Dataset	MATD3BC	MACQL	ICQ	OMAR	OMIGA	MADiff	MAC-Flow	MoSDOT
Spread	Expert	$108.3$ $\pm\,3.3$	$98.2$ $\pm\,5.2$	$114.9$ $\pm\,2.6$	$104.0$ $\pm\,3.4$	$80.8$ $\pm\,13.8$	$95.0$ $\pm\,5.3$	$101.7$ $\pm\,10.9$	$104.40$ $\pm\,19.66$
	Medium	$29.3$ $\pm\,4.8$	$34.1$ $\pm\,7.2$	$47.9$ $\pm\,18.9$	$29.3$ $\pm\,5.5$	$30.1$ $\pm\,16.9$	$64.9$ $\pm\,7.7$	$80.1$ $\pm\,20.6$	$95.97$ $\pm\,12.69$
	M-R	$15.4$ $\pm\,5.6$	$20.0$ $\pm\,8.4$	$37.9$ $\pm\,12.3$	$13.6$ $\pm\,5.7$	$5.4$ $\pm\,11.0$	$30.3$ $\pm\,2.5$	$50.4$ $\pm\,33.2$	$54.12$ $\pm\,15.29$
	Random	$9.8$ $\pm\,4.9$	$24.0$ $\pm\,9.8$	$34.4$ $\pm\,5.3$	$6.3$ $\pm\,3.5$	$-3.8$ $\pm\,12.3$	$6.9$ $\pm\,3.1$	$31.1$ $\pm\,6.8$	$75.92$ $\pm\,23.81$
Average rewards		$40.7$	$44.1$	$58.8$	$38.3$	$28.1$	$49.2$	$65.8$	$82.6$

Table 3: Performance on continuous-action offline MARL benchmarks (MPE Simple Spread). M-R denotes Medium-Replay.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Preprocessing (s)	Training (min)	Total (min)
Standard flow matching	70	$37.75\pm 0.01$	38.9
MoSDOT	279	$38.88\pm 0.29$	43.5

Appendix

Table 4: Landmark wall-clock comparison. Training uses

40{,}000

steps; training times are mean

\pm

one standard deviation over four seeds. Total time includes one preprocessing pass and the mean training time.

Agents	Dataset	Retained rows $M$	Occupied support $K$	$K/M$	Cache time (s)
3	SMACv1 3m / Poor	200,000	6,625	3.3%	10.5
3	MPE Spread / Medium	5,000,000	72,509	1.5%	122.4
5	SMACv1 5m_vs_6m / Good	200,000	20,454	10.2%	10.6
5	SMACv2 zerg_5_vs_5 / Replay	200,000	2,380	1.2%	38.6
8	SMACv1 8m / Poor	200,000	32,653	16.3%	67.1
10	SMACv2 terran_10_vs_10 / Replay	200,000	88,729	44.4%	147.7

Appendix

Table 5: Occupied replay support and one-time cache cost.

M

and

K

are totals over the retained dataset, not counts of all possible joint actions. Times measure cache construction before training.

Figure 7: Joint mode-tuple support, teacher vs. student. Bars show the fraction of

32{,}768

samples (fan included in the denominator) falling into each of the 8 joint mode tuples on the 3-agent landmark task; the dashed line marks the zero-fan uniform ideal

1/8

. Top row: centralized teacher (joint flow ODE). Bottom row: one-step factorized student distilled from each teacher under matched architecture and budget. MoSDOT preserves both validity and balance through distillation (95.0%

\to

91.7% valid), while IMLE loses half its valid mass (97.3%

\to

47.1%) and Drift collapses entirely to the fan region (86.3%

\to

1.0%). This is the per-mode visualization of the aggregate scores in Tab. 6 and the empirical counterpart of the projection result in Appendix A.1 .

Method	Fan-free $\uparrow$	Success $\uparrow$	Balance $\uparrow$
Flow matching	0.849 $\pm 0.004$	0.846 $\pm 0.038$	0.949 $\pm 0.024$
IMLE	0.784 $\pm 0.011$	0.558 $\pm 0.072$	0.621 $\pm 0.104$
Drifting	0.364 $\pm 0.078$	0.685 $\pm 0.071$	0.328 $\pm 0.417$
MoSDOT (ours)	0.971 $\pm 0.003$	0.971 $\pm 0.004$	0.963 $\pm 0.019$

Appendix

Table 6: One-step student landmark diagnostics. Mean

\pm

one standard deviation over 4 seeds; same

8192

noises

\times\,3

agents per seed as the teacher table. Best and second-best in sage and powder blue . Higher is better for all metrics.

Method	Sample-level				MLP-level
Method	$H(m_{i}\!\mid\!z_{i}\,\text{cell})$	Mix rate	KNN flip	AUC excess	Jac. ratio	CF flip
	(lower is better)				(lower is better)
Flow	0.019 $\pm$ 0.024	0.006 $\pm$ 0.010	0.001 $\pm$ 0.000	0.020 $\pm$ 0.008	0.066 $\pm$ 0.005	0.000 $\pm$ 0.000
IMLE	0.319 $\pm$ 0.109	0.098 $\pm$ 0.036	0.135 $\pm$ 0.051	0.105 $\pm$ 0.042	0.292 $\pm$ 0.055	0.129 $\pm$ 0.050
Drift	0.666 $\pm$ 0.252	0.244 $\pm$ 0.117	0.315 $\pm$ 0.134	0.290 $\pm$ 0.123	0.772 $\pm$ 0.148	0.307 $\pm$ 0.128
Ours	0.064 $\pm$ 0.024	0.020 $\pm$ 0.010	0.011 $\pm$ 0.003	0.012 $\pm$ 0.003	0.046 $\pm$ 0.015	0.002 $\pm$ 0.002

Appendix

Table 7: Full route-locality diagnostics. Sample-level metrics use 12 points per method (4 seeds

\times

3 agents).

H

denotes own-noise cell entropy with 10 bins, Mix denotes own-cell mix rate, KNN flip uses

k{=}20

, and AUC excess reports the maximum off-diagonal LDA AUC minus

0.5

. MLP-level metrics report 4 points for the Jacobian ratio (4 seeds, agent-averaged) and 12 points for CF flip (4 seeds

\times

3 agents). All values are mean

\pm

std.

Method	$H,\,b{=}5$	$H,\,b{=}10$	$H,\,b{=}20$	KNN $k{=}10$	KNN $k{=}20$	KNN $k{=}50$
Flow	0.109 $\pm$ 0.061	0.021 $\pm$ 0.024	0.005 $\pm$ 0.006	0.000 $\pm$ 0.000	0.001 $\pm$ 0.000	0.001 $\pm$ 0.001
IMLE	0.360 $\pm$ 0.101	0.319 $\pm$ 0.109	0.291 $\pm$ 0.106	0.134 $\pm$ 0.051	0.135 $\pm$ 0.051	0.137 $\pm$ 0.051
Drift	0.690 $\pm$ 0.240	0.666 $\pm$ 0.252	0.634 $\pm$ 0.252	0.314 $\pm$ 0.134	0.315 $\pm$ 0.134	0.316 $\pm$ 0.133
Ours	0.181 $\pm$ 0.015	0.064 $\pm$ 0.024	0.032 $\pm$ 0.010	0.008 $\pm$ 0.003	0.011 $\pm$ 0.003	0.019 $\pm$ 0.003

Appendix

Table 8: Route-locality sensitivity. Each entry reports 12 points per method (4 seeds

\times

3 agents), mean

\pm

std. The relative ranking is preserved across all bin counts and neighborhood sizes.

Explore similar work

CoDiMAD: Diffusion-Based Privileged Distillation for Communication-Free Multi-Robot Coordination
Jul 10, 2026Jiyue Tao, Shunheng Xin, Tongsheng Shen +2Multi-Robot Systems Aware Heterogeneous Multi-Teacher Multimodal On-Policy Distillation
Robotics and Control Laboratory, School of Advanced Manufacturing and Robotics, and the State Key Laboratory of Turbulence and Complex Systems, Peking University, Beijing, 100871, China · National Innovation Institute of Defense Technology, Beijing 100071, China
CODA: Coordination via On-Policy Diffusion for Multi-Agent Offline Reinforcement Learning
Apr 25, 2026Marcel Hedman, Kale-ab Abebe Tessera, Juan Claude Formanek +5Multi-Agent Reinforcement Learning Diffusion Policies
Inephany Ltd · Department of Statistics, University of Oxford · The University of Edinburgh +2
MAS-OPD: On-Policy Distillation for Multi-agent Systems
Sep 28, 2026Qiyong Zhong, Mao Zheng, Mingyang Song +5Aware Heterogeneous Multi-Teacher Multimodal On-Policy Distillation Token-Level Supervision
University of Science and Technology of China · Foundation Model Department, Tencent · Zhejiang University +1

Jul 10, 2026, cs.ROJ/K move · Enter open · S save

CoDiMAD: Diffusion-Based Privileged Distillation for Communication-Free Multi-Robot Coordination

Jiyue Tao, Shunheng Xin, Tongsheng Shen, Dexin Zhao+1

Robotics and Control Laboratory, School of Advanced Manufacturing and Robotics, and the State Key Laboratory of Turbulence and Complex Systems, Peking University, Beijing, 100871, China · National Innovation Institute of Defense Technology, Beijing 100071, China

Multi-Robot Systems Aware Heterogeneous Multi-Teacher Multimodal On-Policy Distillation Coordination Cooperation Diffusion-Based Dataset Distillation

Open paper AskPDF arXiv

Decentralized multi-robot coordination under partial observability remains challenging, especially in communication-free settings where agents must act solely from local sensor observations. Privileged policy distillation provides a promising approach by transferring knowledge from a globally informed oracle to sensor-constrained students. However, in multi-agent systems, the same local observation may correspond to multiple global configurations requiring qualitatively different cooperative actions, making the conditional action distribution inherently multi-modal. Standard deterministic distillation collapses these modes to their mean, often yielding invalid or hesitant actions. To address this issue, we propose CoDiMAD, a three-stage framework that trains a privileged oracle with MAPPO, constructs an offline dataset of local-observation-oracle-action pairs, and distills the oracle into decentralized students parameterized as conditional denoising diffusion probabilistic models. By approximating the conditional oracle-action distribution through the diffusion reverse process, CoDiMAD samples decisive actions from coherent coordination modes rather than averaging across them. Theoretical analysis characterizes the mode-averaging failure of deterministic distillation and the distributional recovery property of diffusion-based distillation. Experiments on three cooperative tasks show that CoDiMAD consistently outperforms direct local MARL and deterministic distillation baselines. The source code will be made publicly available upon acceptance.

Abstract

Explore similar work

CoDiMAD: Diffusion-Based Privileged Distillation for Communication-Free Multi-Robot Coordination

CODA: Coordination via On-Policy Diffusion for Multi-Agent Offline Reinforcement Learning

MAS-OPD: On-Policy Distillation for Multi-agent Systems