qLowerConfidenceBound (pytorch#2517)

Summary:
Pull Request resolved: pytorch#2517

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

Reviewed By: SebastianAment

Differential Revision: D60624931

fbshipit-source-id: 928c116c970f263aa225e171c6be469a0179f001

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
@@ -50,6 +50,7 @@
 )
 from botorch.acquisition.monte_carlo import (
     qExpectedImprovement,
+    qLowerConfidenceBound,
     qNoisyExpectedImprovement,
     qProbabilityOfImprovement,
     qSimpleRegret,
@@ -767,13 +768,15 @@ def construct_inputs_qPI(
     }
 
 
-@acqf_input_constructor(qUpperConfidenceBound)
+@acqf_input_constructor(qLowerConfidenceBound, qUpperConfidenceBound)
 def construct_inputs_qUCB(
     model: Model,
     objective: Optional[MCAcquisitionObjective] = None,
     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 +791,30 @@ 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.")
+        if objective is None:
+            objective = IdentityMCObjective()
+        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_input_constructors.py

@@ -61,6 +61,7 @@
 )
 from botorch.acquisition.monte_carlo import (
     qExpectedImprovement,
+    qLowerConfidenceBound,
     qNoisyExpectedImprovement,
     qProbabilityOfImprovement,
     qSimpleRegret,
@@ -86,6 +87,7 @@
 from botorch.acquisition.multi_objective.utils import get_default_partitioning_alpha
 from botorch.acquisition.objective import (
     ConstrainedMCObjective,
+    IdentityMCObjective,
     LinearMCObjective,
     ScalarizedPosteriorTransform,
 )
@@ -754,34 +756,86 @@ def test_construct_inputs_qPI(self) -> None:
         self.assertIs(acqf.model, mock_model)
         self.assertIs(acqf.objective, objective)
 
-    def test_construct_inputs_qUCB(self) -> None:
-        c = get_acqf_input_constructor(qUpperConfidenceBound)
+
+class TestQUpperConfidenceBoundInputConstructor(InputConstructorBaseTestCase):
+    acqf_class = qUpperConfidenceBound
+
+    def setUp(self, suppress_input_warnings: bool = True) -> None:
+        super().setUp(suppress_input_warnings=suppress_input_warnings)
+        self.c = get_acqf_input_constructor(self.acqf_class)
+
+    def test_confidence_bound(self) -> None:
         mock_model = self.mock_model
-        kwargs = c(model=mock_model, training_data=self.blockX_blockY)
+        kwargs = self.c(model=mock_model, training_data=self.blockX_blockY)
         self.assertEqual(kwargs["model"], mock_model)
-        self.assertIsNone(kwargs["objective"])
         self.assertIsNone(kwargs["X_pending"])
         self.assertIsNone(kwargs["sampler"])
         self.assertEqual(kwargs["beta"], 0.2)
-        acqf = qUpperConfidenceBound(**kwargs)
+        acqf = self.acqf_class(**kwargs)
         self.assertIs(acqf.model, mock_model)
 
+    def test_confidence_bound_with_objective(self) -> None:
         X_pending = torch.rand(2, 2)
         objective = LinearMCObjective(torch.rand(2))
-        kwargs = c(
-            model=mock_model,
+        kwargs = self.c(
+            model=self.mock_model,
             training_data=self.blockX_blockY,
             objective=objective,
             X_pending=X_pending,
             beta=0.1,
         )
-        self.assertEqual(kwargs["model"], mock_model)
+        self.assertEqual(kwargs["model"], self.mock_model)
         self.assertTrue(torch.equal(kwargs["objective"].weights, objective.weights))
         self.assertTrue(torch.equal(kwargs["X_pending"], X_pending))
         self.assertIsNone(kwargs["sampler"])
         self.assertEqual(kwargs["beta"], 0.1)
-        acqf = qUpperConfidenceBound(**kwargs)
-        self.assertIs(acqf.model, mock_model)
+        acqf = self.acqf_class(**kwargs)
+        self.assertIs(acqf.model, self.mock_model)
+
+    def test_confidence_bound_with_constraints_error(self) -> None:
+        with self.assertRaisesRegex(ValueError, "Constraints require an X_baseline."):
+            self.c(
+                model=self.mock_model,
+                training_data=self.blockX_blockY,
+                constraints=torch.rand(2, 2),
+            )
+
+    def test_confidence_bound_with_constraints(self) -> None:
+        # these are needed for computing the infeasible cost
+        self.mock_model._posterior._mean = torch.zeros(2, 2)
+        self.mock_model._posterior._variance = torch.ones(2, 2)
+
+        X_baseline = torch.rand(2, 2)
+        outcome_constraints = (torch.tensor([[0.0, 1.0]]), torch.tensor([[0.5]]))
+        constraints = get_outcome_constraint_transforms(
+            outcome_constraints=outcome_constraints
+        )
+        for objective in (LinearMCObjective(torch.rand(2)), None):
+            with self.subTest(objective=objective):
+                kwargs = self.c(
+                    model=self.mock_model,
+                    training_data=self.blockX_blockY,
+                    objective=objective,
+                    constraints=constraints,
+                    X_baseline=X_baseline,
+                )
+                final_objective = kwargs["objective"]
+                self.assertIsInstance(final_objective, ConstrainedMCObjective)
+                if objective is None:
+                    self.assertIsInstance(
+                        final_objective.objective, IdentityMCObjective
+                    )
+                else:
+                    self.assertIs(final_objective.objective, objective)
+                self.assertIs(final_objective.constraints, constraints)
+                # test that we can construct the acquisition function
+                self.acqf_class(**kwargs)
+
+
+class TestQLowerConfidenceBoundInputConstructor(
+    TestQUpperConfidenceBoundInputConstructor
+):
+    acqf_class = qLowerConfidenceBound
 
 
 class TestMultiObjectiveAcquisitionFunctionInputConstructors(
@@ -1452,7 +1506,12 @@ def setUp(self, suppress_input_warnings: bool = True) -> None:
         # {key: (list of acquisition functions, arguments they accept)}
         self.cases = {
             "PosteriorMean-type": (
-                [PosteriorMean, UpperConfidenceBound, qUpperConfidenceBound],
+                [
+                    PosteriorMean,
+                    UpperConfidenceBound,
+                    qUpperConfidenceBound,
+                    qLowerConfidenceBound,
+                ],
                 {"model": self.mock_model},
             ),
         }

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):