Agents

Configs and factory:

banditry.agents.factory

Declarative agent configuration and the build_agent factory.

This module is the main entry point of banditry: pick a config dataclass (OFUGPConfig for GP-UCB / optimism-in-the-face-of-uncertainty agents, TSConfig for Thompson-sampling agents), fill in its fields, and pass it to build_agent together with a DesignSpace to obtain a ready-to-use agent that drives the suggest -> evaluate -> observe loop.

Sampler behaviour for Thompson-sampling agents is controlled through TSConfig.sampler_config; the defaults live in DEFAULT_NUTS_CONFIG and DEFAULT_LANGEVIN_CONFIG and can be copied and overridden key by key.

DEFAULT_NUTS_CONFIG = {'num_samples': 16, 'warmup_steps': 64, 'target_accept_prob': 0.6, 'prior_std': 0.5, 'obs_noise_prior_loc': -3.0, 'obs_noise_prior_scale': 0.5, 'max_obs_noise': 1.0, 'max_tree_depth': 6, 'disable_progbar': False} module-attribute

Default TSConfig.sampler_config for sampler="nuts".

The dict is splatted into the NUTSSampler constructor (NUTSSampler(**sampler_config)), so any constructor keyword argument is a valid key. Keys set here:

• num_samples (16): posterior draws kept after warmup; one draw is then selected uniformly at random as the sampled value function.
• warmup_steps (64): NUTS adaptation iterations, discarded before sampling.
• target_accept_prob (0.6): target acceptance probability for NUTS step-size adaptation.
• prior_std (0.5): standard deviation of the zero-mean Gaussian prior placed on every network weight.
• obs_noise_prior_loc (-3.0): location of the prior over the observation-noise std, interpreted in log-noise space (clamped to [log(min_obs_noise), log(max_obs_noise)]).
• obs_noise_prior_scale (0.5): scale of that prior in the unconstrained (pre-sigmoid) parameterisation of the noise.
• max_obs_noise (1.0): upper bound of the observation-noise std (the noise is a sigmoid-squashed parameter living in [min_obs_noise, max_obs_noise]).
• max_tree_depth (6): maximum doubling depth of the NUTS trajectory tree.
• disable_progbar (False): show Pyro's MCMC progress bar (the constructor default hides it).

Other accepted keys include adapt_step_size, adapt_mass_matrix, use_multinomial_sampling, init_obs_noise, min_obs_noise, jit_compile, ignore_jit_warnings, and generator.

DEFAULT_LANGEVIN_CONFIG = {'batch_size': 512, 'num_epochs': 128, 'burn_in': 64, 'temperature': 0.2, 'step_size': welling_teh_schedule(a=0.005, b=1.0, gamma=0.55, lr_floor=0.001), 'max_obs_noise': 10.0, 'precondition': True} module-attribute

Default TSConfig.sampler_config for sampler="langevin".

The dict is splatted into the LangevinSampler constructor (LangevinSampler(**sampler_config)), so any constructor keyword argument is a valid key. Keys set here:

• batch_size (512): SGLD minibatch size (capped at the dataset size).
• num_epochs (128): passes over the data; the total number of updates is at least burn_in + 1 epochs and min_batches updates, whichever is larger.
• burn_in (64): epochs discarded before a post-burn-in state is reservoir-sampled uniformly as the posterior draw.
• temperature (0.2): SGLD noise temperature; 1.0 targets the exact posterior, smaller values sample a colder (sharper) posterior.
• step_size: constant float or callable schedule (t, n_burn_in) -> lr; here a Welling-Teh polynomial decay a * (b + t) ** -gamma that is frozen after burn-in and floored at lr_floor.
• max_obs_noise (10.0): upper clamp of the learned observation-noise std.
• precondition (True): apply RMSProp-style diagonal preconditioning to the SGLD updates.

Other accepted keys include prior_precision, precond_alpha, precond_eps, min_batches, init_obs_noise, min_obs_noise, and generator.

AgentConfig(rand_sample=4) dataclass

Base configuration shared by every agent built by build_agent.

Concrete agents are configured through the subclasses OFUGPConfig and TSConfig; passing a plain AgentConfig to build_agent raises TypeError.

rand_sample = 4 class-attribute instance-attribute

Number of quasi-random (scrambled Sobol) warmup suggestions served before the agent switches to model-based suggestions. Values below 2 are floored to 2 by the agent. Default: 4.

OFUGPConfig(rand_sample=4, surrogate='svgp', frequentist=False, rkhs_norm=None, noise_std_proxy=None, model_config_overrides=dict()) dataclass

Bases: AgentConfig

Configuration for the GP-UCB / optimism-in-the-face-of-uncertainty agent (OFUGPAgent).

Although noise_std_proxy defaults to None (to keep the dataclass keyword-friendly), it is effectively required: build_agent forwards it to OFUGPAgent, whose constructor raises ValueError when it is None.

surrogate = 'svgp' class-attribute instance-attribute

Surrogate model family: "gp" (exact Gaussian process; accurate but cubic in the number of observations) or "svgp" (sparse variational GP; inducing-point approximation that scales to larger datasets). Must be a ModelEnum member name. Default: "svgp".

frequentist = False class-attribute instance-attribute

If True, use the frequentist Chowdhury-Gopalan (2017) confidence width beta_t (requires rkhs_norm) instead of the default Bayesian GP-UCB width. Default: False.

rkhs_norm = None class-attribute instance-attribute

Bound B on the RKHS norm of the reward function, used by the frequentist beta_t. Required when frequentist=True (the agent raises ValueError at suggestion time otherwise); ignored in the Bayesian setting. Default: None.

noise_std_proxy = None class-attribute instance-attribute

Sub-Gaussian / GP observation-noise scale R. Effectively required: building an OFU-GP agent raises ValueError if this is left as None. Used by the frequentist beta_t and by the realised information gain accumulated on every observe. Default: None.

model_config_overrides = field(default_factory=dict) class-attribute instance-attribute

Overrides merged on top of the per-surrogate default model config (see ModelEnum.model_config) and passed through to the surrogate constructor. Default: {}.

TSConfig(rand_sample=4, sampler='nuts', feel_good=False, fg_lambda=1.0, fg_bound=1.0, model_config=dict(), sampler_config=None, should_warm_start=True, latent_dimension=None) dataclass

Bases: AgentConfig

Configuration for the Thompson-sampling agent (TSAgent).

sampler = 'nuts' class-attribute instance-attribute

Posterior-sampling oracle for the value-function weights: "nuts" (Pyro NUTS; requires the optional banditry[nuts] extra) or "langevin" (SGLD; no extra dependencies). Default: "nuts".

feel_good = False class-attribute instance-attribute

If True, enable Feel-Good Thompson sampling (Zhang, 2021): the Gaussian likelihood is reweighted with an exploration-boosting term built from fg_lambda and fg_bound. Default: False.

fg_lambda = 1.0 class-attribute instance-attribute

Strength lambda of the Feel-Good reweighting term (only used when feel_good=True); 0 disables the term. Default: 1.0.

fg_bound = 1.0 class-attribute instance-attribute

Upper truncation of the Feel-Good optimistic-value term (the cap b in min(b, f*)); must be non-negative. Only used when feel_good=True. Default: 1.0.

model_config = field(default_factory=dict) class-attribute instance-attribute

Keyword arguments forwarded to the neural ValueFunction surrogate (e.g. num_uniqs, emb_sizes, a custom feature extractor fe). num_uniqs is filled in automatically for spaces with categorical parameters. Default: {}.

sampler_config = None class-attribute instance-attribute

Keyword arguments splatted into the sampler constructor (NUTSSampler(**sampler_config) or LangevinSampler(**sampler_config)), so any constructor kwarg is a valid key. None (default) selects DEFAULT_NUTS_CONFIG or DEFAULT_LANGEVIN_CONFIG according to sampler. An explicit dict replaces the default rather than being merged into it, so copy the default and override keys.

should_warm_start = True class-attribute instance-attribute

If True (default), initialise each round's MCMC run from the previous round's sampled model instead of a freshly initialised network.

latent_dimension = None class-attribute instance-attribute

Latent dimension of the sampled parameter vector, recorded for labelling/diagnostics. None (default) resolves to TSAgent.num_samplable_params() at build time.

build_agent(config, space)

Build a ready-to-use agent from a declarative config.

Parameters:

Name	Type	Description	Default
config	AgentConfig	An OFUGPConfig (GP-UCB / OFU agent) or a TSConfig (Thompson-sampling agent).	required
space	DesignSpace	The design space the agent optimises over.	required

Returns:

Type	Description
AbstractAgent	A fully constructed OFUGPAgent or TSAgent.

Raises:

Type	Description
TypeError	If config is neither an OFUGPConfig nor a TSConfig.
ValueError	If TSConfig.sampler is not "nuts" or "langevin", or if OFUGPConfig.noise_std_proxy is None (raised by OFUGPAgent).
KeyError	If OFUGPConfig.surrogate is not a ModelEnum member name ("gp" or "svgp").
ImportError	If sampler="nuts" and the optional pyro dependency (pip install "banditry[nuts]") is not installed.

Example

Build an OFU-GP agent on a two-parameter space and run one round:

import numpy as np
from banditry import DesignSpace, OFUGPConfig, build_agent

space = DesignSpace.parse([
    {"name": "x0", "type": "num", "lb": -1, "ub": 1},
    {"name": "x1", "type": "num", "lb": -1, "ub": 1},
])
agent = build_agent(OFUGPConfig(rand_sample=4, surrogate="gp", noise_std_proxy=1.0), space)

rec = agent.suggest(1)  # DataFrame of suggestions
y = ((rec["x0"] - 0.3) ** 2 + (rec["x1"] + 0.2) ** 2).to_numpy()
agent.observe(rec, y)  # agents minimise y

Source code in src/banditry/agents/factory.py

def build_agent(config: AgentConfig, space: DesignSpace) -> AbstractAgent:
    """Build a ready-to-use agent from a declarative config.

    Args:
        config: An `OFUGPConfig` (GP-UCB / OFU agent) or a `TSConfig`
            (Thompson-sampling agent).
        space: The design space the agent optimises over.

    Returns:
        A fully constructed `OFUGPAgent` or `TSAgent`.

    Raises:
        TypeError: If ``config`` is neither an `OFUGPConfig` nor a `TSConfig`.
        ValueError: If ``TSConfig.sampler`` is not ``"nuts"`` or ``"langevin"``, or if
            ``OFUGPConfig.noise_std_proxy`` is ``None`` (raised by `OFUGPAgent`).
        KeyError: If ``OFUGPConfig.surrogate`` is not a `ModelEnum` member name
            (``"gp"`` or ``"svgp"``).
        ImportError: If ``sampler="nuts"`` and the optional pyro dependency
            (``pip install "banditry[nuts]"``) is not installed.

    Example:
        Build an OFU-GP agent on a two-parameter space and run one round:

        ```python
        import numpy as np
        from banditry import DesignSpace, OFUGPConfig, build_agent

        space = DesignSpace.parse([
            {"name": "x0", "type": "num", "lb": -1, "ub": 1},
            {"name": "x1", "type": "num", "lb": -1, "ub": 1},
        ])
        agent = build_agent(OFUGPConfig(rand_sample=4, surrogate="gp", noise_std_proxy=1.0), space)

        rec = agent.suggest(1)  # DataFrame of suggestions
        y = ((rec["x0"] - 0.3) ** 2 + (rec["x1"] + 0.2) ** 2).to_numpy()
        agent.observe(rec, y)  # agents minimise y
        ```
    """
    if isinstance(config, OFUGPConfig):
        model = ModelEnum[config.surrogate]
        return OFUGPAgent(
            space,
            surrogate=model,
            rand_sample=config.rand_sample,
            model_config=model.model_config(space, config.model_config_overrides),
            frequentist=config.frequentist,
            rkhs_norm=config.rkhs_norm,
            noise_std_proxy=config.noise_std_proxy,
        )
    if isinstance(config, TSConfig):
        nll = FeelGoodNLL(fg_lambda=config.fg_lambda, fg_bound=config.fg_bound) if config.feel_good else None
        if config.sampler == "nuts":
            from banditry.sampling_oracles.nuts_sampler import NUTSSampler

            sampler_cls = NUTSSampler
            sampler_config = config.sampler_config if config.sampler_config is not None else DEFAULT_NUTS_CONFIG
        elif config.sampler == "langevin":
            sampler_cls = LangevinSampler
            sampler_config = config.sampler_config if config.sampler_config is not None else DEFAULT_LANGEVIN_CONFIG
        else:
            raise ValueError(f"Unknown sampler: {config.sampler!r}")
        return TSAgent(
            space,
            rand_sample=config.rand_sample,
            sampler_cls=sampler_cls,
            sampler_config=sampler_config,
            nll=nll,
            model_config=config.model_config,
            should_warm_start=config.should_warm_start,
            latent_dimension=config.latent_dimension,
        )
    raise TypeError(f"Unknown config type: {type(config).__name__}")

The agent contract:

banditry.agents.agent.AbstractAgent(space, rand_sample=None)

Bases: ABC

Base class for bandit agents driving a suggest -> evaluate -> observe loop.

Usage contract:

• Call suggest to obtain a batch of candidate configurations (a DataFrame with one column per parameter of the design space), evaluate them externally, then feed the results back with observe. Repeat.
• The library minimises the observed values y; negate rewards if you need maximisation.
• The first rand_sample suggestions are quasi-random (scrambled Sobol) warmup draws; once at least rand_sample observations have been recorded, suggestions come from the subclass's model (get_model + pick_action).
• fix_input turns the loop into a contextual bandit: any subset of parameters can be pinned to observed context values each round, and the remaining parameters are optimised conditionally.

Parameters:

Name	Type	Description	Default
space	DesignSpace	The design space to optimise over.	required
rand_sample	int | None	Sobol warmup budget. None (default) resolves to 1 + space.num_paras; explicit values are floored at 2.	None

Attributes:

Name	Type	Description
space		The design space being optimised.
X		DataFrame of all observed inputs, one column per parameter.
y		(n, 1) float array of observed objective values (lower is better).
rand_sample		Resolved warmup budget (see above).
support_parallel_opt		Whether the agent supports n_suggestions > 1.
support_constraint		Whether constrained optimisation is supported.
support_multi_objective		Whether multi-objective observations are supported.
support_combinatorial		Whether categorical/integer parameters are supported.
support_contextual		Whether fix_input contexts are supported.

Source code in src/banditry/agents/agent.py

def __init__(
    self,
    space: DesignSpace,
    rand_sample: int | None = None,
):
    self.space = space
    self.X = pd.DataFrame(columns=self.space.para_names)
    self.y = np.zeros((0, 1))
    self.rand_sample = 1 + self.space.num_paras if rand_sample is None else max(2, rand_sample)
    # SobolEngine is constructed lazily on first use so its scramble
    # is captured from torch's *post-seed* global state (set by
    # Experiment.run's seed_everything), not whatever torch state
    # happens to be live at factory time.
    self._sobol: SobolEngine | None = None

observe(X, y)

Record evaluated suggestions in the agent's history.

Rows whose y is non-finite (NaN or infinite) are dropped; the remaining rows are appended to X and y.

Parameters:

Name	Type	Description	Default
X	DataFrame	Evaluated inputs, one column per parameter (as returned by suggest).	required
y	ndarray	Observed objective values, one per row of X (reshaped to (n, 1)). Lower is better — the agent minimises.	required

Source code in src/banditry/agents/agent.py

def observe(self, X: pd.DataFrame, y: np.ndarray) -> None:
    """Record evaluated suggestions in the agent's history.

    Rows whose ``y`` is non-finite (NaN or infinite) are dropped; the
    remaining rows are appended to `X` and `y`.

    Args:
        X: Evaluated inputs, one column per parameter (as returned by `suggest`).
        y: Observed objective values, one per row of ``X`` (reshaped to
            ``(n, 1)``). Lower is better — the agent minimises.
    """
    valid_id = np.where(np.isfinite(y.reshape(-1)))[0].tolist()
    XX = X.iloc[valid_id]
    yy = y[valid_id].reshape(-1, 1)
    self.X = pd.concat([self.X, XX], axis=0, ignore_index=True)
    self.y = np.vstack([self.y, yy])

quasi_sample(n, fix_input=None)

Draw quasi-random samples from the design space.

Samples come from a scrambled Sobol sequence in the transformed space and are mapped back to the original parameter domain; numeric parameters that are discrete after transformation are rounded. Used for warmup rounds and to pad incomplete suggestion batches.

Parameters:

Name	Type	Description	Default
n	int	Number of samples to draw.	required
fix_input	dict[str, Any] | None	Optional mapping of parameter name to pinned value; the corresponding columns are overwritten after sampling.	None

Returns:

Type	Description
DataFrame	DataFrame with n rows, one column per parameter, in the
DataFrame	original (untransformed) domain.

Source code in src/banditry/agents/agent.py

def quasi_sample(self, n: int, fix_input: dict[str, Any] | None = None) -> pd.DataFrame:
    """Draw quasi-random samples from the design space.

    Samples come from a scrambled Sobol sequence in the transformed space
    and are mapped back to the original parameter domain; numeric
    parameters that are discrete after transformation are rounded. Used
    for warmup rounds and to pad incomplete suggestion batches.

    Args:
        n: Number of samples to draw.
        fix_input: Optional mapping of parameter name to pinned value; the
            corresponding columns are overwritten after sampling.

    Returns:
        DataFrame with ``n`` rows, one column per parameter, in the
        original (untransformed) domain.
    """
    fix_input = fix_input or {}
    samp = self.sobol.draw(n)
    samp = samp * (self.space.opt_ub - self.space.opt_lb) + self.space.opt_lb
    x = samp[:, : self.space.num_numeric]
    xe = samp[:, self.space.num_numeric :]
    for i, name in enumerate(self.space.numeric_names):
        if self.space.paras[name].is_discrete_after_transform:
            x[:, i] = x[:, i].round()
    df_samp = self.space.inverse_transform(x, xe)
    for k, v in fix_input.items():
        df_samp[k] = v
    return df_samp

get_best_id(fix_input=None)

Return the positional index of the best (lowest-y) observation.

Parameters:

Name	Type	Description	Default
fix_input	dict[str, Any] | None	Optional context filter. When given, only observations whose pinned columns match the given values are considered (fully numeric columns are compared with a tight absolute tolerance, other columns by equality). If no row matches, the method falls back to the global argmin.	None

Returns:

Type	Description
int	Positional index into X/y of the best matching observation.

Source code in src/banditry/agents/agent.py

def get_best_id(self, fix_input: dict[str, Any] | None = None) -> int:
    """Return the positional index of the best (lowest-``y``) observation.

    Args:
        fix_input: Optional context filter. When given, only observations
            whose pinned columns match the given values are considered
            (fully numeric columns are compared with a tight absolute
            tolerance, other columns by equality). If no row matches, the
            method falls back to the global argmin.

    Returns:
        Positional index into `X`/`y` of the best matching observation.
    """
    if not fix_input:
        return int(np.argmin(self.y.reshape(-1)))
    X = self.X.copy()
    y = self.y.copy()
    for k, v in fix_input.items():
        col = X[k]
        col_numeric = pd.to_numeric(col, errors="coerce")
        if col_numeric.notna().all() and np.isscalar(v):
            crit = ~np.isclose(
                col_numeric.to_numpy(dtype=float),
                float(v),
                rtol=0.0,
                atol=np.finfo(np.float64).eps,
            )
        else:
            crit = (col != v).to_numpy()
        y[crit] = np.inf
    if np.isfinite(y).any():
        return int(np.argmin(y.reshape(-1)))
    return int(np.argmin(self.y.reshape(-1)))

check_unique(rec)

Flag which rows of rec are new.

Parameters:

Name	Type	Description	Default
rec	DataFrame	Candidate suggestions to check.	required

Returns:

Type	Description
list[bool]	One boolean per row of rec: True if the row duplicates
list[bool]	neither a past observation nor an earlier row of rec itself.

Source code in src/banditry/agents/agent.py

def check_unique(self, rec: pd.DataFrame) -> list[bool]:
    """Flag which rows of ``rec`` are new.

    Args:
        rec: Candidate suggestions to check.

    Returns:
        One boolean per row of ``rec``: ``True`` if the row duplicates
        neither a past observation nor an earlier row of ``rec`` itself.
    """
    return (~pd.concat([self.X, rec], axis=0).duplicated().tail(rec.shape[0]).values).tolist()

prepare_data()

Transform the observation history into model-ready tensors.

Returns:

Type	Description
	Tuple (Xc, Xe, y): transformed numeric features, transformed
	categorical features, and a float tensor copy of the observed values.

Source code in src/banditry/agents/agent.py

def prepare_data(self):
    """Transform the observation history into model-ready tensors.

    Returns:
        Tuple ``(Xc, Xe, y)``: transformed numeric features, transformed
        categorical features, and a float tensor copy of the observed values.
    """
    X, Xe = self.space.transform(self.X)
    y = torch.FloatTensor(self.y).clone()
    return X, Xe, y

get_model(Xc, Xe, y) abstractmethod

Build the decision model from transformed data.

Subclasses must fit (or posterior-sample) their surrogate here.

Parameters:

Name	Type	Description	Default
Xc		Transformed numeric features.	required
Xe		Transformed categorical features.	required
y		Observed values as a float tensor.	required

Returns:

Type	Description
	A model consumable by pick_action.

Source code in src/banditry/agents/agent.py

@abstractmethod
def get_model(self, Xc, Xe, y):
    """Build the decision model from transformed data.

    Subclasses must fit (or posterior-sample) their surrogate here.

    Args:
        Xc: Transformed numeric features.
        Xe: Transformed categorical features.
        y: Observed values as a float tensor.

    Returns:
        A model consumable by `pick_action`.
    """
    pass

pick_action(model, fix_input, n_suggestions=1) abstractmethod

Choose the next suggestions given a fitted or sampled model.

Subclasses must optimise their acquisition here, honouring fix_input.

Parameters:

Name	Type	Description	Default
model		The model returned by get_model.	required
fix_input		Mapping of parameter name to pinned context value (may be empty).	required
n_suggestions		Number of suggestions to return.	1

Returns:

Type	Description
	DataFrame of n_suggestions suggested configurations.

Source code in src/banditry/agents/agent.py

@abstractmethod
def pick_action(self, model, fix_input, n_suggestions=1):
    """Choose the next suggestions given a fitted or sampled model.

    Subclasses must optimise their acquisition here, honouring ``fix_input``.

    Args:
        model: The model returned by `get_model`.
        fix_input: Mapping of parameter name to pinned context value (may be empty).
        n_suggestions: Number of suggestions to return.

    Returns:
        DataFrame of ``n_suggestions`` suggested configurations.
    """
    pass

suggest(n_suggestions=1, fix_input=None)

Propose the next configurations to evaluate.

While fewer than rand_sample observations have been recorded, this returns quasi-random Sobol samples; afterwards it fits or samples the model (get_model) and optimises the acquisition (pick_action).

Parameters:

Name	Type	Description	Default
n_suggestions	int	Number of configurations to return.	1
fix_input	dict[str, Any] | None	Optional context — parameter values to pin this round (contextual bandit).	None

Returns:

Type	Description
DataFrame	DataFrame with n_suggestions rows, one column per parameter.

Source code in src/banditry/agents/agent.py

def suggest(self, n_suggestions: int = 1, fix_input: dict[str, Any] | None = None) -> pd.DataFrame:
    """Propose the next configurations to evaluate.

    While fewer than `rand_sample` observations have been recorded, this
    returns quasi-random Sobol samples; afterwards it fits or samples the
    model (`get_model`) and optimises the acquisition (`pick_action`).

    Args:
        n_suggestions: Number of configurations to return.
        fix_input: Optional context — parameter values to pin this round
            (contextual bandit).

    Returns:
        DataFrame with ``n_suggestions`` rows, one column per parameter.
    """
    fix_input = fix_input or {}
    if self.X.shape[0] < self.rand_sample:
        return self.quasi_sample(n_suggestions, fix_input)

    Xc, Xe, y = self.prepare_data()

    model = self.get_model(Xc, Xe, y)

    return self.pick_action(model, fix_input, n_suggestions)

n_plays()

Return the number of observations recorded so far.

Source code in src/banditry/agents/agent.py

def n_plays(self) -> int:
    """Return the number of observations recorded so far."""
    return len(self.y)

Agent implementations:

banditry.agents.ofugpagent

GP-UCB / optimism-in-the-face-of-uncertainty agent and its surrogate registry.

ModelEnum

Bases: Enum

GP surrogate registry for OFUGPAgent.

Members map surrogate names to classes: gp is an exact Gaussian process (accurate, but cubic cost in the number of observations) and svgp is a sparse variational GP (inducing-point approximation that scales to larger datasets). model_config builds the per-surrogate default configuration.

model_config(space, configs_to_override=None)

Build the default constructor config for this surrogate.

Parameters:

Name	Type	Description	Default
space	DesignSpace	Design space (used to derive categorical cardinalities).	required
configs_to_override		Optional overrides merged on top of the defaults.	None

Returns:

Type	Description
	Keyword-argument dict for the surrogate class constructor. For
	svgp: batch_size, num_inducing, use_ngd. For gp:
	lr, num_epochs, verbose, noise_lb, pred_likeli,
	optimizer. When the space has categorical parameters,
	num_uniqs (per-parameter cardinalities) is added.

Source code in src/banditry/agents/ofugpagent.py

def model_config(self, space: DesignSpace, configs_to_override=None):
    """Build the default constructor config for this surrogate.

    Args:
        space: Design space (used to derive categorical cardinalities).
        configs_to_override: Optional overrides merged on top of the defaults.

    Returns:
        Keyword-argument dict for the surrogate class constructor. For
        ``svgp``: ``batch_size``, ``num_inducing``, ``use_ngd``. For ``gp``:
        ``lr``, ``num_epochs``, ``verbose``, ``noise_lb``, ``pred_likeli``,
        ``optimizer``. When the space has categorical parameters,
        ``num_uniqs`` (per-parameter cardinalities) is added.
    """
    cfg = {}
    if self == ModelEnum.svgp:
        cfg = {"batch_size": 128, "num_inducing": 256, "use_ngd": False}
    if self == ModelEnum.gp:
        cfg = {
            "lr": 0.01,
            "num_epochs": 100,
            "verbose": False,
            "noise_lb": 8e-4,
            "pred_likeli": True,
            "optimizer": "adam",
        }
    cfg.update(configs_to_override or {})
    if space.num_categorical > 0:
        cfg["num_uniqs"] = [len(space.paras[name].categories) for name in space.enum_names]
    return cfg

OFUGPAgent(space, noise_std_proxy, surrogate=ModelEnum.svgp, rand_sample=None, acq_cls=LCB, model_config=None, frequentist=False, delta=0.01, kappa_fn=None, rkhs_norm=None)

Bases: AbstractAgent

Optimism-in-the-face-of-uncertainty (GP-UCB) agent.

Fits a GP surrogate to the observations and suggests the minimiser of a lower-confidence-bound acquisition mu - kappa * sigma (or of the MACE multi-objective ensemble), optimised with an evolutionary optimiser over the design space, with context parameters pinned via fix_input.

The confidence multiplier kappa is Bayesian by default; with frequentist=True it becomes the Chowdhury-Gopalan (2017, "On Kernelized Multi-armed Bandits") beta_t, which requires the RKHS-norm bound B (rkhs_norm) and the sub-Gaussian noise scale R (noise_std_proxy).

Parameters:

Name	Type	Description	Default
space	DesignSpace	The design space to optimise over.	required
noise_std_proxy	float	Sub-Gaussian / GP observation-noise scale R. Required — the constructor raises ValueError if None. Used by the frequentist beta_t and by the realised information gain accumulated on every observe.	required
surrogate		ModelEnum member selecting the GP surrogate (default: ModelEnum.svgp).	svgp
rand_sample		Sobol warmup budget; see AbstractAgent.	None
acq_cls		Acquisition class — LCB (default) or MACE (MACE is required for n_suggestions > 1).	LCB
model_config		Overrides merged into the surrogate's default config (see ModelEnum.model_config).	None
frequentist	bool	Use the Chowdhury-Gopalan beta_t instead of the Bayesian confidence width (default False).	False
delta		Confidence level delta appearing in both width formulas (default 0.01).	0.01
kappa_fn	Callable[[OFUGPAgent, int], float] | None	Optional callable (agent, n_suggestions) -> float overriding kappa entirely.	None
rkhs_norm		RKHS-norm bound B; required when frequentist=True.	None

Source code in src/banditry/agents/ofugpagent.py

def __init__(
    self,
    space: DesignSpace,
    noise_std_proxy: float,
    surrogate=ModelEnum.svgp,
    rand_sample=None,
    acq_cls=LCB,
    model_config=None,
    frequentist: bool = False,
    delta=0.01,
    kappa_fn: Callable[["OFUGPAgent", int], float] | None = None,
    rkhs_norm=None,
):
    super().__init__(space, rand_sample=rand_sample)
    if noise_std_proxy is None:
        raise ValueError(
            "noise_std_proxy is required (sub-Gaussian / GP noise scale used by "
            "the frequentist β_t and the realised information gain)"
        )
    self.surrogate = surrogate
    self.acq_cls = acq_cls
    self.delta = delta
    self.model_config = model_config
    self.frequentist = frequentist
    self._kappa_fn = kappa_fn
    self._rkhs_norm = rkhs_norm
    self.noise_std_proxy = noise_std_proxy
    self._realised_information_gain = 0.0
    self._pending_var_t: np.ndarray | None = None

get_model(X, Xe, y)

Instantiate the configured GP surrogate and fit it to the data.

Source code in src/banditry/agents/ofugpagent.py

def get_model(self, X, Xe, y):
    """Instantiate the configured GP surrogate and fit it to the data."""
    model = self.surrogate.value(
        self.space.num_numeric,
        self.space.num_categorical,
        1,
        **self.surrogate.model_config(self.space, self.model_config),
    )
    model.fit(X, Xe, y)
    return model

kappa(n_suggestions)

Confidence multiplier for the LCB acquisition mu - kappa * sigma.

If a kappa_fn override was supplied, it is delegated to. Otherwise, with t = max(1, n_plays() // n_suggestions) rounds played and input dimension d:

• Bayesian (default): beta_t = (2 + d/2) * ln(t) + ln(pi^2 / delta) and sqrt(beta_t) is returned.
• Frequentist (frequentist=True): the Chowdhury-Gopalan width beta_t = B + 4 * R * sqrt(gamma_t + 1 + ln(1/delta)), where gamma_t is the realised information gain accumulated across observe calls; returned as is (already on the sqrt scale).

Parameters:

Name	Type	Description	Default
n_suggestions		Batch size of the current round (scales the round counter).	required

Returns:

Type	Description
	The multiplier applied to the predictive standard deviation.

Raises:

Type	Description
ValueError	If frequentist=True and no rkhs_norm was provided.

Source code in src/banditry/agents/ofugpagent.py

def kappa(self, n_suggestions):
    """Confidence multiplier for the LCB acquisition ``mu - kappa * sigma``.

    If a ``kappa_fn`` override was supplied, it is delegated to. Otherwise,
    with ``t = max(1, n_plays() // n_suggestions)`` rounds played and input
    dimension ``d``:

    - Bayesian (default): ``beta_t = (2 + d/2) * ln(t) + ln(pi^2 / delta)``
      and ``sqrt(beta_t)`` is returned.
    - Frequentist (``frequentist=True``): the Chowdhury-Gopalan width
      ``beta_t = B + 4 * R * sqrt(gamma_t + 1 + ln(1/delta))``, where
      ``gamma_t`` is the realised information gain accumulated across
      `observe` calls; returned as is (already on the sqrt scale).

    Args:
        n_suggestions: Batch size of the current round (scales the round counter).

    Returns:
        The multiplier applied to the predictive standard deviation.

    Raises:
        ValueError: If ``frequentist=True`` and no ``rkhs_norm`` was provided.
    """
    # TODO: Rework
    if self._kappa_fn is not None:
        return self._kappa_fn(self, n_suggestions)

    # The benefit of this is arguable
    t = max(1, self.n_plays() // n_suggestions)
    d = self.X.shape[1]
    delta = self.delta

    if self.frequentist:
        if self._rkhs_norm is None:
            raise ValueError("rkhs_norm must be provided for the frequentist setting")
        beta_t = self._rkhs_norm + 4 * self.noise_std_proxy * np.sqrt(
            self._realised_information_gain + 1 + np.log(1 / delta)
        )
        # Already square rooted
        return beta_t
    else:
        beta_t = (2.0 + d / 2.0) * np.log(t) + np.log(PI_SQUARED / delta)
        return np.sqrt(beta_t)

pick_action(model, fix_input, n_suggestions=1)

Optimise the acquisition and select n_suggestions candidates.

Runs an evolutionary optimiser seeded at the incumbent, pads with quasi-random samples if needed, and — for batches larger than 2 — reserves slots for the most uncertain and the best-predicted candidates. The selected candidates' predictive variances are cached for the information-gain update in observe.

Raises:

Type	Description
RuntimeError	If n_suggestions > 1 with a non-MACE acquisition.

Source code in src/banditry/agents/ofugpagent.py

def pick_action(self, model, fix_input, n_suggestions=1):
    """Optimise the acquisition and select ``n_suggestions`` candidates.

    Runs an evolutionary optimiser seeded at the incumbent, pads with
    quasi-random samples if needed, and — for batches larger than 2 —
    reserves slots for the most uncertain and the best-predicted
    candidates. The selected candidates' predictive variances are cached
    for the information-gain update in `observe`.

    Raises:
        RuntimeError: If ``n_suggestions > 1`` with a non-MACE acquisition.
    """
    if self.acq_cls != MACE and n_suggestions != 1:
        raise RuntimeError("Parallel optimization is supported only for MACE acquisition")

    best_id = self.get_best_id(fix_input)
    best_x = self.X.iloc[[best_id]]

    py_best, _ = model.predict(*self.space.transform(best_x))
    py_best = py_best.detach().numpy().squeeze()

    kappa = self.kappa(n_suggestions)

    acq = self.acq_cls(model, best_y=py_best, kappa=kappa)
    opt = EvolutionOpt(self.space, pop=100, max_iters=100, verbose=False)
    prob = ContextualProblem(acq, self.space, fix_input)
    rec = opt.optimise(prob, initial_suggest=best_x)
    rec = self._fill_suggestions(rec, n_suggestions, fix_input, max_retries=4)
    select_id = np.random.choice(rec.shape[0], n_suggestions, replace=False).tolist()

    prev_pred_likeli = model.pred_likeli
    model.pred_likeli = False
    with torch.no_grad():
        py_t, ps2_t = model.predict(*self.space.transform(rec))
        py_all = py_t.reshape(-1).cpu().numpy()
        ps2_all = ps2_t.reshape(-1).cpu().numpy()
        best_pred_id = int(np.argmin(py_all))
        best_unce_id = int(np.argmax(ps2_all))
        if best_unce_id not in select_id and n_suggestions > 2:
            select_id[0] = best_unce_id
        if best_pred_id not in select_id and n_suggestions > 2:
            select_id[1] = best_pred_id
        rec_selected = rec.iloc[select_id].copy()
        self._pending_var_t = ps2_all[select_id]
    model.pred_likeli = prev_pred_likeli
    return rec_selected

observe(X, y)

Record observations and update the realised information gain.

Uses the predictive variances cached by pick_action when available (otherwise refits the GP on the past data, falling back to the prior variance on failure) to accumulate 0.5 * sum(log(1 + var / R^2)), then delegates to AbstractAgent.observe.

Source code in src/banditry/agents/ofugpagent.py

def observe(self, X, y):
    """Record observations and update the realised information gain.

    Uses the predictive variances cached by `pick_action` when available
    (otherwise refits the GP on the past data, falling back to the prior
    variance on failure) to accumulate ``0.5 * sum(log(1 + var / R^2))``,
    then delegates to `AbstractAgent.observe`.
    """
    if self._pending_var_t is not None:
        var_t = self._pending_var_t
        self._pending_var_t = None
    elif len(self.y) > 0:
        try:
            Xc_prev, Xe_prev, y_prev = self.prepare_data()
            model = self.get_model(Xc_prev, Xe_prev, y_prev)
            model.pred_likeli = False  # latent σ_f² for the info gain
            xc_new, xe_new = self.space.transform(X)
            with torch.no_grad():
                _, var_new = model.predict(xc_new, xe_new)
            var_t = var_new.detach().cpu().numpy().reshape(-1)
        except Exception:
            import banditry.logging_utils as log

            log.debug("OFUGPAgent.observe: GP fit failed; falling back to prior variance")
            var_t = np.ones(len(X))
    else:
        var_t = np.ones(len(X))
    self._realised_information_gain += 0.5 * float(np.log1p(var_t / self.noise_std_proxy**2).sum())
    super().observe(X, y)

banditry.agents.tsagent.TSAgent(space, rand_sample=None, model_config=None, sampler_cls=LangevinSampler, sampler_config=None, nll=None, should_warm_start=False, latent_dimension=None)

Bases: AbstractAgent

Thompson-sampling agent with a neural value function.

Each round, the weights of a neural ValueFunction are sampled from their posterior by an MCMC sampling oracle (LangevinSampler / SGLD, or NUTSSampler / Pyro NUTS), and the sampled network is minimised over the design space with an evolutionary optimiser (context parameters pinned via fix_input) — i.e. the agent plays the argmin of one posterior draw.

Optional Feel-Good Thompson sampling (Zhang, 2021, "Feel-Good Thompson Sampling for Contextual Bandits and Reinforcement Learning") reweights the likelihood with an exploration-boosting term; enable it by passing nll=FeelGoodNLL(...).

Parameters:

Name	Type	Description	Default
space	DesignSpace	The design space to optimise over.	required
rand_sample		Sobol warmup budget; see AbstractAgent.	None
model_config	dict | None	Keyword arguments for the ValueFunction surrogate (e.g. num_uniqs, emb_sizes, a custom feature extractor fe). num_uniqs is filled in automatically for spaces with categorical parameters.	None
sampler_cls	type[Sampler]	Sampling-oracle class (default: LangevinSampler).	LangevinSampler
sampler_config	dict | None	Keyword arguments splatted into sampler_cls; None uses the sampler's own constructor defaults.	None
nll	NLL | None	Optional NLL factory replacing the plain Gaussian likelihood, e.g. FeelGoodNLL for Feel-Good Thompson sampling. It is bound to the current context and observation history on every suggest.	None
should_warm_start	bool	If True, reuse the previous round's sampled model as the MCMC initialisation of the next round (default False here; the factory TSConfig enables it).	False
latent_dimension	float | None	Latent dimension of the sampled parameter vector, recorded for labelling/diagnostics; None resolves to num_samplable_params.	None

Source code in src/banditry/agents/tsagent.py

def __init__(
    self,
    space: DesignSpace,
    rand_sample=None,
    model_config: dict | None = None,
    sampler_cls: type[Sampler] = LangevinSampler,
    sampler_config: dict | None = None,
    nll: NLL | None = None,
    should_warm_start: bool = False,
    latent_dimension: float | None = None,
):
    super().__init__(space, rand_sample=rand_sample)
    self.model_config = {} if model_config is None else dict(model_config)

    self.sampler_cls = sampler_cls
    self.sampler = self.sampler_cls(**({} if sampler_config is None else dict(sampler_config)))
    self.nll = nll
    self.should_warm_start = should_warm_start
    self.warm_start_model: ValueFunction | None = None
    self._num_samplable_params: int | None = None
    self.latent_dimension = latent_dimension if latent_dimension is not None else self.num_samplable_params()

get_model(Xc, Xe, y, nll=None)

Draw one posterior sample of the value function.

Starts from the previous round's sampled model when warm-starting is enabled (and available), otherwise from a freshly initialised ValueFunction, and runs the sampling oracle on the data.

Parameters:

Name	Type	Description	Default
Xc		Transformed numeric features.	required
Xe		Transformed categorical features.	required
y		Observed values as a float tensor.	required
nll		Optional bound NLL callable overriding the Gaussian likelihood.	None

Returns:

Type	Description
ValueFunction	The sampled ValueFunction (also stored for warm-starting).

Source code in src/banditry/agents/tsagent.py

def get_model(self, Xc, Xe, y, nll=None) -> ValueFunction:
    """Draw one posterior sample of the value function.

    Starts from the previous round's sampled model when warm-starting is
    enabled (and available), otherwise from a freshly initialised
    ``ValueFunction``, and runs the sampling oracle on the data.

    Args:
        Xc: Transformed numeric features.
        Xe: Transformed categorical features.
        y: Observed values as a float tensor.
        nll: Optional bound NLL callable overriding the Gaussian likelihood.

    Returns:
        The sampled ``ValueFunction`` (also stored for warm-starting).
    """
    if self.should_warm_start and self.warm_start_model is not None:
        initial_model = self.warm_start_model
    else:
        initial_model = ValueFunction(
            self.space.num_numeric,
            self.space.num_categorical,
            1,
            **self._value_function_config(),
        )
    sampled_model = self.sampler.sample(initial_model, Xc, Xe, y, nll=nll)
    if self.should_warm_start:
        self.warm_start_model = sampled_model
    return sampled_model

suggest(n_suggestions=1, fix_input=None)

Propose the next configurations to evaluate (see AbstractAgent.suggest).

Identical to the base implementation, except that the nll factory, when present, is bound to the current context (fix_input) and the observation history before posterior sampling.

Source code in src/banditry/agents/tsagent.py

def suggest(self, n_suggestions=1, fix_input=None):
    """Propose the next configurations to evaluate (see `AbstractAgent.suggest`).

    Identical to the base implementation, except that the ``nll`` factory,
    when present, is bound to the current context (``fix_input``) and the
    observation history before posterior sampling.
    """
    fix_input = fix_input or {}
    if self.X.shape[0] < self.rand_sample:
        return self.quasi_sample(n_suggestions, fix_input)

    Xc, Xe, y = self.prepare_data()
    nll = self.nll(fix_input, self.space, self.X) if self.nll is not None else None
    model = self.get_model(Xc, Xe, y, nll=nll)
    return self.pick_action(model, fix_input, n_suggestions)

pick_action(model, fix_input, n_suggestions=1)

Minimise the sampled value function over the design space.

Runs an evolutionary optimiser seeded at the incumbent, keeps the final population, filters out already-observed rows, pads with quasi-random samples if needed, and returns the top n_suggestions rows.

Source code in src/banditry/agents/tsagent.py

def pick_action(self, model, fix_input, n_suggestions=1):
    """Minimise the sampled value function over the design space.

    Runs an evolutionary optimiser seeded at the incumbent, keeps the
    final population, filters out already-observed rows, pads with
    quasi-random samples if needed, and returns the top ``n_suggestions``
    rows.
    """
    best_id = self.get_best_id(fix_input)
    best_x = self.X.iloc[[best_id]]

    opt = EvolutionOpt(self.space, pop=100, max_iters=100, verbose=False)
    prob = ContextualProblem(ThompsonObjective(model), self.space, fix_input)
    pop_rec = opt.optimise(
        prob,
        initial_suggest=best_x,
        return_pop=True,
    )
    mask = self.check_unique(pop_rec)
    log.debug(f"Leftover solutions after filtration (in parts): {np.mean(mask)}")
    rec = pop_rec[mask].reset_index(drop=True)
    rec = self._fill_suggestions(rec, n_suggestions, fix_input, max_retries=5)

    return rec.head(n_suggestions).copy()

num_samplable_params()

Count the ValueFunction parameters with requires_grad=True.

This is the exact set both LangevinSampler and NUTSSampler iterate over. It is determined by architecture alone (network builder + feature extractor + scalers), so a throwaway ValueFunction instance suffices to read it off. Cached after the first call.

Returns:

Type	Description
int	Total number of samplable scalar parameters.

Source code in src/banditry/agents/tsagent.py

def num_samplable_params(self) -> int:
    """Count the ``ValueFunction`` parameters with ``requires_grad=True``.

    This is the exact set both `LangevinSampler` and ``NUTSSampler``
    iterate over. It is determined by architecture alone (network builder
    + feature extractor + scalers), so a throwaway ``ValueFunction``
    instance suffices to read it off. Cached after the first call.

    Returns:
        Total number of samplable scalar parameters.
    """
    if self._num_samplable_params is None:
        vf = ValueFunction(
            self.space.num_numeric,
            self.space.num_categorical,
            1,
            **self._value_function_config(),
        )
        self._num_samplable_params = sum(p.numel() for p in vf.parameters() if p.requires_grad)
    return self._num_samplable_params