Source code for botorch.acquisition.proximal

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

r"""
A wrapper around AcquisitionFunctions to add proximal weighting of the
acquisition function.
"""

from __future__ import annotations

import torch
from botorch.acquisition import AcquisitionFunction
from botorch.exceptions.errors import UnsupportedError
from botorch.utils import t_batch_mode_transform
from torch import Tensor
from torch.nn import Module


[docs]class ProximalAcquisitionFunction(AcquisitionFunction):
    """A wrapper around AcquisitionFunctions to add proximal weighting of the
    acquisition function. Acquisition function is weighted via a squared exponential
    centered at the last training point, with varying lengthscales corresponding to
    `proximal_weights`. Can only be used with acquisition functions based on single
    batch models.

    Small values of `proximal_weights` corresponds to strong biasing towards recently
    observed points, which smoothes optimization with a small potential decrese in
    convergence rate.

    Example:
        >>> model = SingleTaskGP(train_X, train_Y)
        >>> EI = ExpectedImprovement(model, best_f=0.0)
        >>> proximal_weights = torch.ones(d)
        >>> EI_proximal = ProximalAcquisitionFunction(EI, proximal_weights)
        >>> eip = EI_proximal(test_X)
    """

    def __init__(
        self,
        acq_function: AcquisitionFunction,
        proximal_weights: Tensor,
    ) -> None:
        r"""Derived Acquisition Function weighted by proximity to recently
        observed point.

        Args:
            acq_function: The base acquisition function, operating on input tensors
                of feature dimension `d`.
            proximal_weights: A `d` dim tensor used to bias locality
                along each axis.
        """
        Module.__init__(self)

        self.acq_func = acq_function

        if hasattr(acq_function, "X_pending"):
            if acq_function.X_pending is not None:
                raise UnsupportedError(
                    "Proximal acquisition function requires `X_pending` to be None."
                )
            self.X_pending = acq_function.X_pending

        self.register_buffer("proximal_weights", proximal_weights)

        # check model for train_inputs and single batch
        if not hasattr(self.acq_func.model, "train_inputs"):
            raise UnsupportedError(
                "Acquisition function model must have " "`train_inputs`."
            )

        if (
            self.acq_func.model.batch_shape != torch.Size([])
            and self.acq_func.model.train_inputs[0].shape[1] != 1
        ):
            raise UnsupportedError(
                "Proximal acquisition function requires a single batch model"
            )

        # check to make sure that weights match the training data shape
        if (
            len(self.proximal_weights.shape) != 1
            or self.proximal_weights.shape[0]
            != self.acq_func.model.train_inputs[0][-1].shape[-1]
        ):
            raise ValueError(
                "`proximal_weights` must be a one dimensional tensor with "
                "same feature dimension as model."
            )

[docs]    @t_batch_mode_transform(expected_q=1, assert_output_shape=False)
    def forward(self, X: Tensor) -> Tensor:
        r"""Evaluate base acquisition function with proximal weighting.

        Args:
            X: Input tensor of feature dimension `d` .

        Returns:
            Base acquisition function evaluated on tensor `X` multiplied by proximal
            weighting.
        """
        last_X = self.acq_func.model.train_inputs[0][-1].reshape(1, 1, -1)
        diff = X - last_X

        M = torch.linalg.norm(diff / self.proximal_weights, dim=-1) ** 2

        proximal_acq_weight = torch.exp(-0.5 * M)
        return self.acq_func(X) * proximal_acq_weight.flatten()