0.2.1

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:

  1. Constrained optimization
  2. Interfaces for easily modifying optimization goals
  3. 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.

In [1]:
from ax import Data, Metric, OptimizationConfig, Objective, OutcomeConstraint, ComparisonOp, json_load
from ax.modelbridge.cross_validation import cross_validate
from ax.modelbridge.factory import get_GPEI
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

init_notebook_plotting()

[INFO 08-10 23:52:32] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
In [2]:
experiment = json_load.load_experiment('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).

In [3]:
experiment.trials[0]
Out[3]:
BatchTrial(experiment_name='human_in_the_loop_tutorial', index=0, status=TrialStatus.COMPLETED)
In [4]:
experiment.trials[0].time_created
Out[4]:
datetime.datetime(2019, 3, 29, 18, 10, 6)
In [5]:
# Number of arms in first experiment, including status_quo
len(experiment.trials[0].arms)
Out[5]:
65
In [6]:
# Sample arm configuration
experiment.trials[0].arms[0]
Out[6]:
Arm(name='0_0', parameters={'x_excellent': 0.9715802669525146, 'x_good': 0.8615524768829346, 'x_moderate': 0.7668091654777527, 'x_poor': 0.34871453046798706, 'x_unknown': 0.7675797343254089, 'y_excellent': 2.900710028409958, 'y_good': 1.5137152910232545, 'y_moderate': 0.6775947093963622, 'y_poor': 0.4974367544054985, 'y_unknown': 1.0852564811706542, 'z_excellent': 517803.49761247635, 'z_good': 607874.5171427727, 'z_moderate': 1151881.2023103237, 'z_poor': 2927449.2621421814, 'z_unknown': 2068407.6935052872})

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.

In [7]:
experiment.status_quo
Out[7]:
Arm(name='status_quo', parameters={'x_excellent': 0, 'x_good': 0, 'x_moderate': 0, 'x_poor': 0, 'x_unknown': 0, 'y_excellent': 1, 'y_good': 1, 'y_moderate': 1, 'y_poor': 1, 'y_unknown': 1, 'z_excellent': 1000000, 'z_good': 1000000, 'z_moderate': 1000000, 'z_poor': 1000000, 'z_unknown': 1000000})
In [8]:
objective_metric = Metric(name="metric_1")
constraint_metric = Metric(name="metric_2")

experiment.optimization_config = OptimizationConfig(
    objective=Objective(objective_metric),
    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.

In [9]:
data = Data(pd.read_json('hitl_data.json'))
data.df.head()
Out[9]:
arm_name metric_name mean sem trial_index start_time end_time n
0 0_1 metric_1 495.763048 2.621641 0 2019-03-30 2019-04-03 1599994
1 0_23 metric_1 524.367712 2.731647 0 2019-03-30 2019-04-03 1596356
2 0_14 metric_2 21.460244 0.069457 0 2019-03-30 2019-04-03 1600182
3 0_53 metric_2 21.437433 0.069941 0 2019-03-30 2019-04-03 1601081
4 0_53 metric_1 548.387691 2.893486 0 2019-03-30 2019-04-03 1601081
In [10]:
data.df['arm_name'].unique()
Out[10]:
array(['0_1', '0_23', '0_14', '0_53', '0_0', '0_54', '0_55', '0_56',
       '0_27', '0_57', '0_58', '0_13', '0_59', '0_6', '0_60', '0_61',
       '0_62', '0_63', '0_7', '0_28', '0_15', '0_16', '0_17', '0_18',
       '0_19', '0_29', '0_2', '0_20', '0_21', '0_22', '0_3', '0_30',
       '0_8', '0_10', '0_31', '0_24', '0_32', '0_33', '0_34', '0_35',
       '0_36', '0_37', '0_38', '0_9', '0_39', '0_4', '0_25', '0_11',
       '0_40', '0_41', '0_42', '0_43', '0_44', '0_45', 'status_quo',
       '0_46', '0_47', '0_48', '0_26', '0_49', '0_12', '0_5', '0_50',
       '0_51', '0_52'], dtype=object)
In [11]:
data.df['metric_name'].unique()
Out[11]:
array(['metric_1', 'metric_2'], dtype=object)

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.

In [12]:
experiment.search_space.parameters
Out[12]:
{'x_excellent': RangeParameter(name='x_excellent', parameter_type=FLOAT, range=[0.0, 1.0]),
 'x_good': RangeParameter(name='x_good', parameter_type=FLOAT, range=[0.0, 1.0]),
 'x_moderate': RangeParameter(name='x_moderate', parameter_type=FLOAT, range=[0.0, 1.0]),
 'x_poor': RangeParameter(name='x_poor', parameter_type=FLOAT, range=[0.0, 1.0]),
 'x_unknown': RangeParameter(name='x_unknown', parameter_type=FLOAT, range=[0.0, 1.0]),
 'y_excellent': RangeParameter(name='y_excellent', parameter_type=FLOAT, range=[0.1, 3.0]),
 'y_good': RangeParameter(name='y_good', parameter_type=FLOAT, range=[0.1, 3.0]),
 'y_moderate': RangeParameter(name='y_moderate', parameter_type=FLOAT, range=[0.1, 3.0]),
 'y_poor': RangeParameter(name='y_poor', parameter_type=FLOAT, range=[0.1, 3.0]),
 'y_unknown': RangeParameter(name='y_unknown', parameter_type=FLOAT, range=[0.1, 3.0]),
 'z_excellent': RangeParameter(name='z_excellent', parameter_type=FLOAT, range=[50000.0, 5000000.0]),
 'z_good': RangeParameter(name='z_good', parameter_type=FLOAT, range=[50000.0, 5000000.0]),
 'z_moderate': RangeParameter(name='z_moderate', parameter_type=FLOAT, range=[50000.0, 5000000.0]),
 'z_poor': RangeParameter(name='z_poor', parameter_type=FLOAT, range=[50000.0, 5000000.0]),
 'z_unknown': RangeParameter(name='z_unknown', parameter_type=FLOAT, range=[50000.0, 5000000.0])}
In [13]:
experiment.search_space.parameter_constraints
Out[13]:
[OrderConstraint(x_poor <= x_moderate),
 OrderConstraint(x_moderate <= x_good),
 OrderConstraint(x_good <= x_excellent),
 OrderConstraint(y_poor <= y_moderate),
 OrderConstraint(y_moderate <= y_good),
 OrderConstraint(y_good <= y_excellent)]

Model Fit

Fitting BoTorch's GPEI will allow us to predict new candidates based on our first Sobol batch. Here, we make use of the default settings for GP-EI defined in the ModelBridge factory.

In [14]:
gp = get_GPEI(
    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:

  1. Increase the amount of traffic in this experiment, to reduce the measurement noise.
  2. Increase the number of points run in the random batch, to assist the GP in covering the space.
  3. 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.

In [15]:
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.

In [16]:
render(tile_fitted(gp, rel=True))
In [17]:
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.

In [18]:
unconstrained = gp.gen(
    n=3,
    optimization_config=OptimizationConfig(
        objective=Objective(objective_metric),
    )
)

Let's plot the tradeoffs again, but with our new arms.

In [19]:
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.

In [20]:
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),
        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.

In [21]:
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%.

In [22]:
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),
        outcome_constraints=[constraint_1],
    )
)
In [23]:
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.

In [24]:
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.

In [25]:
# 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])
Out[25]:
BatchTrial(experiment_name='human_in_the_loop_tutorial', index=1, status=TrialStatus.CANDIDATE)

The arms are combined into a single trial, along with the status_quo arm. Their generator can be accessed from the trial as well.

In [26]:
experiment.trials[1].arms
Out[26]:
[Arm(name='1_0', parameters={'x_excellent': 0.6796566123079969, 'x_good': 0.30962911757954, 'x_moderate': 0.22702905229621942, 'x_poor': 0.10294187920100518, 'x_unknown': 0.7553312216512191, 'y_excellent': 2.5545337349582433, 'y_good': 1.7638011725543632, 'y_moderate': 0.8273623333681144, 'y_poor': 0.29717809340106716, 'y_unknown': 1.9016630932270007, 'z_excellent': 4808947.071517051, 'z_good': 3662812.9173141825, 'z_moderate': 2520117.0827065865, 'z_poor': 1423024.1898883623, 'z_unknown': 2420216.2405132814}),
 Arm(name='1_1', parameters={'x_excellent': 0.7235540200089434, 'x_good': 0.3708887245851172, 'x_moderate': 0.32222151425542855, 'x_poor': 0.12534663432399573, 'x_unknown': 0.567433705987285, 'y_excellent': 2.592113018564256, 'y_good': 1.9637059835472426, 'y_moderate': 1.0396611971693528, 'y_poor': 0.31257602169040466, 'y_unknown': 1.1489585406758143, 'z_excellent': 4907780.741431984, 'z_good': 2958948.6428023507, 'z_moderate': 3030674.263508936, 'z_poor': 2546312.040983533, 'z_unknown': 1684560.421822009}),
 Arm(name='1_2', parameters={'x_excellent': 0.718849722059773, 'x_good': 0.258081712381127, 'x_moderate': 0.17138807708169676, 'x_poor': 0.10782203085845236, 'x_unknown': 0.9344894804617083, 'y_excellent': 2.4338296739206378, 'y_good': 1.6640353720609986, 'y_moderate': 0.5348532216310654, 'y_poor': 0.3106162972150125, 'y_unknown': 2.2102484146455503, 'z_excellent': 4463290.132798549, 'z_good': 4178304.8754716087, 'z_moderate': 2029768.0925392308, 'z_poor': 708440.7163803072, 'z_unknown': 3203836.189079878}),
 Arm(name='1_3', parameters={'x_excellent': 0.7894922345488506, 'x_good': 0.5349675173677081, 'x_moderate': 0.3275643686804765, 'x_poor': 0.10251142653184422, 'x_unknown': 0.6817951058664702, 'y_excellent': 2.646912633003732, 'y_good': 1.828546799056513, 'y_moderate': 0.9020190930376933, 'y_poor': 0.2957295595388464, 'y_unknown': 1.6879667438656503, 'z_excellent': 4938964.148179267, 'z_good': 3437565.289352808, 'z_moderate': 2987236.880797905, 'z_poor': 1802487.108063311, 'z_unknown': 2409795.8177433996}),
 Arm(name='1_4', parameters={'x_excellent': 0.6658291291191759, 'x_good': 0.5045439054837153, 'x_moderate': 0.42339666241973867, 'x_poor': 0.06398851337926287, 'x_unknown': 0.5159864638659586, 'y_excellent': 2.504015030029951, 'y_good': 1.7992472372147768, 'y_moderate': 0.9808999091085837, 'y_poor': 0.2845679956943306, 'y_unknown': 1.2823295125553815, 'z_excellent': 4287716.305387594, 'z_good': 2715346.8331379113, 'z_moderate': 2704380.3822467397, 'z_poor': 3156798.372122934, 'z_unknown': 1715217.4559332486}),
 Arm(name='1_5', parameters={'x_excellent': 0.811970791182185, 'x_good': 0.48252435398580285, 'x_moderate': 0.394942462237296, 'x_poor': 0.24961361099544127, 'x_unknown': 0.6907134322343044, 'y_excellent': 2.7041152396619803, 'y_good': 2.17308912658401, 'y_moderate': 0.9116818039362982, 'y_poor': 0.5139324282695359, 'y_unknown': 1.6895096341289553, 'z_excellent': 4609067.8683422385, 'z_good': 4223065.508602824, 'z_moderate': 3938936.7496620137, 'z_poor': 2336138.045313484, 'z_unknown': 3410619.2094468265}),
 Arm(name='1_6', parameters={'x_excellent': 0.7906691335321648, 'x_good': 0.6883888475706466, 'x_moderate': 0.4751459521667211, 'x_poor': 0.04663667070797014, 'x_unknown': 0.7264766199033215, 'y_excellent': 2.618491746265104, 'y_good': 1.82172121023409, 'y_moderate': 0.7911401090439719, 'y_poor': 0.13779894751032995, 'y_unknown': 1.395865302300796, 'z_excellent': 4436012.890843508, 'z_good': 2766295.6580096805, 'z_moderate': 1906321.5790297666, 'z_poor': 2838818.9951091483, 'z_unknown': 1581159.068344447})]

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).

In [27]:
experiment.trials[1]._generator_run_structs
Out[27]:
[GeneratorRunStruct(generator_run=GeneratorRun(3 arms, total weight 3.0), weight=1.0),
 GeneratorRunStruct(generator_run=GeneratorRun(3 arms, total weight 3.0), weight=1.0),
 GeneratorRunStruct(generator_run=GeneratorRun(1 arms, total weight 1.0), weight=1.0)]

Here, we can see the unconstrained set-up used for our first set of candidates.

In [28]:
experiment.trials[1]._generator_run_structs[0].generator_run.optimization_config
Out[28]:
OptimizationConfig(objective=Objective(metric_name="metric_1", minimize=False), outcome_constraints=[])
Download Tutorial Archive (.tar.gz): Notebook, Source and Data
Download Tutorial Jupyter Notebook
Download Tutorial Source Code

Total runtime of script: 2 minutes, 17.06 seconds.