evaluate_acquisition_function does not work correctly about ax

[sdaulton]

Hi @lsassen, Can you share your code including setting up AxClient and the GenerationStrategy for repro? Thanks!

[lsassen]

Hi @sdaulton. Here is how I set my AxClient:

1. Set up the generation startegy:

gs = GenerationStrategy(
    steps=[
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1,
            model_kwargs={
                "surrogate": Surrogate(FixedNoiseGP),
                "botorch_acqf_class": UpperConfidenceBound,
                "acquisition_options":  {"beta": 1.96**2,}
            },
        ),
    ]
)

2. Set up the experiment

search_space = SearchSpace(parameters=[RangeParameter(name="x", parameter_type=ParameterType.FLOAT, lower=-10*np.pi, upper=10*np.pi)])

optimization_config = OptimizationConfig(objective=Objective(metric=Metric(name="f(x)",lower_is_better=False,),minimize=False,))

# Here, we have a dummy runner that does nothing.
class MyRunner(Runner):
    def run(self, trial):
        trial_metadata = {"name": str(trial.index)}
        return trial_metadata
    
experiment =Experiment(
        name="my first experiment",
        search_space=search_space,
        optimization_config=optimization_config,
        tracking_metrics=None,
        runner=MyRunner(),
        status_quo=None,
        description=None,
        is_test=False,
        experiment_type=None,
        properties=None,
        default_data_type=None,
    )

3. Setup the Runner

ax_client = AxClient(
    generation_strategy=gs,
    verbose_logging=True,
    early_stopping_strategy=None,
)
ax_client._set_experiment(
    experiment=experiment, overwrite_existing_experiment=True
)

I am using ax-platform==0.3.5 version. If you need more input from my side to provide feedback please let me know.
Thanks you

[cheeseheist]

FYI in case it helps in identifying the issue - I am using the developer API, so not using the AxClient at all and have a similar issue. In my BO loop, I'm generating trials one by one. During each iteration, I added a call to evaluate_acquisition_function() for the suggested trial. If I complete the trial, fetch data, and retrain the model every iteration before getting another suggestion, I have no issues. If I try to make more than one trial (and evaluate the acquisition function for each) before fetching the data and retraining, every new trial has all the inputs clamped to an upper bound of 1 (see output messages in screenshot below). It is as if the evaluation of the acquisition function normalizes everything including the bounds and forgets to undo the transform on the parameter bounds or something. From that point on, everything is essentially clamped to 1 until the model is trained fresh again. If I remove the evaluate_acquisition_function() call, the problem goes away. Hope this helps!

[Balandat]

@cheeseheist thanks for this keen observation, this is very helpful. Seems like we were transforming the search space in-place when calling evaluate_acquisition_function(), which then leads to an issue during the next generation step if using the same search space object.

I put up #2386 which I hope will fix this. Would you mind giving this a try to see if this indeed fixes your issue?