Multi-fidelity optimization with KG and service API #475
Description
Hi,
First off, thanks for the great work! I've been trying to run multi-fidelity optimization using the one-shot-KG method from botorch with the service API. In principle, everything seems to run fine. However, I noticed that the optimizer keeps running in the lowest fidelity and does not explore in this parameter.
To see I've this is an issue with my problem I tried to reproduce the MFKG example from the botorch documentation where the fidelity is clearly changed during the optimization, but I observe the same behaviour.
This is the code I ran to reproduce the botorch example.
from ax.service.ax_client import AxClient
from botorch.test_functions.multi_fidelity import AugmentedHartmann
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
import torch
problem = AugmentedHartmann(negate=True)
def objective(parameters):
# x7 is the fidelity
x = torch.tensor([parameters.get(f"x{i+1}") for i in range(7)])
return {'f':(problem(x),0.0)}
gs = GenerationStrategy(
steps=[GenerationStep(model=Models.SOBOL,num_trials = 16),
GenerationStep(model=Models.GPKG,num_trials=-1),
])
ax_client = AxClient(generation_strategy=gs)
ax_client.create_experiment(
name="hartmann_mf_experiment",
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x3",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x4",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x5",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x6",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x7",
"type": "range",
"bounds": [0.0, 1.0],
"is_fidelity": True,
"target_value": 1.
},
],
objective_name="f",
)
# Initial sobol samples
for i in range(16):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=trial_index, raw_data=objective(parameters))
# KGBO
for i in range(6):
q_p, q_t = [], []
# Simulate batches
for q in range(4):
parameters, trial_index = ax_client.get_next_trial()
q_p.append(parameters)
q_t.append(trial_index)
for q in range(4):
pi = q_p[q]
ti = q_t[q]
ax_client.complete_trial(trial_index=ti, raw_data=objective(pi))After the initial samples the fidelity stays at 0 except for two trial where it is very close to zero.
Are there any parameters that need to be specified for the acquisition function or is MFKG not yet supported with the service API?
Thanks!