Issue with non linear constraints and qMultifidelityKnowledgeGradient

[Adalia456]

I'm currently working on a project that involves multi-fidelity optimization using BoTorch. I'm trying to optimize an acquisition function, but I keep encountering a RuntimeError when I run the optimization loop. Here's the error message:

RuntimeError                              Traceback (most recent call last)
Cell In[46], line 23
     20 mf_acqf = get_mf(model)
     22 # Optimize the acquisition function and get new candidate observations, manually entered after each experiment 
---> 23 new_x,cost = optimize_mf_and_get_observation(mf_acqf)
     24 print(new_x)
     25 new_x_np = new_x.numpy()

Cell In[45], line 136, in optimize_mf_and_get_observation(acq_func)
    134 def optimize_mf_and_get_observation(acq_func):
    135     torch.set_grad_enabled(True)
--> 136     candidates, _ = optimize_acqf_mixed(
    137         acq_function=acq_func,
    138         bounds=bounds,
    139         fixed_features_list=[{10: 0.3}, {10: 1.0}],
    140         q=BATCH_SIZE,
    141         num_restarts=NUM_RESTARTS,
    142         raw_samples=RAW_SAMPLES,
    143         options={"batch_limit": 1, "maxiter": 200},
    144         nonlinear_inequality_constraints = [(constraint,True)],
    145         batch_initial_conditions = batch_initial_conditions)
    147     # Calculate the cost
    148     cost = cost_model(candidates).sum()

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:985, in optimize_acqf_mixed(acq_function, bounds, q, num_restarts, fixed_features_list, raw_samples, options, inequality_constraints, equality_constraints, nonlinear_inequality_constraints, post_processing_func, batch_initial_conditions, ic_generator, ic_gen_kwargs, **kwargs)
    982 candidates = torch.tensor([], device=bounds.device, dtype=bounds.dtype)
    984 for _ in range(q):
--> 985     candidate, acq_value = optimize_acqf_mixed(
    986         acq_function=acq_function,
    987         bounds=bounds,
    988         q=1,
    989         num_restarts=num_restarts,
    990         raw_samples=raw_samples,
    991         fixed_features_list=fixed_features_list,
    992         options=options or {},
    993         inequality_constraints=inequality_constraints,
    994         equality_constraints=equality_constraints,
    995         nonlinear_inequality_constraints=nonlinear_inequality_constraints,
    996         post_processing_func=post_processing_func,
    997         batch_initial_conditions=batch_initial_conditions,
    998         ic_generator=ic_generator,
    999         ic_gen_kwargs=ic_gen_kwargs,
   1000     )
   1001     candidates = torch.cat([candidates, candidate], dim=-2)
   1002     acq_function.set_X_pending(
   1003         torch.cat([base_X_pending, candidates], dim=-2)
   1004         if base_X_pending is not None
   1005         else candidates
   1006     )

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:954, in optimize_acqf_mixed(acq_function, bounds, q, num_restarts, fixed_features_list, raw_samples, options, inequality_constraints, equality_constraints, nonlinear_inequality_constraints, post_processing_func, batch_initial_conditions, ic_generator, ic_gen_kwargs, **kwargs)
    952 ff_candidate_list, ff_acq_value_list = [], []
    953 for fixed_features in fixed_features_list:
--> 954     candidate, acq_value = optimize_acqf(
    955         acq_function=acq_function,
    956         bounds=bounds,
    957         q=q,
    958         num_restarts=num_restarts,
    959         raw_samples=raw_samples,
    960         options=options or {},
    961         inequality_constraints=inequality_constraints,
    962         equality_constraints=equality_constraints,
    963         nonlinear_inequality_constraints=nonlinear_inequality_constraints,
    964         fixed_features=fixed_features,
    965         post_processing_func=post_processing_func,
    966         batch_initial_conditions=batch_initial_conditions,
    967         ic_generator=ic_generator,
    968         return_best_only=True,
    969         **ic_gen_kwargs,
    970     )
    971     ff_candidate_list.append(candidate)
    972     ff_acq_value_list.append(acq_value)

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:567, in optimize_acqf(acq_function, bounds, q, num_restarts, raw_samples, options, inequality_constraints, equality_constraints, nonlinear_inequality_constraints, fixed_features, post_processing_func, batch_initial_conditions, return_best_only, gen_candidates, sequential, ic_generator, timeout_sec, return_full_tree, retry_on_optimization_warning, **ic_gen_kwargs)
    544     gen_candidates = gen_candidates_scipy
    545 opt_acqf_inputs = OptimizeAcqfInputs(
    546     acq_function=acq_function,
    547     bounds=bounds,
   (...)
    565     ic_gen_kwargs=ic_gen_kwargs,
    566 )
--> 567 return _optimize_acqf(opt_acqf_inputs)

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:588, in _optimize_acqf(opt_inputs)
    585     return _optimize_acqf_sequential_q(opt_inputs=opt_inputs)
    587 # Batch optimization (including the case q=1)
--> 588 return _optimize_acqf_batch(opt_inputs=opt_inputs)

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:352, in _optimize_acqf_batch(opt_inputs)
    349         batch_acq_values = torch.cat(batch_acq_values_list).flatten()
    350     return batch_candidates, batch_acq_values, opt_warnings
--> 352 batch_candidates, batch_acq_values, ws = _optimize_batch_candidates()
    354 optimization_warning_raised = any(
    355     (issubclass(w.category, OptimizationWarning) for w in ws)
    356 )
    357 if optimization_warning_raised and opt_inputs.retry_on_optimization_warning:

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\optimize.py:336, in _optimize_acqf_batch.<locals>._optimize_batch_candidates()
    331 with warnings.catch_warnings(record=True) as ws:
    332     warnings.simplefilter("always", category=OptimizationWarning)
    333     (
    334         batch_candidates_curr,
    335         batch_acq_values_curr,
--> 336     ) = opt_inputs.gen_candidates(
    337         batched_ics_, opt_inputs.acq_function, **filtered_gen_kwargs
    338     )
    339 opt_warnings += ws
    340 batch_candidates_list.append(batch_candidates_curr)

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\generation\gen.py:160, in gen_candidates_scipy(initial_conditions, acquisition_function, lower_bounds, upper_bounds, inequality_constraints, equality_constraints, nonlinear_inequality_constraints, options, fixed_features, timeout_sec)
    149 _no_fixed_features = _remove_fixed_features_from_optimization(
    150     fixed_features=fixed_features,
    151     acquisition_function=acquisition_function,
   (...)
    157     nonlinear_inequality_constraints=nonlinear_inequality_constraints,
    158 )
    159 # call the routine with no fixed_features
--> 160 clamped_candidates, batch_acquisition = gen_candidates_scipy(
    161     initial_conditions=_no_fixed_features.initial_conditions,
    162     acquisition_function=_no_fixed_features.acquisition_function,
    163     lower_bounds=_no_fixed_features.lower_bounds,
    164     upper_bounds=_no_fixed_features.upper_bounds,
    165     inequality_constraints=_no_fixed_features.inequality_constraints,
    166     equality_constraints=_no_fixed_features.equality_constraints,
    167     nonlinear_inequality_constraints=_no_fixed_features.nonlinear_inequality_constraints,  # noqa: E501
    168     options=options,
    169     fixed_features=None,
    170     timeout_sec=timeout_sec,
    171 )
    172 clamped_candidates = _no_fixed_features.acquisition_function._construct_X_full(
    173     clamped_candidates
    174 )
    175 return clamped_candidates, batch_acquisition

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\generation\gen.py:252, in gen_candidates_scipy(initial_conditions, acquisition_function, lower_bounds, upper_bounds, inequality_constraints, equality_constraints, nonlinear_inequality_constraints, options, fixed_features, timeout_sec)
    249 def f(x):
    250     return -acquisition_function(x)
--> 252 res = minimize_with_timeout(
    253     fun=f_np_wrapper,
    254     args=(f,),
    255     x0=x0,
    256     method=options.get("method", "SLSQP" if constraints else "L-BFGS-B"),
    257     jac=with_grad,
    258     bounds=bounds,
    259     constraints=constraints,
    260     callback=options.get("callback", None),
    261     options={
    262         k: v
    263         for k, v in options.items()
    264         if k not in ["method", "callback", "with_grad"]
    265     },
    266     timeout_sec=timeout_sec,
    267 )
    268 _process_scipy_result(res=res, options=options)
    270 candidates = fix_features(
    271     X=torch.from_numpy(res.x).to(initial_conditions).reshape(shapeX),
    272     fixed_features=fixed_features,
    273 )

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\utils\timeout.py:82, in minimize_with_timeout(fun, x0, args, method, jac, hess, hessp, bounds, constraints, tol, callback, options, timeout_sec)
     80 try:
     81     warnings.filterwarnings("error", message="Method .* cannot handle")
---> 82     return optimize.minimize(
     83         fun=fun,
     84         x0=x0,
     85         args=args,
     86         method=method,
     87         jac=jac,
     88         hess=hess,
     89         hessp=hessp,
     90         bounds=bounds,
     91         constraints=constraints,
     92         tol=tol,
     93         callback=wrapped_callback,
     94         options=options,
     95     )
     96 except OptimizationTimeoutError as e:
     97     msg = f"Optimization timed out after {e.runtime} seconds."

File ~\anaconda31\envs\myenv\Lib\site-packages\scipy\optimize\_minimize.py:722, in minimize(fun, x0, args, method, jac, hess, hessp, bounds, constraints, tol, callback, options)
    719     res = _minimize_cobyla(fun, x0, args, constraints, callback=callback,
    720                            bounds=bounds, **options)
    721 elif meth == 'slsqp':
--> 722     res = _minimize_slsqp(fun, x0, args, jac, bounds,
    723                           constraints, callback=callback, **options)
    724 elif meth == 'trust-constr':
    725     res = _minimize_trustregion_constr(fun, x0, args, jac, hess, hessp,
    726                                        bounds, constraints,
    727                                        callback=callback, **options)

File ~\anaconda31\envs\myenv\Lib\site-packages\scipy\optimize\_slsqp_py.py:338, in _minimize_slsqp(func, x0, args, jac, bounds, constraints, maxiter, ftol, iprint, disp, eps, callback, finite_diff_rel_step, **unknown_options)
    334 # Set the parameters that SLSQP will need
    335 # meq, mieq: number of equality and inequality constraints
    336 meq = sum(map(len, [atleast_1d(c['fun'](x, *c['args']))
    337           for c in cons['eq']]))
--> 338 mieq = sum(map(len, [atleast_1d(c['fun'](x, *c['args']))
    339            for c in cons['ineq']]))
    340 # m = The total number of constraints
    341 m = meq + mieq

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\parameter_constraints.py:497, in _make_f_and_grad_nonlinear_inequality_constraints.<locals>.f_obj(X)
    495 if X_c is None or not np.array_equal(X_c, X):
    496     cache["X"] = X.copy()
--> 497     cache["obj"], cache["grad"] = f_obj_and_grad(X)
    498 return cache["obj"]

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\optim\parameter_constraints.py:488, in _make_f_and_grad_nonlinear_inequality_constraints.<locals>.f_obj_and_grad(x)
    487 def f_obj_and_grad(x):
--> 488     obj, grad = f_np_wrapper(x, f=nlc)
    489     return obj, grad

File ~\anaconda31\envs\myenv\Lib\site-packages\botorch\generation\gen.py:211, in gen_candidates_scipy.<locals>.f_np_wrapper(x, f)
    209 loss = f(X_fix).sum()
    210 # compute gradient w.r.t. the inputs (does not accumulate in leaves)
--> 211 gradf = _arrayify(torch.autograd.grad(loss, X)[0].contiguous().view(-1))
    212 if np.isnan(gradf).any():
    213     msg = (
    214         f"{np.isnan(gradf).sum()} elements of the {x.size} element "
    215         "gradient array `gradf` are NaN. "
    216         "This often indicates numerical issues."
    217     )

File ~\anaconda31\envs\myenv\Lib\site-packages\torch\autograd\__init__.py:412, in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched, materialize_grads)
    408     result = _vmap_internals._vmap(vjp, 0, 0, allow_none_pass_through=True)(
    409         grad_outputs_
    410     )
    411 else:
--> 412     result = _engine_run_backward(
    413         t_outputs,
    414         grad_outputs_,
    415         retain_graph,
    416         create_graph,
    417         inputs,
    418         allow_unused,
    419         accumulate_grad=False,
    420     )
    421 if materialize_grads:
    422     if any(
    423         result[i] is None and not is_tensor_like(inputs[i])
    424         for i in range(len(inputs))
    425     ):

File ~\anaconda31\envs\myenv\Lib\site-packages\torch\autograd\graph.py:744, in _engine_run_backward(t_outputs, *args, **kwargs)
    742     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    743 try:
--> 744     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    745         t_outputs, *args, **kwargs
    746     )  # Calls into the C++ engine to run the backward pass
    747 finally:
    748     if attach_logging_hooks:

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I tried using torch.set_grad_enabled(True) but it did not work. Here is a snippet of my code to show how I programmed the non linear constraint :

from botorch.optim.optimize import optimize_acqf_mixed
from botorch.optim.optimize import optimize_acqf_cyclic

NUM_RESTARTS = 5
RAW_SAMPLES = 128
BATCH_SIZE = 2

def initial_conditions(NUM_RESTARTS, BATCH_SIZE, bounds):
    """
    Generate initial conditions for Bayesian optimization.
    
    Args:
        NUM_RESTARTS (int): Number of restarts, i.e., the number of different initial conditions.
        BATCH_SIZE (int): The size of each batch.
        bounds (torch.Tensor): A 2x11 tensor specifying the lower and upper bounds for each dimension.
        
    Returns:
        torch.Tensor: A NUM_RESTARTS x BATCH_SIZE x 11 tensor with initial conditions.
    """
    torch.set_grad_enabled(True)
    d = 11  # Dimension of the problem
    lower_bound, upper_bound = bounds[0, :], bounds[1, :]
    
    # Initialize the tensor to store initial conditions
    initial_conditions_tensor = torch.zeros(NUM_RESTARTS, BATCH_SIZE, d)
    
    for i in range(NUM_RESTARTS):
        for j in range(BATCH_SIZE):
            # Randomly select 3 columns from the first 8 columns
            selected_columns = torch.randperm(7)[:3]
            
            # Fill these 3 columns with values within the bounds
            for col in selected_columns:
                initial_conditions_tensor[i, j, col] = torch.FloatTensor(1).uniform_(lower_bound[col], upper_bound[col])
            
            # Fill columns 7 to 9 with random values within bounds
            for col in range(8, 10):
                initial_conditions_tensor[i, j, col] = torch.FloatTensor(1).uniform_(lower_bound[col], upper_bound[col])
            
            # Set the last column (column 10) to either 0.3 or 1 randomly
            initial_conditions_tensor[i, j, 10] = 0.3 if torch.rand(1).item() < 0.5 else 1.0
    
    return initial_conditions_tensor

batch_initial_conditions = initial_conditions(NUM_RESTARTS, BATCH_SIZE, bounds)


def constraint(X):
    torch.set_grad_enabled(True)
    """
    Vérifie les contraintes sur le tenseur X et retourne un tenseur de résultats.

    Args:
        X (Tensor): Un tenseur de forme (d,) ou (batch_size, d) ou (num_batches, batch_size, d).

    Returns:
        Tensor: Un tenseur avec des valeurs 1 ou -1 indiquant si chaque sous-tenseur respecte les contraintes.
    """
    if X.ndim == 1:
        # X est de forme (d,), convertir en (1, 1, d) pour homogénéiser le traitement
        X = X.unsqueeze(0).unsqueeze(0)
    elif X.ndim == 2:
        # X est de forme (batch_size, d), convertir en (1, batch_size, d)
        X = X.unsqueeze(0)
    elif X.ndim == 3:
        # X est déjà de forme (num_batches, batch_size, d), aucun changement nécessaire
        pass
    else:
        raise ValueError(f"Dimension du tenseur non supportée: {X.ndim}")
    X = torch.tensor(X, requires_grad=True)
    num_batches, batch_size, d = X.shape
    results = torch.full((num_batches, batch_size), -1, dtype=torch.float,requires_grad=True)  # Initialiser avec -1

    for i in range(num_batches):
        for j in range(batch_size):
            row = X[i, j, :7]  # Prendre les 8 premières colonnes
            colonnes_non_nulles = (row > 0).sum().item()  # Compte le nombre de colonnes avec des valeurs > 0
            nb_colonnes_sup_0 = (colonnes_non_nulles > 0)  # Condition 1
            
            somme_colonnes = row.sum().item()  # Somme des valeurs dans les 8 premières colonnes
            condition_2 = somme_colonnes <= 75
            
            if nb_colonnes_sup_0 <= 4 and somme_colonnes <= 75:
                return torch.tensor(1.0,dtype=torch.int)  # Return a tensor
            else:
                return torch.tensor(1.0,dtype=torch.int)  # Return a tensor

    return results

def get_mf(model):

    curr_val_acqf = FixedFeatureAcquisitionFunction(
        acq_function=PosteriorMean(model),
        d=11,
        columns=[10],
        values=[1],
    )

    _, current_value = optimize_acqf(
        acq_function=curr_val_acqf,
        bounds=bounds[:, :-1],
        q=1,
        num_restarts=10,
        raw_samples=1024,
        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)


    
def optimize_mf_and_get_observation(acq_func):
    torch.set_grad_enabled(True)
    candidates, _ = optimize_acqf_mixed(
        acq_function=acq_func,
        bounds=bounds,
        fixed_features_list=[{10: 0.3}, {10: 1.0}],
        q=BATCH_SIZE,
        num_restarts=NUM_RESTARTS,
        raw_samples=RAW_SAMPLES,
        options={"batch_limit": 1, "maxiter": 200},
        nonlinear_inequality_constraints = [(constraint,True)],
        batch_initial_conditions = batch_initial_conditions)
    
    # Calculate the cost
    cost = cost_model(candidates).sum()
    new_x = candidates.detach()
    
    return new_x, cost

I would be extremely grateful for any help, I really don't have any idea on how to proceed...

[Balandat]

The issue is that in your constraint definition constraint() you're losing all the gradients by doing things like using .item(), which returns a python scalar and not a pytorch tensor. You'll have to implement this in a differentiable fashion (or at least using a differentiable approximation as it looks like your constraint may be non-differentiable) if you want to use it in that form. Also, it seems that you're misinterpreting how the constraint is used - the constraint is considered feasible if constraint(X) >= 0 so constraint(X) should output some continuous (differentiable) value.

A word of warning though: The constellation of using qMultifidelityKnowledgeGradient with nonlinear parameter constraints is not exactly well-tested, and I could potentially see additional issues creep up if this is fixed. I recommend as a first step using a dummy constraint that just implements some simple bound on some dimension of the input (e.g. constraint = lambda X: X[0] - bound_1) to see if this setup throws some other errors before spending a lot of time on figuring out the differentiable approximation of the constraint.