Test-time Multi-agent Coordination by Decomposed Value Gradient Flow

2026-10-01

AI Summary

In offline multi-agent reinforcement learning, generative policies struggle to identify high-value regions, while value optimization often induces mode collapse and disrupts coordination. This work proposes SCOUT, a novel framework that pioneers the decoupling of flow-matching behavioral priors from decomposed value functions. It leverages Stein variational gradient descent for test-time action refinement and provides theoretical guarantees via a KL bound under the Individual-Global-Max (IGM) principle. The proposed method effectively balances scalability with coordination, achieving state-of-the-art average performance across both discrete and continuous offline benchmarks. Furthermore, it yields substantial improvements in all offline-to-online configurations.

Abstract

Offline multi-agent reinforcement learning (MARL) faces a persistent trade-off. Expressive generative policies can represent multi-modal coordination in the data, but cannot distinguish high-value regions, while value-optimized policies exploit the learned Q-function but collapse the multi-modal into a single dominant mode. A single agent's mode collapse can break joint coordination, and simultaneous drift across agents can push the joint policy into unseen regions of the action space. We propose scalable coordination via optimal unified transport (SCOUT), the first offline MARL framework to combine a generative foundation model with a learned value function through test-time action refinement. SCOUT trains two decoupled components: a flow-matching behavioral prior and a decomposed value function. At test-time, it transports behavioral samples toward high-value regions via Stein variational gradient descent. The number of transport steps controls adaptive test-time scaling, replacing a fixed regularization coefficient. Under the individual-global-max (IGM) principle, we prove a single-term KL bound on the joint soft-value gap that vanishes as transport converges, with an irreducible additive residual proportional to the IGM violation. Empirically, SCOUT achieves the best average performance across discrete and continuous offline MARL benchmarks and yields performance improvements in all offline-to-online configurations.

Problem

• offline multi-agent reinforcement learning
• mode collapse
• generative policy
• value optimization
• multi-modal coordination

Innovation

• Offline Multi-agent Reinforcement Learning
• Test-time Scaling
• Flow Matching
• Stein Variational Gradient Descent
• Value Decomposition