Multi-Agent Imitation by Learning and Sampling from Factorized Soft Q-Function

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MAFIS is a Multi-Agent Imitation Learning (MAIL) framework that extends the single-agent IQ-Learn method to cooperative multi-agent settings by factorizing the soft Q-function (QMIX-style) and sampling from the optimal policy as an Energy-Based Model via Stochastic Gradient Langevin Dynamics (SGLD). This eliminates adversarial optimization entirely, yielding a stable, non-adversarial objective for both discrete and continuous multi-agent control. MAFIS supports both online and offline imitation and decentralized execution.

Problem

Existing MAIL methods fail in complementary ways: Behavioral Cloning (BC) suffers from compounding error under distributional shift, while Adversarial Imitation Learning (GAIL/MA-GAIL) is training-unstable due to min-max optimization. The single-agent IQ-Learn framework avoids adversarial training by deriving a closed-form policy from the soft Q-function, but extending it to multi-agent settings is non-trivial: treating agents as a single entity prevents decentralized execution, and the required logsumexp over a joint action space is exponentially intractable in continuous control.

Method

Soft Q-function factorization: MAFIS represents the global soft Q-function as a weighted sum of per-agent Q-functions — Q_tot(τ, a) = Σ k^i(s) Q^i(τ^i, a^i) — using a QMIX-style mixing network. This factorization decouples the imitation objective (Equation 7) into a per-agent sum, enabling scalable training and avoiding exponential joint-action enumeration.

EBM sampling for continuous control: For continuous action spaces, computing the logsumexp normalization constant is intractable. MAFIS observes that the per-agent optimal policy π_{Q^i}(a^i|τ^i) is an Energy-Based Model with energy −(1/α) k^i(s) Q^i(τ^i, a^i). Gradient estimation of the imitation objective (Equation 8) only requires sampling from this EBM, which is achieved via Stochastic Gradient Langevin Dynamics (SGLD): a^{i,k+1} ← a^{i,k} + (ε²/2α) ∇_{a^{i,k}} k^i(s) Q^i(τ, a^{i,k}) + εω. This replaces adversarial policy training entirely.

Decentralized execution: During execution, global state s (required for mixing weights k^i(s)) is unavailable. Since k^i(s) > 0 always, the maximizer of the local policy π̃_{Q^i}(a|τ) ∝ exp{(1/α) Q^i(τ, a)} — using only local observation history — is identical to the full policy's maximizer. The action with highest Q-value among N SGLD samples is selected.

Key Contributions

• Non-adversarial MAIL objective for multi-agent settings by combining IQ-Learn with value decomposition (QMIX mixing network).
• First MAIL method to treat the optimal policy as an EBM and use SGLD sampling to handle continuous action spaces, bypassing logsumexp intractability.
• Unified online/offline framework covering both discrete (SMACv2, MPE, Gold Miner) and continuous (MaMuJoCo) multi-agent control.
• Decentralized execution without global state at inference time via local Q-function argmax.
• Theoretical derivation showing the MAFIS objective differs from the concurrent MIFQ method in optimal soft value function representation.

Results

Offline imitation (Table 1, mean over 5 seeds): - SMACv2 terran_5_vs_5: MAFIS 18.35 vs. BC 7.64, MIFQ 15.57 (Expert: 17.55) - SMACv2 terran_10_vs_11: MAFIS 16.27 vs. BC 5.21, MIFQ 12.27 (Expert: 18.05) - MPE simple_spread: MAFIS −18.77 vs. BC −20.27, MIFQ −39.95 (Expert: −11.44) - Gold Miner Easy: MAFIS 4004.02 vs. BC 2395.90, MIFQ 3252.17 (Expert: 4810.23) - MaMuJoCo Ant(2x4): MAFIS 3694.13 vs. BC 1547.32, MIFQ N/A (Expert: 5687.83) - MaMuJoCo Walker2d(2x3): MAFIS 3679.56 vs. BC 1743.18 (Expert: 6481.51)

Online learning (Figure 2): MAFIS consistently outperforms MA-GAIL and MIFQ across all discrete benchmarks; MA-GAIL underperforms BC on MaMuJoCo continuous tasks.

Limitations

• Still lags significantly behind expert performance, especially in continuous control (e.g., Ant(2x4): 3694 vs. expert 5688).
• SGLD sampling introduces computational overhead; N=20 parallel samples required per gradient step for stable estimation.
• Entropy parameter α is sensitive: performance degrades sharply at α=1.5 on Walker2d(2x3) due to overly exploratory sampling.
• Cooperative-only setting (fully cooperative Markov game); not evaluated on competitive or mixed-motive multi-agent tasks.
• Decentralized execution fix (dropping k^i(s)) loses the mixing weights, which may reduce expressivity relative to centralized execution.

Relevance to Agentic AI / LLM Agents

MAFIS addresses a core challenge in training multi-agent systems from demonstrations rather than hand-crafted rewards — a paradigm directly relevant to agentic pipelines where reward specification is impractical but expert rollouts (human demonstrations, stronger model traces) are available. The factorized, decentralized execution model maps naturally onto agent architectures where each agent acts on local context, mirroring how individual LLM-based agents in a multi-agent system receive partial observations. The EBM-SGLD approach for continuous action sampling is also conceptually related to diffusion-style action generation emerging in robot/agent learning. This work strengthens the theoretical and empirical foundation for scalable imitation in cooperative multi-agent settings beyond toy tasks.