From 08ead8f641acb96a680aeeaad5efbee76b3e36ee Mon Sep 17 00:00:00 2001 From: Sam Daulton Date: Mon, 9 Sep 2024 14:36:24 -0700 Subject: [PATCH] qLowerConfidenceBound (#2517) Summary: Pull Request resolved: https://github.com/pytorch/botorch/pull/2517 Implement a qLowerConfidence acquisition function for more confident/risk-averse candidate selection. Reviewed By: SebastianAment Differential Revision: D60624931 --- botorch/acquisition/input_constructors.py | 26 +++++++- botorch/acquisition/monte_carlo.py | 19 +++++- test/acquisition/test_input_constructors.py | 74 ++++++++++++++++++--- test/acquisition/test_monte_carlo.py | 39 +++++++++-- 4 files changed, 138 insertions(+), 20 deletions(-) diff --git a/botorch/acquisition/input_constructors.py b/botorch/acquisition/input_constructors.py index dba1d55f29..6a1f898967 100644 --- a/botorch/acquisition/input_constructors.py +++ b/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, diff --git a/botorch/acquisition/monte_carlo.py b/botorch/acquisition/monte_carlo.py index 766aec144e..b97fb6957e 100644 --- a/botorch/acquisition/monte_carlo.py +++ b/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) diff --git a/test/acquisition/test_input_constructors.py b/test/acquisition/test_input_constructors.py index a8e876eaf3..31d3c34f78 100644 --- a/test/acquisition/test_input_constructors.py +++ b/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( diff --git a/test/acquisition/test_monte_carlo.py b/test/acquisition/test_monte_carlo.py index d877a2f5c5..9e819201db 100644 --- a/test/acquisition/test_monte_carlo.py +++ b/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):