PaperlayerView exact revision on arXiv

Source-linked AI summary

MA-WAM: Multi-Agent World-Action Model for Test-Time Planning

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

arXiv:2609.31281v1cs.AI

TL;DR

Cooperative multi-agent planning must account for how simultaneous actions affect teammates’ later observations and responses. MA-WAM uses a frozen flow policy and routed world model to rank candidate joint futures. Across 30 offline MARL settings, it reports mean relative gains of 22.0% over direct execution and 25.6% over uniform selection.

  • Problem

    Independent single-agent predictions omit dependencies among simultaneous actions, while multi-agent planning needs an efficient evaluator that preserves those interactions when ranking candidate futures.

    Source: §1 · p. 2

  • Method

    MA-WAM uses a frozen flow policy to propose candidate joint-action sequences and a Routed World Model to predict their futures and team returns for receding-horizon selection.

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

  • Results

    Across 30 offline MARL settings on MAMuJoCo, SMAC, and MPE, MA-WAM achieves a mean relative gain of 22.0% over direct execution and 25.6% over uniform selection.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported results support using RWM to rank policy-proposed joint futures while preserving cross-agent dependencies during prediction.

    Source: §6 · p. 9

    The experiments evaluate CTCE only; CTDE execution would require locally factorized proposers and scorers using per-agent information.

    Source: §3 · p. 4, §6 · p. 9

Abstract

from arXiv · show

Multi-agent cooperative tasks require different agents to execute a joint action simultaneously, and each agent's action affects both the observations and responses of the other agents. Hence, a world model is needed to predict the team return resulting from the joint actions of all agents. A naive extension directly applies a single-agent world model to each agent's action when predicting the team return step by step. However, such an extension fails to capture the dependencies among the simultaneous actions of multiple agents. We propose Multi-Agent World-Action Model (MA-WAM), a test-time planning framework that enables a frozen multi-agent flow policy to evaluate futures of candidate joint actions. To our knowledge, MA-WAM is the first test-time world-model planner for multi-agent flow policies. MA-WAM predicts the consequences of each joint action according to cross-agent dependencies and enables efficient candidate scoring. Across 30 offline multi-agent reinforcement learning (MARL) settings on MAMuJoCo, SMAC, and MPE, MA-WAM achieves mean relative gains of 22.0% over direct execution and 25.6% over uniform action selection. Under the standard evaluation protocol on an A100 GPU, MA-WAM adds 12.1 ms, accounting for 2.5% of the measured generation-and-scoring time.

1 Introduction

Simultaneous actions shape teammates’ later observations and responses, so evaluating cooperative futures requires preserving cross-agent dependencies. MA-WAM ranks candidate joint-action sequences with a routed world model, and improves mean returns across the reported offline MARL settings.

1 Introduction

  • Offline MARL deploys policies learned from fixed datasets without further environment interaction, while generative policies can produce multiple candidate sequences.MA-WAM frames deployment as test-time foresight: comparing candidate joint futures by predicted team outcomes before execution.Source: §1 · p. 2
  • Independent per-agent world-model predictions omit how simultaneous actions change teammates’ observations and responses, making candidate rankings unreliable.A suitable evaluator must preserve joint-context interactions while scoring multiple candidate futures efficiently.Source: §1 · p. 2
  • MA-WAM uses a frozen generative policy to propose joint-action sequences and a Routed World Model to predict their futures and cumulative team returns.It executes the first action of the highest-scoring candidate, then replans from the next real observation; selection stays within the policy’s samples.Source: §1 · p. 2
  • 22.0% mean relative gain over direct execution and 25.6% over uniform selection across 30 offline MARL settings on MAMuJoCo, SMAC, and MPE.With eight candidates fixed, RWM ranking has higher mean return on 23 settings; MA-WAM has higher mean return than direct execution on 26.Source: §1 · p. 2

2 Related Work

The related work distinguishes policy learning from deployment-time evaluation of candidate joint futures. MA-WAM applies policy-constrained sample-based MPC to offline MARL deployment.

2 Related Work

  • Offline MARL and generative policies improve the learned policy distribution, but their deployment interface remains reactive and executes one policy output at each decision step.The cited approaches include imitation and value-constrained methods, joint sequence models, and diffusion- or flow-based policies.Source: §2 · p. 2
  • MA-WAM adapts receding-horizon MPC by restricting its search to joint futures proposed by a frozen generative MARL policy.The paper characterizes this as policy-constrained sample-based MPC for offline MARL.Source: §2 · p. 3
  • Prior multi-agent world models mainly use learned dynamics during policy training, policy improvement, or direct model-based control; MA-WAM studies fixed-policy offline deployment.Related approaches include decentralized, diffusion-based, interaction-latent, and mixture-of-experts architectures.Source: §2 · p. 3

3 Preliminaries

The reported experiments use centralized test-time coordination execution (CTCE); a decentralized-execution variant would require locally factorized components.

3 Preliminaries

  • Agent indices use ordinary superscripts such as i, while candidate indices use parenthesized superscripts such as (m).The candidate index m identifies candidate samples.Source: §3 · p. 3
  • The reported evaluation uses CTCE, while CTDE execution would require both proposer and scorer to be factorized over agents and use local information.The planner receives the joint observation and outputs a joint action in the reported setup.Source: §3 · pp. 3–4

4 Method

MA-WAM combines a frozen flow-policy proposer with a routed world-model scorer, then replans after executing the first action of the highest-scoring candidate.

4.1 Framework Overview

  • At deployment, a frozen policy proposes M length-H joint-action sequences, RWM scores predicted returns, and the planner executes only the winner’s first action before replanning.The proposer and world model are both trained from the same offline dataset and remain frozen during deployment.Source: §4.1 · p. 4

4.2 Coordinated Few-Step Flow Policy

  • The coordinated few-step flow proposer maps Gaussian noise samples, conditioned on joint-observation history and target return, to candidate trajectories and decodes them into executable actions.A shared inverse-dynamics head converts consecutive predicted observations into actions; the policy is trained with trajectory-flow and action losses.Source: §4.2 · p. 4

4.3 Routed World Model

  • RWM uses context-conditioned routing over shared experts to predict the next joint observation, while a reward branch predicts per-agent rewards.The branches are trained by supervised prediction on offline transitions; reward inputs use a detached predicted next observation.Source: §4.3 · pp. 4–5

4.4 Test-Time Planning

  • At each decision step, RWM scores each candidate by predicted cumulative return over a finite horizon, and the planner executes its first joint action before replanning.The deployment ranking score uses an undiscounted short horizon; the policy and world model remain frozen throughout deployment.Source: §4.4 · p. 5

5 Experiments

The experiments compare frozen-policy execution with candidate planning across offline MARL benchmarks, then examine when planning helps, how reliably RWM ranks futures, and its inference cost.

5.1 Setup

  • The evaluation covers MAMuJoCo, SMAC, and MPE under fixed-denoising comparisons using five seeds and three denoising steps per candidate.Reactive uses M=1, while Random and Planning use M=8; Table 1 compares deployed offline MARL methods, and the controlled comparison uses a frozen policy.Source: §5.1 · pp. 5–7

5.2 Planning vs. Reactive Execution (RQ1)

  • 26 of 30 settings improve with Planning over Reactive execution; the largest gains reach +31 normalized points on MPE and +254 raw return on MA-MuJoCo.Figure 3 shows the per-setting comparison and tallies 26 positive, 0 zero, and 4 negative differences.Source: Figure 3 · p. 7, §5.2 · p. 8

5.3 When Foresight Helps (RQ2)

  • The largest mean relative gains occur when both reactive-policy headroom and H=8 ranking reliability are above their global medians.Both groups below median headroom remain closest to zero; with the same M=8 candidate budget, RWM selection exceeds uniform Random selection on 23 of 30 settings.Source: §5.3 · p. 8

5.4 Model Accuracy and Cost (RQ3)

  • At H=8, mean Spearman correlation is 0.81, 8-way Top-2 accuracy is 0.78, pairwise accuracy is 0.82, and normalized regret is 0.13.At H=16, Spearman correlation falls to 0.65 and normalized regret rises to 0.22; the requested rollout cap is eight, with replanning after each executed action.Source: §5.4 · p. 9
  • RWM scoring takes 12.1 ms on an A100, or 2.5% of the 473.7 ms generation-and-scoring time.Separately, RWM has the highest mean return on all 10 settings in the alternative-scorer comparison; diagnostics also show context-dependent expert assignment and broad expert use.Source: §5.4 · p. 9

6 Conclusion and Limitations

MA-WAM uses a frozen world model to rank candidate joint futures from a frozen flow policy, reporting gains across 30 offline MARL settings without policy retraining. The experiments are limited to CTCE, and use M=8 with a requested horizon cap of eight.

6 Conclusion and Limitations

  • Across 30 offline MARL settings, MA-WAM gains 22.0% over direct execution, while RWM ranking gains 25.6% over uniform selection at the same M=8 candidate budget.On an A100, one batched RWM scoring pass takes 12.1 ms and accounts for 2.5% of measured generation-and-scoring time.Source: §6 · p. 9
  • Experiments evaluate CTCE only; CTDE would require locally factorized proposers and scorers using per-agent information during execution.All benchmarks use M=8 and a requested horizon cap of eight; adaptive budgets, horizons, and broader scorer comparisons are left for future study.Source: §6 · p. 9

A.1 Deployment Information and Execution Setting

  • Both the proposer and RWM remain frozen at deployment, and candidate scores use model-predicted futures rather than real future transitions.The planner executes the selected sequence’s first joint action, receives the next real observation, and then starts a new planning cycle.Source: §A.1 · p. 13

A.2 Baseline Comparison and Hyperparameters

  • MA-WAM ranks first or second against published baselines on 22 of 30 settings; planning improves mean return on 26 settings in the controlled fixed-denoising comparison.The effective horizon is eight steps on MPE, MA-MuJoCo, and SMAC 2s3z, and three steps on the other SMAC maps because their proposer checkpoints generate three-step sequences.Source: §A.2 · p. 14

A.3 Supplementary Results for RQ1: Planning vs. Reactive Execution

The fixed-denoising comparison finds positive Planning–Reactive gains across most settings, with gains also positive when aggregated by benchmark family. Supplementary win-rate and trajectory results show how selected behavior changes across cooperative tasks.

A.3 Supplementary Results for RQ1: Planning vs. Reactive Execution

  • 22.0% mean setting-wise relative gain over Reactive, with positive ∆ on 26 of 30 settings, no ties, and negative ∆ on four.The two-sided exact sign test gives p = 5.95×10−5.Source: Table 2 · p. 15, §A.3 · p. 15
  • Aggregate relative gains are +14.42% on MPE, +2.75% on SMAC, and +7.68% on MAMuJoCo.The benchmark-level aggregation is positive across all three families; the conclusion notes that gains are smaller on SMAC, where many settings are near the task ceiling.Source: §A.3 · pp. 15, 17, Table 3 · p. 17
  • Representative trajectories show planning-selected changes in spatial coverage, pursuit, coordinated locomotion, and focus fire.Figure 7 includes MPE, MAMuJoCo, and SMAC examples; the trajectories are qualitative complements to return tables.Source: §A.3 · p. 17, Figure 7 · p. 17
  • Planning improves SMAC win rate on 8 of 12 settings and ties at the floor on two.The largest reported change is 8m-Medium, rising from 0.74 to 0.88; win rate decreases slightly on two settings.Source: §A.3 · pp. 17–18, Table 4 · p. 18

A.4 Supplementary Results for RQ2: When Does Foresight Help

With the candidate budget held fixed, RWM ranking outperforms uniform selection across the evaluated settings. Gains are largest where reactive-policy headroom and ranking reliability are high, and generally rise as offline dataset quality falls; the RQ3 evaluation examines prediction accuracy and inference cost across horizons.

A.4 Supplementary Results for RQ2: When Does Foresight Help

  • 25.6% mean relative gain over uniform selection, with RWM higher on 23 of 30 equal-budget settings.Both methods use M=8 candidates; the two-sided exact sign test gives p = 0.00522.Source: §A.4 · p. 18, Table 5 · p. 19
  • The largest mean planning gains occur when both reactive-policy headroom and H=8 ranking reliability are above their cross-setting medians.Both low-headroom groups have mean gains closest to zero; dataset quality provides a related headroom signal.Source: §A.4 · p. 19, Figure 8 · p. 20
  • Mean relative gain rises from +2.6% in the best dataset-quality tier to +14.9% in the middle and +50.8% in the worst.Gains increase from better to worse datasets in 8 of 9 tasks, with 5m_vs_6m the exception; World-Random’s +223% relative gain corresponds to +3.1 normalized points.Source: §A.4 · p. 19, Figure 9 · p. 20

A.5 Supplementary Results for RQ3: Model Accuracy and Inference Cost

  • The RQ3 evaluation compares predicted and realized cumulative returns across horizons, and also measures next-observation rollout error and inference latency.These measures examine how prediction error changes with rollout horizon and whether scoring is fast enough for each decision step.Source: §A.5 · p. 21

Horizon-Wise Score Reliability Metrics

The section evaluates whether world-model scores preserve return rankings across horizons and in same-state counterfactual selection. At H=8, ranking remains reliable across benchmark families, although reliability varies with horizon and setting.

Horizon-Wise Score Reliability Metrics

  • The diagnostics include top-ranked selection, pairwise ordering, and regret; normalized regret scales realized-return loss by the return range within a group or candidate set.For same-state metrics with M=8, random Top-1, Top-2, and pairwise accuracy are 0.125, 0.25, and 0.5, respectively.Source: Horizon-Wise Score Reliability Metrics · pp. 21–23
  • Snippet-based proxy metrics assess score ordering across recorded trajectories that generally start from different states, whereas the same-state diagnostic compares policy-proposed futures from one restored state.Same-state testing restores each state before executing every candidate, directly evaluating selection among alternatives available to the planner.Source: Horizon-Wise Score Reliability Metrics · p. 22
  • Pairwise ranking accuracy exceeds chance on all five same-state settings (0.58–0.74), and RWM’s mean normalized regret is 0.37 versus 0.52 for Random.Top-1 exceeds its 0.125 chance level on four settings, but is at chance on 4Ant-Medium and 0.250 on Tag-Medium.Source: Table 6 · p. 24, Horizon-Wise Score Reliability Metrics · pp. 24–25

Full Per-Setting Score Calibration

  • At H=8, Spearman correlation is positive across all 30 settings: medians are 0.85 on MPE, 0.83 on MAMuJoCo, and 0.69 on SMAC.Planning gain also depends on Reactive performance and the realized-return spread among same-state candidates, which vary across settings.Source: Figure 10 · p. 24, Full Per-Setting Score Calibration · p. 26

Per-Setting Score Reliability across Horizons

  • At H=8, median Spearman correlation is about 0.80–0.84, pairwise ranking accuracy about 0.81–0.82, and normalized regret 0.10–0.16 across the three benchmark families.Ranking deteriorates by H=16; the evaluation uses an eight-step scoring horizon and replans after each executed action.Source: Figure 11 · p. 25, Per-Setting Score Reliability across… · p. 26
  • MAMuJoCo reliability is highest at H=1, while MPE and SMAC peak around H=4; all three families remain comparable at H=8.Median Spearman correlation is 0.96 for MAMuJoCo at H=1, and 0.84 and 0.90 for MPE and SMAC at H=4, respectively.Source: Figure 11 · p. 25, Per-Setting Score Reliability across… · p. 26

State Rollout Error Diagnostic

The section measures multi-step rollout error and inference cost, then reports diagnostics of the routed dynamics branch. Error increases with rollout length, while batched RWM scoring is a small portion of measured planning time.

State Rollout Error Diagnostic

  • At H=8, median next-observation MSE is 0.037 on MPE, 0.057 on SMAC, and 0.199 on MAMuJoCo; by H=16, the medians rise to 0.089, 0.187, and 0.268, respectively.The rollout repeatedly feeds predicted observations into the next model transition; the shorter horizon accompanies replanning after each executed action.Source: State Rollout Error Diagnostic · p. 26

Inference-Time Profiling

  • At M=8, batched RWM scoring takes 13.4 ms on an RTX 3090 and 12.1 ms on an A100, or 3.1% and 2.5% of total generation-and-scoring time.The corresponding total times are 428.0 ms and 473.7 ms, respectively.Source: Table 7 · p. 27, Inference-Time Profiling · p. 27
  • Batching the eight candidates reduces candidate-generation latency by 4.2–6.0× across six settings on both GPUs while retaining the same candidates and RWM selection rule.The optimization changes the generation schedule by folding candidates into one pass.Source: Table 8 · p. 27, Inference-Time Profiling · pp. 27–28
  • Different agents assign different routing mass to shared experts, and removing individual experts produces unequal next-observation prediction degradation.The diagnostics indicate non-uniform expert use and different contributions by individual experts to prediction.Source: Figure 13 · p. 28, Inference-Time Profiling · p. 28

A.6 Diagnostics and Role of Routed Experts

The diagnostics examine how RWM routes shared experts across contexts and agents, and compare routing structure with predicted cross-agent action influence. Across settings, routing is more associated with transition context than fixed agent identity, and representative routing graphs align with influence patterns.

A.6 Diagnostics and Role of Routed Experts

  • RWM uses shared experts with mixtures conditioned on joint observation–action tokens, allowing agents to use different expert mixtures across contexts.Soft dispatch mixes agents’ tokens into expert slots, and agent-specific combine weights use the processed slots to predict each agent’s observation delta.Source: §A.6 · pp. 28–29
  • Routing has stronger context than agent-identity association on 26 of 30 settings, with a mean corrected-NMI difference of 0.076; transition-based routing variation is larger on 23 settings, with a mean JS-divergence difference of 0.047.The router uses an average of 7.51 effective experts out of eight, and the dominant expert changes across 38.1% of consecutive transition pairs.Source: §A.6 · p. 29, Figure 14 · p. 30
  • In representative MPE Tag-Medium and SMAC 8m-Poor cases, routing and action-perturbation influence graphs correlate at 0.964 and 0.926, respectively.Their normalized mean elementwise absolute differences are 0.053 and 0.146; the authors interpret the patterns as capturing cross-agent predictive dependencies.Source: Figure 15 · p. 30, §A.6 · pp. 30–31

A.7 Comparison with Alternative Candidate Selectors

The selector comparison uses a shared candidate-to-score interface to assess RWM against alternative scorers. RWM has the highest mean return in all 10 representative settings, with comparisons interpreted within each setting.

A.7 Comparison with Alternative Candidate Selectors

  • The shared interface maps the current joint-observation context and each proposed joint-action sequence to a score, then executes the first joint action of the highest-scoring sequence and replans.RWM rolls joint dynamics forward and sums predicted rewards; alternatives include a monolithic scorer, a current-action Q-ranker, and a direct trajectory-return predictor.Source: §A.7 · p. 31
  • The alternatives use the same frozen proposer, eight-candidate budget, denoising budget, and horizon protocol; Independent-WM is a separate single-run reference rather than a matched-seed ablation.Because benchmark families use different score scales, return comparisons are made within each setting rather than aggregated across families.Source: §A.7 · pp. 31–32
  • RWM achieves the highest mean return in all 10 representative settings spanning MPE, SMAC, and MA-MuJoCo.The strongest alternative or reference is the direct-return predictor in four settings, Independent-WM in three, the Q-ranker in two, and the monolithic scorer in one.Source: §A.7 · p. 32

A.8 Candidate-Space Constraint and Selection Reliability

RWM selects only among the frozen proposer’s current sampled candidates and executes the winner’s first action before replanning; its rankings improve over uniform selection at the chosen horizon.

A.8 Candidate-Space Constraint and Selection Reliability

  • The planner selects from the frozen proposer’s current sampled candidate set, executes only the selected sequence’s first joint action, then replans.Stochastic samples need not duplicate offline-dataset actions; the constraint is the candidates sampled at that step.Source: §A.8 · p. 32
  • The evaluation compared episode return under RWM ranking and uniform selection with M=8, alongside horizon-wise proxy diagnostics and same-state counterfactual tests.The same-state test covered five MPE and MA-MuJoCo Medium settings.Source: §A.8 · p. 32
  • The result analysis reports that RWM had the higher mean return on 23/30 settings; the supplied excerpt says the test was two-sided but does not provide an exact p-value.Source: §A.8 · p. 32
  • At H=8, median Spearman correlation was about 0.80–0.84, median pairwise ranking accuracy about 0.81–0.82, and median normalized regret 0.10–0.16 across benchmark families.These are the Figure 11 proxy-ranking diagnostics.Source: §A.8 · p. 32
  • In same-state tests on five settings, pairwise accuracy exceeded chance and ranged from 0.58 to 0.74; mean normalized selection regret was 0.37 for RWM versus 0.52 for Random.The test restored 64 anchor states and evaluated all eight candidates for H=8 real steps.Source: §A.8 · p. 32

B Theoretical Analysis and Formulation

This section formalizes cooperative offline MARL and presents MA-WAM, a test-time planner that samples and scores candidate joint-action sequences using a frozen world model with cross-agent information. At deployment, it executes the first action of the highest-scoring candidate and replans from each new observation.

B.1 Preliminaries

  • The setting is cooperative offline MARL formalized as a Dec-POMDP, with a fixed dataset collected by unknown behavior policies and no further interaction for policy learning.The section distinguishes privileged latent state from the agents’ local observations and defines a shared team reward.Source: §B.1 · pp. 32–33

B.2 Detailed Test-Time Planning Procedure

  • The planner samples M joint action sequences from the frozen proposer, scores each by undiscounted predicted reward over H steps, and selects the highest-scoring candidate.The policy objective remains discounted infinite-horizon return; the undiscounted criterion is used for finite-horizon candidate ranking.Source: §B.2 · p. 33
  • Agents execute only the selected sequence’s first joint action, then re-plan from the next real joint observation while keeping the policy and world model frozen.Resetting the imagined trajectory at each step confines model-error accumulation to the current short-horizon score.Source: §B.2 · p. 33

B.3 Frozen Stochastic Trajectory Proposer

  • The frozen proposer denoises joint observation trajectories with the agent axis intact, using cross-agent attention to exchange information among agents.A shared inverse-dynamics head maps consecutive predicted observations to candidate joint actions.Source: §B.3 · pp. 34–35
  • The proposer uses a stop-gradient temporal-consistency surrogate during training and combines conditional and unconditional velocities for target-return guidance at deployment.The same noisy trajectory is evaluated at two time labels; gradients flow only through the reference-time pass, and experiments use guidance weight ω ≈1.2.Source: §B.3 · pp. 34–35

B.4 World-Model and Planning Formulation

  • RWM dynamics soft-route agent observation-action tokens through shared slots and experts, then combine slot outputs for each agent’s next-observation prediction.Dense dispatch lets every agent contribute to every slot, while dense combination lets every processed slot contribute to every agent prediction.Source: §B.4 · pp. 36–37
  • The sparse-routed reward model scores transitions from current joint observations, joint actions, and stop-gradient predicted next observations.It maps each agent’s reward token to a distribution over reward experts.Source: §B.4 · pp. 37–38
  • Reward prediction uses sparse top-k routing: selected expert outputs receive renormalized mixture weights, while all other expert outputs receive zero weight.Source: §B.4 · p. 38
  • The world model is trained by supervised prediction on offline transitions, with a stop-gradient on predicted next observations isolating dynamics optimization from reward-loss gradients.The dynamics loss trains prediction of short-horizon observation changes, while the reward predictor uses the detached predicted next observation as input.Source: §B.4 · pp. 38–39
  • For replicated shared team rewards, summing per-agent reward-head outputs gives an unnormalized score n times the per-agent mean, preserving within-task candidate ordering.Source: §B.4 · p. 39
  • A load-balancing auxiliary term discourages the reward router from concentrating both expert-selection frequency and router probability mass on the same experts.Source: §B.4 · p. 39
  • The planner scores M length-H candidate joint-action sequences through imagined rollouts and executes the first joint action of the highest-scoring candidate.Candidates share a fixed target return and differ in sampled proposer noise; each rollout starts from the current joint observation, accumulates predicted rewards, and yields a score for ranking.Source: §B.4 · pp. 39–40
  • Each planning cycle executes only the first joint action, then samples and scores a fresh candidate set from the next real observation, limiting model-error accumulation to one H-step rollout.Source: §B.4 · p. 40
  • The random-selection ablation uses the same M-candidate proposals and seed schedule as Planning but selects uniformly, isolating the effect of candidate selection.Uniform selection gives a candidate the same marginal proposal distribution as one policy sample; the reactive baseline is M=1 and executes its sole candidate.Source: §B.4 · p. 40

Related papers

Loading 2609.31281v1…