Technical Summary: Learning to Orchestrate Agents under Uncertainty

Problem Definition

The paper addresses the challenge of adaptive orchestration of heterogeneous agents in environments characterized by uncertainty and evolving agent behavior. In real-world multi-agent systems, agents (e.g., specialized AI models or human experts) exhibit varying reliability, operational costs, and output qualities. While prior work in agent orchestration often focuses on expected utility or offline accuracy, it frequently fails to explicitly model uncertainty in agent reliability and output distributions at the orchestration level.

The core problem is formulated as a sequential decision-making process where a meta-controller must decide which agent to delegate a task to. This decision must balance:

1. Reward: The expected performance of the agent.
2. Cost: The operational cost of invoking the agent.
3. Uncertainty: The variability in agent outputs and the alignment between an agent's output distribution and the task's required distribution.

The authors note that classical frameworks like Dec-POMDPs often assume known system models and lack scalability, while standard Multi-Armed Bandit (MAB) approaches typically treat agents as interchangeable up to mean reward, ignoring query costs and distributional heterogeneity.

Methodology: BOT-Orch

The authors propose BOT-Orch, a lightweight framework that recasts the orchestration problem as a stochastic bandit problem regularized by Optimal Transport (OT) distances.

Key Components:

1. OT-Regularized Rewards:
   The framework defines a net utility for selecting an agent $i$ at time $t$ as:
   $$U_t(i) = R_t(i) - \lambda W_t(i)$$
   Where $R_t(i)$ is the reward and $W_t(i)$ is the alignment cost. The alignment cost is quantified using the Wasserstein distance ($W_c$) between the agent's output distribution ($\mu_i$) and a task-specific reference distribution ($\nu_t$). This allows the system to penalize agents whose output distributions are geometrically distant from the desired task outcome, even if their mean rewards are similar.

2. Survival-Based Rewards:
   To model reliability, latency, and time-to-success, the framework incorporates survival analysis with latent frailty variables. Agent performance is modeled via time-to-completion random variables with right-censoring, allowing the system to handle incomplete feedback and heterogeneous reliability profiles.

3. Algorithmic Implementation:

   • i.i.d. Setting: Uses a Boltzmann (softmax) policy that selects agents based on exponentially smoothed rewards penalized by OT misalignment.
   • Non-i.i.d. Setting: Extends the algorithm to handle history-dependent task distributions and non-stationarity by incorporating a correction term that encodes memory effects and temporal correlations.

Key Contributions

The paper makes three primary contributions:

1. Bandit-Based Orchestration with OT Alignment:
   The authors cast delegation as a stochastic bandit problem regularized by OT distances. This enables uncertainty-aware alignment, allowing the orchestrator to adaptively select agents based on the geometric discrepancy between their output distributions and task requirements, rather than just scalar mean rewards.

2. Theoretical Guarantees:
   Under standard assumptions (bounded costs, sub-exponential frailty, Lipschitz continuity), the authors prove:

   • Sublinear Regret: The cumulative OT-regularized regret is $O(\sqrt{T})$, ensuring asymptotic convergence to oracle performance.
   • Structural Optimality: The framework provably induces a preference ordering among agents with identical mean rewards but differing distributional alignment.
   • Robustness: The system is robust to noisy alignment costs (margin robustness under Gaussian noise).
   • Convergence: The orchestration weights converge almost surely to an invariant point of the mean-field ODE.

3. Empirical Validation:
   The framework is evaluated in both synthetic and semi-synthetic settings, including a human-AI triage scenario under deployment shift. Results demonstrate that BOT-Orch consistently outperforms standard bandit baselines (UCB1) and heuristic approaches (Random, No-OT) in both i.i.d. and non-i.i.d. environments.

Experimental Results

The evaluation covers synthetic environments with heterogeneous, non-i.i.d. agent behaviors (e.g., piecewise-stationary variance, sinusoidal drift, Brownian bridge processes) and a semi-synthetic human-AI triage task using the Breast Cancer Wisconsin dataset.

• Performance: BOT-Orch achieves the highest cumulative net utility and lowest oracle regret across all tested environments.
• Adaptation to Shift: In the non-i.i.d. setting (deployment shift), BOT-Orch successfully adapts by increasing the escalation rate to the human expert when distributional mismatch occurs, whereas baselines like No-OT fail to adjust effectively.
• Alignment Cost: The framework significantly reduces cumulative alignment costs (by an order of magnitude in some cases) compared to baselines that ignore distributional alignment.
• Ablation: The study confirms that the OT term is the primary driver of improvement; removing it (No-OT) leads to higher regret and alignment costs.

Significance and Claims

The paper claims that incorporating distributional structure directly into orchestration policies substantially improves robustness and adaptability in modern multi-agent AI systems.

• Theoretical Insight: It establishes that modeling agents via their output distributions and using Optimal Transport to measure alignment provides a principled way to handle uncertainty, offering theoretical guarantees of sublinear regret and convergence.
• Practical Utility: The framework is particularly relevant for emerging ecosystems of specialized models and tool-using agents where performance, risk, and cost must be jointly managed under uncertainty.
• Limitations: The authors modestly acknowledge limitations, including the computational expense of Wasserstein distances in high-dimensional settings, the need for calibration of the alignment parameter $\lambda$, and the requirement for access to suitable task-specific reference distributions.

The work suggests that moving beyond mean-reward optimization to distributional-aware orchestration is a critical step for scalable and robust deployment of heterogeneous agent teams in stochastic, partially observable environments.