Pending observations for botorch MCAcquisitionFunctions (#176)

Summary:
Pull Request resolved: #176

This reverts past changes that resulted in removing the ability to
use pending observations with all MCAcquisitionFunctions.  This adds back
in the previous functionality which is needed in asynchronous optimization
settings.

This changes MCAcquisitionFunction to always return the total acquisition value
of pending points plus new points.

Reviewed By: Balandat

Differential Revision: D15218481

fbshipit-source-id: a3a4049e7768673e676d6e31534725f67e3cc693

Author: bkarrer

botorch/acquisition/acquisition.py

@@ -6,11 +6,14 @@
 Abstract base module for all botorch acquisition functions.
 """
 
+import warnings
 from abc import ABC, abstractmethod
+from typing import Optional
 
 from torch import Tensor
 from torch.nn import Module
 
+from ..exceptions import BotorchWarning
 from ..models.model import Model
 
 
@@ -26,6 +29,24 @@ def __init__(self, model: Model) -> None:
         super().__init__()
         self.add_module("model", model)
 
+    def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
+        r"""Informs the acquisition function about pending design points.
+
+        Args:
+            X_pending: `m x d` Tensor with `m` `d`-dim design points that have
+                been submitted for evaluation but have not yet been evaluated.
+        """
+        if X_pending is not None:
+            if X_pending.requires_grad:
+                warnings.warn(
+                    "Pending points require a gradient but the acquisition function"
+                    " will not provide a gradient to these points.",
+                    BotorchWarning,
+                )
+            self.X_pending = X_pending.clone().detach()
+        else:
+            self.X_pending = X_pending
+
     @abstractmethod
     def forward(self, X: Tensor) -> Tensor:
         r"""Evaluate the acquisition function on the candidate set X.

botorch/acquisition/analytic.py

@@ -39,6 +39,11 @@ def _validate_single_output_posterior(self, posterior: Posterior) -> None:
                 f" function of type {self.__class__.__name__}"
             )
 
+    def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
+        raise UnsupportedError(
+            f"Analytic acquisition functions do not account for X_pending yet."
+        )
+
 
 class ExpectedImprovement(AnalyticAcquisitionFunction):
     r"""Single-outcome Expected Improvement (analytic).

botorch/acquisition/monte_carlo.py

@@ -38,15 +38,19 @@ def __init__(
         model: Model,
         sampler: Optional[MCSampler] = None,
         objective: Optional[MCAcquisitionObjective] = None,
+        X_pending: Optional[Tensor] = None,
     ) -> None:
         r"""Constructor for the MCAcquisitionFunction base class.
 
         Args:
             model: A fitted model.
             sampler: The sampler used to draw base samples. Defaults to
                 `SobolQMCNormalSampler(num_samples=512, collapse_batch_dims=True)`.
-            objective: THe MCAcquisitionObjective under which the samples are
+            objective: The MCAcquisitionObjective under which the samples are
                 evaluated. Defaults to `IdentityMCObjective()`.
+            X_pending:  A `m x d`-dim Tensor of `m` design points that have
+                points that have been submitted for function evaluation
+                but have not yet been evaluated.
         """
         super().__init__(model=model)
         if sampler is None:
@@ -55,12 +59,15 @@ def __init__(
         if objective is None:
             objective = IdentityMCObjective()
         self.add_module("objective", objective)
+        self.set_X_pending(X_pending)
 
     @abstractmethod
     def forward(self, X: Tensor) -> Tensor:
-        r"""Takes in a `(b) x q x d` X Tensor of `b` t-batches with `q` `d`-dim
-        design points each, expands and concatenates `self.X_pending` and
-        returns a one-dimensional Tensor with `b` elements."""
+        r"""Takes in a `(b) x q x d` X Tensor of `(b)` t-batches with `q` `d`-dim
+        design points each, and returns a one-dimensional Tensor with
+        `(b)` elements.  Should utilize the result of set_X_pending as needed
+        to account for pending function evaluations.
+        """
         pass  # pragma: no cover
 
 
@@ -89,6 +96,7 @@ def __init__(
         best_f: Union[float, Tensor],
         sampler: Optional[MCSampler] = None,
         objective: Optional[MCAcquisitionObjective] = None,
+        X_pending: Optional[Tensor] = None,
     ) -> None:
         r"""q-Expected Improvement.
 
@@ -100,8 +108,14 @@ def __init__(
                 `SobolQMCNormalSampler(num_samples=500, collapse_batch_dims=True)`
             objective: The MCAcquisitionObjective under which the samples are
                 evaluated. Defaults to `IdentityMCObjective()`.
+            X_pending:  A `m x d`-dim Tensor of `m` design points that have
+                points that have been submitted for function evaluation
+                but have not yet been evaluated.  Concatenated into X upon
+                forward call.  Copied and set to have no gradient.
         """
-        super().__init__(model=model, sampler=sampler, objective=objective)
+        super().__init__(
+            model=model, sampler=sampler, objective=objective, X_pending=X_pending
+        )
         if not torch.is_tensor(best_f):
             best_f = torch.tensor(float(best_f))
         self.register_buffer("best_f", best_f)
@@ -118,6 +132,8 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b)`-dim Tensor of Expected Improvement values at the given
             design points `X`.
         """
+        if self.X_pending is not None:
+            X = torch.cat([X, match_batch_shape(self.X_pending, X)], dim=-2)
         posterior = self.model.posterior(X)
         samples = self.sampler(posterior)
         obj = self.objective(samples)
@@ -150,20 +166,28 @@ def __init__(
         X_baseline: Tensor,
         sampler: Optional[MCSampler] = None,
         objective: Optional[MCAcquisitionObjective] = None,
+        X_pending: Optional[Tensor] = None,
     ) -> None:
         r"""q-Noisy Expected Improvement.
 
         Args:
             model: A fitted model.
-            X_baseline: A `m x d`-dim Tensor of `m` design points that have
-                either already been observed or whose evaluation is pending.
-                These points are considered as the potential best design point.
+            X_baseline: A `r x d`-dim Tensor of `r` design points that have
+                already been observed. These points are considered as the
+                potential best design point.
             sampler: The sampler used to draw base samples. Defaults to
                 `SobolQMCNormalSampler(num_samples=500, collapse_batch_dims=True)`.
             objective: The MCAcquisitionObjective under which the samples are
                 evaluated. Defaults to `IdentityMCObjective()`.
+            X_pending:  A `m x d`-dim Tensor of `m` design points that have
+                points that have been submitted for function evaluation
+                but have not yet been evaluated.  Concatenated into X upon
+                forward call.  Copied and set to have no gradient.
+
         """
-        super().__init__(model=model, sampler=sampler, objective=objective)
+        super().__init__(
+            model=model, sampler=sampler, objective=objective, X_pending=X_pending
+        )
         self.register_buffer("X_baseline", X_baseline)
 
     @t_batch_mode_transform()
@@ -178,6 +202,8 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b)`-dim Tensor of Noisy Expected Improvement values at the given
             design points `X`.
         """
+        if self.X_pending is not None:
+            X = torch.cat([X, match_batch_shape(self.X_pending, X)], dim=-2)
         q = X.shape[-2]
         X_full = torch.cat([X, match_batch_shape(self.X_baseline, X)], dim=-2)
         # TODO (T41248036): Implement more efficient way to compute posterior
@@ -214,6 +240,7 @@ def __init__(
         best_f: Union[float, Tensor],
         sampler: Optional[MCSampler] = None,
         objective: Optional[MCAcquisitionObjective] = None,
+        X_pending: Optional[Tensor] = None,
         tau: float = 1e-3,
     ) -> None:
         r"""q-Probability of Improvement.
@@ -226,12 +253,18 @@ def __init__(
                 `SobolQMCNormalSampler(num_samples=500, collapse_batch_dims=True)`
             objective: The MCAcquisitionObjective under which the samples are
                 evaluated. Defaults to `IdentityMCObjective()`.
+            X_pending:  A `m x d`-dim Tensor of `m` design points that have
+                points that have been submitted for function evaluation
+                but have not yet been evaluated.  Concatenated into X upon
+                forward call.  Copied and set to have no gradient.
             tau: The temperature parameter used in the sigmoid approximation
                 of the step function. Smaller values yield more accurate
                 approximations of the function, but result in gradients
                 estimates with higher variance.
         """
-        super().__init__(model=model, sampler=sampler, objective=objective)
+        super().__init__(
+            model=model, sampler=sampler, objective=objective, X_pending=X_pending
+        )
         if not torch.is_tensor(best_f):
             best_f = torch.tensor(float(best_f))
         self.register_buffer("best_f", best_f)
@@ -251,6 +284,8 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b)`-dim Tensor of Probability of Improvement values at the given
             design points `X`.
         """
+        if self.X_pending is not None:
+            X = torch.cat([X, match_batch_shape(self.X_pending, X)], dim=-2)
         posterior = self.model.posterior(X)
         samples = self.sampler(posterior)
         obj = self.objective(samples)
@@ -286,6 +321,8 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b)`-dim Tensor of Simple Regret values at the given design
             points `X`.
         """
+        if self.X_pending is not None:
+            X = torch.cat([X, match_batch_shape(self.X_pending, X)], dim=-2)
         posterior = self.model.posterior(X)
         samples = self.sampler(posterior)
         obj = self.objective(samples)
@@ -315,6 +352,7 @@ def __init__(
         beta: float,
         sampler: Optional[MCSampler] = None,
         objective: Optional[MCAcquisitionObjective] = None,
+        X_pending: Optional[Tensor] = None,
     ) -> None:
         r"""q-Upper Confidence Bound.
 
@@ -325,8 +363,14 @@ def __init__(
                 `SobolQMCNormalSampler(num_samples=500, collapse_batch_dims=True)`
             objective: The MCAcquisitionObjective under which the samples are
                 evaluated. Defaults to `IdentityMCObjective()`.
+            X_pending:  A `m x d`-dim Tensor of `m` design points that have
+                points that have been submitted for function evaluation
+                but have not yet been evaluated.  Concatenated into X upon
+                forward call.  Copied and set to have no gradient.
         """
-        super().__init__(model=model, sampler=sampler, objective=objective)
+        super().__init__(
+            model=model, sampler=sampler, objective=objective, X_pending=X_pending
+        )
         self.register_buffer("beta", torch.tensor(float(beta)))
 
     @t_batch_mode_transform()
@@ -341,6 +385,8 @@ def forward(self, X: Tensor) -> Tensor:
             A `(b)`-dim Tensor of Upper Confidence Bound values at the given
             design points `X`.
         """
+        if self.X_pending is not None:
+            X = torch.cat([X, match_batch_shape(self.X_pending, X)], dim=-2)
         posterior = self.model.posterior(X)
         samples = self.sampler(posterior)
         obj = self.objective(samples)

botorch/acquisition/utils.py

@@ -8,7 +8,6 @@
 
 from typing import Callable, Optional
 
-import torch
 from torch import Tensor
 
 from . import analytic, monte_carlo
@@ -64,7 +63,11 @@ def get_acquisition_function(
     if acquisition_function_name == "qEI":
         best_f = objective(model.posterior(X_observed).mean).max().item()
         return monte_carlo.qExpectedImprovement(
-            model=model, best_f=best_f, sampler=sampler, objective=objective
+            model=model,
+            best_f=best_f,
+            sampler=sampler,
+            objective=objective,
+            X_pending=X_pending,
         )
     elif acquisition_function_name == "qPI":
         best_f = objective(model.posterior(X_observed).mean).max().item()
@@ -73,25 +76,30 @@ def get_acquisition_function(
             best_f=best_f,
             sampler=sampler,
             objective=objective,
+            X_pending=X_pending,
             tau=kwargs.get("tau", 1e-3),
         )
     elif acquisition_function_name == "qNEI":
-        if X_pending is None:
-            X_baseline = X_observed
-        else:
-            X_baseline = torch.cat([X_observed, X_pending], dim=-2)
         return monte_carlo.qNoisyExpectedImprovement(
-            model=model, X_baseline=X_baseline, sampler=sampler, objective=objective
+            model=model,
+            X_baseline=X_observed,
+            sampler=sampler,
+            objective=objective,
+            X_pending=X_pending,
         )
     elif acquisition_function_name == "qSR":
         return monte_carlo.qSimpleRegret(
-            model=model, sampler=sampler, objective=objective
+            model=model, sampler=sampler, objective=objective, X_pending=X_pending
         )
     elif acquisition_function_name == "qUCB":
         if "beta" not in kwargs:
             raise ValueError("`beta` must be specified in kwargs for qUCB.")
         return monte_carlo.qUpperConfidenceBound(
-            model=model, beta=kwargs["beta"], sampler=sampler, objective=objective
+            model=model,
+            beta=kwargs["beta"],
+            sampler=sampler,
+            objective=objective,
+            X_pending=X_pending,
         )
     raise NotImplementedError(
         f"Unknown acquisition function {acquisition_function_name}"

botorch/optim/optimize.py

@@ -15,7 +15,7 @@
 from ..acquisition import AcquisitionFunction
 from ..acquisition.analytic import AnalyticAcquisitionFunction
 from ..acquisition.utils import is_nonnegative
-from ..exceptions import BadInitialCandidatesWarning, UnsupportedError
+from ..exceptions import BadInitialCandidatesWarning
 from ..gen import gen_candidates_scipy, get_best_candidates
 from ..utils.sampling import draw_sobol_samples
 from .initializers import initialize_q_batch, initialize_q_batch_nonneg
@@ -36,7 +36,7 @@ def sequential_optimize(
     r"""Generate a set of candidates via sequential multi-start optimization.
 
     Args:
-        acq_function: The qNoisyExpectedImprovement acquisition function.
+        acq_function: An AcquisitionFunction
         bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`.
         q: The number of candidates.
         num_restarts:  Number of starting points for multistart acquisition
@@ -65,14 +65,9 @@ def sequential_optimize(
         >>> bounds = torch.tensor([[0.], [1.]])
         >>> candidates = sequential_optimize(qEI, bounds, 2, 20, 500)
     """
-    if not hasattr(acq_function, "X_baseline"):
-        raise UnsupportedError(  # pyre-ignore: [16]
-            "Sequential Optimization is only supported for acquisition functions "
-            "with an `X_baseline` property."
-        )
     candidate_list = []
     candidates = torch.tensor([])
-    base_X_baseline = acq_function.X_baseline  # pyre-ignore: [16]
+    base_X_pending = acq_function.X_pending  # pyre-ignore: [16]
     for _ in range(q):
         candidate = joint_optimize(
             acq_function=acq_function,
@@ -90,13 +85,13 @@ def sequential_optimize(
             candidate = post_processing_func(candidate.view(-1)).view(*candidate_shape)
         candidate_list.append(candidate)
         candidates = torch.cat(candidate_list, dim=-2)
-        acq_function.X_baseline = (
-            torch.cat([base_X_baseline, candidates], dim=-2)
-            if base_X_baseline is not None
+        acq_function.set_X_pending(
+            torch.cat([base_X_pending, candidates], dim=-2)
+            if base_X_pending is not None
             else candidates
         )
-    # Reset acq_func to previous X_baseline state
-    acq_function.X_baseline = base_X_baseline
+    # Reset acq_func to previous X_pending state
+    acq_function.set_X_pending(base_X_pending)
     return candidates
 
 
@@ -263,7 +258,7 @@ def gen_batch_initial_conditions(
             if factor < max_factor:
                 factor += 1
     warnings.warn(
-        "Unable to find non-zero acquistion function values - initial conditions "
+        "Unable to find non-zero acquisition function values - initial conditions "
         "are being selected randomly.",
         BadInitialCandidatesWarning,
     )

test/acquisition/test_analytic.py

@@ -48,6 +48,8 @@ def test_expected_improvement(self, cuda=False):
             ei = module(X)
             ei_expected = torch.tensor(0.6978, device=device, dtype=dtype)
             self.assertTrue(torch.allclose(ei, ei_expected, atol=1e-4))
+            with self.assertRaises(UnsupportedError):
+                module.set_X_pending(None)
 
     def test_expected_improvement_cuda(self, cuda=False):
         if torch.cuda.is_available():