[Bug] Multi-Fidelity BO with Discrete Fidelities using KG doesn't work on new version

[rimaabaidi]

🐛 Bug

Hello, I updated Botorch to the newest version (0.5.0). Before that I was using v0.4.0. The tutorial "Multi-Fidelity BO with Discrete Fidelities using KG" used to work fine with it (if q=4 then 4 candidates are evaluated in parallel), and accordingly the function I am interested in optimization, regardless of others numerical instabilities when the batch-size is too large. But now, with the new version, I get an error when I run the tutorial:

Traceback (most recent call last):
  File "/home/abaidi/Desktop/Inverse_Stefan_Problem/Code/gpinformationfusion/tutorials/hartmann.py", line 135, in <module>
    new_x, new_obj, cost = optimize_mfkg_and_get_observation(mfkg_acqf)
  File "/home/abaidi/Desktop/Inverse_Stefan_Problem/Code/gpinformationfusion/tutorials/hartmann.py", line 109, in optimize_mfkg_and_get_observation
    candidates, _ = optimize_acqf_mixed(
  File "/home/abaidi/miniconda3/envs/gpif/lib/python3.8/site-packages/botorch/optim/optimize.py", line 512, in optimize_acqf_mixed
    acq_value = acq_function(candidates)  # compute joint acquisition value
  File "/home/abaidi/miniconda3/envs/gpif/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/abaidi/miniconda3/envs/gpif/lib/python3.8/site-packages/botorch/utils/transforms.py", line 220, in decorated
    output = method(acqf, X, *args, **kwargs)
  File "/home/abaidi/miniconda3/envs/gpif/lib/python3.8/site-packages/botorch/acquisition/knowledge_gradient.py", line 416, in forward
    X_actual, X_fantasies = _split_fantasy_points(X=X, n_f=self.num_fantasies)
  File "/home/abaidi/miniconda3/envs/gpif/lib/python3.8/site-packages/botorch/acquisition/knowledge_gradient.py", line 527, in _split_fantasy_points
    raise ValueError(
ValueError: n_f (128) must be less than the q-batch dimension of X (16)

So, I change the num-fantasy to 15 just to try it out. It works, but now the number of candidates in one trial is 16 as you can see here:

candidates: tensor([[0.518, 0.843, 0.053, 0.632, 0.248, 0.113, 0.750], [0.544, 0.916, 0.025, 0.594, 0.347, 0.029, 0.750], [0.368, 0.829, 0.063, 0.659, 0.293, 0.000, 0.750], [0.473, 0.772, 0.183, 0.487, 0.340, 0.010, 0.750], [0.415, 0.913, 0.123, 0.638, 0.328, 0.110, 0.500], [0.573, 0.829, 0.072, 0.588, 0.384, 0.179, 0.500], [0.605, 0.804, 0.000, 0.568, 0.356, 0.003, 0.500], [0.475, 0.827, 0.000, 0.472, 0.320, 0.160, 0.500], [0.558, 0.859, 0.230, 0.676, 0.367, 0.040, 0.500], [0.479, 0.833, 0.213, 0.693, 0.321, 0.000, 0.500], [0.466, 0.876, 0.173, 0.573, 0.337, 0.000, 0.500], [0.574, 0.920, 0.123, 0.490, 0.325, 0.118, 0.500], [0.551, 0.855, 0.319, 0.613, 0.452, 0.016, 0.500], [0.400, 0.845, 0.130, 0.664, 0.379, 0.034, 0.500], [0.448, 0.900, 0.027, 0.631, 0.257, 0.077, 0.500], [0.533, 0.841, 0.081, 0.485, 0.372, 0.000, 0.500]], dtype=torch.float64)

To reproduce

import os
import torch

tkwargs = {
    "dtype": torch.double,
    "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
}
SMOKE_TEST = os.environ.get("SMOKE_TEST")
from botorch.test_functions.multi_fidelity import AugmentedHartmann

problem = AugmentedHartmann(negate=True).to(**tkwargs)
fidelities = torch.tensor([0.5, 0.75, 1.0], **tkwargs)
from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
from botorch.models.transforms.outcome import Standardize
from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
from botorch.utils.transforms import unnormalize, standardize
from botorch.utils.sampling import draw_sobol_samples

def generate_initial_data(n=16):
    # generate training data
    train_x = torch.rand(n, 6, **tkwargs)
    train_f = fidelities[torch.randint(3, (n, 1))]
    train_x_full = torch.cat((train_x, train_f), dim=1)
    train_obj = problem(train_x_full).unsqueeze(-1)  # add output dimension
    return train_x_full, train_obj

def initialize_model(train_x, train_obj):
    # define a surrogate model suited for a "training data"-like fidelity parameter
    # in dimension 6, as in [2]
    model = SingleTaskMultiFidelityGP(
        train_x,
        train_obj,
        outcome_transform=Standardize(m=1),
        data_fidelity=6
    )
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    return mll, model

from botorch import fit_gpytorch_model
from botorch.models.cost import AffineFidelityCostModel
from botorch.acquisition.cost_aware import InverseCostWeightedUtility
from botorch.acquisition import PosteriorMean
from botorch.acquisition.knowledge_gradient import qMultiFidelityKnowledgeGradient
from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction
from botorch.optim.optimize import optimize_acqf
from botorch.acquisition.utils import project_to_target_fidelity

bounds = torch.tensor([[0.0] * problem.dim, [1.0] * problem.dim], **tkwargs)
target_fidelities = {6: 1.0}

cost_model = AffineFidelityCostModel(fidelity_weights={6: 1.0}, fixed_cost=5.0)
cost_aware_utility = InverseCostWeightedUtility(cost_model=cost_model)

def project(X):
    return project_to_target_fidelity(X=X, target_fidelities=target_fidelities)

def get_mfkg(model):
    curr_val_acqf = FixedFeatureAcquisitionFunction(
        acq_function=PosteriorMean(model),
        d=7,
        columns=[6],
        values=[1],
    )

    _, current_value = optimize_acqf(
        acq_function=curr_val_acqf,
        bounds=bounds[:, :-1],
        q=1,
        num_restarts=10 if not SMOKE_TEST else 2,
        raw_samples=1024 if not SMOKE_TEST else 4,
        options={"batch_limit": 10, "maxiter": 200},
    )

    return qMultiFidelityKnowledgeGradient(
        model=model,
        num_fantasies=128 if not SMOKE_TEST else 2,
        current_value=current_value,
        cost_aware_utility=cost_aware_utility,
        project=project,
    )


from botorch.optim.initializers import gen_one_shot_kg_initial_conditions
from botorch.optim.optimize import optimize_acqf_mixed

torch.set_printoptions(precision=3, sci_mode=False)

NUM_RESTARTS = 10 if not SMOKE_TEST else 2
RAW_SAMPLES = 512 if not SMOKE_TEST else 4

def optimize_mfkg_and_get_observation(mfkg_acqf):
    """Optimizes MFKG and returns a new candidate, observation, and cost."""

    X_init = gen_one_shot_kg_initial_conditions(
        acq_function=mfkg_acqf,
        bounds=bounds,
        q=4,
        num_restarts=10,
        raw_samples=512,
    )
    candidates, _ = optimize_acqf_mixed(
        acq_function=mfkg_acqf,
        bounds=bounds,
        fixed_features_list=[{6: 0.5}, {6: 0.75}, {6: 1.0}],
        q=4,
        num_restarts=NUM_RESTARTS,
        raw_samples=RAW_SAMPLES,
        batch_initial_conditions=X_init,
        options={"batch_limit": 5, "maxiter": 200},
    )
    # observe new values
    cost = cost_model(candidates).sum()
    new_x = candidates.detach()
    new_obj = problem(new_x).unsqueeze(-1)
    print(f"candidates:\n{new_x}\n")
    print(f"observations:\n{new_obj}\n\n")
    return new_x, new_obj, cost

train_x, train_obj = generate_initial_data(n=16)
cumulative_cost = 0.0
N_ITER = 3 if not SMOKE_TEST else 1

for _ in range(N_ITER):
    mll, model = initialize_model(train_x, train_obj)
    fit_gpytorch_model(mll)
    mfkg_acqf = get_mfkg(model)
    new_x, new_obj, cost = optimize_mfkg_and_get_observation(mfkg_acqf)
    train_x = torch.cat([train_x, new_x])
    train_obj = torch.cat([train_obj, new_obj])
    cumulative_cost += cost

System information

BoTorch Version (0.5.0) GPyTorch Version (1.5.0) PyTorch Version (1.8.1) Ubuntu

Should I just get back to the older version?

[Balandat]

Thanks for flagging this. I think I know what's going on, it looks like this issue was introduced in #804. Let me see how hard it is to fix this.

[esantorella]

I am not able to reproduce this (a year and a half later).