This is Part 2 of a series on engineering AI systems that learn in production. Part 1 introduced the complexity classifier and three reasoning tiers. This part covers the reinforcement learning layer that replaces static thresholds with a contextual bandit that learns optimal (model, strategy) pairs from production outcomes.
The map is not the terrain
Part 1 gave us a complexity classifier. It scores each request across five dimensions, maps that score to a reasoning tier, and dispatches accordingly. For a static system, it works reasonably well.
But here's the thing about heuristic classifiers — they encode your assumptions at design time and then never update them. Those weights (0.30 for task type, 0.20 for ambiguity, and so on) were chosen based on intuition and early testing. They don't adapt when the system encounters patterns that violate those assumptions.
Consider a technical support agent handling database troubleshooting. The classifier might score most requests as moderate complexity (DRAFT tier). But maybe for this specific agent dealing with this specific domain, DRAFT consistently produces responses that require operator correction. The classifier doesn't know that. It keeps routing to DRAFT because the input signals look moderate. The problem isn't the classification of the input. The problem is that input classification doesn't predict which strategy will produce a good outcome for this particular context.
Think about the difference between a map and a navigator. The classifier is the map: static, useful for orientation, gives you a reasonable default route. What we need is a navigator that drives the actual roads, learns where the potholes are, and adjusts the route based on experience.
That navigator is a contextual bandit.
Thompson Sampling in plain language
Before the math, the intuition.
You moved to a new city and you're trying to find the best lunch spot near your office. Nine restaurants to choose from (three cuisines times three price points, say). You don't know which one you'll like best. How do you figure it out?
One approach: always go to whatever place seems best based on what you know so far. Problem is you'll settle on the first decent restaurant and never discover the amazing one two blocks further. You exploit what you know without exploring what you don't.
Another approach: go to a random restaurant every day. Problem is you'll eat bad meals indefinitely because you never concentrate on the places that have proven good.
Thompson Sampling does something more natural. You maintain a mental model of how good you think each restaurant is, but with uncertainty. For the place you've been to ten times and loved, you're pretty confident it's great. For the place you tried once and it was mediocre — you think it's probably mediocre but you're not sure. Maybe the chef was having a bad day.
Each lunchtime, you mentally "sample" from your uncertainty about each restaurant. Sometimes that sample for the once-visited place comes out surprisingly high (your uncertainty is wide, so extreme values are possible). When that happens, you go try it again. Most of the time, your confident favorites produce the highest samples, so you go there. But occasionally, your uncertainty about less-explored options produces a high enough sample to trigger exploration.
Over time, you visit everything enough to form confident beliefs. Restaurants you keep going back to are genuinely the best ones. You explored efficiently because you explored in proportion to your uncertainty, not randomly.
That's Thompson Sampling. Now replace "restaurants" with "(model, strategy) pairs" and "how much you liked the meal" with "how good the AI response was, accounting for cost and latency."
The action space
The bandit selects from the Cartesian product of available models and reasoning strategies.
Say the system has access to three models (gpt-4o, o3, deepseek-r1) and three strategies (NONE, DRAFT, FULL_COT). That's nine possible arms:
| Arm | Model | Strategy |
|---|---|---|
| 1 | gpt-4o | NONE |
| 2 | gpt-4o | DRAFT |
| 3 | gpt-4o | FULL_COT |
| 4 | o3 | NONE |
| 5 | o3 | DRAFT |
| 6 | o3 | FULL_COT |
| 7 | deepseek-r1 | NONE |
| 8 | deepseek-r1 | DRAFT |
| 9 | deepseek-r1 | FULL_COT |
Not all combinations make equal sense. You wouldn't typically run FULL_COT on a model without native thinking-token support. But the bandit can learn that. If arm 4 (o3 with no reasoning) consistently underperforms arm 6 (o3 with full reasoning), the posterior for arm 4 will shrink and it'll rarely get selected. You don't need to manually prune bad combinations — the learning handles it.
Per-agent policies define which models and strategies are available. An agent optimized for cost might only have gpt-4o in its allowlist. An agent handling high-stakes analytical work might have all three models available. The bandit operates within whatever action space the policy permits.
Context: the bandit isn't blind
A plain multi-armed bandit treats every request the same. Pull the arm with the best track record, regardless of what the request looks like.
That's too coarse. The optimal (model, strategy) pair for "summarize this email" is probably different from the optimal pair for "diagnose why our deployment is failing under load." We need the bandit to condition its selection on the request context.
The system encodes each request into a context bucket — a composite key built from:
- Task complexity (from the classifier): low / medium / high
- Intent label: the detected user intent (scheduling, billing, analysis, debugging, etc.)
- Conversation length bucket: short (1-3 turns), medium (4-10), long (11+)
- Tenant plan tier: free, professional, enterprise
- Time bucket: business hours, off-hours, weekend
The full context key looks something like: high | debugging | medium_conv | enterprise | business_hours
Each unique context bucket gets its own set of Beta distributions, one per arm. So the system maintains separate beliefs about "how good is (gpt-4o, DRAFT) for high-complexity debugging requests from enterprise tenants during business hours" versus "how good is (gpt-4o, DRAFT) for low-complexity scheduling requests from free-tier tenants on weekends."
Why buckets instead of continuous features? Three reasons.
First, manageability: maintaining a Beta distribution per arm per context bucket works. Second, interpretability: an operator can look at the bucket high | analysis | long_conv | enterprise | business_hours and understand exactly what context it represents. Third, cold-start sharing: a new agent handling similar request types can inherit priors from existing agents with the same context patterns, instead of starting from zero.
The selection algorithm
Here's the core of Thompson Sampling with the Beta-Bernoulli model.
Each arm a_i in context bucket b has parameters (α_i, β_i) representing the system's belief about that arm's reward probability. Both start at (1, 1), which is a uniform prior. ("I know nothing. This arm could be anywhere from terrible to excellent.")
When a request arrives with context bucket b:
For each arm a_i in the agent's allowlist:
Sample θ_i ~ Beta(α_i^b, β_i^b)
Select a* = argmax_i(θ_i)
Execute the request using (model, strategy) = a*
That's it. Sample from your current beliefs, pick whichever sample came out highest, use that arm. Arms you're uncertain about (low α + β, meaning wide distributions) will occasionally produce high samples and get explored. Arms you're confident about (high α + β, narrow distributions) will consistently produce samples near their true mean, so good arms get exploited.
The warm-up period. When the system hasn't seen many requests for a given context, every arm has wide distributions and selection is nearly random. We define a warm-up threshold: if the selected arm has fewer than 50 observations in this context, the selection gets flagged as low-confidence. During warm-up (fewer than 50 total observations per arm in the context), exploration rate starts at 30%. With exploration active, the system bypasses Thompson Sampling entirely and selects a random arm uniformly. This guarantees that even arms the sampling might overlook get a fair number of trials early on.
Exploration decay. That 30% exploration rate decays by 1% per week, settling at a 2% floor. Even in steady state, the system occasionally tries a random arm. This handles non-stationarity: if a model provider improves their API quality, or if token costs change, the system can discover that a previously inferior arm has become competitive.
The reward signal
After every request, the system observes three outcomes:
- Quality: how good was the response? (Evaluated by automated quality metrics, sentence coherence scoring, factual grounding checks, or downstream task completion.)
- Cost: how many tokens were consumed? (Translated to monetary cost using the model's pricing.)
- Latency: how long did the user wait?
These three measurements get normalized and combined into a single scalar reward:
R = w_q · Q_norm + w_c · (1 - C_norm) + w_l · (1 - L_norm)
The default weights: w_q = 0.5, w_c = 0.3, w_l = 0.2. Quality matters most, cost matters second, latency matters third. These weights are configurable per agent, so cost-sensitive deployments can raise w_c and latency-sensitive ones can raise w_l.
Normalization is relative. Q_norm, C_norm, and L_norm are each scaled to [0, 1] based on the observed min/max across all arms in the model pool over a trailing 7-day window. If the cheapest model costs $0.001 per request and the most expensive costs $0.05, then a request costing $0.025 gets C_norm = 0.49. Cost and latency are inverted (lower is better), which is why they appear as (1 - C_norm) and (1 - L_norm).
Edge case: when min equals max (no variance observed, maybe the system just started), the normalized value defaults to 0.5. System assumes middling performance until it has enough data to differentiate.
A concrete example. An enterprise tenant asks a debugging question. The bandit selects (o3, FULL_COT). Response takes 3.2 seconds, costs $0.038, and scores 0.92 on quality evaluation. Over the trailing 7 days, quality ranges [0.6, 0.95], cost ranges [$0.001, $0.05], latency ranges [0.2s, 5.0s].
Q_norm = (0.92 - 0.6) / (0.95 - 0.6) = 0.914C_norm = (0.038 - 0.001) / (0.05 - 0.001) = 0.755-
L_norm = (3.2 - 0.2) / (5.0 - 0.2) = 0.625
R = 0.5(0.914) + 0.3(1 - 0.755) + 0.2(1 - 0.625)
= 0.457 + 0.074 + 0.075
= 0.606
A reward of 0.606. Decent but not outstanding. Quality was high, but cost and latency pulled it down. If a cheaper arm could have produced similar quality, the bandit will learn to prefer it.
Updating beliefs
After computing the reward, the system updates the Beta distribution for the arm that was selected, in the context bucket where it was selected:
α_new = α_old + R
β_new = β_old + (1 - R)
With our example reward of 0.606: if the arm previously had α = 12.4, β = 8.2, it becomes α = 13.006, β = 8.594. Distribution shifts slightly toward higher expected reward. Over hundreds of observations, the distribution narrows and the expected value converges on the arm's true performance in that context.
Why does this work? Beta distribution's mean is α / (α + β). Adding R to α and (1 - R) to β pushes the mean toward R while gradually reducing variance (as α + β grows, the distribution gets tighter). High rewards increase α faster, shifting the mean upward. Low rewards increase β faster, shifting the mean downward. Update is proportional to the reward magnitude, so a reward of 0.9 shifts the distribution more than a reward of 0.55.
Batch processing. Rewards don't update the posterior immediately after each request. They accumulate in a Redis queue and get processed every 5 minutes. This is a pragmatic choice: it reduces database write pressure, allows the system to handle bursts without locking contention, and makes the update process idempotent (if a batch fails, it retries the same set of rewards).
Persistence. Every 15 minutes, current state of all Beta distributions gets persisted to PostgreSQL with a SHA-256 checksum. On startup, the system loads the latest valid checkpoint. If the checksum doesn't match (corrupted state from a partial write or a weird failure mode), the system falls back to Beta(1, 1) priors (uniform, "I know nothing") and rebuilds from the last 7 days of routing decisions stored in the audit log.
Worst-case data loss is 7 days of learning, not a catastrophic reset to zero knowledge. And because the system maintains the raw routing decisions alongside the computed posteriors, it can always reconstruct its beliefs from history.
What makes this work in practice
A few design decisions that matter for production but don't show up in the math:
Per-agent isolation. Each agent has its own set of posteriors. The billing agent's beliefs about (gpt-4o, DRAFT) are completely separate from the analytics agent's beliefs about the same arm. What works for billing queries doesn't necessarily transfer to analytical queries, even in the same context bucket.
Policy constraints. Bandit can only select from the agent's configured allowlist. An enterprise-tier agent might have all nine arms available. A cost-constrained free-tier agent might only have three (gpt-4o with each strategy). The bandit optimizes within whatever boundaries the policy sets.
The classifier as a prior. The complexity classifier from Part 1 doesn't disappear. It provides the initial routing decision for new contexts where the bandit has insufficient data (the warm-up period). As the bandit accumulates observations, its learned posteriors gradually override the classifier's heuristic. In contexts where the bandit has 200+ observations per arm, the classifier's initial score barely influences the final selection.
Non-stationarity handling. Model providers change their APIs. Pricing shifts. Models get updated. The 2% exploration floor ensures the system can detect these changes, but slowly. For abrupt changes (a model provider announces a price cut), operators can manually reset the posteriors for affected arms, triggering a new warm-up period that quickly re-learns the new performance landscape.
The diagram
Here's the full bandit decision loop:
A request arrives with context features (complexity, intent, conversation depth, tenant tier, time). Context encoder maps these to a bucket string, which determines which set of Beta distributions to load, one per arm. System checks the exploration rate; if exploring, it picks a random arm. If exploiting, it samples from each arm's Beta distribution and picks the arm whose sample is highest. Selected arm executes, system observes quality/cost/latency, computes a normalized reward, enqueues it to Redis. Every 5 minutes a batch process applies the Bayesian update to the posterior. Loop closes. Next time this context appears, the distributions reflect what was learned.
What's still missing
The reward signal so far is fully automated. Quality is measured by an eval function. Cost and latency are measured by instrumentation. System learns from these numbers.
But numbers miss something. An automated quality scorer might rate a response highly because it's coherent, grammatically correct, and addresses the topic. But the human reading it might find it unhelpful, off-target, or missing the point in a way that's hard for an automated metric to capture.
What happens when an operator reads the response and corrects it? What happens when a customer gives a thumbs-down? What happens when a human takes over the conversation entirely because the agent's response was so inadequate that no correction could salvage it?
These are human preference signals, and they carry information that automated metrics can't replicate. Part 3 covers how to collect them, how to convert them into reward adjustments, and how to feed them back into the bandit so that it learns not just from numbers but from human judgment.
We'll also cover what happens when things go wrong at the infrastructure level: model timeouts, rate limits, budget exhaustion. System needs to degrade gracefully, and that graceful degradation is itself a learned behavior.
Next in the series: Part 3 explores how human preference signals enrich the learning loop, and how the system survives production failure modes.
Top comments (0)