Source-linked AI summary

G2MAF: Test-Time Gradient Guidance for Multi-Agent Flow Policies

Guowei Zou, Haitao Wang, Guoxin Wang, Zhiquan Chen, Beiwen Zhang, Guojie Wang, Hejun Wu

arXiv:2609.31286v1cs.AI

TL;DR

Frozen offline multi-agent policies execute a single joint-action proposal, which may leave better nearby, behavior-consistent alternatives unexplored. G2MAF refines proposals with a globally normalized, projected centralized-critic gradient; across MPE and SMAC settings, it improves most frozen settings with a modest inference-latency increase.

  • Problem

    A frozen policy’s single joint-action sample does not certify a well-coordinated decision, even when better nearby alternatives remain consistent with behavior data.

    Source: Introduction · p. 2
  • Method

    G2MAF keeps the joint policy and centralized critic frozen, applying one globally normalized critic-gradient step to the proposed joint action and projecting it onto the feasible set.

    Source: Introduction · p. 2
  • Results

    20 frozen settings improve, with mean relative gains of 9.2% on MPE and 8.9% on SMAC; model inference latency increases about 6%.

    Source: Introduction · p. 2
  • Takeaways & Limitations

    The canonical action-space update closes recoverable local gaps without retraining, using one backward pass and no world model or rollouts.

    Source: Conclusion and Limitations · p. 9

    G2MAF assumes centralized access to joint observations and actions and reliable critic ranking of nearby actions; evaluation covers two benchmark families and selected operating points.

    Source: Conclusion and Limitations · p. 9

Abstract

from arXiv · show

Offline multi-agent reinforcement learning (MARL) learns cooperative policies from fixed datasets without further environment interaction and a learned policy is frozen at deployment. Such a frozen policy typically proposes a single joint action and executes it directly at deployment time. However, this one-shot deployment often commits to a suboptimal proposal, even when better nearby alternatives remain consistent with the behavior data. To address this issue, we propose Gradient Guided Multi Agent Flow (G2MAF), a refinement framework for optimizing joint policies at test-time. G2MAF applies one globally normalized, projected critic gradient to guide and coordinate all agents' corrections while keeping the action both feasible and close to the frozen policy proposal. Across 24 MPE and SMAC settings, its canonical variant improves 20 frozen settings, with mean relative gains of 9.2% on MPE and 8.9% on SMAC, with model inference latency increased by about 6% only.

Introduction

A frozen policy’s single joint-action sample may leave nearby, behavior-consistent alternatives unexplored. G2MAF uses centralized critic guidance to refine that action at test time, improving most evaluated settings.

Introduction

  • A single sample from a generative policy does not certify a well-coordinated joint decision, motivating the coordination-gap perspective.The coordination gap is the critic-predicted score difference between the frozen proposal and a higher-scoring nearby joint action.Source: Introduction · p. 2
  • The framework studies whether a local critic-guided joint update can improve a frozen cooperative policy without retraining, a world model, or candidate rollouts.For few-step generators, guidance can be injected before decoding or after the executable action is produced.Source: Introduction · p. 2
  • G2MAF applies one globally normalized gradient of the centralized critic to the joint action, then projects the result onto the feasible set.Each agent’s correction depends on teammates’ choices, while shared normalization gives the team one correction budget.Source: Introduction · p. 2
  • The canonical variant improves 20 frozen settings, with mean relative gains of 9.2% on MPE and 8.9% on SMAC; the critic step increases model inference latency by about 6% on average.Source: Introduction · p. 2
  • The contributions include trajectory-space injection and masked-logit refinement for discrete actions, alongside the canonical action-space update.Under explicit conditions, the paper reports a local critic-ascent guarantee.Source: Introduction · p. 2

Related Work

The paper situates G2MAF among offline MARL, generative policies, and test-time policy improvement. Its distinction is centralized test-time editing of a proposed joint action while both policy and critic remain frozen.

Related Work

  • Offline MARL includes value decomposition and centralized-training methods, while G2MAF keeps a centralized joint critic available to refine the complete action at test time.The cited background includes QMIX [30] and MADDPG [22].Source: Related Work · p. 2
  • Multi-agent generative policies include diffusion [37] and flow methods [16] [19] [26] [38], with CoFlow serving as this paper’s coordinated generative backbone.The passage also identifies flow-policy guidance [2] [31] and one-step or shortcut flows [6] [7] in single-agent work.Source: Related Work · pp. 2–3
  • Unlike approaches that fine-tune policy parameters or update them using logged transitions, G2MAF edits only the proposed joint action with both networks frozen.World-action-model planning is described as an orthogonal route to test-time deliberation.Source: Related Work · p. 3

Preliminaries

The preliminaries define a cooperative offline MARL setting with joint observations and actions, a fixed logged dataset, and a frozen return-conditioned generative policy scored by a centralized critic.

Preliminaries

  • Each agent has a local observation and action, and the joint observation and action collect those local quantities across agents.The passage denotes the joint quantities as tuples over the agents’ local spaces.Source: Preliminaries · p. 3
  • A fixed dataset contains observations, actions, shared rewards, next observations, and continuation indicators, collected by an unknown joint behavior policy..Source: Preliminaries · p. 3
  • The frozen generative policy samples a joint action conditioned on observation and target return, while a centralized critic scores the joint observation-action pair.The reported execution uses one centralized coordinator to refine and dispatch all action blocks; neither model is updated online.Source: Preliminaries · p. 3

Method

G2MAF locally edits a frozen joint-policy proposal using a centralized behavior critic, then executes a feasible refined action. The paper also characterizes its flow prior and value-guidance formulation.

Framework Overview

  • At each decision, the frozen policy proposes a complete joint action, the critic differentiates team value with respect to it, and G2MAF applies one bounded update before execution.The policy supplies a behavior-like starting point, and neither model is updated at deployment.Source: Framework Overview · p. 3

Coordinated Flow Prior

  • The frozen coordinated flow policy transports a normalized joint observation trajectory from noise toward a data trajectory, conditioned on observation and target return.Cross-agent attention makes the velocity a joint function, and a shared inverse-dynamics head decodes the final trajectory into an action.Source: Coordinated Flow Prior · p. 3

Centralized Critic as a Coordination Signal

  • The centralized critic scores the complete joint action, so each action-block gradient depends on teammates’ actions and supplies coordinated corrections.The critic is trained once by behavior fitted critic evaluation and then frozen; the framework figure shows global gradient normalization and feasibility projection.Source: Centralized Critic as a Coordination… · pp. 3–4, Figure 2 · p. 4

Joint Value Guidance at Test Time

  • The ideal policy target raises predicted team value while staying close to the frozen joint prior; G2MAF uses the value-gradient term to refine a decoded action directly.For continuous densities, the value tilt adds the centralized value gradient to the prior score; G2MAF refines one sample rather than evaluating that score or fitting the tilted policy.Source: Joint Value Guidance at Test Time · pp. 4–5

G2MAF: Refinement after Generation

G2MAF refines the generated joint action with one globally normalized critic-gradient step and projects it onto the feasible set. The section also describes trajectory-space injection and masked-logit control for discrete actions.

G2MAF: Refinement after Generation

  • One globally normalized critic-gradient step refines the executable joint action, then Euclidean projection keeps it within the feasible joint-action box.The complete joint action is differentiated, so each agent’s correction depends on teammates’ current actions; shared normalization gives the team one correction budget.Source: G2MAF: Refinement after Generation · p. 5

Where to Inject the Value Signal

  • Trajectory-space variants inject the value signal before decoding, propagating the critic direction through the decoder Jacobian; Post applies the update after decoding in action space.All updates after every reverse Euler step, while Final updates only at k = 1; a nonlinear decoder generally makes these edits nonequivalent.Source: Where to Inject the Value Signal · p. 5

Discrete Joint Actions

  • For SMAC, per-agent legal-action masks constrain softmax logits, and masking again before arg max prevents an illegal action from being executed.The concatenated logits share one normalization, with ∥˜ℓ−ℓ∥2 ≤η.Source: Discrete Joint Actions · p. 5

Experiments

Across 24 MPE and SMAC settings, experiments assess gains, injection locations, discrete control, and deployment cost; G2MAF improves 20/24 settings.

Setup

  • Evaluation uses frozen CGB checkpoints across MPE navigation and predator-prey tasks and four SMAC maps, spanning Poor-to-Expert datasets and five-seed operating-point sweeps.Published baseline columns retain their source-specific benchmark splits, score scales, and uncertainty conventions.Source: Setup · pp. 6–7

Q1: Gains Track Local Headroom

  • G2MAF improves 20/24 settings, with largest gains on several lower-quality datasets and little change near saturation.SMAC success rises by up to +21.1 percentage points on several Medium/Poor maps while Good-map performance remains strong.Source: Q1: Gains Track Local Headroom · p. 7, Figure 3 · p. 8
  • LOTO step transfer is positive on 18/24 settings, critic ranking exceeds chance on 23/24, and continuous updates remain local to empirical action distances.Gains span multiple observed headroom–reliability cells rather than concentrating in the high/high cell.Source: Q1: Gains Track Local Headroom · p. 7

Q2: The Injection Point Matters

  • On MPE, All and Final trajectory variants are positive on 12/12 settings, versus 10/12 for Post; they are largest on 8, 1, and 3 settings, respectively.These variant evaluations use separate batches; Post is deployed because it refines executable actions using the behavior critic.Source: Q2: The Injection Point Matters · p. 7
  • The paired full-grid rerun is positive on 22/24 settings, with critic value increasing on all 24; comparison against the factorized update is mixed on continuous MPE.G2MAF exceeds the best of 16 candidates on 18/24 settings and the one-agent update on 19/24.Source: Q2: The Injection Point Matters · p. 7

Q3: Discrete Actions and Deployment Cost

  • Masked-logit discrete sweeps yield positive frontier gains on 10/11 nonsaturated settings; a critic gradient costs 3.85 ms and total inference latency rises from 157.86 to 167.11 ms (1.06×).G2MAF uses one proposal and one backward pass without a world model or candidate rollouts.Source: Q3: Discrete Actions and Deployment Cost · p. 7, Figure 5 · p. 9

From Local Ascent to Policy Improvement

  • A sufficiently small step guarantees centralized critic-value improvement, and paired rollouts confirm policy improvement across both benchmarks.The transfer test also reports improved mean return on all three tested frozen behavior-cloning MLP settings.Source: From Local Ascent to Policy Improvement · p. 8

Conclusion and Limitations

G2MAF refines frozen offline multi-agent policies with centralized value guidance and no policy retraining. Its benefits depend on reliable nearby-action ranking and are evaluated on two benchmark families.

Conclusion

  • G2MAF applies one normalized, projected joint-action gradient step with a local critic-ascent guarantee, improving the equal-weighted mean setting-wise relative return by 9.1% over the frozen backbone at 1.06× average model-side latency.The post-generation variant refines decoded actions with the behavior critic and avoids the separate normalized-coordinate critic required for trajectory-space injection.Source: Conclusion and Limitations · p. 9

Limitations

  • The method assumes centralized access to joint observations and actions plus reliable nearby-action ranking; little local headroom limits benefit, and evaluation remains limited to two benchmark families and selected operating points.Larger teams, stronger distribution shifts, and decentralized execution remain future work.Source: Conclusion and Limitations · p. 9

Appendix section

The appendix contains experimental supplements and algorithmic and mathematical derivations.

Appendix section

  • Appendix materials cover supplementary experiments, algorithm and mathematical derivations, the formal setting and coordination gap, and value-injection variants.Source: Appendix section · p. 12

Appendix : Experimental Supplement

The supplement organizes quantitative tests by question and fixes an operating-point protocol; it also specifies the step-size sweep and centralized critic training setup.

Appendix : Experimental Supplement

  • The supplement separates detailed results for Q1–Q3 from tests of step-size transfer, mechanism and structure controls, perturbation controls, and benchmark context.Each quantitative experiment states its question, comparison protocol, metrics, observed result, and limits of its conclusion.Source: Appendix : Experimental Supplement · p. 12
  • The evaluation sweeps step size independently by setting, reports each grid’s best value, and uses benchmark-specific signed and relative gains.Continuous η ranges from 0.01 to 0.1 and discrete η from 2 to 50; relative gain is compared only within a benchmark because reward origins differ.Source: Appendix : Experimental Supplement · p. 12
  • Each setting’s centralized behavior critic is a four-layer, 512-unit MLP trained for 20k TD-regression steps on the policy’s offline dataset.For discrete domains, per-agent one-hot actions keep the critic differentiable in relaxed action probabilities; G2MAF differentiates the online critic, while the target network forms TD targets only.Source: Appendix : Experimental Supplement · p. 12

Main-Table Reporting Conventions

This section presents G2MAF’s benchmark results and diagnostics alongside perturbation controls and qualitative context. It also describes the centralized test-time refinement method, its critic and coordinated generative backbone, and value-injection variants, including how refinement constrains updates and preserves action feasibility.

Main-Table Reporting Conventions

Refinement Locality Relative to Empirical Support

  • A single G2MAF step stayed below the offline data’s characteristic nearest-neighbor action scale and exhibited local value ascent across 12 continuous MPE settings.Source: Refinement Locality Relative to… · p. 13

Direct Critic Reliability and a Two-Factor Headroom Test

  • Critic pairwise accuracy exceeded chance in all 12 MPE settings and 11 of 12 SMAC settings, except near-saturated 3m-Good.Mean standardized gains across low-headroom/low-reliability, low-headroom/high-reliability, high-headroom/low-reliability, and high-headroom/high-reliability cells were −0.45, +0.59, −0.75, and +0.77, respectively; positive-setting counts were 2/2, 4/4, 2/4, and 2/2. The largest gains were distributed across all four cells.Source: Direct Critic Reliability and a… · p. 14, Figure 6 · p. 14

Q2: G2MAF Value-Injection Locations in Few-Step Multi-Agent Flow

  • All-step trajectory guidance had the largest gain on 8 of 12 MPE settings and was positive on all; final-step guidance was positive on all, while post-generation guidance was largest on 3 and positive on 10.Post-generation action-space G2MAF is the deployment choice because it acts on executable actions and reuses the behavior critic; trajectory-space variants require a separate normalized-coordinate critic.Source: Q2: G2MAF Value-Injection Locations in… · p. 15

Mechanism and Structure Controls

  • G2MAF gains were positive on 22/24 settings, exceeding Best-N on 18/24 and the one-agent gain on 19/24; comparison with Factorized was mixed on MPE.The critic’s predicted value change was positive on all 24 settings, continuous displacements followed setting-specific radii, and SMAC updates remained legal after re-masking.Source: Table 4 · p. 16, Mechanism and Structure Controls · p. 17

Non-Oracle Step-Size Transfer

  • A domain-transferred step produced positive held-out gains on 11/12 MPE settings and 7/12 SMAC settings, or 18/24 overall.Canonical gains were positive on 10/12 settings in each benchmark; the two evaluation batches are reported separately.Source: Non-Oracle Step-Size Transfer · pp. 17–18, Table 5 · p. 18

Perturbation Controls

  • Perturbation controls replaced the critic gradient with a random direction or a gradient from mismatched observations, using the same nominal pre-projection radius; canonical G2MAF was positive on 20/24 settings and zero on 3m-Good.Source: Perturbation Controls · pp. 18–19
  • Separate-batch controls show random movement and mismatched conditioning have a different ordering from the complete method, while several control rows have positive directions.The table compares critic-guided effects with two matched nominal-radius perturbations.Source: Perturbation Controls · p. 19

Formal Multi-Agent Setting and Coordination Gap

  • A centralized coordinator uses a frozen joint policy and critic to refine all agents’ actions jointly, then dispatches them for simultaneous execution without environment interaction or policy updates.Under support matching and local critic consistency, a feasible refinement within the stated radius reduces the samplewise coordination gap by the critic-value increase.Source: Formal Multi-Agent Setting and… · p. 20

Centralized Behavior Critic Training

  • The centralized critic is fit offline by one-step fitted-Q evaluation of the empirical behavior policy, using retained logged pairs with successor-action bootstrapping.Because episode-final tuples are discarded, it is a behavior-continuation surrogate on retained data-supported pairs, not an exact finite-horizon estimator at omitted endpoints.Source: Centralized Behavior Critic Training · p. 20

The G2MAF Algorithm

  • Each G2MAF decision applies one globally normalized joint critic gradient, then projects continuous actions or executes a masked-logit discrete update using one critic backward pass and no rollout.The shared normalization bounds the pre-projection joint displacement by η.Source: The G2MAF Algorithm · p. 21

Discrete Joint Actions

  • For discrete actions, G2MAF differentiates through a masked-softmax relaxation of per-agent logits, then re-masks and discretizes the actions to preserve legality.The critic is trained on concatenated one-hot actions, while the relaxation supplies the joint-logit gradient.Source: Discrete Joint Actions · p. 21, The G2MAF Algorithm · p. 21

Derivation: G2MAF as a Local Centralized Value Tilt

  • Concatenating all agents’ gradient blocks before normalization preserves the critic’s relative correction magnitudes while assigning the team one shared displacement budget.For continuous actions, projection keeps the final vector within legal bounds; under the stated smoothness and step-size conditions, a nonzero projected update strictly improves critic value.Source: Derivation: G2MAF as a Local… · p. 22

The Coordinated Generative Backbone

  • The coordinated generative backbone couples agents through cross-agent velocity attention before inverse-dynamics decoding, so its induced joint policy need not factor into independent per-agent policies.The canonical action-space variant edits the decoded executable action; the text notes that this short flow makes intermediate-trajectory perturbations a distinct design choice.Source: The Coordinated Generative Backbone · pp. 22–24

G2MAF Value-Injection Variants for Coordinated Generators

  • G2MAF guidance can modify every denoising step, only the final trajectory step, or the decoded action after generation, with trajectory-space gradients passing through the decoder.Because a nonlinear decoder can rotate or rescale trajectory-space directions, its intervention need not match an action-space update.Source: G2MAF Value-Injection Variants for… · p. 24
Loading 2609.31286v1…