Skip to content

LEMON: Learning Executable Multi-Agent Orchestration via Counterfactual Reinforcement Learning

🕒 Published (v1): 2026-05-14 07:24 UTC · Source: Arxiv · Venue: NEURIPS 2026 · link

Why this paper was selected

NEURIPS; RL learns multi-agent orchestration topology outperforming hand-designed workflows

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LEMON is an LLM-based orchestrator that jointly generates a full executable multi-agent specification—roles, capacity levels, and dependency graph—as a single YAML artifact, trained with GRPO augmented by a localized counterfactual signal that assigns credit to individual orchestration fields rather than whole specifications. It achieves state-of-the-art average accuracy (90.72%) across six reasoning and coding benchmarks while sitting on or near the Pareto frontier for accuracy vs. worker-token cost.

Problem

Existing MAS orchestration methods optimize roles, topology, or capacity separately and sequentially, so inter-decision coupling is never captured. All current RL training signals are attached to end-to-end execution outcomes, making credit assignment too sparse to distinguish which role, capacity level, or dependency edge caused a success or failure.

Method

Compositional executable orchestration generation. Given task \(x\), policy \(\pi_\theta\) generates a YAML specification \(y\) encoding for each agent: identity/base-role/customized-duty triple \(R_v = (a_v, b_v, d_v)\), capacity level \(L_v \in \mathcal{L}\) (maps to a worker LLM tier), and dependency references that define edges \(E\). The YAML is validated, compiled into an orchestration graph \(G = \mathcal{C}(y)\), and executed in dependency order.

Two-stage training: 1. SFT warm-start. A teacher model generates valid YAML examples; filtered to parse-/schema-valid instances; trained with cross-entropy on the token sequence. 2. GRPO + localized counterfactual RL. The global orchestration reward is $\(R_\text{orch}(x, y_i) = R_\text{task} + \lambda_\text{tok} B_\text{tok} - \lambda_\text{graph} C_\text{graph}\)$ (task correctness + token efficiency bonus − structural cost penalty). GRPO uses group-relative advantage \(\hat{A}_i = (R_i - \bar{R}_x)/(\sigma_x + \epsilon)\) over \(B\) samples.

Localized counterfactual (CF) objective. For each valid specification, one of three mutation families \(\mathcal{M} = \{m_E, m_R, m_L\}\) is sampled (adaptive probability via running contrast estimates): \(m_E\) removes a dependency reference, \(m_R\) rolls back a customized duty to the base role, \(m_L\) downgrades one capacity level. The mutated spec \(\tilde{y}\) is executed to obtain \(\Delta_c = R_\text{orch}(x,y) - R_\text{orch}(x,\tilde{y})\). The CF objective scores only the edited token spans \((S, \tilde{S})\): $\(J_\text{cf}(\theta) = \mathbb{E}\!\left[w(\Delta_c)\log\sigma\!\left(\beta_\text{cf}\, b_c\left(s_\theta^\text{orig} - s_\theta^\text{cf}\right)\right)\right]\)$ where \(w(\Delta_c) = \min(|\Delta_c|, \delta_\text{max})/\delta_\text{max}\) downweights small/saturated contrasts and \(b_c = \text{sign}(\Delta_c)\). Tokens outside the edited span provide conditioning context but receive no gradient from \(J_\text{cf}\).

Key Contributions

  • Joint compositional generation of role specifications, capacity levels, and dependency structure as a single validated, directly executable YAML artifact.
  • Localized counterfactual RL: converts execution-level reward differences into span-level supervision, attributing credit to the specific orchestration field (edge, role, or capacity) whose edit caused the contrast.
  • Adaptive mutation sampling: per-family probability updated from running contrast magnitudes to focus exploration on the most informative mutation type.
  • Empirical demonstration across 6 benchmarks that joint generation + localized credit outperforms decomposed topology, workflow, and routing baselines on both accuracy and token efficiency.

Results

  • Average accuracy 90.72% across MMLU, GSM8K, AQuA, MultiArith, SVAMP, HumanEval.
  • Outperforms SC-CoT (single-agent) by +3.96 pp, best fixed-structure MAS (Tree-MAS) by +2.19 pp, AFlow (adaptive workflow) by +5.35 pp, best topology-design baseline (AgentPrune) by +2.59 pp.
  • Best on 5/6 benchmarks; second on MMLU (84.32 vs ARG-Designer 85.30, gap 0.98 pp).
  • Token efficiency: on MMLU reaches competitive accuracy at ~\(0.2\times10^4\) worker tokens vs ~\(0.3\times10^4\) (ARG-Designer) and >\(0.6\times10^4\) (several topology baselines).
  • Ablation: adding CF over Global RL Only improves MMLU/GSM8K/AQuA accuracy from 83.08/93.40/85.43 → 84.32/94.32/86.22 while reducing tokens from 2543/3743/3013 → 2148/3362/2056.
  • Removing SFT warm-start catastrophically degrades accuracy to 48.53/75.82/64.96.
  • Each CF mutation family (dependency, role, capacity) contributes independently: ablating any single family reduces accuracy and/or increases token cost.

Limitations

  • Counterfactual mutations require executing the mutated specification, doubling execution cost per training step for each mutation sample.
  • Only three mutation families are considered; compound edits or inter-field interactions are not explored.
  • Evaluated on reasoning/coding benchmarks; generalization to tool-use, open-ended generation, or long-horizon planning tasks is unverified.
  • The SFT initialization relies on a "stronger teacher model" generating valid YAML; quality and availability of this teacher is a dependency not fully analyzed.
  • MMLU evaluation uses only 1,531 validation examples (not the full test set), limiting direct comparison to other published results.

Relevance to Harnesses / Meta-Harnesses

LEMON is directly a meta-harness: rather than hand-coding a multi-agent pipeline, it learns to generate the harness specification (roles, capacities, wiring) on demand for each task input. The YAML orchestration spec is an executable meta-artifact—analogous to a harness config—that a downstream executor interprets without further human design. The counterfactual RL training mechanism is a principled approach to automatically optimizing harness structure, addressing the core challenge of credit assignment that any harness-auto-configuration system faces. For researchers building meta-harnesses, the localized CF signal offers a reusable training technique for attributing outcome quality to individual configuration decisions rather than to the harness as a whole.