[Bug] Issues with qKG.evaluate()

[saitcakmak]

🐛 Bug

i) This line raises an AttributeError with models that do not have an _input_batch_shape attribute. (#578 as an example)

botorch/botorch/optim/initializers.py

Line 328 in 7e2a404

batch_shape = acq_function.model._input_batch_shape

ii) Since it is a subclass of qKnowledgeGradient, qMultiFidelityKnowledgeGradient has evaluate method implemented. However, the current implementation ignores the project operator in the inner optimization, and may produce buggy output.

Expected Behavior

i) The batch shape could be inferred from model.train_inputs batch shape. My observation is _input_batch_shape = model.train_inputs[0].shape[:-2], but I am not sure if this is always true. If this is true, a simple try/except should do it.

ii) It should either raise a NotImplementedError, or be modified to accommodate the project operator.

[saitcakmak]

To add to this, gen_one_shot_kg_initial_conditions ignores the project operator as well.

[Balandat]

Thanks for raising this.

The batch shape could be inferred from model.train_inputs batch shape. My observation is _input_batch_shape = model.train_inputs[0].shape[:-2], but I am not sure if this is always true. If this is true, a simple try/except should do it.

I think rather than doing something ad-hoc in the initializer, the proper thing to do here is to add a batch_shape property to GPyTorchModel (or even Model), similar to num_outputs. Then we should be able to use that transparently everywhere.

It should either raise a NotImplementedError, or be modified to accommodate the project operator.

Good point. I think the most logical thing to do would be to somehow incorporate the projection operation as part of the value_function - that way we don't have to manually project the X whenever we need to evaluate the value function.

Let me see what I can come up with here.

[saitcakmak]

I see that you put up #588 for the batch shape. For the project issue, I was thinking a wrapper like this may work:

class ProjectedValueFunction(AcquisitionFunction):
    def __init__(
            self,
            base_value_function: AcquisitionFunction,
            project: Callable[[Tensor], Tensor]
    ):
        self.base_value_function = base_value_function
        self.project = project

    def forward(self, X: Tensor) -> Tensor:
        return self.base_value_function(self.project(X))


def _get_value_function(
    model: Model,
    objective: Optional[Union[MCAcquisitionObjective, ScalarizedObjective]] = None,
    sampler: Optional[MCSampler] = None,
    project: Optional[Callable[[Tensor], Tensor]] = None,
) -> AcquisitionFunction:
    r"""Construct value function (i.e. inner acquisition function)."""
    if isinstance(objective, MCAcquisitionObjective):
        base_value_function = qSimpleRegret(
            model=model, sampler=sampler, objective=objective)
    else:
        base_value_function = PosteriorMean(model=model, objective=objective)
    if project is None:
        return base_value_function
    else:
        return ProjectedValueFunction(base_value_function, project)

[Balandat]

Yeah this seems quite reasonable to me. @danielrjiang any thoughts?

[danielrjiang]

Thanks @saitcakmak, seems pretty useful to me too.

[saitcakmak]

Cool, I’ll put up a PR soon.