Allow passing Y_samples directly in MARS.set_baseline_Y (#1364)

Summary:
Pull Request resolved: #1364

We currently sample from the model to generate `Y` to compute the `MVaR` of, and set the Pareto subset of MVaR as the `baseline_Y`. This adds an option to skip the sampling bit and provide the `Y` directly. This saves us compute in cases where we may have access to pre-computed samples and may not want to re-sample.

Reviewed By: esantorella

Differential Revision: D38950165

fbshipit-source-id: 8eb24b086a03ad0035e4018f54975d057eaa1c95

Author: saitcakmak

botorch/acquisition/multi_objective/multi_output_risk_measures.py

@@ -38,6 +38,7 @@
 )
 from botorch.acquisition.risk_measures import CVaR, RiskMeasureMCObjective, VaR
 from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.warnings import BotorchWarning
 from botorch.models.model import Model
 from botorch.utils.multi_objective.pareto import is_non_dominated
 from botorch.utils.transforms import normalize
@@ -568,17 +569,34 @@ def __init__(
         self.mvar = MVaR(n_w=self.n_w, alpha=self.alpha)
         self._chebyshev_objective = None
 
-    def set_baseline_Y(self, model: Model, X_baseline: Tensor) -> None:
+    def set_baseline_Y(
+        self,
+        model: Optional[Model],
+        X_baseline: Optional[Tensor],
+        Y_samples: Optional[Tensor] = None,
+    ) -> None:
         r"""Set the `baseline_Y` based on the MVaR predictions of the `model`
         for `X_baseline`.
 
         Args:
             model: The model being used for MARS optimization. Must have a compatible
-                `InputPerturbation` transform attached.
+                `InputPerturbation` transform attached. Ignored if `Y_samples` is given.
             X_baseline: An `n x d`-dim tensor of previously evaluated points.
+                Ignored if `Y_samples` is given.
+            Y_samples: An optional `(n * n_w) x d`-dim tensor of predictions. If given,
+                instead of sampling from the model, these are used.
         """
-        with torch.no_grad():
-            Y = model.posterior(X_baseline.unsqueeze(-2)).mean.squeeze(-2)
+        if Y_samples is None:
+            with torch.no_grad():
+                Y = model.posterior(X_baseline.unsqueeze(-2)).mean.squeeze(-2)
+        else:
+            if model is not None or X_baseline is not None:
+                warnings.warn(
+                    "`model` and `X_baseline` are ignored when `Y_samples` is "
+                    "provided to `MARS.set_baseline_Y`.",
+                    BotorchWarning,
+                )
+            Y = Y_samples
         Y = self.preprocessing_function(Y)
         Y = self.mvar(Y).view(-1, Y.shape[-1])
         Y = Y[is_non_dominated(Y)]

test/acquisition/multi_objective/test_multi_output_risk_measures.py

@@ -4,9 +4,11 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import warnings
 from typing import Optional
 
 import torch
+from botorch import settings
 from botorch.acquisition.multi_objective.multi_output_risk_measures import (
     IndependentCVaR,
     IndependentVaR,
@@ -18,6 +20,7 @@
 )
 from botorch.acquisition.multi_objective.objective import IdentityMCMultiOutputObjective
 from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.warnings import BotorchWarning
 from botorch.models.deterministic import GenericDeterministicModel
 from botorch.models.transforms.input import InputPerturbation
 from botorch.utils.multi_objective.pareto import is_non_dominated
@@ -475,20 +478,24 @@ def test_set_baseline_Y(self):
         )
         model = GenericDeterministicModel(f=lambda X: X, num_outputs=2)
         model.input_transform = perturbation
-        mars.set_baseline_Y(
-            model=model, X_baseline=torch.tensor([[0.0, 0.0], [1.0, 1.0]])
-        )
+        X_baseline = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
+        mars.set_baseline_Y(model=model, X_baseline=X_baseline)
+        self.assertTrue(torch.equal(mars.baseline_Y, torch.tensor([[1.5, 1.5]])))
+        # With Y_samples.
+        mars._baseline_Y = None
+        Y_samples = model.posterior(X_baseline).mean
+        with warnings.catch_warnings(record=True) as ws, settings.debug(True):
+            mars.set_baseline_Y(model=model, X_baseline=X_baseline, Y_samples=Y_samples)
         self.assertTrue(torch.equal(mars.baseline_Y, torch.tensor([[1.5, 1.5]])))
+        self.assertTrue(any(w.category == BotorchWarning for w in ws))
         # With pre-processing function.
         mars = MARS(
             alpha=0.5,
             n_w=3,
             chebyshev_weights=[0.5, 0.5],
             preprocessing_function=lambda Y: -Y,
         )
-        mars.set_baseline_Y(
-            model=model, X_baseline=torch.tensor([[0.0, 0.0], [1.0, 1.0]])
-        )
+        mars.set_baseline_Y(model=model, X_baseline=X_baseline)
         self.assertTrue(torch.equal(mars.baseline_Y, torch.tensor([[-0.5, -0.5]])))
 
     def test_get_Y_normalization_bounds(self):