Botorch model_bridge changes after calling evaluate_acquisition function

[amgonzalezd]

Hi, I came up with some strange behavior using the service API and the model_bridge.evaluate_acquisition_function. I'm not sure if this behavior is normal and I'm missing something or if something is happening here. I'd like to use the model_bridge to keep track on the evaluation of the acquisition function. Here is the code I'm using:

test_function = AdaptedBranin()

def evaluation_function(parameters):
    x = np.array([parameters.get(f"x{i+1}") for i in range(test_function.dim)])
    tf_botorch = tf_from_botorch(botorch_synthetic_function=test_function)
    return {"AdaptedBranin": (tf_botorch(x), None)}

generation_strategy = GenerationStrategy(
    name='SOBOL+BotorchModularDefaults',
    steps=[
        GenerationStep(
            model=Models.SOBOL,
            num_trials=5
        ),
        GenerationStep(
            model=Models.BOTORCH_MODULAR,
            num_trials=-1
        )
    ]
)

ax_client = AxClient(random_seed=42, generation_strategy=generation_strategy)
ax_client.create_experiment(
    name="adapted_branin",
    parameters=[
        {"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
        {"name": "x2", "type": "range", "bounds": [0.0, 15.0]}
    ],
    objectives={"AdaptedBranin": ObjectiveProperties(minimize=True)}
)

for i in range(20):
    parameters, trial_index = ax_client.get_next_trial()
    model_bridge = ax_client.generation_strategy.model
    
    ax_client.complete_trial(trial_index=trial_index, raw_data=evaluation_function(parameters))

After training, I use the last model_bridge to plot a slice of the model

render(plot_slice(model_bridge, "x2", "AdaptedBranin"))

This looks fine. Afterwards if I call the evaluate_acquisition_function() on the model and plot it again, I get the following plot.

obs_feat = ObservationFeatures(parameters=parameters) # this are the last parameters of the loop
value = model_bridge.evaluate_acquisition_function([obs_feat])

render(plot_slice(model_bridge, "x2", "AdaptedBranin"))

Why does calling evaluate_acquisition_function() is changing the behavior of my model? It seems that the model is not transforming backwards since the x-axis is now between 0-1 and not 0-15.

[Balandat]

Hmm yeah this does not seem right; it does appear that there is something going on with the transforms being modified in some way that they shouldn't be.

[smao-astro]

Hi,

I faced a similar issue when I using Service API with GenerationStrategy. The plots generated with plot_slice are different before and after the following line is executed

acq_values = ax_client.generation_strategy.model.evaluate_acquisition_function(
        observation_features=[ax.core.observation.ObservationFeatures({"x": x}) for x in xs],
        # search_space=ax_client.experiment.search_space,
    )

As you have said, this looks not correct.

In addition, what I found is that the plot won't be changed unexpectedly if I specify the parameter search_space,

acq_values = ax_client.generation_strategy.model.evaluate_acquisition_function(
        observation_features=[ax.core.observation.ObservationFeatures({"x": x}) for x in xs],
        search_space=ax_client.experiment.search_space,
    )

However, I am not 100% sure whether this is the correct way to do this. So this is just a walkaround to prevent the slice plot being changed. Therefore, I want to report this and hopefully, this would be helpful for you guys to track down the bug.

[saitcakmak]

The issue does not reproduce anymore. I am guessing it has been resolved as part of other efforts. Feel free to open a new issue if you run into similar problems!