Wrapping qKG with FixedFeature: full_tree evaluates to True s.t. fantasy candidates aren't masked

[jakobzeitler]

I suspect I am ignorant to some bigger design issue, but I am currently observing the following behaviour for my code:

Because I am fixing the fidelity in the qKG via the FixedFeatureAcquisitionFunction, return_full_tree evaluates to True (https://github.com/pytorch/botorch/blob/9a45872462ab4a6917bf147946111ad3f0bbc86e/botorch/optim/optimize.py#L94) s.t. _optimize_acqf (https://github.com/pytorch/botorch/blob/9a45872462ab4a6917bf147946111ad3f0bbc86e/botorch/optim/optimize.py#L574) evaluates to False and therefore extract_candidates is not run.

Code
`# Construct qKG
qKG = qKnowledgeGradient(
model=model,
num_fantasies=128 if not SMOKE_TEST else 2,
current_value=current_value,
)

    batch_initial_conditions = gen_one_shot_kg_initial_conditions(
        acq_function=qKG,
        bounds=bounds,
        q=1,
        fixed_features={1:0}, # fidelity is in column 1
        num_restarts=NUM_RESTARTS,
        raw_samples=RAW_SAMPLES,
    )

    # throw away the fidelity dimension
    ff_batch_initial_conditions = batch_initial_conditions[..., :-1]

    # Optimise qKG with fixed IS
    fixed_qKG = FixedFeatureAcquisitionFunction(
        acq_function=qKG,
        d=2,
        columns=[1],  # 1 is atm our IS column
        values=[l],
    )

    candidate, value = optimize_acqf(
        acq_function=fixed_qKG,
        bounds=bounds[:, :-1],
        q=1,
        num_restarts=10 if not SMOKE_TEST else 2,
        raw_samples=1024 if not SMOKE_TEST else 4,
        batch_initial_conditions=ff_batch_initial_conditions,
        options={"batch_limit": 10, "maxiter": 200},
    )`

What am I missing?

Thanks!

[hvarfner]

Hi @jakobzeitler !

I'll look into it and get back to you ASAP.

[jakobzeitler]

My fix is to run

candidate = qKG.extract_candidates(candidate)

but I am worried I am violating design conventions.

Thanks in advance for any help or guidance.

[hvarfner]

Hi @jakobzeitler ,

I don't see any issues with this solution. If anything, we should probably change the full_tree property to account for this case and others like it, where you would want to wrap the acquisition function with something like a FixedFeatureAcqusitionFunction.

However, it seems to me like you would accomplish the same thing by skipping the FixedFeatureAcqusitionFunction altogether and instead fixing the fidelity in optimize_acqf:

candidate, value = optimize_acqf(
    acq_function=qKG,
    bounds=bounds,
    q=1,
    num_restarts=10 if not SMOKE_TEST else 2,
    raw_samples=1024 if not SMOKE_TEST else 4,
    fixed_features={1: 0},
    options={"batch_limit": 10, "maxiter": 200},
)
candidate = candidate[:, :-1] # And extract all but the fidelity.

Have you tried this option?

[jakobzeitler]

Thanks so much for the quick reply!

Your suggested solution makes sense, I was looking for that in the optimize_acqf API yesterday, but seem to have overlooked that, thanks for pointing me to it.

From your experience, will these two approaches (FixedFeatureAcquisitionFunction vs. fixed_features option in optimize_acqf) behave similarly or identical? Any obvious pitfalls?

Thanks a lot in advance!

[hvarfner]

Not that I can think of. I am fairly confident that the machinery is the same, but the imputation of fixed values occur in different places in the codebase (inside the acqf for FixedFeatureAcquisitionFunction, inside optimize_acqf for the fixed_features options. Running them side by side for 20ish rounds each, they yielded similar candidates.

In Ax, we use the latter, so that is what I would personally recommend based on convention.