Dr. MAS: Stable Reinforcement Learning for Multi-Agent LLM Systems

Lang Feng; Longtao Zheng; Shuo He; Fuxiang Zhang; Bo An

Multi-agent reinforcement learning

Dr. MASStable Reinforcement Learning
for Multi-Agent LLM Systems

NeurIPS 2026 MALGAI @ ICLR 2026

Lang Feng Longtao Zheng Shuo He Fuxiang Zhang Bo An^*

Nanyang Technological University, Singapore

Read the paper Code 162 Cite this work

arXiv: 2602.08847 Models on Hugging Face

Figure 1 from the paper. Left: GRPO uses a global advantage baseline for all agents and exhibits gradient spikes, whereas Dr. MAS uses an agent-specific baseline and stable gradients. Right: existing single-actor frameworks are contrasted with Dr. MAS co-training multiple LLM actors. — Figure 1 A simple change in the learning signal, paired with a framework built for multiple LLM actors.

+5.6%

Better math reasoning

Relative avg@16 improvement

+15.2%

Stronger multi-turn search

Relative avg@16 improvement

Relative performance gains over multi-agent GRPO, as reported in the paper. See experimental results for settings.

01 / Overview

Collaboration needs
a better learning signal.

Abstract

Multi-agent LLM systems improve reasoning and tool use via role specialization, yet reinforcement learning (RL) post-training for such systems remains unstable and underexplored. We theoretically pinpoint a key source of instability when extending group-based RL to cooperative multi-agent LLM systems: under GRPO-style optimization, a global normalization baseline can mismatch heterogeneous agents' reward distributions, inducing gradient-norm instability.

Based on this finding, we propose Dr. MAS, a simple and stable RL recipe for cooperative multi-agent LLM systems. Algorithmically, Dr. MAS normalizes advantages per agent using each agent's own reward statistics, which calibrates gradient scales and dramatically stabilizes the training. Systemically, Dr. MAS provides an end-to-end multi-agent RL framework with scalable orchestration, flexible per-agent LLM serving and optimization, and shared resource scheduling of actor backends.

Read the full abstract

Unlike single-actor frameworks such as veRL, Dr. MAS natively supports multiple heterogeneous LLMs over a unified GPU pool. Lifecycle-aware backend management and dynamic dispatch release inactive models' GPU memory and schedule active agents on demand, enabling hardware-efficient co-training. We evaluate Dr. MAS on multi-agent math reasoning and multi-turn search with Qwen2.5 and Qwen3. Dr. MAS achieves clear gains over vanilla GRPO (e.g., +5.6% avg@16 and +4.6% pass@16 on math, and +15.2% avg@16 and +13.1% pass@16 on search) while largely eliminating gradient spikes. It also remains highly effective under heterogeneous agent-model assignments while improving efficiency.

Understand the instability

A theoretical analysis connects mismatched global reward statistics to inflated per-agent gradient norms.

Calibrate every agent

Agent-wise advantage normalization aligns update scales with each role's own reward distribution.

Train as one system

Flexible orchestration, per-model optimization, and lifecycle-aware scheduling in a unified GPU pool.

02 / Theory & Method

Find the instability.
Then remove the mismatch.

A theoretical diagnosis.
A one-line change to normalization.

4.1 / Risk of Gradient Norm Explosion

The instability has a precise scaling law.

Proof in Appendix B

Following the paper's notation, the score function and unclipped gradient contribution for an active agent are:

z_{i,t}^{(k)} \triangleq \nabla_{\theta_k}\rho_{\theta_k}(\bm{a}_t^i)

\tilde{g}_k^{\mathrm{global}} \triangleq \frac{R^i-\mu}{\sigma}\,z_{i,t}^{(k)}

Expectations sample outputs uniformly from $\mathcal{Y}_k$ , the steps where agent $k$ is active. Assumption 4.1 bounds the score second moment: $0\leq\mathbb{E}[\|z_{i,t}^{(k)}\|^2]\leq C_k\lt\infty$ .

Lemma 4.2

The gradient second moment

Agent-active reward statistics determine the multiplicative scale; $\Delta_k$ retains score-reward dependence.

\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}\!\left[\|\tilde{g}_k^{\mathrm{global}}\|^2\right]=\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}\!\left[\|z_{i,t}^{(k)}\|^2\right]\,\frac{\sigma_k^2+(\mu_k-\mu)^2}{\sigma^2}+\Delta_k

See the three-step proof

Appendix B.1 expands the reward about the agent-active mean, then uses the covariance identity. No independence assumption is needed.

Center the reward on the active agent.
Rewrite the trajectory reward using the same $R^i,\mu,\mu_k$ as the paper.
$R^i-\mu=(R^i-\mu_k)+(\mu_k-\mu)$
The cross term vanishes.
Under uniform sampling from $\mathcal{Y}_k$ , the centered reward has zero mean. Its second moment is:
$\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}\!\left[(R^i-\mu)^2\right]=\sigma_k^2+(\mu_k-\mu)^2$
Keep the covariance correction.
Factor the score and normalized-reward second moments, retaining their covariance. Substitution gives Lemma 4.2 exactly.
$\Delta_k=\operatorname{Cov}\!\Big(\|z_{i,t}^{(k)}\|^2,\frac{(R^i-\mu)^2}{\sigma^2}\Big)$

The identity requires finite moments and a nonzero normalization standard deviation. It concerns the unclipped, per-active-output gradient contribution, not the norm of an averaged minibatch update.

Proposition 4.3 / Gradient-Norm Inflation

Large mean offsets or variance ratios can amplify gradients.

A larger mean offset or variance ratio increases the multiplicative term. A smaller global standard deviation magnifies both.

\frac{\sigma_k^2+(\mu_k-\mu)^2}{\sigma^2}=\underbrace{\frac{\sigma_k^2}{\sigma^2}}_{\text{variance mismatch}}+\underbrace{\frac{(\mu_k-\mu)^2}{\sigma^2}}_{\text{mean mismatch}}

Along training iterations $m$ , if the score and covariance do not cancel this growth:

\frac{\sigma_{k,m}^2+(\mu_{k,m}-\mu_m)^2}{\sigma_m^2}\to\infty\quad\Longrightarrow\quad\mathbb{E}[\|\tilde{g}_m^{\mathrm{global}}\|^2]\to\infty

Here $\tilde{g}_m^{\mathrm{global}}$ stacks all agents' gradients. One diverging component is enough.

As qualified in Appendix B.2, divergence follows when the growing factor is not cancelled by the score second moment or covariance correction. A growing factor alone is not an unconditional guarantee of gradient divergence.

Interactive theory lab

Gradient second moment.

Adjust the reward statistics. Isolate how normalization changes the gradient scale.

Lemma 4.2 Proposition 4.3

01 / Set the statistics

Global statistics Mean

\mu

0.00 Std. dev.

\sigma

0.50

Agent statistics Mean

\mu_k

1.00 Std. dev.

\sigma_k

1.00

The one-line remedy

(\mu,\sigma)\longrightarrow(\mu_k,\sigma_k)

Same rewards. Only the statistics change.

02 / Observe the effectGlobal normalization

8.00×

+700.00%vs. agent-wise normalization

Active gradient second moment

\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}[\|\tilde g_k^{\mathrm{global}}\|^2]

\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}[\|\tilde g_k^{\mathrm{agent}}\|^2]

Agent-wise reference: 1.00×

As the global scale shrinks

Theoretical response

Global Agent-wise reference Selected scale

Sweep $\sigma$ with $\mu,\mu_k,\sigma_k$ held fixed. The dot follows the selected normalization; the axis stays fixed as you move $\sigma$ .

Inside the global multiplier

Two sources of mismatch.

Amplification

\displaystyle\frac{\sigma_k^2}{\sigma^2}

4.00Variance ratio

\displaystyle\frac{(\mu_k-\mu)^2}{\sigma^2}

4.00Mean mismatch

Global8.00×Second-moment multiplier

Controlled comparison, not training data

\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}[\|z_{i,t}^{(k)}\|^2]=1,\qquad\Delta_k=\Delta_k^{\mathrm{agent}}=0

These fixed quantities isolate the normalization effect. Reward moments alone do not determine an observed training gradient.

The independently adjustable moments are a teaching model, not a sampled rollout. Changing the other moments recomputes the curve and its axis; the global comparison remains visible when the remedy is active.

4.2 / Agent-Wise Remedy

Change the statistics.
Keep the objective.

Replace $(\mu,\sigma)$ with $(\mu_k,\sigma_k)$ . The multiplicative term becomes exactly 1, giving Eq. (6):

\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}\!\left[\|\tilde{g}_k^{\mathrm{agent}}\|^2\right]=\mathbb{E}_{\bm{a}_t^i\sim\mathcal{Y}_k}\!\left[\|z_{i,t}^{(k)}\|^2\right]+\Delta_k^{\mathrm{agent}}

This removes the reward-statistics multiplier, not every possible source of gradient instability.

GRPOGlobal baseline

The same statistics for every role.

A_{\mathrm{global}}^{i} = \frac{R^{i} - \mu}{\sigma}

All agents inherit the same normalized advantage, even when their active steps follow different reward distributions.

Baseline mismatch can inflate gradient norms.

Dr. MASAgent-wise baseline

Each role learns on its own scale.

A_{\mathrm{agent}}^{i,k} = \frac{R^{i} - \color{#8c569b}{\mu_k}}{\color{#8c569b}{\sigma_k}}

Normalize each agent's advantage with the reward statistics of its own active steps, keeping update scales better calibrated.

Agent-wise normalization stabilizes co-training.

Statistics that follow the agent.

For the set of outputs $\mathcal{Y}_k$ where agent $k$ is active:

\mu_k \triangleq \frac{1}{|\mathcal{Y}_k|} \sum_{\bm{a}_t^i \in \mathcal{Y}_k} R^i

\sigma_k^2 \triangleq \frac{1}{|\mathcal{Y}_k|} \sum_{\bm{a}_t^i \in \mathcal{Y}_k} (R^i - \mu_k)^2

These are action-weighted statistics: a trajectory's reward contributes once for every step in which the agent is active. Section 4.2

Inside the update A shared objective, with agent-specific advantages

As illustrated in Figure 1, the clipped group-relative objective is retained. Dr. MAS replaces the global advantage with $A_{\mathrm{agent}}^{i,k}$ for each agent.

\mathcal{J}_k(\theta_k) = \mathbb{E}_{x \sim p(x)}\!\left[ \frac{1}{|\mathcal{Y}_k|} \sum_{\bm{a}_t^i \in \mathcal{Y}_k} \min\!\left( \rho_{\theta_k}(\bm{a}_t^i) A_{\mathrm{agent}}^{i,k},\; \operatorname{clip}\!\left(\rho_{\theta_k}(\bm{a}_t^i), 1 \pm \epsilon\right) A_{\mathrm{agent}}^{i,k} \right) \right]

Scroll horizontally to view the full objective.

The importance-sampling ratio is $\rho_{\theta_k}(\bm{a}_t^i) = \pi_{\theta_k}(\bm{a}_t^i \mid \bm{s}_t^i) / \pi_{\theta_k^{\mathrm{old}}}(\bm{a}_t^i \mid \bm{s}_t^i)$ . As in the paper's presentation, the KL regularization term is omitted here for clarity.

Training dynamics

Less noise. More learning.

Qwen2.5-3B / Non-sharing

Original paper plots compare training accuracy and verifier, answer, and search gradient norms. Dr. MAS substantially reduces the large gradient spikes seen with GRPO. — Figure 3 from the paper. Agent-wise normalization smooths gradient norms across the three-agent search system.

Appendix F.2: a more severe failure case at 7B

Under Qwen2.5-7B non-sharing, the paper reports a search-agent gradient spike above 80 followed by NaN. The missing orange curve is not a return to zero.

Figure 7: original training accuracy and verifier, answer and search gradient plots. Under Qwen2.5-7B non-sharing, GRPO search-agent gradients spike above 80 and subsequently become NaN; Dr. MAS shows smoother curves. — Figure 7 / Additional stability evidence, Qwen2.5-7B, non-sharing. Appendix F.2

03 / Framework

Framework for
Multi-Agent LLM RL

An end-to-end RL post-training pipeline for cooperative multi-agent LLM systems, from flexible orchestration to per-model updates over a shared GPU pool.

Section 4.3 Appendix C.1

Figure 2. Environment workers interact with the multi-agent orchestrator. Logical agents map to shared or separate ActorRollout worker groups. Lifecycle-aware backends use a shared GPU pool, and the RL trainer splits collected trajectories into per-model micro-batches for optimization. — Figure 2 / End-to-end architecture From distributed rollout to per-worker-group policy updates.

Multi-Agent Orchestration

A multi-agent trajectory collector coordinates interactions between agents and the environment. A pluggable orchestra defines agent roles and execution flows, including sequential collaboration, conditional routing, and per-agent active masking. At each step, it selects the active agent based on the current state and previous agent outputs.

Figure 5. Left: a solver proposes a solution and a verifier either approves it or requests another refinement. Right: a verifier repeatedly routes to a searcher while more evidence is needed, then invokes an answer agent. — Figure 5 / Two evaluated workflows. Appendix C.1

Math: a refinement loop: A solver proposes or revises a solution; a verifier accepts it or requests another round.
Search: a verifier-led hierarchy: The verifier invokes search for more evidence, then calls the answer agent when ready.

Agent-Model Assignment

Logical agents are mapped to physical LLM worker groups (wg_id). In non-shared settings, each agent has a distinct worker group. In shared settings, agents configured with the same model use one worker group, reusing its weights for joint training and inference.

Shared LLM Optional

Same-model agents reuse one set of weights.

Separate LLMs Optional

Independent weights; model choices can differ.

Per-Agent Configuration

Agent-specific training hyperparameters, such as actor.optim.lr, are injected into each agent's configuration and attached to its corresponding LLM worker group.

Shared GPU Resource Pooling and Scheduling

Dr. MAS co-locates LLM worker groups in a unified GPU resource pool. ActorRollout backends are scheduled through Ray placement groups and served by SGLang. During rollout, agent_to_wg_mapping routes requests to the corresponding actor_rollout_wg[wg_id] backend.

Route
Dispatch each active agent's generation request to its assigned worker group.
resume()
Load the active backend's model weights and KV caches onto GPU.
offload()
Release inactive backends' GPU memory; models need not all remain resident.

Per-Worker-Group Optimization

After rollout collection, the trainer partitions the aggregated batch by worker-group ID:

\mathcal{B}_{g}=\{\,\tau\in\mathcal{B}\mid\texttt{wg\_id}(\tau)=g\,\}

Policy updates are performed for each worker group, so gradients from an agent's trajectories update only its designated LLM backend.

04 / Experiments

Stable training.
Measurable gains.

Two workflows. Four model sizes.
Shared and independent LLM weights.

Model

Agent weights

Metric

Multi-agent GRPO57.8%

Dr. MAS62.3%

+4.5 pp

Average across benchmarks

Score differences

Every benchmark, from a shared zero.

GRPO higher Dr. MAS higher

Math reasoning with Qwen3-8B, shared weights, avg@16
Benchmark	Dr. MAS - GRPO (pp)	GRPO	Dr. MAS	Gain (pp)

Scores are percentages. Gains are percentage-point (pp) differences between Dr. MAS and multi-agent GRPO in the same setting. Displayed values are transcribed from Tables 1 and 2; deltas use the rounded scores.

What do avg@16 and pass@16 mean?

avg@16 measures average correctness over 16 sampled attempts. pass@16 measures whether at least one of those 16 attempts is correct. Shared weights means the logical agents use one shared LLM; separate weights means each agent has its own LLM parameters.

Heterogeneous by design

Different models.
One effective team.

Keep Qwen2.5-7B as the verifier. Use Llama-3.2-3B-Instruct for search and answering. The heterogeneous assignment retains nearly the same benchmark performance as the all-7B team.

Explore heterogeneous assignments

42.0%avg@16

All-7B baseline: 42.5%

57.5%pass@16

All-7B baseline: 57.7%

Figure 4: three-agent multi-turn search, heterogeneous versus all-Qwen2.5-7B model assignment.

Scope & limitations. Evaluated on two representative workflows. Dr. MAS does not address every source of multi-agent RL instability, including credit assignment across agents and turns. Generalization to broader architectures and larger models remains open.

05 / Citation

Build on
this work.

If you find Dr. MAS useful in your research, please consider citing our paper.

Download .bib

@article{feng2026dr,
  title={{Dr. MAS}: Stable reinforcement learning for multi-agent {LLM} systems},
  author={Feng, Lang and Zheng, Longtao and He, Shuo and Zhang, Fuxiang and An, Bo},
  journal={arXiv preprint arXiv:2602.08847},
  year={2026}
}

Dr. MASStable Reinforcement Learning
for Multi-Agent LLM Systems

Better math reasoning

Stronger multi-turn search