Revamp qKG constructor to avoid issue with missing objective (#351)

Summary:
Pull Request resolved: #351

D18739840 added a check to `MCAcquisitionFunction` that would raise if no objective was specified for multi-output models. This doesn't play well with the current qKG constructor, something that was overlooked in testing.

This modifies the qKG constructor to avoid this explicit check.

Reviewed By: 2timesjay

Differential Revision: D19352464

fbshipit-source-id: eb99329d393be68f5a9c652e1244c63efc2b7d87

Author: Balandat

botorch/acquisition/knowledge_gradient.py

@@ -103,14 +103,20 @@ def __init__(
                 )
         else:
             num_fantasies = sampler.sample_shape[0]
-        super().__init__(model=model, sampler=sampler, X_pending=X_pending)
+        super(MCAcquisitionFunction, self).__init__(model=model)
         # if not explicitly specified, we use the posterior mean for linear objs
         if isinstance(objective, MCAcquisitionObjective) and inner_sampler is None:
             inner_sampler = SobolQMCNormalSampler(
                 num_samples=128, resample=False, collapse_batch_dims=True
             )
-        self.inner_sampler = inner_sampler
+        if objective is None and model.num_outputs != 1:
+            raise UnsupportedError(
+                "Must specify an objective when using a multi-output model."
+            )
+        self.sampler = sampler
         self.objective = objective
+        self.set_X_pending(X_pending)
+        self.inner_sampler = inner_sampler
         self.num_fantasies = num_fantasies
         self.current_value = current_value
 

test/acquisition/test_knowledge_gradient.py

@@ -16,10 +16,12 @@
     qMultiFidelityKnowledgeGradient,
 )
 from botorch.acquisition.monte_carlo import qSimpleRegret
-from botorch.acquisition.objective import GenericMCObjective
+from botorch.acquisition.objective import GenericMCObjective, ScalarizedObjective
 from botorch.exceptions.errors import UnsupportedError
+from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.sampling.samplers import IIDNormalSampler, SobolQMCNormalSampler
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
+from gpytorch.distributions import MultitaskMultivariateNormal
 
 
 NO = "botorch.utils.testing.MockModel.num_outputs"
@@ -92,6 +94,20 @@ def test_initialize_q_knowledge_gradient(self):
             self.assertIsNone(qKG.X_pending)
             self.assertTrue(torch.equal(qKG.current_value, current_value))
             self.assertEqual(qKG.get_augmented_q_batch_size(q=3), 8 + 3)
+            # test construction with non-MC objective (ScalarizedObjective)
+            qKG_s = qKnowledgeGradient(
+                model=mm,
+                num_fantasies=16,
+                sampler=sampler,
+                objective=ScalarizedObjective(weights=torch.rand(2)),
+            )
+            self.assertIsNone(qKG_s.inner_sampler)
+            self.assertIsInstance(qKG_s.objective, ScalarizedObjective)
+            # test error if no objective and multi-output model
+            mean2 = torch.zeros(1, 2, device=self.device, dtype=dtype)
+            mm2 = MockModel(MockPosterior(mean=mean2))
+            with self.assertRaises(UnsupportedError):
+                qKnowledgeGradient(model=mm2)
 
     def test_evaluate_q_knowledge_gradient(self):
         for dtype in (torch.float, torch.double):
@@ -172,6 +188,30 @@ def test_evaluate_q_knowledge_gradient(self):
                     self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1]))
             self.assertTrue(torch.allclose(val, objective(samples).mean(), atol=1e-4))
             self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
+            # test non-MC objective (ScalarizedObjective)
+            weights = torch.rand(2, device=self.device, dtype=dtype)
+            objective = ScalarizedObjective(weights=weights)
+            mean = torch.tensor([1.0, 0.5], device=self.device, dtype=dtype).expand(
+                n_f, 1, 2
+            )
+            cov = torch.tensor(
+                [[1.0, 0.1], [0.1, 0.5]], device=self.device, dtype=dtype
+            ).expand(n_f, 2, 2)
+            posterior = GPyTorchPosterior(MultitaskMultivariateNormal(mean, cov))
+            mfm = MockModel(posterior)
+            with mock.patch.object(MockModel, "fantasize", return_value=mfm) as patch_f:
+                with mock.patch(NO, new_callable=mock.PropertyMock) as mock_num_outputs:
+                    mock_num_outputs.return_value = 2
+                    mm = MockModel(None)
+                    qKG = qKnowledgeGradient(
+                        model=mm, num_fantasies=n_f, objective=objective
+                    )
+                    val = qKG(X)
+                    patch_f.assert_called_once()
+                    cargs, ckwargs = patch_f.call_args
+                    self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1]))
+                    val_expected = (mean * weights).sum(-1).mean(0)
+                    self.assertTrue(torch.allclose(val, val_expected))
 
 
 class TestQMultiFidelityKnowledgeGradient(BotorchTestCase):