Ax: Abandoned arms and trials are not excluded from modeling data

[XRen615]

hey guys, thanks for offering such a great lib!

In case I found the data of an already completed arm is invalid, how can I remove it to make sure the model prediction is not effected by this invalid arm?

I tried to abandon this arm like
experiment.trials[0].mark_arm_abandoned(arm_name='0_0')
but seems the abandoned arm is still taken into account in model fit process.

Thanks in advance :)

[ldworkin]

Hi @XRen615 ! Marking the arm abandoned is indeed the right way to do this. When you say "it seems the abandoned arm is still taken into account in model fit process", what makes you think that?

[XRen615]

Hi @XRen615 ! Marking the arm abandoned is indeed the right way to do this. When you say "it seems the abandoned arm is still taken into account in model fit process", what makes you think that?

well, I made up a case like this

def fetch_trial_data(trial):
    records = []
    for arm_name, arm in trial.arms_by_name.items():
        params = arm.parameters
        records.append({
            "arm_name": arm_name,
            "metric_name": 'obj',
            # to make arm 0_0 an outlier
            "mean": 9999999 if arm_name == '0_0' else params["x1"] ** 2 + params["x2"] ** 2 + params["x3"] ** 2,
            "sem": None,
            "trial_index": trial.index,
        })
    return Data(df=pd.DataFrame.from_records(records))

As every x_i is defined between (0,10), manually set arm 0_0 to be 9999999 will make it a sufficient large outlier.

Later, after attach_data, I abandoned this arm 0_0 by
experiment.trials[0].mark_arm_abandoned(arm_name='0_0')

Finally, I plot the contour by
render(interact_contour(model=model, metric_name='obj'))
and get a figure like this

apparently the large outlier I abandoned affected the prediction, otherwise I was expecting something like below as the metrics is actually a sphere

Maybe I misused something?

Thanks!

[Balandat]

Maybe I misused something?

I don't think so, it seems that we don't actually properly filter out abandoned arms if they have data associated with them:
https://github.com/facebook/Ax/blob/master/ax/modelbridge/base.py#L136-L140

We should add a check that filters out data associated with abandoned arms.

Ideally we'd do something smarter about handling systematic outlier observations, but doing this generically for a broad range of problems without requiring domain knowledge is tough. For now ignoring abandoned arms is probably the best bet in the short term.

[lena-kashtelyan]

@XRen615, we'd like to understand your usage pattern a little bit better to ensure we fully support it (and fix the filtering of abandoned arms).

A few questions that would help us understand your use case better:
1) In the code snipped you provided (defining fetch_trial_data), what is the context for the function –– is the it fetch_trial_data method of a Metric object?
2) How are you creating the model object in your code snippet? What data are you passing to the model (e.g., are you using experiment.fetch_data)?
3) Are you using Trial-s or BatchTrial-s in your experiment?

[XRen615]

@XRen615, we'd like to understand your usage pattern a little bit better to ensure we fully support it (and fix the filtering of abandoned arms).

A few questions that would help us understand your use case better:

1. In the code snipped you provided (defining fetch_trial_data), what is the context for the function –– is the it fetch_trial_data method of a Metric object?
2. How are you creating the model object in your code snippet? What data are you passing to the model (e.g., are you using experiment.fetch_data)?
3. Are you using Trial-s or BatchTrial-s in your experiment?

Thanks guys!

• I use the function fetch_trial_data to produce the Data object that been attach to the experiment later on

trial_data = fetch_trial_data(experiment.trials[0])
experiment.attach_data(trial_data)

• I create model object with Data from methodexperiment.fetch_data(), I think that's where the problem lies.
  Models.BOTORCH(experiment=experiment, data=experiment.fetch_data())
• I'm using BatchTrial in this case.

I attach the whole script below which can reproduce the problem for your reference

import pandas as pd
from ax import *
from ax.storage.metric_registry import register_metric
from ax.storage.runner_registry import register_runner
import numpy as np


class MyRunner(Runner):
    def run(self, trial):
        """
        push to production logic
        async return insert to here
        """
        print('push trial to production...')
        arms = []
        for arm in trial.arms:
            arms.append({'name': arm.name, 'parameters': arm.parameters})
        print(arms)
        return {"name": str(trial.index)}


def _construct_exp(search_space_list, objective_dict, constraint_list):
    # construct search space
    parameters = []
    for params_info in search_space_list:
        type = params_info['type']
        if type in ['INT', 'FLOAT']:
            parameters.append(
                RangeParameter(name=params_info['name'], lower=params_info['lower'], upper=params_info['upper'],
                               parameter_type=ParameterType.__getattr__(params_info['type'])))
        else:
            parameters.append(
                ChoiceParameter(name=params_info['name'], values=params_info['values'],
                                parameter_type=ParameterType.__getattr__(params_info['type'])))
    search_space = SearchSpace(parameters)

    # construct optimization config = objective + constraints
    objective = Objective(
        metric=Metric(name=objective_dict['name']),
        minimize=objective_dict['minimize'],
    )

    outcome_constraints = []
    for constraint_info in constraint_list:
        outcome_constraints.append(OutcomeConstraint(
            metric=Metric(constraint_info['name']),
            op=ComparisonOp.__getattr__(constraint_info['op']),
            bound=constraint_info['bound'],
            relative=False
        ))

    optimization_config = OptimizationConfig(
        objective=objective,
        outcome_constraints=outcome_constraints
    )

    # construct experiment objective
    experiment = Experiment(
        name="foo",
        search_space=search_space,
        optimization_config=optimization_config,
        runner=MyRunner()
    )
    register_runner(MyRunner)

    return experiment


def fetch_trial_data(trial):
    """
    replace with reporting logic in real case
    """
    records = []
    for arm_name, arm in trial.arms_by_name.items():
        params = arm.parameters

        records.append({
            "arm_name": arm_name,
            "metric_name": 'obj',
            # to make arm 0_0 an outlier
            "mean": 9999999 if arm_name == '0_0' else params["x1"] ** 2 + params["x2"] ** 2 + params["x3"] ** 2,
            "sem": None,
            "trial_index": trial.index,
        })
        records.append({
            "arm_name": arm_name,
            "metric_name": 'constraint_1',
            # to make arm 0_0 an outlier
            "mean": 9999999 if arm_name == '0_0' else params["x1"] + params["x2"] + params["x3"],
            "sem": None,
            "trial_index": trial.index,
        })
    return Data(df=pd.DataFrame.from_records(records))


# get from backend
search_space_list = [{'name': 'x1', 'type': 'FLOAT', 'lower': 0.0, 'upper': 10.0},
                     {'name': 'x2', 'type': 'INT', 'lower': 0, 'upper': 10},
                     {'name': 'x3', 'type': 'FLOAT', 'lower': 0.0, 'upper': 10.0}]

objective_dict = {'name': 'obj', 'minimize': False}

constraint_list = [{'name': 'constraint_1', 'op': 'LEQ', 'bound': 6}]

# construct a experiment
experiment = _construct_exp(search_space_list=search_space_list, objective_dict=objective_dict,
                            constraint_list=constraint_list)
# cold start
# insert user defined start position(s)
user_defined_starts = [{'name': 'udp_1', 'parameters': {'x1': 1.0, 'x2': 1, 'x3': 1.0}}]

sobol_num = 3 - len(user_defined_starts)
sobol = Models.SOBOL(search_space=experiment.search_space)
generator_run = sobol.gen(sobol_num)
experiment.new_batch_trial(generator_run=generator_run)

# add user inserted points
for user_defined_start in user_defined_starts:
    experiment.trials[0].add_arm(Arm(name=user_defined_start['name'], parameters=user_defined_start['parameters']))

experiment.trials[0].run()

# report metrics
trial_data = fetch_trial_data(experiment.trials[0])
experiment.attach_data(trial_data)
experiment.trials[0].mark_completed()

print(experiment.fetch_data().df)

# update search space
range_param1 = RangeParameter(name="x1", lower=1.0, upper=10.0, parameter_type=ParameterType.FLOAT)
range_param2 = RangeParameter(name="x2", lower=1.0, upper=10.0, parameter_type=ParameterType.INT)
range_param3 = RangeParameter(name="x3", lower=1.0, upper=10.0, parameter_type=ParameterType.FLOAT)

experiment.search_space = SearchSpace(
    parameters=[range_param1, range_param2, range_param3],
)
# abandon trial
experiment.trials[0].mark_arm_abandoned(arm_name='0_0')

save(experiment, "./foo.json")

# start optimization loop
ROUND = 2
for _ in range(1, ROUND + 1):
    print('now start round {_}'.format(_=_))

    # load experiment from json
    experiment = load("./foo.json")

    gpei = Models.BOTORCH(experiment=experiment, data=experiment.fetch_data())
    generator_run = gpei.gen(2)
    experiment.new_batch_trial(generator_run=generator_run)
    experiment.trials[_].run()

    trial_data = fetch_trial_data(experiment.trials[_])
    experiment.attach_data(trial_data)
    experiment.trials[_].mark_completed()

    print('************')
    save(experiment, "foo.json")

trial_df = experiment.fetch_data().df
print(trial_df)

# visualizations
from ax.plot.slice import plot_slice
from ax.utils.notebook.plotting import render
from ax.modelbridge.cross_validation import cross_validate
from ax.modelbridge.registry import Models
from ax.plot.contour import interact_contour
from ax.plot.diagnostic import interact_cross_validation
from ax.plot.scatter import (
    plot_objective_vs_constraints,
)

experiment = load("./foo.json")

model = Models.BOTORCH(experiment=experiment, data=experiment.fetch_data())
# slice
render(plot_slice(model, "x1", "obj"))

# interact contour
render(interact_contour(model=model, metric_name='obj'))
render(interact_contour(model=model, metric_name='constraint_1'))

# trade-off
render(plot_objective_vs_constraints(model, 'obj', rel=False))

# cross-validation
cv_results = cross_validate(model, folds=-1)  # folds -1 = leave-one-out, arm level
render(interact_cross_validation(cv_results))

[lena-kashtelyan]

@XRen615, thank you very much, this is super helpful! A few things:

1) You are exactly right about the problem coming from the fact that you pass data containing the abandoned arms to Models.BOTORCH –– the model works with whatever data you pass into it, and that is actually by design, as we wanted to keep models and experiments 'decoupled' from each other. Usually it's systems within Ax that are responsible for data-fetching (either via Metric-s implementations in our Developer API or via built-in utilities in Service and Loop APIs), so we get to filter the data however we need to before passing it to the model, which includes filtering out data for abandoned arms or trials.

2) Given that in your setup, you are controlling what data you are attaching to the experiment and passing to the model, can you just not include the outlier arms into the data you are attaching to the model? That should be an easy workaround for you for now.

3) What is the reason why you are using BatchTrial-s for your experiment? In general, batch trials have a pretty specific purpose, see this paragraph in the docs: https://ax.dev/docs/core.html#trial-vs-batch-trial. If you don't in fact need batch trials, you could use our Service API, which is more user-friendly (you can still evaluate multiple arms in parallel, they will just each be their own trial). Tutorial lives here: https://ax.dev/tutorials/gpei_hartmann_service.html.

4) With 1) said, there is currently a bug with the filtering out of abandoned arms' data currently, which we will investigate and fix : ) Currently abandoned arms or trials are not filtered out from the data passed to model, so you are likely best off with my suggestion in 2) while we fix the filtering. Note that if you were to go with the Service API, you still could apply that suggestion –– just don't pass the data for the arms/tirials you would like to abandon to ax_client.complete_trial.

Let me know if this makes sense and whether the suggested workaround helps!

[XRen615]

@lena-kashtelyan thanks for the help!

• As for the usage of BatchTrial: in my case the evaluation is done via an external online A/B system, on which multiple slots may be available so I can make parallelism to save some time. e.g. There is an experiment with 3 variants (slots) and I decide to update every 1 hour, then in each round I will generate a BatchTrial with 3 arms, deploy them, after 1 hour I query the metrics and complete the them in the same time. The arm evaluations are practically simultaneous but independent —— is this a correct use case of BatchTrial? For me personally I didn't find much difference between BatchTrial with multiple arms and multiple Trial with one arm in this scenario.

• As for the Service API: I tried it, it's indeed easier to use but seems lack of some agility I need. e.g. In some cases I may want a parallelism more than ax_client.get_max_parallelism() could provide and/or I want to manually insert some points to explore in the next round. I didn't find a proper way to do those in Service API paradigm so I tried the script I pasted above.

• For the quick fix, I simply added a function to manually filtered out abandoned arms in Data object passed to model like below:

def _filter_abandonded_arms(experiment):
    """filter out abandoned arm(s) then reconstruct the data"""
    # name of abandoned arms
    abandonded_arms = []
    for trial_index, trial in experiment.trials.items():
        abandonded_arms += [arm.name for arm in trial.abandoned_arms]
    df = experiment.fetch_data().df
    # filter out arms in abandonded_arms
    df_filtered = df[~(df.arm_name.isin(abandonded_arms))]
    # reconstruct Data and return
    return Data(df_filtered)

data = _filter_abandonded_arms(experiment) if experiment.num_abandoned_arms else experiment.fetch_data() 
gpei = Models.BOTORCH(experiment=experiment, data=data)

I hope this will work?

Thanks!

[lena-kashtelyan]

1) Regarding BatchTrial usage: 1-arm Trial-s would be correct for your case, if the arm evaluations are fully independent (and if you don't need them running in a batch so they are all deployed at the exact same time together, if there is non-stationarity in the data).
2) For the Service API: the parallelism restrictions can actually be bypassed, see this section of the Service API tutorial. For manually inserting points, you can use ax_client.attach_trial (discussed in this section) : ) Obviously you don't have to use the Service API, just letting you know that you can and that it might be more convenient, since Ax is a complex system and Service API is harder to misuse.
3) Your fix sounds exactly right and should work –– let me know how it goes for you!

[lena-kashtelyan]

The fix for this is now on master and will be included in the next stable version release we do (should be within a week or two).