Source-linked AI summary
Dr. MAS: Stable Reinforcement Learning for Multi-Agent LLM Systems
Lang Feng, Longtao Zheng, Shuo He, Fuxiang Zhang, Bo An
TL;DR
Reliable RL post-training for multi-agent LLM systems is difficult because global GRPO normalization can mismatch agents’ reward distributions and destabilize gradients. Dr. MAS uses per-agent reward normalization within an end-to-end training framework, yielding consistent gains over vanilla GRPO across math and search while largely eliminating gradient spikes. It remains effective with heterogeneous agent-model assignments and improves efficiency.
Problem
RL post-training for multi-agent LLM systems remains underexplored and unstable when single-agent group-based methods use a global baseline across agents with different reward distributions.
Method
Dr. MAS normalizes each agent’s advantages using its own reward mean and variance and provides orchestration, per-agent configuration, model-sharing, and resource-pooling support.
Results
+5.6% avg@16 and +4.6% pass@16 on math, and +15.2% avg@16 and +13.1% pass@16 on search, were achieved over vanilla GRPO while gradient spikes were largely eliminated.
Takeaways & Limitations
Per-agent normalization provides more stable multi-agent RL across sharing and non-sharing settings, including heterogeneous agent-model assignments and improved efficiency.
Takeaways & Limitations
Evaluation covers a two-agent math loop and a three-agent multi-turn search pipeline, so the demonstrated scope is limited to these settings.
Abstract
from arXiv · showhide
Multi-agent LLM systems enable advanced reasoning and tool use via role specialization, yet reliable reinforcement learning (RL) post-training for such systems remains difficult. In this work, we theoretically pinpoint a key reason for training instability when extending group-based RL to multi-agent LLM systems. We show that under GRPO-style optimization, a global normalization baseline may deviate from diverse agents' reward distributions, which ultimately leads to gradient-norm instability. Based on this finding, we propose Dr. MAS, a simple and stable RL training recipe for multi-agent LLM systems. Dr. MAS uses an agent-wise remedy: normalizing advantages per agent using each agent's own reward statistics, which calibrates gradient scales and dramatically stabilizes training, both theoretically and empirically. Beyond the algorithm, Dr. MAS provides an end-to-end RL training framework for multi-agent LLM systems, supporting scalable orchestration, flexible per-agent LLM serving and optimization configs, and shared resource scheduling of LLM actor backends. We evaluate Dr. MAS on multi-agent math reasoning and multi-turn search benchmarks using Qwen2.5 and Qwen3 series models. 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. Moreover, it remains highly effective under heterogeneous agent-model assignments while improving efficiency.
1 Introduction
Multi-agent LLM RL remains underexplored and unstable when extending single-agent group-based methods. Dr. MAS addresses this with agent-wise normalization and an end-to-end framework, improving performance and stability across tasks and configurations.
- Multi-agent LLM applications require agents to interact and coordinate in information retrieval, software engineering, and device control.
- RL extensions to multi-agent LLM systems face algorithmic instability and limited support for heterogeneous orchestration, co-training, scheduling, and resource sharing.
- Dr. MAS normalizes each agent’s rewards using its own mean and variance, balancing per-agent gradients and reducing policy-gradient variance.
- +5.6% avg@16 and +4.6% pass@16 on math, and +15.2% avg@16 and +13.1% pass@16 on search, were achieved over vanilla GRPO.
- Dr. MAS largely eliminates gradient-norm spikes and remains effective with heterogeneous agent-model assignments while improving efficiency.
2 Related Work
Related work spans RL methods for LLM post-training, multi-agent coordination, and scalable RL infrastructure. Existing systems increasingly support multi-turn and tool-integrated training but retain limitations for heterogeneous model assignments.
- RL for LLMs has progressed from alignment-focused RLHF toward RLVR and group-based methods such as GRPO and RLOO.
- Multi-agent RL work studies role-specialized, multi-turn coordination through self-play and related interfaces, but often remains confined to two-agent settings.
- General-purpose RL stacks provide modular pipeline abstractions and distributed execution to improve utilization during heavy sampling.
- Agentic RL frameworks increasingly emphasize multi-turn rollout and tool integration, while support for heterogeneous model assignment remains limited.
3 Preliminaries
The paper models cooperative multi-agent LLM systems as dynamically interacting policies that produce shared trajectory rewards. GRPO normalizes group rewards and propagates the resulting advantage across each agent’s outputs.
- A cooperative system contains K distinct LLM agents, each with policy πθk, jointly solving tasks sampled from a distribution.
- A trajectory records conversational states, text actions, and the active agent at each step, allowing agent identity to change dynamically.
- Agents may share parameters with role-specific prompts or maintain distinct parameters, depending on the system design.
- 3.2 Group Relative Policy Optimization: GRPO compares multiple rollouts from one instruction and normalizes their scalar terminal rewards using the group mean and standard deviation.
- 3.2 Group Relative Policy Optimization: The normalized trajectory advantage is propagated to every output contributed by each agent, whose invocation frequency determines its sample count.
- 3.2 Group Relative Policy Optimization: The policy objective is defined separately for each agent using its outputs and importance-sampling ratios, with KL regularization omitted here.
4 Methodology
Dr. MAS addresses gradient instability in multi-agent GRPO by replacing a potentially misaligned global reward baseline with agent-wise normalization, while providing an end-to-end framework for orchestration, model assignment, configuration, and resource scheduling.
- Gradient instability: Different specialized agents can have distinct reward distributions, making vanilla GRPO’s shared global baseline suboptimal and potentially destabilizing training.Agents may remain systematically above or below the global mean, shifting their effective advantages and inflating gradient-estimator variance.
- Gradient instability: Under bounded score-function second moments, the per-agent gradient’s second moment decomposes into a dominant scaling factor and a residual score-reward covariance correction.The analysis uses the unclipped GRPO gradient and separates the contribution of reward statistics from covariance effects.
- Gradient instability: Global normalization can inflate gradient norms when an agent’s reward mean differs from the global mean or its conditional variance exceeds the global variance.The resulting inflation can produce severe gradient spikes and destabilize the entire multi-agent training process.
- End-to-end framework: The framework combines distributed orchestration, logical-to-worker-group assignment with optional model sharing, per-agent optimization settings, and shared resource scheduling.It supports conditional agent execution and separates logical assignments from physical resource placement for efficient multi-agent rollouts.
- Agent-wise remedy: Dr. MAS normalizes each agent’s rewards with statistics computed only on the steps where that agent is active, replacing the global baseline (µ, σ) with (µk, σk).This makes the dominant normalization factor equal to 1 and bounds each agent’s gradient second moment by its own score statistics.
- End-to-end framework: The framework instantiates a two-agent solver-verifier loop for math and a hierarchical three-agent verifier-search-answer pipeline for multi-turn search.The search workflow selectively invokes retrieval or answer generation based on the verifier’s decision.
5 Experiment
Dr. MAS is evaluated on multi-agent math reasoning and multi-turn search under shared and non-shared LLM settings, showing more reliable performance and training dynamics than vanilla GRPO. Ablations and heterogeneous assignments further examine normalization choices, gradient stability, and deployment efficiency.
- Math Orchestration: Dr. MAS improves math performance over vanilla GRPO under both LLM-sharing and non-sharing settings, with an overall increase of 5.6% in avg@16 and 4.6% in pass@16.The evaluation uses Qwen3-4B/8B agents on competitive mathematical reasoning benchmarks.
- Search Orchestration: 15.2% avg@16 and 13.1% pass@16 are the overall gains reported for Dr. MAS on multi-turn search.The search workflow uses verifier, search, and answer agents under shared and non-shared LLM settings.
- Gradient-Norm Instability: Vanilla GRPO produces frequent, high-magnitude per-agent gradient spikes, whereas Dr. MAS keeps gradient norms smoother and lower during three-agent search training.The search agent exhibits the largest spikes, with additional peaks for the answer and verifier agents.
- Ablation Study: Fully per-agent normalization using both mean and standard deviation achieves the best ablation results on the search task.Using either per-agent mean or per-agent standard deviation already improves over global normalization, while combining both provides the most reliable learning signal.
6 Conclusions and Limitations
The paper concludes that Dr. MAS improves multi-agent RL stability and performance, while leaving credit assignment and large-agent scalability unresolved.
- Conclusions: Dr. MAS normalizes advantages per agent and improves stability over vanilla GRPO across math and multi-turn search settings.The framework supports sharing and non-sharing configurations.
- Limitations: Dr. MAS does not resolve all instability sources, including credit assignment across agents and turns.
- Limitations: The framework has not been evaluated with much larger numbers of agents.Such settings may raise resource-allocation and asynchronous-execution issues.
A.1 Proof of Lemma 4.2
The lemma decomposes each agent’s gradient norm into a dominant reward-statistics scaling factor and a residual score–reward covariance term.
- Per-agent statistics: For agent k, µk and σk^2 are the mean and variance of rewards sampled when agent k is active.The active time steps are sampled uniformly from Yk.
- Gradient decomposition: The gradient expression separates a dominant scaling component from a residual score–reward covariance correction.The correction term is denoted ∆k.
- Gradient decomposition: In large-scale LLM training, low-dimensional final-quality rewards and local token-level policy stochasticity make their covariance an empirical consideration.
A.2 Proof of Proposition 4.3
The proposition explains gradient-norm inflation when an agent’s active reward statistics differ from global statistics, with either variance mismatch or mean misalignment increasing gradient fluctuations.
- Blow-up argument: The proof applies the gradient expression at successive training iterations and concludes divergence when the stated growth conditions hold.
- Gradient inflation: A global-normalized gradient for agent k is amplified when active-subset reward statistics differ from global reward statistics.
- Gradient inflation: The inflation factor contains variance mismatch and mean misalignment terms relative to the global variance.Variance mismatch reflects different reward spread; mean misalignment reflects a shifted active-agent reward mean.
- Gradient inflation: Either mismatch increases the gradient’s second moment proportionally and therefore produces larger gradient fluctuations.
B Experimental Details
The experiments use role-specialized two-agent math and three-agent search systems with controlled hyperparameters, configurable agent prompts, and Dr. MAS training infrastructure; Dr. MAS produces smoother gradients than GRPO.
- Math task: The math setup uses Solver and Verifier agents with up to two loops, 8192-token prompts, 4096-token responses, group size 8, and binary rewards.Each agent uses a 1 × 10^-6 actor learning rate.
- Search task: The search setup uses Verifier, Search, and Answer agents, up to four turns, group size 5, and a 1 × 10^-6 actor learning rate.
- Prompt construction: Multi-agent prompts combine an environment observation, accumulated team context, and a role-specific instruction.
- Training framework: The framework supports configurable agent–model assignments, optional LLM sharing, worker-group mapping, distributed rollouts, and per-worker-group updates.
- Stability results: Dr. MAS keeps math gradient norms smoother than GRPO and leads to steadier training-accuracy improvement.This comparison is reported for the two-agent Qwen3-4B non-sharing orchestration.
- Stability results: On multi-turn search, GRPO’s search-agent gradient norm exceeds 80 and reaches NaN, whereas Dr. MAS maintains relatively low norms across agents.
E.3 Training Curves of Ablation Study
Figure 8 compares training curves for alternative advantage-normalization variants in the ablation study. The global baseline is contrasted with increasingly agent-wise normalization choices.
- Figure 8 presents training curves for different advantage-normalization variants in the ablation study.
F.1 Multi-Turn Search Task: Hierarchical Coordination
The search-task materials trace a process of identifying Mark Dismore’s birthplace, locating Greenfield in Hancock County, and concluding that Hancock County answers the question.
- Mark Dismore’s birthplace is identified as Greenfield, Indiana.
- Greenfield is identified as the county seat of Hancock County, Indiana.
- The reasoning concludes that Mark Dismore’s birthplace is located in Hancock County.
F.2 Math Task: Iterative Coordination
The math-task materials describe an expected-region calculation for two fixed diameters and 25 additional random segments, while highlighting that infinite-line formulas do not directly apply to line segments.
- The setup contains 27 segments: two perpendicular diameters plus 25 additional segments joining perimeter points in different quadrants.
- Infinite-line region formulas are not directly applicable because line segments may fail to intersect and obey disk-boundary constraints.
- A separate solver conclusion reports 379 regions, but the accompanying critique identifies missing probability and expectation derivations.
- The expected-region approach requires calculating intersection probabilities under the quadrant constraints.
- One listed calculation totals 176 expected intersections and reports E[R] = 204.