Iteration Scheduler
A research loop without a scheduler is a queue with delusions. The scheduler decides what the loop stops exploring, and that decision is the whole game.
Learning Objectives
- Model a research workflow as a hypothesis queue feeding parallel experiment slots whose results fan back in.
- Run multiple experiments concurrently with asyncio so the scheduler can keep all slots busy.
- Score each hypothesis branch with UCB so the scheduler can prune low-yield branches without abandoning exploration.
- Fan out finished results to a paper-write stage and a re-queue stage so a high-yield branch spawns follow-up hypotheses.
- Surface a per-iteration trace with slot occupancy, per-branch mean reward and run counts, and pruning decisions.
Why a scheduler, not a worklist
A worklist runs jobs in submission order, which is fine when each job is independent. A scheduler, as I use the word here, reads each result as it lands and reorders the queue. Research needs that, because a finding from experiment three changes the priority of experiments four and five, and a scheduler that reads the result fan-in gets more useful work done per unit of compute.
The design choice that matters is the scoring rule. A greedy scorer picks the current leader and never explores. A uniform scorer never exploits. UCB (upper confidence bound) sits between them: it exploits the leader while reserving capacity for branches that have been tried less.
The system shape
The queue holds hypotheses. When a slot frees, the scheduler picks the highest-UCB hypothesis and runs it there as an asynchronous experiment. Finished experiments put their result onto the bus. The bus updates UCB statistics on the originating branch, and when a branch's yield crosses a threshold it also fans out to the paper-write stage.
The Hypothesis shape
A hypothesis is one trial within a research direction, and branch names that direction. The Hypothesis record carries only id, branch, and payload. branch is the key for UCB statistics, and multiple hypotheses may share a branch. The statistics live on a separate BranchStats record per branch: runs is the count of completed experiments for that branch, and reward_sum is the cumulative reward. UCB reads both.
UCB scoring
UCB1 scores a branch as its mean reward plus an exploration bonus. This lesson uses the classic form.
ucb(branch) = mean_reward(branch) + c * sqrt( ln(total_runs) / runs(branch) )total_runs is the count of all experiments completed across all branches. c is the exploration weight, and the lesson defaults it to sqrt(2). A branch with zero runs gets +inf, so untried branches are scheduled first. A branch with high mean reward keeps a high score until other branches catch up, and a branch that runs many times without much reward gets eclipsed by less-run alternatives.
The pruning gate is separate from the picker. It removes a branch from future scheduling when its mean reward falls below an absolute floor (default 0.2) after at least prune_after_runs trials (default 3). That keeps the queue bounded.
Parallel slots with asyncio
The scheduler drives experiments with asyncio.create_task. Each task runs the experiment runner (an async def callable) and returns a Result. The main loop waits on the set of in-flight tasks with asyncio.wait(..., return_when=asyncio.FIRST_COMPLETED) and fires the scoring update on each completion.
Three slots run concurrently, and the main loop does not block on any single experiment. It starts a new task as soon as a slot frees, and it stops once the queue is empty and no tasks are in flight.
Fan-out: paper triggers
A paper trigger is a paper.trigger event. The scheduler fans one onto an output list when a branch's mean reward reaches paper_threshold (default 0.7) and that branch has not yet produced a paper. There is no minimum run count: after one run the mean is that run's reward, so a single lucky result can trigger a paper. Downstream, the paper writer from lesson fifty-four would pick it up. Here the trigger is captured as a list so tests can assert it.
Fan-out: follow-up hypotheses
A follow-up hypothesis is new work on the same branch, produced, when a high-yield result lands, by the user-supplied expander if one is set. The expander is a pure function from Result to list[Hypothesis]. The scheduler calls it only when an expander is set and the result's reward is at or above paper_threshold. The lesson ships a deterministic expander that always returns two follow-ups, so every qualifying result queues two new hypotheses, recorded as an expand event.
Budgets
Two budgets protect the scheduler from runaway loops.
max_experiments : total count of experiments dispatched across all branches
max_seconds : wall-clock cap (monotonic time)The scheduler stops dispatching new tasks as soon as either budget is hit. It also checks both after each batch of completions, and when one fires it stops the main loop, awaits the in-flight tasks, and returns a SchedulerReport. Its stop_reason field is "max_experiments" when the experiment budget stopped the run, "deadline" when the wall-clock cap did, and "queue_empty" when the queue and the in-flight set both ran dry. experiments_run counts dispatched tasks, including any still in flight when the budget hit. Those in-flight tasks are drained: their results update branch stats and can still fire a paper trigger, but no pruning or expansion runs for them.
The Trace and final report
The trace is the ordered event log of the run. Each scheduling decision emits an event: dispatch (with the slot count after dispatch), result (with the branch mean and run count), prune, error for a failed experiment, and result.drain for results collected after a budget stop. A drain event carries only the branch and the reward, with no branch mean or run count, and when a drained result fires a paper trigger, the paper.trigger event is logged just before its result.drain event. A drained task that raised is skipped without an error event. Fan-out is two events: paper.trigger and expand. The final report summarises per-branch stats, experiments_run, wall_seconds, the stop_reason, and the paper triggers fired. The next lesson, the end-to-end demo, reads this report to drive the paper writer.
I'd keep the pruning gate separate from the picker. UCB ranks the queue, and pruning decides when a branch stops being a candidate at all. I'd rather read those as two rules than one blended score.
My bet is that prune_after_runs is the knob to tune first, because the floor of 0.2 only applies once a branch has at least prune_after_runs runs (default 3) behind its mean. I'd also set both budgets, max_experiments and max_seconds, from the first run, so a runaway loop is capped before it costs anything.
How to read the code
code/main.py defines Hypothesis, Result, BranchStats, IterationScheduler, and a make_deterministic_runner factory that returns an asyncio experiment runner with predictable rewards. The runner sleeps for a fixed delay_ms (default 5ms), which makes the concurrency observable.
code/tests/test_scheduler.py covers: UCB picks untried branches first, parallel slot occupancy, paper triggers when threshold is crossed, branch pruning after low-yield trials, fan-out follow-up hypotheses, and budget exit (both experiment count and wall clock).
Going further
Three extensions a real implementation will want. First, persistent UCB stats across sessions: the current statistics live in memory, and a real scheduler would checkpoint them so a restart preserves the exploration budget already spent. Second, multi-objective scoring: instead of a scalar reward, each result emits a vector and UCB becomes a Pareto-style picker. Third, contextual bandits: the picker conditions on hypothesis features (length, complexity) so similar hypotheses share exploration.
The scheduler is where research becomes more than a worklist. Once UCB is wired and the slots run in parallel, other improvements compose on top.
Code
- main.py (python, 361 lines)
- tests/test_scheduler.py (python, 220 lines)