UnityMAS-O: A General RL Optimization Framework for LLM-Based Multi-Agent Systems¶

🕒 Published (v1): 2026-05-26 07:30 UTC · Source: Arxiv · link

Why this paper was selected

Unified RL interface optimizing all roles in LLM multi-agent systems end-to-end

Ask a follow-up

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

TL;DR¶

UnityMAS-O is a general RL optimization framework that extends the single-policy verl post-training stack to handle arbitrary LLM-based multi-agent workflows. It introduces four first-class abstractions—logical agent roles, workflow graphs, agent-to-model mappings, and role-specific reward functions—and backs them with a Ray star-topology runtime for distributed PPO-style training. Evaluated on retrieval-augmented QA (Natural Questions, HotpotQA) and reflective code generation, MARL training improves upon manually-specified baselines, with the largest gains on smaller models and strict all-passed code metrics.

Problem¶

Existing LLM multi-agent systems are manually orchestrated: workflows are fixed via prompts and routing rules, and RL post-training frameworks (TRL, OpenRLHF, verl) target single-policy optimization. There is no reusable substrate that jointly represents the workflow graph, role identities, agent-to-model parameter mappings, and multi-level reward attribution needed to train an entire multi-agent pipeline end-to-end via RL.

Method¶

UnityMAS-O represents a multi-agent workflow as a directed computation graph $G = (V, E)$ where each node $v \in V$ is a logical agent role and edges encode information/control flow. Four first-class objects are defined:

Logical agent roles $v_i = \langle P_{v_i}, X_{v_i}, Y_{v_i}, T_{v_i}, C_{v_i} \rangle$ (prompt template, I/O spaces, tools, constraints).
Agent–LLM mapping $\phi: V \to M$ decouples logical roles from physical model instances, supporting full sharing ($\forall v, \phi(v) = m_1$), full separation ($\phi(v_i) = m_i$), or partial sharing (subgroups share).
Graph-structured trajectory $\tau = (S, A, R, C)$ linearized as $s_0, k_1, a_1, s_1, \ldots, k_T, a_T, s_T$, where actions are sampled as $a_t \sim \pi_{\theta_{\phi(k_t)}}(\cdot \mid s_t, k_t)$.
Agent-specific rewards $r_v = R_v(s, y_v, s', \tau)$ combining node-level, turn-level, and trajectory-level signals with optional weights $r_v = \lambda^{\text{node}}_v r_v^{\text{node}} + \lambda^{\text{turn}}_v r_v^{\text{turn}} + \lambda^{\text{traj}}_v r_v^{\text{traj}}$.

The optimization objective is: $$\max_\Theta \mathbb{E}_{x \sim \mathcal{D},\, \tau \sim p_\Theta(\cdot | x, G, \phi)} \!\left[ \sum_{v \in V} \sum_{t \in I_v(\tau)} \gamma^{t-1} r_{v,t} \right]$$

At the system level, a central controller executes workflow scheduling, tool invocation, reward assembly, and trajectory recording. Per-model Ray worker groups $W = \{W_1, \ldots, W_K\}$ handle async inference, tensor buffering, advantage computation, and PPO updates. The controller routes role invocations to the correct worker group via $\phi$, and worker groups return only lightweight outputs ("thin data") to the controller while retaining full rollout tensors ("fat data") locally.

Key Contributions¶

Unified optimization interface that extends verl from single-policy post-training to full MARL over user-defined multi-agent workflow graphs.
First-class abstractions for logical agent roles, explicit agent-to-model mappings, graph-structured trajectories, and multi-level reward functions.
Native support for full/partial/independent parameter sharing without changing the logical workflow specification.
Ray + verl star-topology distributed runtime separating workflow-level control from model-local training.
Instantiations on three workflow families: retrieval-augmented QA, iterative agentic search, and reflective code generation; open-source release.

Results¶

Natural Questions / HotpotQA (QA/search workflows): MARL training improves all three search workflow variants relative to their manually-specified (pre-training) baselines; gains are "especially large for smaller models."
Reflective code generation (plan–code–verify–reflect, 3-round): MARL training improves strict all-passed metrics; the trained workflow also uses fewer verification turns on average (improved execution efficiency).
Parameter-shared training (all roles sharing one LLM) remains competitive with fully separated training across experiments.
Specific numbers are not reported in the provided text excerpt.

Limitations¶

No ablation isolating the contribution of each reward level (node/turn/trajectory) from the overall MARL gain is reported in the excerpt.
Evaluation is limited to two task families (QA/search, code); generalization to other workflow types (e.g., multi-agent planning, negotiation) is asserted but not demonstrated.
The star-topology runtime introduces a central controller bottleneck; scalability under high agent counts or very long trajectories is not characterized.
Credit assignment heuristics (the $\lambda$ weights) require task-specific tuning; no automated credit-assignment method is proposed.
Relies on verl internals; portability to other RL backends (TRL, OpenRLHF) is not addressed.

Relevance to Harnesses / Meta-Harnesses¶

UnityMAS-O is a training-time meta-harness: it wraps arbitrary, user-defined multi-agent workflows and exposes them as a unified RL optimization substrate, exactly the pattern that harness/meta-harness researchers study at the inference and orchestration level but now extended into the training loop. The decoupling of logical workflow graphs from physical model parameters and the pluggable reward interface mirror the separation-of-concerns principles central to meta-harness design (e.g., separating orchestration logic from execution backends). The framework's ability to swap parameter-sharing regimes without rewriting the workflow specification is directly analogous to how meta-harnesses swap underlying models or tools without changing the orchestration layer. Anyone building or studying harnesses that need to go beyond prompt-engineering into end-to-end trainability will find UnityMAS-O a concrete existence proof and reference architecture.