Multi-Agent Coordination via Support-Preserving Distillation

AI summary

Overview

Research area: Offline multi-agent reinforcement learning (MARL), generative policy learning (flow matching, diffusion), and policy distillation under centralized training with decentralized execution (CTDE).

Technical level: Advanced. The paper builds on semi-discrete optimal transport, flow-matching velocity-field regression, Dec-POMDP formalism, and IGM-constrained value factorization. The motivating problem is stated accessibly, but the method requires familiarity with generative modeling and multi-agent RL.

Scope: The paper identifies and repairs a teacher-side failure mode in generative offline MARL distillation by aligning source noise with discrete coordination modes via conditional semi-discrete optimal transport, and separately probes the residual factorization error that remains at decentralized execution.

What This Paper Is About

In offline MARL, a centralized generative "teacher" is trained on joint actions and then distilled into one-step decentralized "students" that each act from their own observation and noise. The paper shows that standard flow-based teachers pair noise samples with replay targets independently, so nearby noise points can be routed toward conflicting coordination modes; the teacher then emits samples sitting between valid modes, and the ℓ2 distillation loss regresses each local actor onto the conditional mean, propagating rather than absorbing that error. The goal is to remove this teacher-side artifact — while being explicit about the second, separate error that comes from independent local sampling at execution time.

Key Contributions

1. Diagnosis of a teacher-side failure mode. The paper isolates a stage before distillation: independent noise–target pairing in standard flow-based teacher training entangles nearby source samples with incompatible coordination modes, producing off-support "fan" mass between valid modes. Because the distillation loss regresses each local actor onto the conditional mean of the teacher's output given local input, this teacher error is not absorbed but propagated to students.

2. MoSDOT (Mode-Support Semi-Discrete Optimal Transport). A method that summarizes multimodal replay into a finite mode support (a discrete set of representative joint actions per context) with prescribed capacities, then partitions source-noise space via conditional semi-discrete optimal transport into Laguerre cells, each assigned to a single mode, before teacher training. Capacities are set to b_k(x) = |C_k(x)| / M_x, matching the empirical replay frequency of each mode. Source assignment operates on the joint mode y_k(x) as a whole, not on a product of per-agent action modes.

3. Formalization of the strict-product gap. The paper defines the residual mismatch between what a centralized teacher can represent and what a strict-product actor can represent, using the anti-correlated XOR target {(-1,+1), (+1,-1)} as the canonical illustration: a strict-product student matching the local marginals also assigns mass to the unsupported tuples (+1,+1) and (-1,-1).

4. A shared-randomness variant to isolate that gap. A policy form where agents read a shared execution-time signal h ~ ρ_h but still act from their own observations and local noise ε_i, with no exchange of observations, actions, or messages. This separates teacher-side assignment error from product-projection error.

Main Findings

• Endpoint accuracy does not imply distillable assignments. On the landmark diagnostic, Flow matching reaches a fan-free score of 0.857 ± 0.010 and success of 0.915 ± 0.029 but shows samples between modes. IMLE achieves the best fan-free score (0.991 ± 0.007) with sharp endpoints but sacrifices success (0.848 ± 0.064) and balance (0.761 ± 0.029). Drifting shows the mismatch most strongly: high success (0.985 ± 0.002) but the lowest route consistency (0.653 ± 0.023) and the largest Jacobian off-diagonal ratio, meaning its per-agent output is strongly shaped by other agents' noise — information unavailable to the local actor.

• MoSDOT leads on routing while remaining competitive on endpoints. MoSDOT records route consistency of 0.972 ± 0.001 (highest across methods), fan-free 0.983 ± 0.002, success 0.991 ± 0.002, and balance 0.961 ± 0.011 (also the best). Its Jacobian off-diagonal ratio is reported as comparable to flow matching and well below IMLE and Drifting. Table 1 reports means ± one standard deviation over 4 seeds.

• Source-space visualization supports the two-axis claim. In Agent 1's noise plane, Flow leaves a fan between modes, IMLE assigns unbalanced area, and Drifting mixes the two modes throughout, while MoSDOT yields balanced, fan-less, disentangled regions. The teacher's joint mode coverage is reported as the highest across methods.

• Cleaner teacher assignments survive decentralized distillation. Students distilled from Flow Matching, IMLE, Drifting, and MoSDOT teachers under the same architecture, seed list, aggregation rule, and training budget show that the MoSDOT student largely follows the teacher's ray structure, while baselines show fan mass between modes. The MoSDOT-distilled student retains higher fan-free, success, and balance scores than baselines, whereas IMLE and Drifting students lose substantial endpoint quality through distillation.

• Shared randomness isolates the strict-product gap. On XOR, the unrouted flow teacher places bridge mass between the two valid tuples and a strict-product student spreads across all four quadrants; MAC-Flow shows bridge mass across all quadrants under strict-product execution. MoSDOT concentrates teacher mass on the two valid tuples, yet independent local sampling still mixes the modes after distillation. Adding the shared signal with the same MoSDOT teacher concentrates the student on the valid tuples. Because only the execution-time signal changes between those last two stages, the strict-product gap is isolated from teacher-side assignment failure.

• Benchmark results. On SMACv1 (6 seeds, ±2σ), average rewards are: BC 12.2, MABCQ 5.5, MACQL 13.1, Diffusion BC 13.0, MADiff 13.8, DoF 15.6, Flow BC 13.4, MAC-Flow 15.6, MoSDOT 16.7. On SMACv2, averages are: BC 7.2, MABCQ 12.3, MACQL 11.3, Diffusion BC 7.6, MADiff 12.4, DoF 14.0, Flow BC 6.2, MAC-Flow 13.1, MoSDOT 11.9. The paper describes MoSDOT as competitive with existing generative baselines, particularly on datasets containing diverse joint behaviors. The MPE Simple Spread table is truncated in the provided content; only MATD3BC (108.3 ± 3.3), MACQL (98.2 ± 5.2), and ICQ (114.9 ± 2.6) on the Expert dataset are visible, and MoSDOT's MPE numbers are not reported in the available text.

Methodology in Plain English

The framework is standard CTDE. A centralized teacher takes the full joint observation context and a joint noise vector and outputs a joint action; the students are one-step decentralized actors, each seeing only its own observation and its own noise block, and their joint policy is the product of the per-agent policies. A centralized critic satisfying the Individual-Global-Max principle supports value-guided actor extraction.

The distinctive step happens before teacher training. For each context, replay joint actions are grouped into a small number of clusters by an aggregation rule, and each cluster is summarized by a representative joint action called a mode. The number of modes is set by the observed targets and that rule, not by enumerating the full combinatorial joint-action space. Identity support is used for discrete-action benchmarks and simple joint-action quantization for continuous ones. Each mode receives a capacity equal to the fraction of replay targets in its cluster.

Source noise is then partitioned so that each cell carries the prescribed mass and maps to exactly one mode. This is a semi-discrete optimal transport problem: the target is a discrete measure placing mass b_k(x) on each mode, the cost is squared Euclidean in the space defining the mode support, and the primal constraint requires each cell to carry its assigned mass. The dual is a finite-dimensional concave problem solvable stochastically with Monte-Carlo samples from the noise law. The optimal dual weights define Laguerre cells, and the assignment is the argmin of cost minus dual weight. The teacher then learns a velocity field by regressing onto the displacement from the source sample to its assigned mode along a linear interpolation path — the standard flow-matching objective, but with the assignment replacing independent pairing.

Distillation regresses each local actor onto the corresponding component of the centralized teacher rollout, optionally combined with a negative centralized Q term weighted against the teacher regularization by a coefficient α. The assignment is used only to train the teacher and is never given to the actors at execution. The shared-randomness variant adds a signal h read by all agents at execution; without h it reduces to strict-product execution.

Why This Matters

Impact on research. The paper reframes a known problem in offline MARL generative distillation: rather than treating distillation as the sole source of mode collapse, it shows the teacher's source-to-mode routing is an upstream cause that pointwise student losses then amplify. By separating teacher-side routing artifacts from the residual product-projection gap and providing measured interfaces for each, it gives the CTDE community a cleaner decomposition of where multimodal coordination capability is lost. It also connects offline MARL to semi-discrete optimal transport techniques previously used in image generation.

Real-world applications:

• Multi-robot fleets (warehouse, logistics) where robots must commit jointly to one of several valid coordination patterns rather than averaging between them.
• Cooperative autonomous driving, where distinct valid negotiation maneuvers must not blend into an unsafe in-between trajectory.
• Swarm and multi-drone control, where individual per-agent plausibility can combine into unseen joint actions.
• Game AI and simulated unit control (the SMACv1/SMACv2 benchmarks derive from StarCraft-style micromanagement).

Industry relevance. Offline RL learns from fixed logged data without further environment interaction, which matters where live exploration is expensive, slow, or unsafe — robotics, industrial control, and simulation-driven product teams. One-step decentralized actors are also cheaper to deploy than multi-step diffusion samplers, which the paper identifies as a motivation for this whole teacher–student line of work. The shared-randomness variant, however, presumes a shared execution-time signal, so deploying it requires an available coordination channel; the paper notes agents exchange no observations, actions, or messages but still correlate mode choice through that signal.

Future Directions

• How to obtain shared randomness in practice. The shared-randomness variant closes the strict-product gap in the XOR diagnostic, but the paper does not specify how such a shared signal would be produced or synchronized in a real decentralized deployment.
• Choosing the aggregation rule and support size. The support size K_x is determined by the observed targets and the aggregation rule, and the rule fixes the granularity of the mode support. How sensitive results are to that choice, and how to select it principledly, is left open (identity support for discrete benchmarks, joint-action quantization for continuous ones).
• Reducing the strict-product gap without an execution-time channel. MoSDOT "leaves the product-projection gap explicit" rather than solving it; whether teacher-side construction alone can narrow it is unresolved.
• Scalability and broader validation. The method's behavior with larger agent counts, richer continuous action spaces, and datasets whose modality structure is less clean than the landmark and XOR diagnostics is not established in the reported content. Detailed theory and algorithmic extensions are deferred to appendices (A.5, B.2), and further distillation results to Appendix F.1, which are not included in the truncated text.

Target Audience

Researchers and graduate students working on offline reinforcement learning, multi-agent coordination, and generative policies — particularly those already familiar with CTDE, flow matching, diffusion policies, and value factorization. It is also relevant to practitioners building multi-agent systems from logged data who need one-step decentralized actors and are considering where multimodal coordination quality is actually lost in their pipeline. Readers without a background in optimal transport or generative modeling will find the problem framing accessible but the method section demanding.

Authors’ 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.