PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Dongsu Lee, Haoran Xu, Amy Zhang

arXiv:2610.02554v1cs.LGcs.RO

TL;DR

Offline MARL must learn coordinated behavior from offline data while preserving diverse modes and limiting unsupported actions. SCOUT combines a flow-matching behavioral prior with decomposed value-based test-time refinement. It achieves the best average performance across continuous and discrete domains and improves performance in all offline-to-online experiments.

  • Problem

    Offline MARL must learn coordinated behavior from heterogeneous, suboptimal demonstrations while avoiding policy actions in state-action regions unsupported by the data.

    Source: §1 · p. 2

  • Method

    SCOUT trains a flow-matching behavioral prior and decomposed value function, then refines behavioral samples toward high-value regions through decentralized test-time action transport.

    Source: §1 · p. 2, §4 · p. 5, §4.1 · p. 5

  • Results

    SCOUT achieves the best average performance in continuous and discrete domains and improves performance in all offline-to-online RL experiments.

    Source: §1 · p. 2

  • Takeaways & Limitations

    A test-time refinement budget replaces training-time regularization tuning while separating behavioral modeling from value improvement.

    Source: §1 · p. 2

    Refinement hyperparameters show task- and dataset-dependent sensitivity, and a principled adaptive stopping criterion remains open.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Offline MARL must model diverse coordinated behavior while avoiding policies that drift beyond the offline data. SCOUT preserves a generative behavioral prior and postpones value guidance until test time.

1 Introduction

  • Value optimization can collapse diverse coordination modes, while independently drifting agents can produce joint actions unsupported by the offline data.A fixed regularization coefficient constrains marginal drift but cannot bound its compound effect on the joint policy.Source: §1 · p. 2
  • Test-time value guidance refines samples from a behavioral prior toward high-value regions, with the adjustment extent replacing a fixed regularization coefficient.This approach is motivated by the need to preserve alternatives that training-time mode selection can suppress.Source: §1 · p. 2
  • SCOUT separates behavioral modeling from value improvement, replacing training-time regularization tuning with a test-time refinement budget.It combines flow-matching behavioral priors with decomposed value gradients for test-time action refinement.Source: §1 · p. 2
  • SCOUT achieves the best average performance in continuous and discrete domains and improves performance in all offline-to-online RL experiments.The refinement strengthens coordination without retraining.Source: §1 · p. 2

2 Related Works

Prior offline MARL methods have limited ability to capture multimodal coordination, while expressive generative policies still couple value learning to policy generation. SCOUT instead refines generated actions at test time.

2 Related Works

  • Offline MARL methods use value conservatism, value decomposition, centralized critics, or world models, but often struggle to capture multimodal data with limited policy expressiveness.Generative policies improve expressiveness, yet some rely on BC-like anchoring while others backpropagate value gradients through iterative sampling.Source: §2 · p. 3
  • SCOUT uses Stein variational gradient fields to move flow-generated action candidates toward high-value regions without explicit policy parameterization.This retains generative expressiveness while providing adaptive test-time scaling.Source: §2 · p. 3
  • Single-agent test-time refinement includes candidate rejection, gradient optimization, residual policies, and particle-based optimal-transport methods.The cited approaches remain limited to single-agent settings.Source: §2 · p. 3

3 Preliminaries

The paper frames offline cooperative MARL as learning decentralized coordinated policies from a fixed dataset, using decomposed values and behavioral-policy extraction to address the joint action problem.

Problem setting

  • The setting is a cooperative Dec-POMDP with multiple agents, local observations and actions, a shared transition process, and offline trajectories from unknown behavioral policies.The framework assumes centralized training and decentralized execution.Source: §3 · p. 3

Value decomposition

  • Under individual-global-max (IGM), additive per-agent utilities allow independent action maximization to yield a joint optimum.The global action-value function is decomposed into per-agent utilities.Source: §3 · p. 3

Policy extraction

  • Direct maximization of a learned Q-function can produce out-of-distribution errors, motivating behavior-regularized policy extraction.A fixed behavioral-distance constraint can be expressed with a divergence measure and threshold.Source: §3 · pp. 3–4
  • Mode-seeking extraction can yield Gaussian or distilled policies that collapse modes, while gradients through iterative generative sampling can be expensive and numerically fragile.The cited alternatives include parameterized gradients and one-step distillation.Source: §3 · p. 4
  • Best-of-N samples behavioral actions and selects the highest-valued candidate, preserving dataset strategies that independent mode selection could disrupt.Flow matching models a behavioral distribution by transporting noise to data and generates samples through ODE integration.Source: §3 · p. 4

4 Scalable Coordination via Optimal Unified Transport (SCOUT)

SCOUT combines a flow-matching behavioral prior and decomposed value function, then refines sampled actions independently at test time. Its analysis bounds the joint soft-value gap under approximate IGM, and a toy example contrasts this refinement with training-time value guidance.

4 Scalable Coordination via Optimal Unified Transport (SCOUT)

  • SCOUT constructs an implicit policy by transporting behavioral samples toward a maximum-entropy value target, with the gradient-step count controlling deployment-time coordination strength.When L = 0 the method recovers behavioral-policy sampling; more steps move actions toward Q-maximizing regions.Source: §4 · p. 5

4.1 Why is this hard in multi-agent settings?

  • The multi-agent challenge is to refine actions using only each agent’s local observation and Q-function while keeping the resulting joint action coherent.IGM aligns independent maximization at the greedy argmax, but does not by itself establish alignment throughout iterative transport or restore correlations missing from individual priors.Source: §4.1 · p. 5

4.2 Proposed solution

  • SCOUT trains a flow-matching behavioral prior and additive per-agent critics, then uses local value gradients to refine behavioral action particles at test time.SVGD smooths value gradients across particles and adds repulsion to prevent collapse; discrete actions are refined in logit space.Source: §4.2 · pp. 5–6

4.3 Theoretical analysis

  • Under approximate IGM, the joint soft-value gap has a transport KL term plus an irreducible 4δ residual, and the KL term vanishes as L grows under idealized assumptions.The result assumes population SVGD, exact Q-functions, and suitable step sizes; per-agent transport also tracks joint best-of-P within the IGM gap.Source: §4.3 · p. 6

4.4 Algorithm summary

  • Training alternates between flow-matching prior learning and TD learning for decomposed Q-functions, while test-time transport steps control refinement independently of the training budget.Ltest = 0 recovers best-of-P sampling; Ltest > Ltrain permits additional scaling when the Q-function generalizes beyond data support.Source: §4.4 · p. 7

4.5 Didactic example: Why do we need test-time value guidance, not training time?

  • In the toy coordination game, five SCOUT refinement steps move candidates toward the optimal mode, while MAC-Flow’s training-time compromise collapses near the origin.The two anti-correlated modes have rewards 4 and 2, respectively.Source: §4.5 · pp. 7–8, Figure 2 · p. 8

5 Experiments

SCOUT achieves the highest average score on both evaluated offline MARL benchmarks, and it outperforms MAC-Flow across the reported offline-to-online scenarios. The experiments also examine benchmark coverage and how refinement-step choices affect performance.

5.1 Experimental setup

  • Across seven scenarios and 24 offline datasets, evaluation covers SMACv1’s discrete actions and MA-MuJoCo’s continuous actions.SMACv1 includes four maps paired with Good, Medium, and Poor data-quality levels; MA-MuJoCo factorizes a robot’s joints across agents.Source: §5.1 · p. 8

5.2 Experimental results

  • SCOUT achieves the highest average score on both benchmarks, outperforming all baselines and matching or exceeding the best baseline on most individual configurations.Table 2 covers MA-MuJoCo’s three tasks and 12 datasets.Source: Table 2 · p. 9, §5.2 · p. 9
  • SCOUT outperforms MAC-Flow across all four offline-to-online scenarios and continues improving after online fine-tuning begins.Figure 3 uses eight seeds, with fine-tuning starting at normalized gradient steps 0.5.Source: §5.2 · p. 9, Figure 3 · p. 9
  • Even a small number of test-time refinement steps improves over no refinement, while the optimal Ltest varies across scenarios.A small deployment-time validation sweep can select Ltest without retraining.Source: §5.2 · pp. 9–10, Figure 4 · p. 10

6 Conclusion

SCOUT separates behavioral modeling from value maximization and refines actions at test time using decomposed Stein variational gradient transport. The paper gives a joint soft-value-gap bound under approximate IGM, while identifying task- and dataset-dependent refinement sensitivity as a limitation.

6 Conclusion

  • SCOUT preserves a multimodal behavioral prior and defers value-driven refinement to test-time inference through decomposed Stein variational gradient transport.The framework decouples behavioral modeling from value maximization.Source: §6 · p. 10
  • Under approximate IGM, SCOUT provides a single-term KL bound on the joint soft-value gap and explains how additive TD targets transfer inter-agent dependencies into per-agent gradients.The bound is stated for the joint soft-value gap.Source: §6 · p. 10
  • Refinement hyperparameters are task- and dataset-dependent, and a principled adaptive stopping criterion remains open.Source: §6 · p. 10

A Proofs

The proof appendix establishes notation for joint observation and action spaces, per-agent behavioral priors and Boltzmann targets, and policies obtained through per-agent transport.

A Proofs

  • The standing notation defines O as the product of the agents’ joint observation and action spaces.Source: A Proofs · p. 15
  • Each agent’s behavioral prior is denoted μi(· | oi), and its factorized Boltzmann target is defined from Qi(oi, ai) and α.Source: A Proofs · p. 15
  • The notation also identifies the joint policy obtained after L independent per-agent transport steps.Source: A Proofs · p. 15

A.1 Proof of Theorem 4.2

The proof relates the additive-objective gap to a KL divergence, then accounts for the shift between additive and soft objectives. The resulting bound includes a δ-dependent correction.

A.1 Proof of Theorem 4.2

  • π*fact is the unique maximizer of Jadd on the joint simplex, and a variational identity expresses the additive-objective gap as a KL divergence.The proof then combines this identity with the δ-shift between Jsoft and Jadd.Source: §A.1 · p. 15
  • Adding the two preceding bounds yields the stated combined bound.Source: §A.1 · p. 15
  • The soft-value gap bound includes an additive 4δ term, and its lower bound of zero follows from Jsoft-optimality of π*.Source: §A.1 · p. 15

A.2 Proof of Proposition 4.3

The proof analyzes per-agent transport and its relation to joint action selection. Agent-specific updates preserve independence, while IGM connects per-agent and joint optimality.

A.2 Proof of Proposition 4.3

  • Each agent’s transport depends only on its own particles and Qi, so independent initialization preserves independence across agents.Source: §A.2 · p. 15
  • Under the IGM identity, per-agent argmax choices are jointly optimal.Source: §A.2 · p. 15
  • For the second proof part, per-agent optimality is applied to any joint candidate index, and the argument maximizes over those candidates.Source: §A.2 · p. 15

B Experimental Details

Experiments were run on one server with a specified CPU and GPU, using a CUDA, cuDNN, Python, and JAX software stack.

B Experimental Details

  • All experiments used a single server equipped with an Intel Xeon E5-2650 v3 CPU and an NVIDIA A40 GPU.Source: §B.1 · p. 16
  • The software stack included CUDA 12.2 and cuDNN 8.9.7.Source: §B.1 · p. 16
  • Experiments used Python 3.9 and JAX 0.4.30.Source: §B.1 · p. 16

B.2 Implementation details

The implementation uses a flow-matching behavioral prior, validation-selected transport settings, and online fine-tuning with mixed offline and newly collected data. The appendix also specifies SCOUT and baseline implementation choices.

B.2 Implementation details

  • The Flow BC prior learns the offline behavior distribution with conditional flow matching and generates candidate actions by fixed-step Euler integration at inference.Its samples retain a multimodal behavioral prior and are biased toward modeled dataset support.Source: §B.2 · p. 16
  • The flow actor receives a Fourier-feature embedding of time, using K = 8 exponentially spaced frequencies to represent multiple temporal scales.The passage says these features support coarse trajectory phases and sharp transitions near t = 0 and t = 1.Source: §B.2 · p. 16
  • Online fine-tuning collects rollout trajectories, uses a warm-up step without online updates, then samples a 50:50 offline-to-online data mixture per gradient step.Source: §B.2 · p. 16

B.3 Hyperparameters

  • Ltrain and Ltest are selected from a small grid on held-out offline data without environment interaction, while baselines receive comparable tuning budgets.Source: §B.3 · p. 16
  • Table 3 is identified as the hyperparameter reference for RL optimization.Source: Table 3 · p. 16

B.4 Baselines

  • MADiff-D uses its official decentralized variant and original architecture and hyperparameters, with a Gaussian DDPM scheduler of 200 denoising steps.It trains on centralized joint demonstrations but executes decentrally using local observations and teammate-trajectory predictions.Source: §B.4 · p. 16
  • DoF uses the official DoF-Trajectory implementation with W-concat factorization and follows the original training and deployment configurations.Action sampling uses a Gaussian DDPM scheduler with 200 denoising steps.Source: §B.4 · p. 17
  • MAC-Flow models multimodal distributions with flow matching and Q-guidance, then distills the guided flow into one-step decentralized policies.The experiments retain its original architecture and hyperparameters.Source: §B.4 · p. 17

C Additional Results

SCOUT improves across continuous offline-to-online configurations, and refinement sensitivity varies by dataset quality. The checklist reports details on reproducibility, experimental variability, theory, and limitations.

C Additional Results

  • SCOUT improves after the online transition in all six MA-MuJoCo configurations and across offline-to-online configurations in both benchmarks.The continuous-control results suggest the undeformed flow prior lets the transport budget exploit newly collected interaction data.Source: Figure 5 · p. 17, C Additional Results · p. 17
  • In Medium and Medium-Replay datasets, increasing Ltest from 0 to 3 yields substantial gains, while Ltest = 10 can degrade performance in some cases.Expert and Medium-Expert performance is largely insensitive to Ltest; the authors suggest excessive transport degrades critic accuracy as actions move away from behavioral support.Source: C Additional Results · p. 17, Figure 6 · p. 18

NeurIPS Paper Checklist

  • The checklist says the paper discusses limitations in its Conclusion section and reports the type of GPU/CPU hardware used.Source: NeurIPS Paper Checklist · pp. 18, 21
  • The paper states Assumption 4.1 and reports complete proofs for three theoretical results in Appendix A.Source: NeurIPS Paper Checklist · p. 19
  • The paper reports detailed algorithm pseudocode, hyperparameter tables, implementation details, and baseline descriptions to support reproduction.The checklist also says Section 5 describes testbeds, datasets, and baselines, with comprehensive implementation details in Appendix B.Source: NeurIPS Paper Checklist · pp. 19, 21
  • Tables 1 and 2 report mean ± two standard deviations over 8 random seeds.Source: NeurIPS Paper Checklist · p. 21
Loading 2610.02554v1…