💻 computer science

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

LEMON is a novel LLM-based orchestrator that utilizes counterfactual reinforcement learning to generate executable multi-agent specifications by integrating roles, duties, capacities, and dependencies, achieving state-of-the-art performance across diverse reasoning and coding benchmarks.

Original authors: Xudong Chen, Yixin Liu, Hua Wei, Kaize Ding

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xudong Chen, Yixin Liu, Hua Wei, Kaize Ding

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a team of AI assistants, each with different skills and levels of "brainpower." Some are fast but simple (like a junior intern), while others are slow but brilliant (like a senior professor). The big challenge isn't just having these assistants; it's figuring out how to organize them to solve a specific problem.

If you ask a simple question, you don't need a whole committee. If you ask a complex math problem, you need a specific chain of experts. If you ask a coding question, you need a different team structure entirely.

This paper introduces LEMON, a smart system that acts as a dynamic team manager. Instead of using a fixed team for every job, LEMON learns to build the perfect team structure from scratch for every single task.

Here is how it works, broken down into simple concepts:

1. The Problem: The "One-Size-Fits-All" Trap

Most current systems use a fixed team. Imagine a construction crew that always uses the same blueprint, whether they are building a doghouse or a skyscraper.

The Issue: Sometimes the blueprint is too complicated for a simple job (wasting time and money). Sometimes it's too weak for a hard job (failing the task).
The Old Way: Researchers tried to fix this by tweaking one part at a time—maybe just changing who talks to whom, or just changing who does the work. But the paper argues that you can't fix the team by tweaking just one piece; you have to design the whole team, the roles, and the workflow together as a single package.

2. The Solution: LEMON (The Master Architect)

LEMON is an AI that acts as an Architect. When you give it a task (like "Solve this math problem" or "Write this code"), it doesn't just answer it. Instead, it writes a blueprint (called an "executable orchestration specification").

This blueprint tells the system:

Who to hire: Which specific AI agents to use.
What their job is: A custom description of their duty (e.g., "Check the units," not just "Solver").
How smart they need to be: Assigning a "capacity level" (Small, Medium, or Large brainpower) to each agent.
The flow: Who needs to talk to whom and in what order.

Think of it like a conductor writing a specific score for an orchestra. For a simple song, they might only need a flute and a drum. For a complex symphony, they need the full string section. LEMON writes the score specifically for the song you are playing.

3. The Secret Sauce: "What If?" Training

How does LEMON learn to write these perfect blueprints? It uses a clever training method called Counterfactual Reinforcement Learning.

Imagine you are a teacher grading a student's essay.

The Old Way (Sparse Feedback): You read the whole essay, give it a grade of "B," and say, "Good job, but try again." The student doesn't know which sentence was bad or which paragraph was great. They just know the whole thing got a B.
LEMON's Way (Localized Feedback): LEMON looks at the blueprint it just wrote. Then, it asks a "What if?" question:
- "What if I made this specific agent 'Small' instead of 'Large'? Would the result be worse?"
- "What if I removed this connection between two agents? Did it break the flow?"

It runs these "What if" scenarios instantly.

If changing a specific part makes the result worse, LEMON learns: "Ah, that specific part was crucial!"
If changing a part makes no difference, LEMON learns: "That part was unnecessary; I can save resources next time."

This allows LEMON to learn exactly which decisions in the blueprint matter, rather than just guessing based on the final grade.

4. The Results: Smarter and Cheaper

The paper tested LEMON on six different types of challenges, including hard math problems, logic puzzles, and coding tasks.

Better Performance: LEMON solved more problems correctly than other methods that use fixed teams or single AI agents.
More Efficient: Because LEMON knows exactly how much "brainpower" is needed for each step, it doesn't waste money using a super-computer for a simple task. It uses the right-sized tool for the job.
The Analogy: If other methods are like ordering a massive banquet for a single sandwich, LEMON is like a chef who cooks exactly the right meal for the number of guests.

Summary

LEMON is a system that learns to design its own team structure for every new problem. Instead of using a rigid, pre-made team, it builds a custom workflow, assigns the right level of intelligence to each step, and learns from "What if?" experiments to ensure every part of the plan is necessary and effective. The result is a system that is both smarter at solving problems and cheaper to run.

Technical Summary: LEMON

Problem Definition

Large Language Model (LLM)-based multi-agent systems (MAS) rely heavily on orchestration design—specifically the definition of agent roles, capacity assignments, and dependency structures—to determine solution quality and execution efficiency. Existing approaches face two primary limitations:

Decomposed Design: Current methods often optimize orchestration components (topology, workflow, routing) separately or sequentially. This fails to treat role design, capacity assignment, and dependency construction as mutually dependent decisions within a single executable specification, potentially leading to suboptimal or inefficient systems.
Sparse Credit Assignment: Training signals for orchestrators are typically attached to end-to-end execution outcomes (e.g., final task success). This "sparse" feedback makes it difficult to determine which specific local decisions (e.g., a specific role definition or a single dependency link) contributed to success or failure, leading to conflation of necessary and incidental choices during policy updates.

Methodology: LEMON

The authors propose LEMON (Learning Executable Multi-agent OrchestratioN via Counterfactual Reinforcement Learning), a framework that treats MAS orchestration as the generation of a single, executable specification.

1. Compositional Executable Orchestration Generation

Instead of optimizing isolated components, LEMON's policy ( $\pi_\theta$ ) generates a structured YAML specification for each input task. This specification simultaneously defines:

Agents: The set of workers to instantiate.
Roles: Task-specific duties, inherited base roles, and customized responsibilities.
Capacity: The specific LLM model size (e.g., Small, Medium, Large) assigned to each agent.
Dependencies: The information flow structure (edges) connecting agents.

The specification is validated against a schema, compiled into a directed acyclic graph (DAG), and executed by worker agents.

2. Training Pipeline

LEMON employs a two-stage training process:

Supervised Fine-Tuning (SFT): The policy is initialized using teacher-generated executable YAML specifications. This ensures the model can produce valid, schema-compliant outputs before reinforcement learning begins, reducing the frequency of invalid rollouts.
Reinforcement Learning (RL): The policy is optimized using a combination of global and localized objectives:
- Global Orchestration-Level Reward (GRPO): The system uses Group Relative Policy Optimization (GRPO) based on execution utility. The reward ( $R_{orch}$ ) combines task success, token efficiency (worker token usage), and structural cost (agent/dependency count).
- Localized Counterfactual Credit Assignment: To address sparse credit assignment, LEMON introduces a counterfactual mechanism. For a valid generated specification, the system applies a "mutation" to a single field:
  - Dependency Deletion: Removing a reference to test necessity.
  - Role Rollback: Reverting a customized duty to a base role.
  - Capacity Downgrade: Reducing the model size assigned to an agent.
  The mutated (counterfactual) specification is executed, and the reward difference ( $\Delta_c$ ) is calculated. This reward contrast is applied only to the edited token spans in the original specification. This allows the policy to learn which specific local decisions drove the performance change, rather than reinforcing the entire sequence indiscriminately.

Key Contributions

Compositional Orchestration Generation: The paper formulates MAS orchestration as the generation of a unified executable specification that composes role design, capacity profiles, and dependency structure, rather than optimizing these decisions in isolation.
Localized Counterfactual Reinforcement Learning: The authors introduce a method to convert execution-level reward differences into span-level supervision. By editing specific orchestration fields and comparing rewards, the system assigns credit to the specific decisions responsible for performance variations, mitigating the sparse credit assignment problem.
Empirical Validation: Extensive experiments demonstrate that LEMON outperforms single-agent prompting, fixed-structure MAS, topology-design methods, and adaptive workflow baselines across reasoning and coding benchmarks.

Experimental Results

The authors evaluated LEMON on six benchmarks: MMLU, GSM8K, AQuA, MultiArith, SVAMP, and HumanEval.

Performance: LEMON achieved a state-of-the-art average score of 90.72%, outperforming the strongest single-agent baseline (SC-CoT, 86.76%), the strongest fixed-structure MAS (Tree-MAS, 88.53%), and the strongest topology-design baseline (AgentPrune, 88.13%). It achieved the best results on five of the six benchmarks.
Token Efficiency: LEMON demonstrated superior accuracy-to-token trade-offs. On benchmarks like MMLU and AQuA, it achieved high accuracy with significantly fewer worker tokens compared to high-cost topology baselines, often lying on or near the Pareto frontier.
Ablation Studies:
- Removing the localized counterfactual objective (Global RL Only) resulted in lower accuracy and higher token usage.
- Removing specific generation components (Role, Capacity, or Structure) consistently degraded performance, confirming the necessity of joint optimization.
- Removing the SFT warm start caused a drastic performance drop (e.g., MMLU accuracy fell from 84.32% to 48.53%), highlighting the importance of initialization for valid exploration.
Case Studies: Qualitative analysis showed that LEMON generates task-aligned structures (e.g., specific reasoning steps for math problems) rather than generic, repeated roles often found in other methods.

Significance and Claims

The paper claims that LEMON advances the field by shifting the optimization focus from isolated topology or workflow decisions to the executable specification itself. By jointly generating roles, capacities, and dependencies, and by utilizing localized counterfactual signals to refine these decisions, LEMON produces multi-agent systems that are both more effective (higher accuracy) and more efficient (lower token cost) than existing adaptive orchestration methods. The authors emphasize that this approach provides a more systematic way to construct task-adaptive orchestration specifications, making multi-agent systems easier to study and deploy under controlled resource budgets.

The authors note limitations, including the non-trivial training cost associated with counterfactual evaluation (though inference remains efficient) and the current focus on text-based reasoning and coding benchmarks, leaving broader scenarios like long-horizon tool use for future work.