Using Ax for Human-in-the-loop Experimentation¶
While Ax can be used in as a fully automated service, generating and deploying candidates Ax can be also used in a trial-by-trial fashion, allowing for human oversight.
Typically, human intervention in Ax is necessary when there are clear tradeoffs between multiple metrics of interest. Condensing multiple outcomes of interest into a single scalar quantity can be really challenging. Instead, it can be useful to specify an objective and constraints, and tweak these based on the information from the experiment.
To facilitate this, Ax provides the following key features:
- Constrained optimization
- Interfaces for easily modifying optimization goals
- Utilities for visualizing and deploying new trials composed of multiple optimizations.
In this tutorial, we'll demonstrate how Ax enables users to explore these tradeoffs. With an understanding of the tradeoffs present in our data, we'll then make use of the constrained optimization utilities to generate candidates from multiple different optimization objectives, and create a conglomerate batch, with all of these candidates in together in one trial.
Experiment Setup
For this tutorial, we will assume our experiment has already been created.
import sys
in_colab = 'google.colab' in sys.modules
if in_colab:
%pip install ax-platform
import os
from ax import (
Data,
Metric,
OptimizationConfig,
Objective,
OutcomeConstraint,
ComparisonOp,
json_load,
)
from ax.modelbridge.cross_validation import cross_validate
from ax.modelbridge.registry import Models
from ax.plot.diagnostic import tile_cross_validation
from ax.plot.scatter import plot_multiple_metrics, tile_fitted
from ax.utils.notebook.plotting import render, init_notebook_plotting
import pandas as pd
import plotly.io as pio
init_notebook_plotting()
if in_colab:
pio.renderers.default = "colab"
NOTE: The path below assumes the tutorial is being run either from the root directory of
the Ax package or from the human_in_the_loop directory that this tutorial lives in.
This is needed since the jupyter notebooks may change active directory during runtime,
making it tricky to find the file in a consistent way.
curr_dir = os.getcwd()
if "human_in_the_loop" not in curr_dir:
curr_dir = os.path.join(curr_dir, "tutorials", "human_in_the_loop")
experiment = json_load.load_experiment(os.path.join(curr_dir, "hitl_exp.json"))
Initial Sobol Trial
Bayesian Optimization experiments almost always begin with a set of random points. In
this experiment, these points were chosen via a Sobol sequence, accessible via the
ModelBridge factory.
A collection of points run and analyzed together form a BatchTrial. A Trial object
provides metadata pertaining to the deployment of these points, including details such
as when they were deployed, and the current status of their experiment.
Here, we see an initial experiment has finished running (COMPLETED status).
experiment.trials[0]
experiment.trials[0].time_created
# Number of arms in first experiment, including status_quo
len(experiment.trials[0].arms)
# Sample arm configuration
experiment.trials[0].arms[0]
Experiment Analysis
Optimization Config
An important construct for analyzing an experiment is an OptimizationConfig. An OptimizationConfig contains an objective, and outcome constraints. Experiment's can have a default OptimizationConfig, but models can also take an OptimizationConfig as input independent of the default.
Objective: A metric to optimize, along with a direction to optimize (default: maximize)
Outcome Constraint: A metric to constrain, along with a constraint direction (<= or >=), as well as a bound.
Let's start with a simple OptimizationConfig. By default, our objective metric will be
maximized, but can be minimized by setting the minimize flag. Our outcome constraint
will, by default, be evaluated as a relative percentage change. This percentage change
is computed relative to the experiment's status quo arm.
experiment.status_quo
objective_metric = Metric(name="metric_1")
constraint_metric = Metric(name="metric_2")
experiment.optimization_config = OptimizationConfig(
objective=Objective(objective_metric, minimize=False),
outcome_constraints=[
OutcomeConstraint(metric=constraint_metric, op=ComparisonOp.LEQ, bound=5),
],
)
Data
Another critical piece of analysis is data itself! Ax data follows a standard format, shown below. This format is imposed upon the underlying data structure, which is a Pandas DataFrame.
A key set of fields are required for all data, for use with Ax models.
It's a good idea to double check our data before fitting models -- let's make sure all of our expected metrics and arms are present.
data = Data(pd.read_json(os.path.join(curr_dir, "hitl_data.json")))
data.df.head()
| arm_name | metric_name | mean | sem | trial_index | start_time | end_time | n | |
|---|---|---|---|---|---|---|---|---|
| 0 | 0_1 | metric_1 | 495.763 | 2.62164 | 0 | 2019-03-30 | 2019-04-03 | 1599994 |
| 1 | 0_23 | metric_1 | 524.368 | 2.73165 | 0 | 2019-03-30 | 2019-04-03 | 1596356 |
| 2 | 0_14 | metric_2 | 21.4602 | 0.069457 | 0 | 2019-03-30 | 2019-04-03 | 1600182 |
| 3 | 0_53 | metric_2 | 21.4374 | 0.069941 | 0 | 2019-03-30 | 2019-04-03 | 1601081 |
| 4 | 0_53 | metric_1 | 548.388 | 2.89349 | 0 | 2019-03-30 | 2019-04-03 | 1601081 |
data.df["arm_name"].unique()
data.df["metric_name"].unique()
Search Space
The final component necessary for human-in-the-loop optimization is a SearchSpace. A SearchSpace defines the feasible region for our parameters, as well as their types.
Here, we have both parameters and a set of constraints on those parameters.
Without a SearchSpace, our models are unable to generate new candidates. By default, the models will read the search space off of the experiment, when they are told to generate candidates. SearchSpaces can also be specified by the user at this time. Sometimes, the first round of an experiment is too restrictive--perhaps the experimenter was too cautious when defining their initial ranges for exploration! In this case, it can be useful to generate candidates from new, expanded search spaces, beyond that specified in the experiment.
experiment.search_space.parameters
experiment.search_space.parameter_constraints
Model Fit
Fitting a Modular BoTorch Model will allow us to predict new candidates based on our
first Sobol batch. Here, we make use of the default settings for BOTORCH_MODULAR
defined in the ModelBridge registry (uses BoTorch's SingleTaskGP and
qLogNoisyExpectedImprovement by default for single objective optimization).
gp = Models.BOTORCH_MODULAR(
search_space=experiment.search_space,
experiment=experiment,
data=data,
)
We can validate the model fits using cross validation, shown below for each metric of interest. Here, our model fits leave something to be desired--the tail ends of each metric are hard to model. In this situation, there are three potential actions to take:
- Increase the amount of traffic in this experiment, to reduce the measurement noise.
- Increase the number of points run in the random batch, to assist the GP in covering the space.
- Reduce the number of parameters tuned at one time.
However, away from the tail effects, the fits do show a strong correlations, so we will proceed with candidate generation.
cv_result = cross_validate(gp)
render(tile_cross_validation(cv_result))
The parameters from the initial batch have a wide range of effects on the metrics of interest, as shown from the outcomes from our fitted GP model.
render(tile_fitted(gp, rel=True))
METRIC_X_AXIS = "metric_1"
METRIC_Y_AXIS = "metric_2"
render(
plot_multiple_metrics(
gp,
metric_x=METRIC_X_AXIS,
metric_y=METRIC_Y_AXIS,
)
)
Candidate Generation
With our fitted GPEI model, we can optimize EI (Expected Improvement) based on any optimization config. We can start with our initial optimization config, and aim to simply maximize the playback smoothness, without worrying about the constraint on quality.
unconstrained = gp.gen(
n=3,
optimization_config=OptimizationConfig(
objective=Objective(objective_metric, minimize=False),
),
)
Let's plot the tradeoffs again, but with our new arms.
render(
plot_multiple_metrics(
gp,
metric_x=METRIC_X_AXIS,
metric_y=METRIC_Y_AXIS,
generator_runs_dict={
"unconstrained": unconstrained,
},
)
)
Change Objectives
With our unconstrained optimization, we generate some candidates which are pretty promising with respect to our objective! However, there is a clear regression in our constraint metric, above our initial 5% desired constraint. Let's add that constraint back in.
constraint_5 = OutcomeConstraint(metric=constraint_metric, op=ComparisonOp.LEQ, bound=5)
constraint_5_results = gp.gen(
n=3,
optimization_config=OptimizationConfig(
objective=Objective(objective_metric, minimize=False), outcome_constraints=[constraint_5]
),
)
This yields a GeneratorRun, which contains points according to our specified optimization config, along with metadata about how the points were generated. Let's plot the tradeoffs in these new points.
from ax.plot.scatter import plot_multiple_metrics
render(
plot_multiple_metrics(
gp,
metric_x=METRIC_X_AXIS,
metric_y=METRIC_Y_AXIS,
generator_runs_dict={"constraint_5": constraint_5_results},
)
)
It is important to note that the treatment of constraints in GP EI is probabilistic. The acquisition function weights our objective by the probability that each constraint is feasible. Thus, we may allow points with a very small probability of violating the constraint to be generated, as long as the chance of the points increasing our objective is high enough.
You can see above that the point estimate for each point is significantly below a 5% increase in the constraint metric, but that there is uncertainty in our prediction, and the tail probabilities do include probabilities of small regressions beyond 5%.
constraint_1 = OutcomeConstraint(metric=constraint_metric, op=ComparisonOp.LEQ, bound=1)
constraint_1_results = gp.gen(
n=3,
optimization_config=OptimizationConfig(
objective=Objective(objective_metric, minimize=False),
outcome_constraints=[constraint_1],
),
)
render(
plot_multiple_metrics(
gp,
metric_x=METRIC_X_AXIS,
metric_y=METRIC_Y_AXIS,
generator_runs_dict={
"constraint_1": constraint_1_results,
},
)
)
Finally, let's view all three sets of candidates together.
render(
plot_multiple_metrics(
gp,
metric_x=METRIC_X_AXIS,
metric_y=METRIC_Y_AXIS,
generator_runs_dict={
"unconstrained": unconstrained,
"loose_constraint": constraint_5_results,
"tight_constraint": constraint_1_results,
},
)
)
Creating a New Trial
Having done the analysis and candidate generation for three different optimization
configs, we can easily create a new BatchTrial which combines the candidates from
these three different optimizations. Each set of candidates looks promising -- the point
estimates are higher along both metric values than in the previous batch. However, there
is still a good bit of uncertainty in our predictions. It is hard to choose between the
different constraint settings without reducing this noise, so we choose to run a new
trial with all three constraint settings. However, we're generally convinced that the
tight constraint is too conservative. We'd still like to reduce our uncertainty in that
region, but we'll only take one arm from that set.
# We can add entire generator runs, when constructing a new trial.
trial = (
experiment.new_batch_trial()
.add_generator_run(unconstrained)
.add_generator_run(constraint_5_results)
)
# Or, we can hand-pick arms.
trial.add_arm(constraint_1_results.arms[0])
The arms are combined into a single trial, along with the status_quo arm. Their
generator can be accessed from the trial as well.
experiment.trials[1].arms
The original GeneratorRuns can be accessed from within the trial as well. This is
useful for later analyses, allowing introspection of the OptimizationConfig used for
generation (as well as other information, e.g. SearchSpace used for generation).
experiment.trials[1]._generator_run_structs
Here, we can see the unconstrained set-up used for our first set of candidates.
experiment.trials[1]._generator_run_structs[0].generator_run.optimization_config