IMABO: Infinite Multi-Armed Bandits with Oracles

IMABO is a Python framework for online hyperparameter optimization: choosing the configuration to serve on each request, and learning from the reward it returns. It combines any bandit policy for choosing among the configurations already tried with any oracle for proposing new ones.

This repository contains the imabo package itself, plus every experiment, baseline, and plotting script used to produce the results in the paper: Bandits in Prod: Hyperparameter Optimization at Inference Time.

Installation

pip install -e .

To also install the experiment dependencies (plotting, async server client, TabPFN, etc.):

pip install -e ".[experiments]"

To also install the development tools (tests, linter):

pip install -e ".[dev]"

Quick Start

from imabo import IMABO

search_space = {
    "learning_rate": {"lower": 1e-5, "upper": 1.0, "log": True},
    "n_layers": {"lower": 1, "upper": 8, "int": True},
    "dropout": {"lower": 0.0, "upper": 0.5},
    "optimizer": {"choices": ["adam", "sgd", "adamw"]},
}

optimizer = IMABO(search_space, seed=42)          # IMOSS-TPE

for step in range(200):
    config = optimizer.suggest()
    reward = train_and_evaluate(config)  # your evaluation function
    optimizer.observe(reward)

print(f"Best config: {optimizer.best_config}")

By default, IMABO uses IMOSS(beta=0.5) paired with TPEOracle(), which is IMOSS-TPE.

from imabo import IMABO, IMOSS, RandomOracle, TPEOracle

IMABO(search_space, seed=42)                                  # IMOSS-TPE (the default)
IMABO(search_space, IMOSS(beta=0.5), TPEOracle(), seed=42)     # same

The four proposed instantiations from the paper are also available under their own names:

from imabo import IMOSSRandom, IMOSSTPE, IMOSSMutateKLTPE, IMOSSTabPFN, load_tabpfn

IMOSSRandom(search_space, beta=0.5, seed=42)                          # IMOSS-Random
IMOSSTPE(search_space, beta=0.5, seed=42)                             # IMOSS-TPE
IMOSSMutateKLTPE(search_space, beta=0.5, seed=42)                     # IMOSS-mutate-KLxTPE
IMOSSTabPFN(search_space, beta=0.5, seed=42, model=load_tabpfn())     # IMOSS-TabPFN

Algorithm Overview

IMABO maintains a growing active set of configurations M_t. Each round it makes one binary choice:

for t = 1, 2, ...
    if policy.expand(): x = oracle.suggest(); admit x
    else:               x = policy.select()
    serve x, observe its reward, update x's statistics

The policy decides both when to admit a new arm and which known arm to serve otherwise. The oracle only decides what a new arm should be.

IMOSS is the policy from the paper: it admits a new arm while $|M_t| < t^{\beta}$ and otherwise serves the arm maximizing the anytime MOSS index. The active set therefore grows as $O(t^\beta)$, which keeps the per-round cost bounded while preserving a sublinear cumulative quantile-regret guarantee.

API Reference

IMABO(search_space, policy, oracle, *, seed=None, memory=None)

IMOSS(beta=0.5, alpha=0.1, n_warmup=10, max_pending=20, min_rewards=1, delayed=False)

TPEOracle(n_candidates=24, gamma_func=None, weights_func=None, prior_weight=1.0, multivariate=True, split_index="policy", categorical_distance_func=None)

Other policies and oracles: BudgetedUCB (fixed horizon, not in the paper), RandomOracle, MutateKLTPEOracle, TabPFNOracle, TabFMOracle. Each carries its own documented arguments.

Search Space Format

A space can be given two ways: a dict or a suggestion function.

A dict, one entry per parameter. Any mix of the four kinds:

{
    "param_name": {"lower": float, "upper": float},                    # continuous
    "param_name": {"lower": float, "upper": float, "log": True},      # log-uniform
    "param_name": {"lower": int, "upper": int, "int": True},          # integer
    "param_name": {"choices": ["a", "b", "c"]},                       # categorical
}

from imabo import IMABO

SPACES = {
    "grid": {f"x{i}": {"choices": list(range(6))} for i in range(4)},
    "mixed": {
        "num_layers": {"choices": [1, 2, 3, 4, 5]},
        "batch_size": {"lower": 16, "upper": 512, "int": True, "log": True},
        "learning_rate": {"lower": 1e-4, "upper": 1e-1, "log": True},
        "momentum": {"lower": 0.1, "upper": 0.99},
        "max_dropout": {"lower": 0.0, "upper": 1.0},
    },
}

optimizer = IMABO(SPACES["mixed"], seed=42)

A suggestion function, which asks a Trial for its parameters instead of declaring them as data:

def space(trial):
    trial.suggest_float("learning_rate", 1e-5, 1.0, log=True)
    trial.suggest_int("n_layers", 1, 8)
    trial.suggest_float("dropout", 0.0, 0.5)
    trial.suggest_categorical("optimizer", ["adam", "sgd", "adamw"])

optimizer = IMABO(space, seed=42)

Each suggest_* call does two things: it declares a parameter — name, type, bounds — and it returns a value drawn for it.

Methods

• optimizer.suggest() -> dict — Returns the next configuration to serve.
• optimizer.observe(reward, config=None) — Reports a reward. Defaults to the last suggestion; pass config= when feedback arrives out of order.
• optimizer.best_config -> dict | None — The configuration the policy reports.
• optimizer.run(objective, n_rounds) -> dict — Drives the loop for offline use.
• optimizer.state -> CurrentState — Snapshot of the active set.
• optimizer.propose() -> dict — What the oracle would admit right now, for diagnostics. Advances the oracle's state, so call it on a copy.deepcopy.

Custom Memory

For production use with persistent storage (database, Redis, etc.), implement the Memory interface:

from imabo import Memory, ArmStats, CurrentState

class RedisMemory(Memory):
    def set(self, key, stats): ...
    def get_reward_frequency(self) -> float: ...
    def increment_step_counter(self): ...
    def get_current_state(self) -> CurrentState: ...
    def pull_arm(self, arm_key): ...
    def observe(self, config, reward): ...

optimizer = IMABO(space, memory=RedisMemory(...))

The oracles keep nothing of their own -- whatever they accumulate lives in the memory -- so a restarted process resumes as soon as your Memory reloads. Nothing has to be saved around a crash.

Writing Your Own Instantiation

An oracle is one method. suggest is handed everything any oracle could need -- the active set, the arms that have returned a reward, and how the policy ranks an arm -- and returns a configuration to admit.

from imabo import IMABO, IMOSS, Oracle

class NearestGridOracle(Oracle):
    """Admit the first point of a coarse grid that is not already an arm."""

    def suggest(self, state, rewarded_arms, score):
        for value in (0.25, 0.5, 0.75):
            config = {name: value for name in self.space.names}
            if self.space.encode(config) not in state.arms:
                return config
        return self.space.sample(self.rng)          # P0, the baseline distribution

optimizer = IMABO(search_space, IMOSS(beta=0.5), NearestGridOracle(), seed=0)

A policy is two: when to admit, and which known arm to serve.

from imabo import AllocationPolicy, ArmStats

class RoundRobin(AllocationPolicy):
    def setup(self, space, rng, memory):
        super().setup(space, rng, memory)
        for _ in range(5):                          # an initial design, drawn from P0
            memory.set(space.encode(space.sample(rng)), ArmStats())

    def expand(self, state, rewarded_arms):
        return len(state.arms) < 20

    def select(self, state, rewarded_arms):
        return rewarded_arms[state.nb_steps % len(rewarded_arms)][0]

Both are bound to the run by setup, which gives them the search space and the run's shared RNG; neither may draw from it during setup except to seed an initial design. Override AllocationPolicy.score to change how oracles rank the active set, and AllocationPolicy.best_arm to change what the optimizer reports.

Baselines

The baselines the paper compares against ship with the library and drive the same loop as IMABO — suggest(), then observe(reward), with best_config for the configuration to report:

from imabo import IMABO, QRM2, UCBAIR, RandomSearch

for optimizer in (IMABO(space), QRM2(space), UCBAIR(space), RandomSearch(space)):
    for _ in range(1000):
        config = optimizer.suggest()
        optimizer.observe(serve(config))
    print(optimizer.best_config)

The tree-search algorithms are exported as the generators they have always been — stroquool, stosoo, hoo_t — plus the TimedOptimizer wrapper. They work on [0, 1]**d rather than on a search space, so mapping their coordinates to configurations is the caller's job, exactly as the paper's experiments do it:

from imabo import TimedOptimizer, stosoo

optimizer = TimedOptimizer(stosoo, budget, dim)
while not optimizer.done:
    x = optimizer.suggest()                  # a point of [0, 1]**dim
    optimizer.observe(x, serve(to_config(x)))

Reproducing Paper Experiments

Install the experiment dependencies first:

pip install -e ".[experiments]"

Every experiment follows the same pattern: run the script to (re)generate results under results/ — runs are seed-checkpointed and resumable, so re-running skips completed seeds — then plot from those results (usually the same script, or a companion script in experiments/utils/plots/).

Command-line options are not uniform across the scripts:

coordination_barrier_experiment --plot reproduces Figure 4 (the two-dimensional Gaussian counter-example) and is the only paper figure whose experiment needs no external data at all: the landscapes are synthetic. Note that --quick shortens the run but does not skip imoss_tabpfn_tuned, so it still loads TabPFN.

Example invocations:

python -m experiments.toy_experiment
python -m experiments.rf_arm_distribution_experiment                          # all algorithms x all benchmarks
python -m experiments.rf_arm_distribution_experiment --algorithm IMOSS-TabPFN # single method
python -m experiments.hotpotqa_experiment --algorithm IMOSS-TABPFN            # or IMOSS-TABFM
python -m experiments.hotpotqa_experiment --algorithm IMOSS-TABPFN --plot-only

HPO benchmark server (Docker)

hpo_experiment (and anything using HPOBench, e.g. Logistic Regression / SVM) needs the bundled HPOBench server running:

docker compose up hpo-server

This builds a container with a Python 3.7 conda env and the vendored HPOBench/ source, and exposes it on localhost:8901. A dev service with the same image is also available for interactive use (docker compose run dev).

Tests

pip install -e ".[dev]"
pytest

Citation

If you use this work, please cite the paper:

@misc{abraham2026banditsprodhyperparameteroptimization,
      title={Bandits in Prod: Hyperparameter Optimization at Inference Time},
      author={Louis Abraham and Tuan-Anh Nguyen and Nicolas Devatine},
      year={2026},
      eprint={2609.01335},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.01335},
}