Constrained Multi-Agent Reinforcement Learning in Imagination

Raz Shmueli; Shahaf S. Shperberg

cs.AIOct 15, 2026

Constrained Multi-Agent Reinforcement Learning in Imagination

Authors: Raz Shmueli, Shahaf S. Shperberg

Organizations: Ben-Gurion University of the Negev, Beer-Sheva, Israel

Abstract

Predicting cost in a multi-agent world model moves most constrained policy updates into imagination: the stronger model-free baseline needs $10$ to $20$ times CA-MAMBA's interactions on the median SMAC configuration, and CA-MAMBA earns $2.8$ to $12$ times the baselines' MAMuJoCo reward at equal budget. In controlled tests on one task per benchmark, pricing the model's cost advantages at parity more than halves a cost-blind learner's cost and meets a tight continuous-control limit the learned multiplier misses, suggesting the model's cost estimates are not the bottleneck.

Figures & tables

Figure 1. The two training phases for agent

i

; Sections 3.2 to 3.4 define each component (GRU: gated recurrent unit, the recurrent network

f_{\theta}

; enc: the observation encoder; dec: the decoder). (a) On replay, attention over

\mathcal{N}_{i}

turns the previous actions and stochastic states into a message

m_{t}^{i}

that drives the recurrent update; the posterior combines

h_{t}^{i}

with the encoded observation, and a Kullback–Leibler (KL) penalty ties the prior, which predicts

z_{t}^{i}

from

h_{t}^{i}

alone, to it. The decoder reads

x_{t}^{i}=(h_{t}^{i},z_{t}^{i})

; the reward and cost heads read

z_{t}^{i}

with

h_{t+1}^{i}

, which carries the agents’ actions

a_{t}

, and the per-agent costs are averaged into

\hat{c}_{t}

. (b) From replay posteriors, the actor (policy) and prior roll out

H_{\mathrm{im}}

imagined steps; the critics (value estimators) form

A_{u}^{\beta}

for the actor, while the price

\beta

steps on the measured cost

\bar{C}_{k}

of the latest real episode. Diagram of the two training phases: world-model learning on replay, and actor-critic learning in imagination with a multiplier updated from measured cost.

Figure 2. How many interactions MAPPO-Lagrangian needs to catch up with CA-MAMBA at 100k on SMAC. (a) Among the 23 configurations where CA-MAMBA wins any episode, the fraction on which MAPPO-Lagrangian has matched CA-MAMBA’s 100k mean cost (dashed), mean win rate (dotted), or both (solid) by each checkpoint; the solid curve is the distribution of

m(g)

, defined in the text. (b) Medians over the 24 configurations of casualties per allied unit (red) and win rate (blue): MAPPO-Lagrangian (lines), MACPO at 100k and 5M (squares), CA-MAMBA at 100k (stars). Two panels. Left: fraction of SMAC configurations on which MAPPO-Lagrangian matches CA-MAMBA's 100k cost and win rate, against MAPPO-Lagrangian's interactions. Right: median casualties per allied unit and median win rate against interactions.

Ant $2{\times}4$
Method (interactions)	Reward $\uparrow$	Cost $\downarrow$
MACPO	$268\mathbin{\pm}70$	$0.2\mathbin{\pm}0.1$
MAPPO-Lag	$439\mathbin{\pm}46$	$0.1\mathbin{\pm}0.2$
CA-MAMBA	$\mathbf{1836}\mathbin{\pm}678$	$17.4\mathbin{\pm}15.1$
no comm.	$1491\mathbin{\pm}80$	$19.4\mathbin{\pm}30.3$
Ant $4{\times}2$

Table 1. MAMuJoCo at the SafePO default limit

d{=}25

, every method at a matched budget of 1M interactions. No comm. acts without messages (

\mathcal{N}_{i}=\{i\}

) but learns and imagines with them. Mean

\mathbin{\pm}

sample standard deviation over three seeds (

\ast

: two). Bold: highest mean reward;

\dagger

: mean cost above

d

	At 100k
Multiplier	Cost $\downarrow$	Win rate $\uparrow$	$\beta$
Imagined	$2.49\mathbin{\pm}0.43$	$0.95\mathbin{\pm}0.02$	$0.22\mathbin{\pm}0.01$
Measured	$2.05\mathbin{\pm}0.42$	$0.95\mathbin{\pm}0.02$	$0.54\mathbin{\pm}0.02$
PID	$8.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$8.28\mathbin{\pm}0.01$
$\beta{=}5$	$8.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$5.00\mathbin{\pm}0.00$
$\beta{=}1$	$1.75\mathbin{\pm}0.20$	$0.97\mathbin{\pm}0.02$	$1.00\mathbin{\pm}0.00$

Table 2. Multiplier rules on SMAC 8m at

d=0

. Mean

\mathbin{\pm}

sample standard deviation over three runs (seeds 1–3) of each run’s episodes between 95k and 100k interactions;

\beta

at 100k. Imagined and Measured update

\beta

(

\alpha_{\beta}=10^{-5}

) from imagined or measured cost; PID uses

K_{P}=1

K_{I}=10^{-5}

K_{D}=1

; the

\beta{=}5,1,0

rows hold

\beta

fixed, so

d

does not enter them.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3. Execution: every agent runs its own world model and a parameter-tied actor. Each world model sends its previous action and stochastic latent

a^{i}_{t-1},z^{i}_{t-1}

to an attention layer over the neighborhood

\mathcal{N}_{i}

and receives the message

m^{i}_{t}

, which drives its recurrent update; the actor then reads the latent

x^{i}_{t}=(h^{i}_{t},z^{i}_{t})

. Under

\mathcal{N}_{i}=\mathcal{I}

the attention is all-to-all; setting

\mathcal{N}_{i}=\{i\}

removes the exchange and changes nothing else. No privileged global state is used. Execution diagram: each agent's world model exchanges previous actions and latents through attention and feeds its latent to a shared actor.

Quantity	Value
World-model learning rate	$2\times 10^{-4}$
Actor learning rate (Adam, wd $10^{-5}$ )	$5\times 10^{-4}$
Critic learning rate	$5\times 10^{-4}$
Multiplier step $\alpha_{\beta}$ (applied $E=4$ times per step)	$10^{-5}$
Initial multiplier $\beta_{0}$	$10^{-4}$
PPO clip $\varepsilon$	$0.2$

Appendix

Table 3. Optimization and training settings, read from the training code. Every term of Eq. (1), including the divergence and the three terms inherited from MAMBA (the continuation and action-availability heads and, for discrete actions, the action classifier), enters the world-model loss with weight

1

; KL balancing splits the divergence

0.8

toward the direction that updates the prior and

0.2

toward the posterior.

Component	Specification
Deterministic state $h$	$256$ (GRU)
Stochastic state $z$	$32\!\times\!32$ categ. (DreamerV2)
Feature $x=(h,z)$	$1280$
Observation embedding	$256$
Reward / cost heads	2-layer MLP, width $256$
Continuation head	2-layer MLP, width $256$

Appendix

Table 4. Network architecture, read from the training code. The transition attention is the only inter-agent communication path; the actor is a per-agent MLP over the communication-conditioned latent.

Environment	Steps	$n$	Reward $\uparrow$	Cost $\downarrow$
Ant $2{\times}4$	100k	3	$836\mathbin{\pm}93$	$0.0\mathbin{\pm}0.0$
	500k	3	$1086\mathbin{\pm}412$	$6.8\mathbin{\pm}11.5$
	700k	3	$1614\mathbin{\pm}553$	$53.1\mathbin{\pm}47.1$
	800k	3	$1794\mathbin{\pm}672$	$82.9\mathbin{\pm}72.0$
	900k	3	$1838\mathbin{\pm}583$	$32.7\mathbin{\pm}33.7$
	1M	3	$1836\mathbin{\pm}678$	$17.4\mathbin{\pm}15.1$

Appendix

Table 5. CA-MAMBA at

d{=}25

against interaction budget, with both model-free baselines at the matched 1M budget and at 10M. The baseline rows are the same runs read at two checkpoints. Cost is not monotone in the budget: it exceeds

d

on Ant

2{\times}4

at 700k to 900k and on HalfCheetah

2{\times}3

at 100k and 800k.

		Ant $2{\times}4$ , $d=0.2$		Ant $4{\times}2$ , $d=1.0$		HalfCheetah $2{\times}3$ , $d=5.0$
Method	Steps	Reward $\uparrow$	Cost $\downarrow$	Reward $\uparrow$	Cost $\downarrow$	Reward $\uparrow$	Cost $\downarrow$
CA-MAMBA, $\alpha_{\beta}{=}10^{-5}$	1M	$2268\mathbin{\pm}134$	$24.27\mathbin{\pm}5.62^{\dagger}$	$1241\mathbin{\pm}566$	$2.53\mathbin{\pm}3.11^{\dagger}$	$1973\mathbin{\pm}506$	$18.17\mathbin{\pm}15.95^{\dagger}$
CA-MAMBA, $\alpha_{\beta}{=}10^{-4}$	1M	$1618\mathbin{\pm}641$	$4.07\mathbin{\pm}5.84^{\dagger}$	$1559\mathbin{\pm}359$	$4.20\mathbin{\pm}7.27^{\dagger}$	$1481\mathbin{\pm}98$	$4.56\mathbin{\pm}1.99$
CA-MAMBA, PID	1M	$1289\mathbin{\pm}421$	$0.91\mathbin{\pm}0.54^{\dagger}$	$1417\mathbin{\pm}196$	$2.40\mathbin{\pm}1.60^{\dagger}$	$2271\mathbin{\pm}282$	$5.22\mathbin{\pm}2.87^{\dagger}$
MACPO	1M	$299\mathbin{\pm}42$	$0.19\mathbin{\pm}0.17$	$315\mathbin{\pm}47$	$0.60\mathbin{\pm}0.66$	$181\mathbin{\pm}43$	$0.00\mathbin{\pm}0.00$
MAPPO-Lag	1M	$520\mathbin{\pm}37$	$0.29\mathbin{\pm}0.37^{\dagger}$	$552\mathbin{\pm}58$	$0.92\mathbin{\pm}0.07$	$775\mathbin{\pm}60$	$0.23\mathbin{\pm}0.32$

Appendix

Table 6. Reward and cost on the multi-agent MuJoCo velocity tasks at tight cost limits:

d=0.2

on Ant

2{\times}4

d=1.0

on Ant

4{\times}2

d=5.0

on HalfCheetah

2{\times}3

. CA-MAMBA updates its multiplier

\beta

from the measured episode cost with step size

\alpha_{\beta}

, or sets it with a PID controller on the same signal (

K_{P}=1

K_{I}=10^{-5}

K_{D}=1

), and is read after 1M interactions; MACPO and MAPPO-Lag are read after 1M and 10M. Each run’s reward and cost average its episodes (CA-MAMBA) or logged interval means (MACPO, MAPPO-Lag) in the last 5k interactions before the budget; entries are mean

\mathbin{\pm}

sample standard deviation over three seeds per cell.

\dagger

: mean cost above

d

. The PID multiplier reached its clip value

\beta=100

within the first 1M interactions in 1 of 3 runs on HalfCheetah

2{\times}3

Environment	$d$	$\alpha_{\beta}$	Reward $\uparrow$	Cost $\downarrow$	$\bar{c}\leq d$
Ant $2{\times}4$	0.2	$10^{-5}$	$2268\mathbin{\pm}134$	$24.3\mathbin{\pm}5.6$	$\times$
		$10^{-4}$	$1618\mathbin{\pm}641$	$4.1\mathbin{\pm}5.8$	$\times$
	25	$10^{-5}$	$1836\mathbin{\pm}678$	$17.4\mathbin{\pm}15.1$	✓
		$10^{-4}$	$1552\mathbin{\pm}475$	$22.6\mathbin{\pm}18.3$	✓
Ant $4{\times}2$	1.0	$10^{-5}$	$1241\mathbin{\pm}566$	$2.5\mathbin{\pm}3.1$	$\times$
		$10^{-4}$	$1559\mathbin{\pm}359$	$4.2\mathbin{\pm}7.3$	$\times$

Appendix

Table 7. Multiplier learning rate

\alpha_{\beta}

at 1M interactions, at both cost limits per environment.

	$d=5$		$d=25$
Multiplier	Reward $\uparrow$	Cost $\downarrow$	Reward $\uparrow$	Cost $\downarrow$
Learned	$1974\mathbin{\pm}506$	$18.2^{\dagger}\mathbin{\pm}15.9$ (1/3)	$2527\mathbin{\pm}229$	$16.6\mathbin{\pm}2.0$
Parity start	$2266\mathbin{\pm}480$	$4.0\mathbin{\pm}6.1$ (2/3)	$2605\mathbin{\pm}235$	$19.1\mathbin{\pm}25.5$ (2/3)
Imagined	$1971\mathbin{\pm}310$	$29.1^{\dagger}\mathbin{\pm}48.0$ (2/3)	$1913\mathbin{\pm}395$	$87.0^{\dagger}\mathbin{\pm}75.0$ (1/3)
PID	$2077\mathbin{\pm}45$	$6.1^{\dagger}\mathbin{\pm}4.4$ (1/3)	$1555\mathbin{\pm}22$	$25.3^{\dagger}\mathbin{\pm}9.1$ (2/3)
$\beta{=}1$ (A)	$1954\mathbin{\pm}298$	$0.6\mathbin{\pm}0.3$ (3/3)	$2282\mathbin{\pm}219$	$2.7\mathbin{\pm}1.4$ (3/3)

Appendix

Table 8. Multiplier rules on HalfCheetah

2{\times}3

after 1M interactions. Each run’s reward and cost average its episodes in the last 5k interactions; cells give mean

\pm

sample standard deviation over three runs and, in parentheses, how many runs end within

d

(not available for Learned at

d=25

). Learned updates

\beta

from the measured episode cost (

\alpha_{\beta}=10^{-5}

\beta_{0}=10^{-4}

; the June–July 2026 runs of the main MAMuJoCo table); Parity start starts it at

\beta_{0}=1

; Imagined feeds it the imagined cost return instead; PID sets

\beta

by PID control on the measured cost (

K_{P}=1

K_{I}=10^{-5}

K_{D}=1

). The last three rows hold

\beta

fixed;

\beta=1

weights the standardized cost and reward advantages equally, and (A) and (B) are two independent sets of runs with the same settings. With

\beta

fixed,

d

does not enter training, so the two columns of a fixed row differ only by run-to-run variation.

\dagger

: mean cost above

d

Environment	$d$	Steps	Reward, messages	Reward, none	Cost, messages	Cost, none
Ant $2{\times}4$	25	500k	$1086\mathbin{\pm}412$	$1083\mathbin{\pm}115$	$6.8\mathbin{\pm}11.5$	$0.0\mathbin{\pm}0.0$
		700k	$1614\mathbin{\pm}553$	$1216\mathbin{\pm}244$	$53.1\mathbin{\pm}47.1$	$1.0\mathbin{\pm}1.1$
		800k	$1794\mathbin{\pm}672$	$1221\mathbin{\pm}60$	$82.9\mathbin{\pm}72.0$	$1.3\mathbin{\pm}1.3$
		900k	$1838\mathbin{\pm}583$	$1506\mathbin{\pm}346$	$32.7\mathbin{\pm}33.7$	$4.7\mathbin{\pm}7.6$
		1M	$1836\mathbin{\pm}678$	$1491\mathbin{\pm}80$	$17.4\mathbin{\pm}15.1$	$19.4\mathbin{\pm}30.3$
	0.2	500k	$1410\mathbin{\pm}395$	$816\mathbin{\pm}181$	$2.0\mathbin{\pm}1.8$	$8.6\mathbin{\pm}14.8$

Appendix

Table 9. CA-MAMBA acting with messages (

\mathcal{N}_{i}=\mathcal{I}

) and without them (

\mathcal{N}_{i}=\{i\}

in every collected episode; world-model learning and imagination use

\mathcal{N}_{i}=\mathcal{I}

in both),

\alpha_{\beta}=10^{-5}

, mean

\pm

sample standard deviation over three runs per cell.

Panel of Figure 4 — Figure 4. Training dynamics at $d=0$ . Top row, the easier maps bane_vs_bane and 2m_vs_1z , where MACPO-Discrete catches up only at 5M; bottom row, 8m and 1c3s5z , where it stays at zero win rate. Each panel pairs the full $0$ – $5$ M baseline range with the zoomed $0$ – $100$ k range. CA-MAMBA is green, MACPO-Discrete red. Four panels of SMAC learning curves at d=0 for bane_vs_bane, 2m_vs_1z, 8m and 1c3s5z, each comparing CA-MAMBA (green) with MACPO-Discrete (red) over the full 0 to 5M range and a zoomed 0 to 100k range.

	MACPO (1M)			CA-MAMBA (100k)
Scenario	Score $\uparrow$	Cost $\downarrow$	Win rate $\uparrow$	Score $\uparrow$	Cost $\downarrow$	Win rate $\uparrow$
2s_vs_1sc	$12.23\mathbin{\pm}0.48$	$\mathbf{0.00}\mathbin{\pm}0.00$	$0.11\mathbin{\pm}0.04$	$\mathbf{19.86}\mathbin{\pm}0.22$	$\mathbf{0.00}\mathbin{\pm}0.00$	$\mathbf{0.95}\mathbin{\pm}0.03$
3m	$11.20\mathbin{\pm}0.96$	$0.12\mathbin{\pm}0.03$	$0.34\mathbin{\pm}0.07$	$\mathbf{19.03}\mathbin{\pm}0.91$	$\mathbf{0.02}\mathbin{\pm}0.02$	$\mathbf{0.92}\mathbin{\pm}0.07$
2s3z	$9.24\mathbin{\pm}0.21$	$\mathbf{0.07}\mathbin{\pm}0.02$	$0.00\mathbin{\pm}0.00$	$\mathbf{18.72}\mathbin{\pm}0.44$	$0.16\mathbin{\pm}0.11$	$\mathbf{0.79}\mathbin{\pm}0.07$
3s_vs_3z	$8.32\mathbin{\pm}0.07$	$\mathbf{0.00}\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$\mathbf{19.89}\mathbin{\pm}0.05$	$0.02\mathbin{\pm}0.03$	$\mathbf{0.97}\mathbin{\pm}0.01$
3s_vs_4z	$5.77\mathbin{\pm}0.13$	$\mathbf{0.00}\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$\mathbf{15.49}\mathbin{\pm}3.71$	$0.26\mathbin{\pm}0.23$	$\mathbf{0.06}\mathbin{\pm}0.09$

Appendix

Table 10. Dense collision cost at

d=0

: CA-MAMBA after 100k interactions versus MACPO after 1M. Score is SMAC’s episode return, which credits damage and kills as well as the win (near

20

for a clean win). Mean

\mathbin{\pm}

sample standard deviation over three seeds; bold marks the better mean of the two methods per metric.

Figure 5. Safety-mechanism ablation on 8m (left) and 1c3s5z (right),

d=0

. Constrained CA-MAMBA (green) versus the unconstrained MAMBA learner (red), over cost, return, and win rate. Two panels, 8m and 1c3s5z at d=0, comparing cost, return and win rate of constrained CA-MAMBA (green) and the unconstrained MAMBA learner (red) over training.

Figure 6. HalfCheetah

2{\times}3

d=25

: predicted against realized team cost on training batches. The prediction follows the full rise-and-fall arc. Predicted and realized team cost on HalfCheetah 2x3 at d=25 against training interactions; the prediction tracks the rise and fall of realized cost.

Figure 7. Ant

2{\times}4

d=25

: cost is zero until roughly 0.4M interactions. The head predicts the zero phase and the onset. Predicted and realized team cost on Ant 2x4 at d=25 against training interactions; both stay at zero until about 0.4M interactions, then rise together.

Figure 8. Ant

4{\times}2

d=25

: cost stays near zero throughout and the prediction stays with it, within regression noise. Predicted and realized team cost on Ant 4x2 at d=25 against training interactions; both stay near zero throughout.

Env	$d$	Price	KL per agent	TV per agent	$\hat{B}_{c}$ [ $10^{-4}$ ]	$\hat{B}_{V}$ [ $10^{-3}$ ]	Imagined 14-step cost	Real 14-step cost	Real $-$ imagined
8m	0	Learned	$4.1\mathbin{\pm}0.1$	$0.66\mathbin{\pm}0.01$	$3.4\mathbin{\pm}1.5$	$63.3\mathbin{\pm}2.7$	$0.728\mathbin{\pm}0.017$	$1.083\mathbin{\pm}0.128$	$0.354\mathbin{\pm}0.132$ ( $+33\%$ )
HalfCheetah $2{\times}3$	5	Parity start	$5.6\mathbin{\pm}1.0$	$0.76\mathbin{\pm}0.04$	$-1.0\mathbin{\pm}1.9$	$-0.0\mathbin{\pm}0.2$	$0.014\mathbin{\pm}0.030$	$0.068\mathbin{\pm}0.112$	$0.054\mathbin{\pm}0.084$ ( $+79\%^{\diamond}$ )
	5	Fixed $\beta{=}1$	$5.4\mathbin{\pm}0.6$	$0.72\mathbin{\pm}0.07$	$2.0\mathbin{\pm}2.5$	$0.1\mathbin{\pm}0.2$	$-0.011\mathbin{\pm}0.009$	$0.009\mathbin{\pm}0.004$	$0.020\mathbin{\pm}0.012$ ( $+220\%$ )
	25	Learned	$4.9\mathbin{\pm}0.8$	$0.73\mathbin{\pm}0.04$	$3.6\mathbin{\pm}6.0$	$-0.2\mathbin{\pm}0.5$	$0.037\mathbin{\pm}0.017$	$0.056\mathbin{\pm}0.015$	$0.020\mathbin{\pm}0.020$ ( $+35\%^{\diamond}$ )
	25	Parity start	$4.1\mathbin{\pm}1.0$	$0.68\mathbin{\pm}0.06$	$0.3\mathbin{\pm}6.7$	$0.0\mathbin{\pm}1.2$	$0.092\mathbin{\pm}0.046$	$0.281\mathbin{\pm}0.109$	$0.189\mathbin{\pm}0.093$ ( $+67\%$ )
	25	Fixed $\beta{=}1$	$5.9\mathbin{\pm}0.0$	$0.77\mathbin{\pm}0.01$	$0.1\mathbin{\pm}2.3$	$0.3\mathbin{\pm}0.2$	$0.006\mathbin{\pm}0.028$	$0.035\mathbin{\pm}0.025$	$0.029\mathbin{\pm}0.009$ ( $+83\%$ )

Appendix

Table 11. World-model error on the runs whose logs record every episode’s exact length, three runs (seeds 1–3) per row. Each quantity is a run’s average over its last

10\%

of interactions (100k for 8m , 1M for HalfCheetah), then mean

\mathbin{\pm}

sample standard deviation over runs. Learned is the measured-cost multiplier of the main results; Parity start runs the same update from

\beta_{0}=1

; Fixed

\beta{=}1

holds the price at parity. KL is between the prior and the posterior over one agent’s latent (

32

categoricals), TV a Monte Carlo estimate of their total variation.

\hat{B}_{c}

and

\hat{B}_{V}

are per-step estimates on replay, real minus model as in Theorem L.1 . Imagined and real costs are discounted windows of the

14

costed imagined steps (

\sum_{t<14}0.99^{t}\approx 13.1

): imagined team cost per step, and realized episode cost divided by exact episode length, each times that sum. The cost head is an unconstrained regressor, so imagined cost can fall slightly below zero where real cost is near zero. The percentage is

(\bar{R}-\bar{I})/\bar{R}

from the row means

\bar{R}

and

\bar{I}

, positive when the model under-predicts;

\diamond

marks rows whose mean gap is within one standard deviation of the per-run gaps.

Multiplier	$\hat{B}_{V}$	$\overline{\|\delta\|}$	$\bar{V}_{C}$	$\widehat{\mathrm{TV}}$	$\widehat{\mathrm{TV}}_{\max}$
Imagined	$0.06$ $\pm 0.01$	$0.23$ $\pm 0.01$	$2.39$ $\pm 0.38$	$0.65$ $\pm 0.00$	$0.70$ $\pm 0.00$
Measured	$0.07$ $\pm 0.00$	$0.21$ $\pm 0.01$	$1.85$ $\pm 0.04$	$0.64$ $\pm 0.01$	$0.69$ $\pm 0.01$
PID	$0.01$ $\pm 0.00$	$0.04$ $\pm 0.02$	$0.28$ $\pm 0.14$	$0.45$ $\pm 0.05$	$0.50$ $\pm 0.07$
$\beta{=}5$	$0.01$ $\pm 0.00$	$0.02$ $\pm 0.01$	$0.13$ $\pm 0.07$	$0.63$ $\pm 0.04$	$0.69$ $\pm 0.02$
$\beta{=}1$	$0.04$ $\pm 0.01$	$0.12$ $\pm 0.03$	$0.82$ $\pm 0.29$	$0.52$ $\pm 0.05$	$0.56$ $\pm 0.05$
$\beta{=}0$	$0.06$ $\pm 0.01$	$0.22$ $\pm 0.02$	$2.19$ $\pm 0.19$	$0.66$ $\pm 0.01$	$0.71$ $\pm 0.01$

Appendix

Table 12. World-model transition diagnostics on SMAC 8m at

d=4

, logged on replayed sequences during training and averaged over the second half of training (50k to 100k environment interactions); each entry is a mean over the seeds that logged it, with the sample standard deviation (

\pm

) below it (three per row). For a replay step with deterministic state

h

\delta=V_{C}(h,z^{q})-\frac{1}{K}\sum_{k=1}^{K}V_{C}(h,z^{p,k})

compares the cost critic

V_{C}

at a posterior latent

z^{q}

with its mean over

K

prior latents

z^{p,k}

\hat{B}_{V}

is the mean of

\delta

, an estimate of

B_{V}

in Theorem L.1 (i) with the cost critic as the cost-to-go.

\overline{|\delta|}

is the mean of

|\delta|

\bar{V}_{C}

is the mean critic value at posterior latents, the scale for the first two columns.

\widehat{\mathrm{TV}}

is the Monte Carlo total-variation distance between an agent’s posterior and prior over its whole latent, averaged over agents, which upper-bounds the latent-space per-agent transition error

\epsilon_{P}^{i}

of Proposition L.3 on replay states.

\widehat{\mathrm{TV}}_{\max}

is the largest per-agent average. Rows list the arms whose runs logged these diagnostics.

Figure 9. HalfCheetah

2{\times}3

d=25

\beta

rises while episode cost exceeds the limit, then relaxes once cost returns below it. The multiplier beta and episode cost on HalfCheetah 2x3 at d=25 against training interactions; beta rises while cost exceeds the limit and falls once cost is back below it.

Figure 10. Ant

2{\times}4

d=25

\beta

holds at zero while the constraint is inactive and rises when it activates. The multiplier beta and episode cost on Ant 2x4 at d=25 against training interactions; beta stays at zero while cost is below the limit and rises when the limit becomes active.

Figure 11. Ant

4{\times}2

d=25

: cost never exceeds the limit and

\beta

never leaves zero. The multiplier beta and episode cost on Ant 4x2 at d=25 against training interactions; cost stays below the limit and beta stays at zero.

Figure 12. SMAC at

d=0

with the measured-cost multiplier,

\alpha_{\beta}=10^{-5}

: MMM (left) and 1c3s5z (right). The multiplier

\beta

(top, mean and one standard deviation over three seeds) rises linearly and stays below

1

at 100k interactions, while episode cost (bottom) falls but stays above the limit. Two panels of multiplier and episode cost against environment interactions for MMM and 1c3s5z at a zero casualty limit.

Baseline	CA-MAMBA lower cost	CA-MAMBA higher win rate	Baseline never wins
MACPO (100k)	23/24	22/24	22/24
MACPO (5M)	16/24	17/24	16/24
MAPPO-Lag (100k)	23/24	22/24	22/24
MAPPO-Lag (1M)	14/24	13/24	10/24
MAPPO-Lag (5M)	5/24	5/24	5/24
CA-MAMBA feasible	5/24

Appendix

Table 13. CA-MAMBA after 100k interactions against MACPO and MAPPO-Lag at the listed interaction budgets, over the 24 SMAC scenario–limit configurations all methods ran. Counts compare unrounded seed means at the same configuration with strict inequalities; a baseline never wins when its mean win rate is zero, and CA-MAMBA is feasible when its mean cost is at most

d

MAPPO-Lag budget $m(g)$	Configurations
$1\times$ (100k)	0
$5\times$ (500k)	3
$10\times$ (1M)	7
$20\times$ (2M)	3
$50\times$ (5M)	5
Not within $50\times$	5

Appendix

Table 14. The interaction multiple

m(g)

of a SMAC configuration

g

is the smallest multiple of CA-MAMBA’s 100k interactions, among

1\times

5\times

10\times

20\times

50\times

, at which MAPPO-Lag’s mean cost is at most and its mean win rate at least CA-MAMBA’s at 100k, ties counting as matched. Configurations where CA-MAMBA never wins are excluded, and the median is over the remaining 23, counting unmatched ones as

+\infty

. Resampling runs within each configuration, the median exceeds

10\times

99\%

2000

resamples and equals

20\times

97\%

Scenario	$d$	Score $\uparrow$	Cost $\downarrow$	Win rate $\uparrow$
1c3s5z	0	$19.29\mathbin{\pm}0.14$	$4.23\mathbin{\pm}0.27$	$0.82\mathbin{\pm}0.02$
1c3s5z	4	$19.15\mathbin{\pm}0.52$	$4.83\mathbin{\pm}0.94$	$0.79\mathbin{\pm}0.12$
2m_vs_1z	0	$19.30\mathbin{\pm}0.49$	$0.14\mathbin{\pm}0.08$	$0.95\mathbin{\pm}0.04$
2m_vs_1z	1	$19.34\mathbin{\pm}0.38$	$0.14\mathbin{\pm}0.09$	$0.95\mathbin{\pm}0.03$
2s_vs_1sc	0	$19.95\mathbin{\pm}0.21$	$0.23\mathbin{\pm}0.11$	$0.96\mathbin{\pm}0.03$
2s_vs_1sc	1	$19.56\mathbin{\pm}0.49$	$0.31\mathbin{\pm}0.10$	$0.92\mathbin{\pm}0.06$

Appendix

Table 15. CA-MAMBA results at exactly 100k interactions on the 24 SMAC configurations shared with the baselines, with the multiplier updated from the measured episode cost. Score is SMAC’s episode return, which credits damage and kills as well as the win (near

20

for a clean win). Entries are mean

\mathbin{\pm}

sample standard deviation across seeds; CA-MAMBA uses three valid seeds except the starred configuration, which uses two.

		CA-MAMBA (100k)		MAPPO-Lag (100k)		MAPPO-Lag (1M)		MAPPO-Lag (5M)
Scenario	$d$	Cost $\downarrow$	Win $\uparrow$	Cost $\downarrow$	Win $\uparrow$	Cost $\downarrow$	Win $\uparrow$	Cost $\downarrow$	Win $\uparrow$
1c3s5z	0	$4.23\mathbin{\pm}0.27$	$0.82\mathbin{\pm}0.02$	$9.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$2.22\mathbin{\pm}0.64$	$0.96\mathbin{\pm}0.04$	$0.44\mathbin{\pm}0.11$	$1.00\mathbin{\pm}0.00$
1c3s5z	4	$4.83\mathbin{\pm}0.94$	$0.79\mathbin{\pm}0.12$	$9.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$2.08\mathbin{\pm}0.66$	$0.96\mathbin{\pm}0.04$	$0.37\mathbin{\pm}0.14$	$1.00\mathbin{\pm}0.00$
2m_vs_1z	0	$0.14\mathbin{\pm}0.08$	$0.95\mathbin{\pm}0.04$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$0.01\mathbin{\pm}0.00$	$1.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00^{\dagger}$	$1.00\mathbin{\pm}0.00^{\dagger}$
2m_vs_1z	1	$0.14\mathbin{\pm}0.09$	$0.95\mathbin{\pm}0.03$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$0.01\mathbin{\pm}0.00$	$1.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$1.00\mathbin{\pm}0.00$
2s_vs_1sc	0	$0.23\mathbin{\pm}0.11$	$0.96\mathbin{\pm}0.03$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$0.21\mathbin{\pm}0.01$	$1.00\mathbin{\pm}0.00$	$0.11\mathbin{\pm}0.07$	$1.00\mathbin{\pm}0.00$

Appendix

Table 16. CA-MAMBA after 100k environment interactions against MAPPO-Lag after 100k, 1M and 5M interactions on the 24 SMAC scenario–limit configurations both methods ran. Read each row across: at the same scenario and limit

d

, lower cost and higher win rate are better, and a method is feasible when its mean cost is at most

d

. Entries are mean

\mathbin{\pm}

sample standard deviation across seeds. MAPPO-Lag uses three seeds at every checkpoint; CA-MAMBA uses three valid seeds except the starred configuration, which uses two. Each daggered entry includes one run that stopped before 5M and contributes its last logged window ( 2m_vs_1z at

d=0

, seed 1, stopped at 3.01M; bane_vs_bane at

d=0

, seed 1, stopped at 4.52M).

		CA-MAMBA (100k)		MACPO (100k)		MACPO (5M)
Scenario	$d$	Cost $\downarrow$	Win $\uparrow$	Cost $\downarrow$	Win $\uparrow$	Cost $\downarrow$	Win $\uparrow$
1c3s5z	0	$4.23\mathbin{\pm}0.27$	$0.82\mathbin{\pm}0.02$	$9.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$9.00\mathbin{\pm}0.00^{\dagger}$	$0.00\mathbin{\pm}0.00^{\dagger}$
1c3s5z	4	$4.83\mathbin{\pm}0.94$	$0.79\mathbin{\pm}0.12$	$9.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$9.00\mathbin{\pm}0.00^{\dagger}$	$0.00\mathbin{\pm}0.00^{\dagger}$
2m_vs_1z	0	$0.14\mathbin{\pm}0.08$	$0.95\mathbin{\pm}0.04$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$0.02\mathbin{\pm}0.02$	$0.99\mathbin{\pm}0.01$
2m_vs_1z	1	$0.14\mathbin{\pm}0.09$	$0.95\mathbin{\pm}0.03$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$1.00\mathbin{\pm}0.00$
2s_vs_1sc	0	$0.23\mathbin{\pm}0.11$	$0.96\mathbin{\pm}0.03$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$	$2.00\mathbin{\pm}0.00$	$0.00\mathbin{\pm}0.00$

Appendix

Table 17. CA-MAMBA after 100k environment interactions against MACPO after 100k and 5M interactions on the 24 SMAC scenario–limit configurations all methods ran. Read each row across: at the same scenario and limit

d

, lower cost and higher win rate are better. Entries are mean

\mathbin{\pm}

sample standard deviation over three seeds, or two where starred. Daggered MACPO (5M) entries include runs that stopped before 5M, whose tracked values predate 5M ( 1c3s5z at

d=0

, seed 1, last logged at 2.155M; 1c3s5z at

d=4

, seed 2, last logged at 2.025M; 3s_vs_5z at

d=0

, seed 1, last logged at 4.25M; 3s_vs_5z at

d=0

, seed 2, last logged at 4.28M; 3s_vs_5z at

d=1

, seed 1, last logged at 4.05M; 3s5z_vs_3s6z at

d=0

, seed 2, last logged at 2.55M; 3s5z_vs_3s6z at

d=4

, seed 2, last logged at 2.295M; bane_vs_bane at

d=4

, seed 2, last logged at 0.735M; bane_vs_bane at

d=4

, seed 3, last logged at 1.29M).

Network
Parameter	Default	SMAC override
Hidden size	512	128
Layers	2	1
Activation	ReLU	–
Orthogonal init (gain)	0.01	–
Feature normalization	yes	–

Appendix

Table 18. Hyperparameters of the MACPO baseline on SMAC. SMAC-specific overrides are listed separately; all other parameters use the defaults shown.

Constrained MARL

See all 14 papers

Oct 5, 2026

Who Bears the Burden? Learning Responsibility for Shared Constraints in Multi-Agent Reinforcement Learning

When multiple agents share a cost budget, a common Lagrange multiplier can enforce the aggregate constraint but does not determine how its penalty should be allocated across agents. Uniform penalties ignore heterogeneity in the rewards agents sacrifice, while agent-specific multipliers may still rely on the same aggregate cost signal. We introduce Lagrangian Responsibility Allocation (LiRA), which learns each agent's share of a common multiplier by optimizing social welfare over a finite training horizon. The multiplier enforces the aggregate budget, while responsibility shares redistribute its influence without modifying the original rewards or constraints. For convex games under standard regularity conditions, varying these shares induces a smooth family of normalized generalized Nash equilibria in which active constraints remain at their budgets while welfare varies. To optimize responsibility before convergence, we derive a welfare gradient that accounts for both learning updates and the induced change in data distribution. Across CityLearn, MABIM, Harvest, and MetaDrive, spanning 3 to 400 agents, LiRA improves average social welfare by up to 29% over uniform and agent-specific multiplier baselines. Grid and driving costs remain within budget, inventory violations decrease, and Harvest makes more effective use of available budget.

Xiaoyang Cao, Jingqi Li, Zhe Fu +1

Massachusetts Institute of Technology · The University of Texas at Austin · Stanford University +1

Game Theory Constrained MARL Constrained RL Cooperative MARL Multi-Agent Reinforcement Learning

Ask PDF

Sep 25, 2026

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

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.

Guowei Zou, Haitao Wang, Guoxin Wang +4

Sun Yat-sen University

Multi-Agent Planning Cooperative MARL Action-Conditioned World Models Offline MARL Efficient Inference for World Action Models

Ask PDF

Jun 1, 2026

Coordination Graphs for Constrained Multi-Agent Reinforcement Learning

Constrained Multi-agent reinforcement learning (CMARL) faces two intertwined challenges: the joint action space grows exponentially with the number of agents, and additional requirements couple agents in ways that reward structure alone does not capture. We introduce Coordination Graphs for Constrained Multi-Agent Reinforcement Learning (CG-CMARL), a framework that addresses both challenges by combining coordination graphs with Lagrangian duality. The system decomposes the joint problem into pairwise regions, each served by a set of shared Q-functions, one for the primary objective and one for each of the constraints, so that the number of learned models is independent of the number of agents. At execution time, Max-Sum message passing coordinates actions across the factor graph, while a Lagrangian multiplier controls the objective--constraint tradeoff, allowing a single trained model to trace a Pareto front without retraining. We provide convergence guarantees under mild conditions, together with a compositional error bound that decomposes into separate interpretable sources, each traceable to a specific design choice and independently controllable. Experiments on cooperative navigation tasks (where teams of up to 10 agents must coordinate to reach target positions while satisfying pairwise constraints) show that our method produces Pareto fronts dominating established baselines trained at fixed reward-shaping ratios, while scaling to team sizes where centralized approaches become intractable.

Santiago Amaya-Corredor, Miguel Calvo-Fullana, Anders Jonsson

Department of Engineering, Universitat Pompeu Fabra, Barcelona, Spain

Multi-Objective Reinforcement Learning Multi-Agent System Optimization Constrained RL Constrained Optimization Cooperative MARL Constrained MARL

Ask PDF

Abstract

Constrained MARL

Explore similar work

Who Bears the Burden? Learning Responsibility for Shared Constraints in Multi-Agent Reinforcement Learning

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

Coordination Graphs for Constrained Multi-Agent Reinforcement Learning