🤖 AI

Learning to Assign Prediction Tasks to Agents with Capacity Constraints

This paper introduces a theoretical framework and sequential explore-exploit algorithms for dynamically assigning prediction tasks to human or AI agents with capacity constraints, demonstrating through experiments that these methods significantly outperform non-contextual baselines by learning to optimize agent expertise and task context.

Original authors: Shang Wu, Saatvik Kher, Padhraic Smyth

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shang Wu, Saatvik Kher, Padhraic Smyth

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 run a busy kitchen with a team of chefs. Some chefs are amazing at baking cakes but terrible at grilling steaks. Others are grill masters who burn every dessert they touch. You also have a rule: no single chef can do more than a certain percentage of the total work, or they'll get too tired and quit (or in the case of AI, get too expensive or overloaded).

This paper is about how to automatically figure out who is good at what and assign the right orders to the right chefs in real-time, while making sure no one gets overworked.

Here is the breakdown of the paper's ideas using simple analogies:

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

In many real-world situations (like medical diagnoses or answering customer questions), you have a mix of workers: some are humans, some are AI models.

The Old Way: You might just flip a coin to decide who gets the next task, or you might always send everything to the "best" chef you know.
The Flaw: If you always send everything to your best chef, they burn out. If you flip a coin, you might send a tricky steak order to the baker, and it gets ruined. Also, chefs often have "hidden talents" that only show up with specific types of food (context).

2. The Solution: A Smart Dispatcher

The authors created a system that acts like a smart dispatcher. It has two main jobs:

Learn on the Fly: It doesn't know exactly who is good at what at the start. As tasks come in, it watches the results. If "Chef A" nails a cake order but fails a steak order, the system learns: "Chef A is a baker, not a grill master."
Respect the Limits: It keeps a mental tally (a "queue") for each chef. If Chef A has already done 40% of the work and their limit is 50%, the system starts saying, "Okay, we need to send the next few orders to Chef B, even if Chef A might be slightly better at them, just to keep the balance."

3. The "Shadow Price" Analogy

The paper uses a mathematical concept called a "shadow price," which is a bit like a dynamic toll booth.

Imagine every time you assign a task to a chef, you have to pay a "toll."
If a chef is underworked, the toll is low (free!).
If a chef is getting close to their limit, the toll goes up.
The system looks at the "toll" and the chef's skill. It might choose a slightly less skilled chef because their "toll" is zero, saving the highly skilled (but expensive/limited) chef for the really hard tasks.

4. The Experiments: Testing the System

The researchers tested this idea in three different "kitchens":

Medical Images (Camelyon17): They used AI models trained on data from different hospitals. One model was great at spotting tumors in Hospital A's images but bad at Hospital B's. The system learned to route Hospital A's images to Model A and Hospital B's to Model B, rather than sending them randomly.
Tabular Data (Bank, Credit, etc.): They simulated tasks where one model knew some facts and another knew different facts. The system learned to split the work based on the specific details of the task.
Text & Images (LLMs and Humans): They tested Large Language Models (like the ones you might chat with) and human annotators. They found that different AI models were better at different topics (e.g., one was great at Chemistry, another at Foreign Policy). The system routed the questions accordingly.

5. The Results: Why It Matters

The paper found that:

Context is King: A system that looks at the specific details of the task (the "context") and routes it to the right person performs much better than a system that just picks randomly or picks the "overall best" person.
The Whole is Greater than the Sum of Parts: By splitting the work based on who is good at what, the team's combined accuracy was often higher than the accuracy of the single best individual.
Constraints Help: Surprisingly, having a hard limit on how much work each person can do actually helped the system learn faster. It forced the system to try different people, preventing it from getting stuck thinking one person is the best at everything.

6. The "Batch" Twist

The paper also looked at what happens if you don't assign tasks one by one, but in small groups (like a batch of 10 orders at a time).

Analogy: It's like a chef prepping a whole tray of appetizers at once.
Result: Doing this in small batches sometimes works slightly better because you can balance the workload more evenly across the group, but it takes a tiny bit longer to get the results.

Summary

This paper proves that if you have a team of mixed experts (humans and AIs) with limited energy, you shouldn't just guess who does what. Instead, you should use a smart, learning system that:

Watches who succeeds at what type of task.
Keeps a strict count of who has done how much work.
Routes the next task to the person who is best suited for it right now, without letting anyone get overworked.

The result is a team that works faster, makes fewer mistakes, and keeps everyone happy.

Technical Summary: Learning to Assign Prediction Tasks to Agents with Capacity Constraints

Problem Definition

The paper addresses the problem of online sequential task assignment in environments where tasks arrive sequentially and must be allocated to a set of available agents (human or AI). The core challenge involves two simultaneous constraints:

Heterogeneous, Context-Dependent Expertise: Agents possess varying levels of expertise that depend on the specific context of the task (e.g., a radiologist may excel at specific medical cases but not others; an LLM may perform better on specific query types). This expertise is initially unknown and must be learned online.
Capacity Constraints: Unlike cost-budget constraints where a single cheap agent could theoretically handle all tasks, this setting imposes long-run capacity limits on individual agents. Each agent $a$ has a target fraction $\alpha_a$ of the total workload they can handle. This reflects real-world operational limits such as human fatigue, burnout, or staffing levels, preventing the over-allocation of tasks to a single high-performing agent.

The objective is to learn an assignment policy that maximizes the overall expected reward (e.g., prediction accuracy) while strictly adhering to these long-run capacity constraints, despite the initial uncertainty regarding agent performance.

Methodology

Theoretical Framework

The authors model the problem as a contextual multi-armed bandit (MAB).

Optimal Unconstrained Policy: In an oracle setting where agent expertise $\mu_a(x)$ is known, the optimal policy assigns each context $x$ to the agent with the highest conditional expected reward.
Optimal Constrained Policy: When capacity constraints $\alpha_a$ are introduced, the authors derive the optimal policy using Lagrange multipliers. They prove that the optimal constrained policy takes the form:
$\phi^*(x) = \arg\max_{a} \{ \mu_a(x) - \lambda_a \}$
where $\lambda_a$ acts as a "shadow price" for the capacity constraint. This implies that the optimal strategy is to assign tasks to the agent with the highest adjusted reward, effectively penalizing agents that are approaching their capacity limits.
Two-Agent Analysis: The paper provides specific characterizations for the two-agent case, showing that the optimal policy assigns the agent with capacity $\alpha$ to the $\alpha$ -mass of contexts where the reward gap $\Delta(x) = \mu_1(x) - \mu_2(x)$ is largest. Theoretical analysis demonstrates that contextual policies strictly outperform non-contextual (random or fixed-agent) baselines whenever agent performance varies with context.

Online Learning Algorithm

To handle the practical scenario where $\mu_a(x)$ is unknown, the authors propose a queue-based online learning framework:

Context Modeling: Agent expertise is modeled using either logistic regression (with Laplace approximation for posterior estimation) or tree-based models (Random Forests). The framework supports both Thompson Sampling (for exploration via posterior sampling) and Greedy selection (using point estimates).
Queue-Based Constraint Enforcement: To enforce long-run capacity constraints without knowing the future, the algorithm maintains a virtual queue $Q_{t,a}$ for each agent. The queue tracks the cumulative deviation between realized assignments and the target capacity $\alpha_a$ :
$Q_{t+1,a} = [Q_{t,a} + \mathbb{I}(a_t = a) - \alpha_a]_+$
Decision Rule: At each step $t$ , the agent is selected by maximizing an adjusted score that balances estimated reward and queue pressure:
$a_t = \arg\max_{a} \{ \mu_{a,t}(x_t) - \eta Q_{t,a} \}$
Here, $\eta$ is a penalty parameter controlling the trade-off between reward maximization and constraint satisfaction. The queue length $Q_{t,a}$ serves as a dynamic, time-varying estimate of the shadow price $\lambda_a$ .

The authors prove that this algorithm inherits standard sublinear regret bounds from contextual Thompson Sampling literature, even with the added complexity of virtual queues.

Experimental Results

The framework was evaluated across diverse settings involving tabular, image, and text prediction tasks, using machine learning classifiers, Large Language Models (LLMs), and human annotators (both real and simulated).

Datasets: Experiments included Camelyon17 (medical imaging), MMLU (LLM routing), ImageNet16H (human vs. AI), and four tabular datasets (Bank, Credit, Coupon, Cardio).
Performance: Contextual assignment policies (both logistic and tree-based, using Thompson Sampling or Greedy strategies) systematically outperformed non-contextual baselines across all capacity levels.
- The combined policies often achieved lower error rates than the marginal error rate of the best individual agent, validating the theoretical claim that leveraging complementary expertise improves overall system performance.
- Tree-based models generally outperformed logistic models, suggesting the importance of capturing non-linear patterns in agent expertise.
- Greedy and Thompson Sampling performed similarly, with the constraints themselves providing sufficient implicit exploration.
Mini-Batch Setting: The authors explored a mini-batch setting where assignments are made for a batch of tasks before observing outcomes. They found a U-shaped error curve relative to batch size: small batches allow frequent updates but noisy capacity enforcement, while large batches delay learning. Intermediate batch sizes offered the best trade-off, though at the cost of increased latency.
Sensitivity: Sensitivity analysis on the queue penalty parameter $\eta$ showed that smaller values prioritize reward maximization (approaching unconstrained performance), while larger values enforce capacity more aggressively (approaching non-contextual random assignment).

Key Contributions

Theoretical Characterization: The paper provides a general theoretical framework for task assignment with capacity-constrained agents, characterizing optimal policies in terms of shadow prices and context-dependent expertise differences.
Algorithmic Framework: It introduces a practical online learning algorithm that jointly learns agent expertise and enforces long-run capacity constraints using a queue-based mechanism, applicable to both human and AI agents.
Empirical Validation: Through extensive experiments across multiple modalities (text, image, tabular) and agent types (LLMs, humans, classifiers), the paper demonstrates that contextual capacity-constrained policies yield systematic gains over non-contextual baselines.

Significance and Claims

The paper claims that effective human-AI collaboration and large-scale organizational task allocation depend not only on improving individual agent performance but also on designing allocation policies that learn heterogeneous expertise while respecting real-world capacity limits.

The authors emphasize that their approach addresses a critical gap in prior work (such as learning to defer or LLM routing), which often assumes offline learning or ignores agent-specific workload limits. By modeling the joint problem of learning expertise and managing capacity online, the framework enables systems to adapt to distribution shifts and operational constraints dynamically. The results suggest that ignoring context or capacity constraints leads to suboptimal performance, whereas jointly modeling them allows for the exploitation of agent complementarity, potentially outperforming even the best single agent in the system.

The paper remains modest regarding its scope, noting that evaluations were conducted using proxy agents in controlled settings with relatively low-dimensional contexts. It acknowledges that real-world applications may involve high-dimensional, multimodal contexts, delayed feedback, and evolving agent behaviors, identifying these as directions for future work.