Learning to Orchestrate Agents under Uncertainty

Abstract: Adaptive orchestration of heterogeneous agents requires making sequential delegation decisions under uncertain and evolving agent behaviour, e.g., coordinating specialised AI models with varying reliability, cost, and response quality. While prior work on agent orchestration focuses on performance or cost, uncertainty in agent reliability and output distributions is typically not modelled explicitly at the orchestration level. In this work, we study the problem of adaptive orchestration of heterogeneous agents under uncertainty, where a meta-controller must decide when to delegate to an agent, accounting for reliability, cost, and uncertainty. We propose BOT-Orch, a lightweight framework that recasts orchestration as a bandit problem over agents, regularized by OT distances between agent output distributions and task-specific reference distributions. We show that the regularised orchestration enjoys $\mathcal{O}(\sqrt{T})$ regret under standard assumptions, and provably induces preference ordering among agents with identical mean rewards but differing distributional alignment. Empirically, we demonstrate that BOT-Orch outperforms standard bandit and heuristic baselines in synthetic but adversarial task allocation settings with heterogeneous, non-i.i.d. agent behaviour.

• The paper proposes BOT-Orch, a framework that integrates optimal transport regularization with bandit optimization to achieve sublinear cumulative regret.
• It demonstrates robust agent selection by accounting for distributional alignment and uncertainty in both i.i.d. and non-i.i.d. task settings.
• Empirical evaluations, including synthetic and semi-synthetic tests, confirm BOT-Orch outperforms classical bandits under distribution shifts and heterogeneous agent performance.

Uncertainty-Aware Agent Orchestration via Optimal Transport-Regularized Bandits

Problem Formulation and Motivation

The paper "Learning to Orchestrate Agents under Uncertainty" (2605.27073) addresses adaptive orchestration in heterogeneous agent teams operating under uncertainty and temporal non-stationarity. Traditional orchestration frameworks focus primarily on maximizing mean reward or minimizing cost but ignore the statistical structure of agent outputs and their alignment with task-specific requirements. Furthermore, prior bandit-based delegation methods do not account for heterogeneity in agent reliability, output variance, or explicit query costs. The authors formalize a meta-control problem in which an orchestrator sequentially delegates tasks to agents while explicitly modeling both reward distributions and distributional alignment via optimal transport (OT) distances.

Framework and Theoretical Guarantees

The main contribution is BOT-Orch, a lightweight orchestration methodology that recasts delegation as a bandit optimization problem regularized by OT distances. The framework leverages agent output distributions $\mu_i$ and task-specific reference distributions $\nu_t$ to compute Wasserstein alignment costs, which are integrated into the bandit reward structure. This regularization induces agent selection based not only on mean-reward but on distributional compatibility, supporting robust delegation even when agents have identical average performance but differing output uncertainty or higher-order distributional features.

BOT-Orch operates under both i.i.d. and non-i.i.d. task settings, utilizing softmax or Boltzmann bandit policies informed by alignment-adjusted rewards (expected reward minus weighted OT cost). The non-i.i.d. variant incorporates history-dependent corrections to address temporal drift and non-stationarity.

The theoretical analysis establishes:

• Sublinear OT-regularized regret: BOT-Orch achieves cumulative regret of $O(\sqrt{T})$ against an oracle with full knowledge under standard bandit and survival-based reward assumptions.
• Preference ordering from OT regularization: Agents with identical mean reward but superior alignment (lower OT distance) are strictly preferred; the framework proves that alignment cost differentiates agents beyond traditional reward metrics.
• Robustness under noisy alignment: BOT-Orch's decision margin is tolerant to Gaussian noise in alignment costs, with specific probabilistic guarantees on correct agent ordering persistence.
• Consistency and convergence: Empirical reward estimates and orchestration weights converge almost surely via stochastic approximation dynamics, yielding stable long-term policies.

Synthetic Evaluation: Stationary and Non-Stationary Regimes

BOT-Orch is systematically evaluated on synthetic environments with both stationary and temporally evolving task distributions:

• IID data generation: Tasks are sampled i.i.d. from structured distributions (Fig. 2), with agent rewards exhibiting higher-order heterogeneity (variance, skewness, bimodality), isolating cases where mean reward alone is insufficient for robust orchestration.

  Figure 1: IID tasks and reward distributions illustrating stationary yet diverse agent output statistics.

• Non-IID scenarios: The experiments include piecewise-stationary rewards with changepoints (Fig. 3a-b), smooth mean drift (Fig. 3c), and Brownian-bridge latent processes (Fig. 3d), capturing continuous and abrupt environment variation.

  Figure 2: Non-IID reward evolution including variance shifts, smooth drift, and temporally correlated latent means.

BOT-Orch consistently delivers highest cumulative net utility and lowest oracle regret across all settings, especially in challenging non-stationary regimes. In these environments, BOT-Orch leverages OT alignment regularization to adapt to distributional shifts and agent heterogeneity, outperforming classical bandit (UCB1), uniform random selection, and ablations without OT. Notably, BOT-Orch maintains low regret and robust adaptation in the presence of abrupt and gradual reward distribution changes.

Semi-Synthetic Evaluation: Human-AI Triage under Distributional Shift

The framework is further validated in a clinical triage scenario using the Breast Cancer Wisconsin dataset, simulating deployment shift by perturbing test data distributions. Two agents—a calibrated AI classifier and a simulated human clinician—are orchestrated to maximize diagnostic correctness and minimize alignment cost. The AI performs better in-distribution, while the human outperforms in shifted population, creating complementary strengths.

BOT-Orch achieves top net utility and regret minimization with highly targeted escalation: it routes more shifted patients to the human when the AI's alignment cost increases, adapting to distributional shift without explicit shift-awareness (Fig. 1). Notably, BOT-Orch outperforms all baselines in both IID and non-IID deployments, reducing cumulative alignment cost by an order-of-magnitude relative to conventional policies. The alignment penalty parameter $\lambda$ is shown via ablation to control the exploitation/alignment trade-off, with optimal settings reliably maximizing utility in both regimes.

Figure 3: BOT-Orch learning performance, showing cumulative net utility, oracle regret, and adaptive escalation rates in semi-synthetic triage.

Escalation rate analysis (Fig. 4) confirms BOT-Orch's targeted routing: it increases referrals to the human on shifted cases with low variance, while non-OT bandit variants exhibit unstable or degenerate routing under cold-start. Longitudinal selection trajectories (Fig. 5) demonstrate BOT-Orch's precise adaptation post-distributional shift, while classical bandits and random policies lack this targeting.

Figure 4: Escalation rates for in-distribution (blue) and shifted (red) cases, highlighting BOT-Orch's targeted escalation and low variance.

Figure 5: Agent selection probabilities over time, with BOT-Orch adapting following shift onset to optimally escalate.

Further sensitivity analysis (Fig. 6, Fig. 7) reveals a phase transition: small $\lambda$ values can inhibit adaptation under abrupt shift, while moderate to large $\lambda$ robustly enforce alignment-driven orchestration. Excessively large $\lambda$ induces degeneracy where alignment dominates reward history.

Figure 6: $\lambda$ sensitivity curves for BOT-Orch, demonstrating monotonic improvement up to optimal joint alignment-utility trade-off.

Figure 7: Escalation rate on shifted patients and cumulative alignment cost as functions of $\lambda$, confirming monotonic alignment gains.

Practical and Theoretical Implications

BOT-Orch advances distributional-awareness in agent orchestration, providing a rigorous mechanism to differentiate agents by uncertainty structure and output compatibility, not just mean performance. Its analytic regret bounds and empirical dominance in heterogeneous, non-stationary conditions position it as a robust orchestration baseline for AI agent ecosystems, complex decision networks, and deployment scenarios where reward and utility structure are insufficient.

From a theoretical perspective, the integration of OT into bandit policy design yields entropy-regularized decision flows and margin-based preference ordering, extending stochastic approximation theory to geometry-aware reward-coupled learning. The use of OT barycenters and adaptive reference measures enables seamless adaptation to evolving distributions, a feature critical for modern AI systems exposed to continual drift and shift.

Future Directions

Scalability remains a challenge due to the computational cost of Wasserstein distances in high-dimensional spaces. Practical deployments may require approximate OT solvers or dimensionality-reduction strategies. Calibration of the alignment penalty $\lambda$ is environment-sensitive, motivating meta-learning or adaptive tuning techniques. Finally, extending BOT-Orch to multi-agent coordination, combinatorial orchestration, and richer reward structures—possibly with hierarchical or compositional tasks—is a promising direction.

Conclusion

BOT-Orch introduces a principled, uncertainty-aware orchestration methodology, combining bandit learning and optimal transport regularization to achieve robust, adaptive agent selection in heterogeneous and non-stationary environments. The approach is theoretically grounded, empirically validated, and offers substantial performance advantages in scenarios where reward structure alone fails to inform optimal delegation. As orchestration becomes central to large-scale AI and human-AI ecosystems, distributionally-aware selection mechanisms such as BOT-Orch will be critical for ensuring system robustness, adaptability, and alignment.