qLowerConfidenceBound (#2517)

Summary:
Pull Request resolved: #2517

Implement a qLowerConfidence acquisition function for more confident/risk-averse candidate selection.

Reviewed By: SebastianAment

Differential Revision: D60624931

botorch/acquisition/input_constructors.py

@@ -13,7 +13,7 @@
 
 import inspect
 from collections.abc import Hashable, Iterable, Sequence
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, Callable, List, Optional, TypeVar, Union
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
@@ -774,6 +774,8 @@ def construct_inputs_qUCB(
     posterior_transform: Optional[PosteriorTransform] = None,
     X_pending: Optional[Tensor] = None,
     sampler: Optional[MCSampler] = None,
+    X_baseline: Optional[Tensor] = None,
+    constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
     beta: float = 0.2,
 ) -> dict[str, Any]:
     r"""Construct kwargs for the `qUpperConfidenceBound` constructor.
@@ -788,11 +790,28 @@ def construct_inputs_qUCB(
             Concatenated into X upon forward call.
         sampler: The sampler used to draw base samples. If omitted, uses
             the acquisition functions's default sampler.
+        X_baseline: A `batch_shape x r x d`-dim Tensor of `r` design points
+            that have already been observed. These points are used to
+            compute with infeasible cost when there are constraints.
+        constraints: A list of constraint callables which map a Tensor of posterior
+            samples of dimension `sample_shape x batch-shape x q x m`-dim to a
+            `sample_shape x batch-shape x q`-dim Tensor. The associated constraints
+            are considered satisfied if the output is less than zero.
         beta: Controls tradeoff between mean and standard deviation in UCB.
 
     Returns:
         A dict mapping kwarg names of the constructor to values.
     """
+    if constraints is not None:
+        if X_baseline is None:
+            raise ValueError("Constraints require an X_baseline.")
+        objective = ConstrainedMCObjective(
+            objective=objective,
+            constraints=constraints,
+            infeasible_cost=get_infeasible_cost(
+                X=X_baseline, model=model, objective=objective
+            ),
+        )
     return {
         "model": model,
         "objective": objective,

botorch/acquisition/monte_carlo.py

@@ -856,7 +856,10 @@ def __init__(
             posterior_transform=posterior_transform,
             X_pending=X_pending,
         )
-        self.beta_prime = math.sqrt(beta * math.pi / 2)
+        self.beta_prime = self._get_beta_prime(beta=beta)
+
+    def _get_beta_prime(self, beta: float) -> float:
+        return math.sqrt(beta * math.pi / 2)
 
     def _sample_forward(self, obj: Tensor) -> Tensor:
         r"""Evaluate qUpperConfidenceBound per sample on the candidate set `X`.
@@ -869,3 +872,17 @@ def _sample_forward(self, obj: Tensor) -> Tensor:
         """
         mean = obj.mean(dim=0)
         return mean + self.beta_prime * (obj - mean).abs()
+
+
+class qLowerConfidenceBound(qUpperConfidenceBound):
+    r"""MC-based batched lower confidence bound.
+
+    This acquisition function is useful for confident/risk-averse decision making.
+    This acquisition function is intended to be maximized as with qUpperConfidenceBound,
+    but the qLowerConfidenceBound will be pessimistic in the face of uncertainty and
+    lead to conservative candidates.
+    """
+
+    def _get_beta_prime(self, beta: float) -> float:
+        """Multiply beta prime by -1 to get the lower confidence bound."""
+        return -super()._get_beta_prime(beta=beta)

test/acquisition/test_monte_carlo.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import math
 import warnings
 from copy import deepcopy
 from functools import partial
@@ -17,6 +18,7 @@
 from botorch.acquisition.monte_carlo import (
     MCAcquisitionFunction,
     qExpectedImprovement,
+    qLowerConfidenceBound,
     qNoisyExpectedImprovement,
     qProbabilityOfImprovement,
     qSimpleRegret,
@@ -871,7 +873,9 @@ def test_q_simple_regret_constraints(self):
 
 
 class TestQUpperConfidenceBound(BotorchTestCase):
-    def test_q_upper_confidence_bound(self):
+    acqf_class = qUpperConfidenceBound
+
+    def test_q_confidence_bound(self):
         for dtype in (torch.float, torch.double):
             # the event shape is `b x q x t` = 1 x 1 x 1
             samples = torch.zeros(1, 1, 1, device=self.device, dtype=dtype)
@@ -881,13 +885,13 @@ def test_q_upper_confidence_bound(self):
 
             # basic test
             sampler = IIDNormalSampler(sample_shape=torch.Size([2]))
-            acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler)
+            acqf = self.acqf_class(model=mm, beta=0.5, sampler=sampler)
             res = acqf(X)
             self.assertEqual(res.item(), 0.0)
 
             # basic test
             sampler = IIDNormalSampler(sample_shape=torch.Size([2]), seed=12345)
-            acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler)
+            acqf = self.acqf_class(model=mm, beta=0.5, sampler=sampler)
             res = acqf(X)
             self.assertEqual(res.item(), 0.0)
             self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 1, 1]))
@@ -924,7 +928,7 @@ def test_q_upper_confidence_bound(self):
                     sum(issubclass(w.category, BotorchWarning) for w in ws), 1
                 )
 
-    def test_q_upper_confidence_bound_batch(self):
+    def test_q_confidence_bound_batch(self):
         # TODO: T41739913 Implement tests for all MCAcquisitionFunctions
         for dtype in (torch.float, torch.double):
             samples = torch.zeros(2, 2, 1, device=self.device, dtype=dtype)
@@ -935,14 +939,14 @@ def test_q_upper_confidence_bound_batch(self):
 
             # test batch mode
             sampler = IIDNormalSampler(sample_shape=torch.Size([2]))
-            acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler)
+            acqf = self.acqf_class(model=mm, beta=0.5, sampler=sampler)
             res = acqf(X)
             self.assertEqual(res[0].item(), 1.0)
             self.assertEqual(res[1].item(), 0.0)
 
             # test batch mode
             sampler = IIDNormalSampler(sample_shape=torch.Size([2]), seed=12345)
-            acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler)
+            acqf = self.acqf_class(model=mm, beta=0.5, sampler=sampler)
             res = acqf(X)  # 1-dim batch
             self.assertEqual(res[0].item(), 1.0)
             self.assertEqual(res[1].item(), 0.0)
@@ -961,7 +965,7 @@ def test_q_upper_confidence_bound_batch(self):
 
             # test batch mode, qmc
             sampler = SobolQMCNormalSampler(sample_shape=torch.Size([2]))
-            acqf = qUpperConfidenceBound(model=mm, beta=0.5, sampler=sampler)
+            acqf = self.acqf_class(model=mm, beta=0.5, sampler=sampler)
             res = acqf(X)
             self.assertEqual(res[0].item(), 1.0)
             self.assertEqual(res[1].item(), 0.0)
@@ -991,9 +995,30 @@ def test_q_upper_confidence_bound_batch(self):
                     sum(issubclass(w.category, BotorchWarning) for w in ws), 1
                 )
 
+    def test_beta_prime(self, negate: bool = False) -> None:
+        acqf = self.acqf_class(
+            model=MockModel(
+                posterior=MockPosterior(
+                    samples=torch.zeros(2, 2, 1, device=self.device, dtype=torch.double)
+                )
+            ),
+            beta=1.96,
+        )
+        expected_value = math.sqrt(1.96 * math.pi / 2)
+        if negate:
+            expected_value *= -1
+        self.assertEqual(acqf.beta_prime, expected_value)
+
     # TODO: Test different objectives (incl. constraints)
 
 
+class TestQLowerConfidenceBound(TestQUpperConfidenceBound):
+    acqf_class = qLowerConfidenceBound
+
+    def test_beta_prime(self):
+        super().test_beta_prime(negate=True)
+
+
 class TestMCAcquisitionFunctionWithConstraints(BotorchTestCase):
     def test_mc_acquisition_function_with_constraints(self):
         for dtype in (torch.float, torch.double):