Scalable Constrained Bayesian Optimization (#1257)

Summary:
<!--
Thank you for sending the PR! We appreciate you spending the time to make BoTorch better.

Help us understand your motivation by explaining why you decided to make this change.

You can learn more about contributing to BoTorch here: https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md
-->

## Motivation

I created this PR to add support for Scalable Constrained Bayesian Optimization (SCBO) as described in [1] to BoTorch. Since BoTorch already supports TuRBO, relatively few modifications were needed to add support for SCBO. I have found the SCBO method to be very useful in my research, and I hope that this addition of SCBO to BoTorch will make it easier for others to use SCBO in the future.

[1] David Eriksson and Matthias Poloczek. Scalable constrained Bayesian optimization. In International Conference on Artificial Intelligence and Statistics, pages 730–738. PMLR, 2021.
(https://doi.org/10.48550/arxiv.2002.08526)

To implement SCBO, this PR includes two additions to BoTorch
1. A new sampling class called ConstrainedMaxPosteriorSampling that we use to implement the constrained Thompson sampling described in [1].
2. A tutorial Jupiter notebook that walks the user through a simple example, constrained optimization problem with SCBO.

### Have you read the [Contributing Guidelines on pull requests](https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes

Pull Request resolved: #1257

Test Plan:
I tested the added code in botorch.generation.sampling.ConstrainedMaxPosteriorSampling by running the tutorial notebook (tutorials/scalable_constrained_bo.ipynb) which uses ConstrainedMaxPosteriorSampling.

## Related PRs

(If this PR adds or changes functionality, please take some time to update the docs at https://github.com/pytorch/botorch, and link to your PR here.)

See the SCBO tutorial included in this PR (tutorials/scalable_constrained_bo.ipynb) for a description of the added implementation. Additionally, I am more than happy to add more formal documentation, specifically for the ConstrainedMaxPosteriorSampling class, I just wasn't sure where this documentation should go in within botorch/docs?

jacobrgardner

Reviewed By: Balandat

Differential Revision: D37669540

Pulled By: dme65

fbshipit-source-id: c2c602e9cd2af2730c0df5553246e793499d2f18

Author: nataliemaus

botorch/generation/sampling.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
@@ -29,6 +29,9 @@
 )
 from botorch.generation.utils import _flip_sub_unique
 from botorch.models.model import Model
+
+from botorch.models.model_list_gp_regression import ModelListGP
+from botorch.models.multitask import MultiTaskGP
 from botorch.utils.sampling import batched_multinomial
 from botorch.utils.transforms import standardize
 from torch import Tensor
@@ -122,9 +125,11 @@ def forward(
             observation_noise=observation_noise,
             posterior_transform=self.posterior_transform,
         )
-
         # num_samples x batch_shape x N x m
         samples = posterior.rsample(sample_shape=torch.Size([num_samples]))
+        return self.maximize_samples(X, samples, num_samples)
+
+    def maximize_samples(self, X: Tensor, samples: Tensor, num_samples: int = 1):
         obj = self.objective(samples, X=X)  # num_samples x batch_shape x N
         if self.replacement:
             # if we allow replacement then things are simple(r)
@@ -232,3 +237,111 @@ def forward(self, X: Tensor, num_samples: int = 1) -> Tensor:
         )
         # now do some gathering acrobatics to select the right elements from X
         return torch.gather(X, -2, idcs.unsqueeze(-1).expand(*idcs.shape, X.size(-1)))
+
+
+class ConstrainedMaxPosteriorSampling(MaxPosteriorSampling):
+    r"""Sample from a set of points according to
+    their max posterior value,
+    which also likely meet a set of constraints
+    c1(x) <= 0, c2(x) <= 0, ..., cm(x) <= 0
+    c1, c2, ..., cm are black-box constraint functions
+    Each constraint function is modeled by a seperate
+    surrogate GP constraint model
+    We sample points for which the posterior value
+    for each constraint model <= 0,
+    as described in https://doi.org/10.48550/arxiv.2002.08526
+
+    Example:
+        >>> CMPS = ConstrainedMaxPosteriorSampling(model,
+                    constraint_model=ModelListGP(cmodel1, cmodel2,
+                    ..., cmodelm)  # models w/ feature dim d=3
+        >>> X = torch.rand(2, 100, 3)
+        >>> sampled_X = CMPS(X, num_samples=5)
+    """
+
+    def __init__(
+        self,
+        model: Model,
+        constraint_model: Union[ModelListGP, MultiTaskGP],
+        objective: Optional[MCAcquisitionObjective] = None,
+        posterior_transform: Optional[PosteriorTransform] = None,
+        replacement: bool = True,
+        minimize_constraints_only: bool = False,
+    ) -> None:
+        r"""Constructor for the SamplingStrategy base class.
+
+        Args:
+            model: A fitted model.
+            objective: The MCAcquisitionObjective under
+                which the samples are evaluated.
+                Defaults to `IdentityMCObjective()`.
+            posterior_transform: An optional PosteriorTransform.
+            replacement: If True, sample with replacement.
+            constraint_model: either a ModelListGP where each submodel
+                is a GP model for one constraint function,
+                or a MultiTaskGP model where each task is one
+                constraint function
+                All constraints are of the form c(x) <= 0.
+                In the case when the constraint model predicts
+                that all candidates violate constraints,
+                we pick the candidates with minimum violation.
+            minimize_constraints_only: False by default, if true,
+                we will automatically return the candidates
+                with minimum posterior constraint values,
+                (minimum predicted c(x) summed over all constraints)
+                reguardless of predicted objective values.
+        """
+        super().__init__(
+            model=model,
+            objective=objective,
+            posterior_transform=posterior_transform,
+            replacement=replacement,
+        )
+        self.constraint_model = constraint_model
+        self.minimize_constraints_only = minimize_constraints_only
+
+    def forward(
+        self, X: Tensor, num_samples: int = 1, observation_noise: bool = False
+    ) -> Tensor:
+        r"""Sample from the model posterior.
+
+        Args:
+            X: A `batch_shape x N x d`-dim Tensor
+                from which to sample (in the `N`
+                dimension) according to the maximum
+                posterior value under the objective.
+            num_samples: The number of samples to draw.
+            observation_noise: If True, sample with observation noise.
+
+        Returns:
+            A `batch_shape x num_samples x d`-dim
+            Tensor of samples from `X`, where
+            `X[..., i, :]` is the `i`-th sample.
+        """
+        posterior = self.model.posterior(X, observation_noise=observation_noise)
+        samples = posterior.rsample(sample_shape=torch.Size([num_samples]))
+
+        c_posterior = self.constraint_model.posterior(
+            X, observation_noise=observation_noise
+        )
+        constraint_samples = c_posterior.rsample(sample_shape=torch.Size([num_samples]))
+        valid_samples = constraint_samples <= 0
+        if valid_samples.shape[-1] > 1:  # if more than one constraint
+            valid_samples = torch.all(valid_samples, dim=-1).unsqueeze(-1)
+        if (valid_samples.sum() == 0) or self.minimize_constraints_only:
+            # if none of the samples meet the constraints
+            # we pick the one that minimizes total violation
+            constraint_samples = constraint_samples.sum(dim=-1)
+            idcs = torch.argmin(constraint_samples, dim=-1)
+            if idcs.ndim > 1:
+                idcs = idcs.permute(*range(1, idcs.ndim), 0)
+            idcs = idcs.unsqueeze(-1).expand(*idcs.shape, X.size(-1))
+            Xe = X.expand(*constraint_samples.shape[1:], X.size(-1))
+            return torch.gather(Xe, -2, idcs)
+        # replace all violators with -infinty so it will never choose them
+        replacement_infs = -torch.inf * torch.ones(samples.shape).to(X.device).to(
+            X.dtype
+        )
+        samples = torch.where(valid_samples, samples, replacement_infs)
+
+        return self.maximize_samples(X, samples, num_samples)

test/generation/test_sampling.py

@@ -19,9 +19,12 @@
 )
 from botorch.generation.sampling import (
     BoltzmannSampling,
+    ConstrainedMaxPosteriorSampling,
     MaxPosteriorSampling,
     SamplingStrategy,
 )
+from botorch.models import SingleTaskGP
+from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
 
 
@@ -190,3 +193,130 @@ def test_boltzmann_sampling(self):
             BS = BoltzmannSampling(acqf, eta=10.0)
             samples = BS(X, num_samples=1)
             self.assertTrue(torch.equal(samples, X[max_idx, :]))
+
+
+class TestConstrainedMaxPosteriorSampling(BotorchTestCase):
+    def test_init(self):
+        mm = MockModel(MockPosterior(mean=None))
+        cmms = MockModel(MockPosterior(mean=None))
+        MPS = ConstrainedMaxPosteriorSampling(mm, cmms)
+        self.assertEqual(MPS.model, mm)
+        self.assertTrue(MPS.replacement)
+        self.assertIsInstance(MPS.objective, IdentityMCObjective)
+        obj = LinearMCObjective(torch.rand(2))
+        MPS = ConstrainedMaxPosteriorSampling(
+            mm, cmms, objective=obj, replacement=False
+        )
+        self.assertEqual(MPS.objective, obj)
+        self.assertFalse(MPS.replacement)
+
+    def test_constrained_max_posterior_sampling(self):
+        batch_shapes = (torch.Size(), torch.Size([3]), torch.Size([3, 2]))
+        dtypes = (torch.float, torch.double)
+        for batch_shape, dtype, N, num_samples, d in itertools.product(
+            batch_shapes, dtypes, (5, 6), (1, 2), (1, 2)
+        ):
+            tkwargs = {"device": self.device, "dtype": dtype}
+            # X is `batch_shape x N x d` = batch_shape x N x 1.
+            X = torch.randn(*batch_shape, N, d, **tkwargs)
+            # the event shape is `num_samples x batch_shape x N x m`
+            psamples = torch.zeros(num_samples, *batch_shape, N, 1, **tkwargs)
+            psamples[..., 0, :] = 1.0
+
+            # IdentityMCObjective, with replacement
+            with mock.patch.object(MockPosterior, "rsample", return_value=psamples):
+                mp = MockPosterior(None)
+                with mock.patch.object(MockModel, "posterior", return_value=mp):
+                    mm = MockModel(None)
+                    c_model1 = SingleTaskGP(X, torch.randn(X.shape[0:-1]).unsqueeze(-1))
+                    c_model2 = SingleTaskGP(X, torch.randn(X.shape[0:-1]).unsqueeze(-1))
+                    c_model3 = SingleTaskGP(X, torch.randn(X.shape[0:-1]).unsqueeze(-1))
+                    cmms1 = MockModel(MockPosterior(mean=None))
+                    cmms2 = ModelListGP(c_model1, c_model2)
+                    cmms3 = ModelListGP(c_model1, c_model2, c_model3)
+                    for cmms in [cmms1, cmms2, cmms3]:
+                        MPS = ConstrainedMaxPosteriorSampling(mm, cmms)
+                        s1 = MPS(X, num_samples=num_samples)
+                        # run again with minimize_constraints_only
+                        MPS = ConstrainedMaxPosteriorSampling(
+                            mm, cmms, minimize_constraints_only=True
+                        )
+                        s2 = MPS(X, num_samples=num_samples)
+                        assert s1.shape == s2.shape
+
+            # ScalarizedObjective, with replacement
+            with mock.patch.object(MockPosterior, "rsample", return_value=psamples):
+                mp = MockPosterior(None)
+                with mock.patch.object(MockModel, "posterior", return_value=mp):
+                    mm = MockModel(None)
+                    cmms = MockModel(None)
+                    with mock.patch.object(
+                        ScalarizedObjective, "forward", return_value=mp
+                    ):
+                        obj = ScalarizedObjective(torch.rand(2, **tkwargs))
+                        MPS = ConstrainedMaxPosteriorSampling(mm, cmms, objective=obj)
+                        s = MPS(X, num_samples=num_samples)
+                        self.assertTrue(s.shape[-2] == num_samples)
+
+            # ScalarizedPosteriorTransform w/ replacement
+            with mock.patch.object(MockPosterior, "rsample", return_value=psamples):
+                mp = MockPosterior(None)
+                with mock.patch.object(MockModel, "posterior", return_value=mp):
+                    mm = MockModel(None)
+                    cmms = MockModel(None)
+                    with mock.patch.object(
+                        ScalarizedPosteriorTransform, "forward", return_value=mp
+                    ):
+                        post_tf = ScalarizedPosteriorTransform(torch.rand(2, **tkwargs))
+                        MPS = ConstrainedMaxPosteriorSampling(
+                            mm, cmms, posterior_transform=post_tf
+                        )
+                        s = MPS(X, num_samples=num_samples)
+                        self.assertTrue(s.shape[-2] == num_samples)
+
+            # ScalarizedPosteriorTransform and Scalarized obj
+            mp = MockPosterior(None)
+            mm = MockModel(posterior=mp)
+            mp = MockPosterior(None)
+            cmms = MockModel(posterior=mp)
+            obj = ScalarizedObjective(torch.rand(2, **tkwargs))
+            post_tf = ScalarizedPosteriorTransform(torch.rand(2, **tkwargs))
+            with self.assertRaises(RuntimeError):
+                ConstrainedMaxPosteriorSampling(
+                    mm, cmms, posterior_transform=post_tf, objective=obj
+                )
+
+            # without replacement
+            psamples[..., 1, 0] = 1e-6
+            with mock.patch.object(MockPosterior, "rsample", return_value=psamples):
+                mp = MockPosterior(None)
+                with mock.patch.object(MockModel, "posterior", return_value=mp):
+                    mm = MockModel(None)
+                    cmms = MockModel(None)
+                    MPS = ConstrainedMaxPosteriorSampling(mm, cmms, replacement=False)
+                    if len(batch_shape) > 1:
+                        with self.assertRaises(NotImplementedError):
+                            MPS(X, num_samples=num_samples)
+                    else:
+                        s = MPS(X, num_samples=num_samples)
+                        self.assertTrue(s.shape[-2] == num_samples)
+
+            # ScalarizedMCObjective, without replacement
+            with mock.patch.object(MockPosterior, "rsample", return_value=psamples):
+                mp = MockPosterior(None)
+                with mock.patch.object(MockModel, "posterior", return_value=mp):
+                    mm = MockModel(None)
+                    cmms = MockModel(None)
+                    with mock.patch.object(
+                        ScalarizedObjective, "forward", return_value=mp
+                    ):
+                        obj = ScalarizedObjective(torch.rand(2, **tkwargs))
+                        MPS = ConstrainedMaxPosteriorSampling(
+                            mm, cmms, objective=obj, replacement=False
+                        )
+                        if len(batch_shape) > 1:
+                            with self.assertRaises(NotImplementedError):
+                                MPS(X, num_samples=num_samples)
+                        else:
+                            s = MPS(X, num_samples=num_samples)
+                            self.assertTrue(s.shape[-2] == num_samples)