Backprop wrt candidates in gen_candidates_torch (#766)

Summary:
## Motivation

Bugfix for #765. The current error was occurring when trying to optimize qKnowledgeGradient with gen_candidates_torch. The fix is to backprop only wrt candidates like what occurs in `gen_candidates_scipy`.

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

Yes.

Pull Request resolved: #766

Test Plan:
Not sure how to setup a unit test here.

## Related PRs

N/A.

Reviewed By: Balandat

Differential Revision: D29401945

Pulled By: wjmaddox

fbshipit-source-id: a00e0b93b41206487b8b5ecb39bd19116d533b58

Author: wjmaddox

botorch/generation/gen.py

@@ -276,8 +276,6 @@ def gen_candidates_torch(
     bayes_optimizer = optimizer(
         params=[clamped_candidates], lr=options.get("lr", 0.025)
     )
-    param_trajectory: Dict[str, List[Tensor]] = {"candidates": []}
-    loss_trajectory: List[float] = []
     i = 0
     stop = False
     stopping_criterion = ExpMAStoppingCriterion(
@@ -288,19 +286,18 @@ def gen_candidates_torch(
         loss = -acquisition_function(clamped_candidates).sum()
         if verbose:
             print("Iter: {} - Value: {:.3f}".format(i, -(loss.item())))
-        loss_trajectory.append(loss.item())
-        param_trajectory["candidates"].append(clamped_candidates.clone())
 
         def closure():
             bayes_optimizer.zero_grad()
-            loss.backward()
+            output_grad = torch.autograd.grad(loss, clamped_candidates)[0]
+            clamped_candidates.grad = output_grad
             return loss
 
         bayes_optimizer.step(closure)
         with torch.no_grad():
             clamped_candidates = columnwise_clamp(
                 X=clamped_candidates, lower=lower_bounds, upper=upper_bounds
-            )
+            ).requires_grad_(True)
         stop = stopping_criterion.evaluate(fvals=loss.detach())
     clamped_candidates = columnwise_clamp(
         X=clamped_candidates,

test/generation/test_gen.py

@@ -9,7 +9,7 @@
 from unittest import mock
 
 import torch
-from botorch.acquisition import qExpectedImprovement
+from botorch.acquisition import qExpectedImprovement, qKnowledgeGradient
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.fit import fit_gpytorch_model
 from botorch.generation.gen import (
@@ -72,15 +72,28 @@ def test_gen_candidates(self, gen_candidates=gen_candidates_scipy, options=None)
         options = {**options, "maxiter": 5}
         for double in (True, False):
             self._setUp(double=double)
-            qEI = qExpectedImprovement(self.model, best_f=self.f_best)
-            candidates, _ = gen_candidates(
-                initial_conditions=self.initial_conditions,
-                acquisition_function=qEI,
-                lower_bounds=0,
-                upper_bounds=1,
-                options=options or {},
-            )
-            self.assertTrue(-EPS <= candidates <= 1 + EPS)
+            acqfs = [
+                qExpectedImprovement(self.model, best_f=self.f_best),
+                qKnowledgeGradient(
+                    self.model, num_fantasies=4, current_value=self.f_best
+                ),
+            ]
+            for acqf in acqfs:
+                ics = self.initial_conditions
+                if isinstance(acqf, qKnowledgeGradient):
+                    ics = ics.repeat(5, 1)
+
+                candidates, _ = gen_candidates(
+                    initial_conditions=ics,
+                    acquisition_function=acqf,
+                    lower_bounds=0,
+                    upper_bounds=1,
+                    options=options or {},
+                )
+                if isinstance(acqf, qKnowledgeGradient):
+                    candidates = acqf.extract_candidates(candidates)
+
+                self.assertTrue(-EPS <= candidates <= 1 + EPS)
 
     def test_gen_candidates_torch(self):
         self.test_gen_candidates(
@@ -96,18 +109,30 @@ def test_gen_candidates_with_none_fixed_features(
         options = {**options, "maxiter": 5}
         for double in (True, False):
             self._setUp(double=double, expand=True)
-            qEI = qExpectedImprovement(self.model, best_f=self.f_best)
-            candidates, _ = gen_candidates(
-                initial_conditions=self.initial_conditions,
-                acquisition_function=qEI,
-                lower_bounds=0,
-                upper_bounds=1,
-                fixed_features={1: None},
-                options=options or {},
-            )
-            candidates = candidates.squeeze(0)
-            self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
-            self.assertTrue(candidates[1].item() == 1.0)
+            acqfs = [
+                qExpectedImprovement(self.model, best_f=self.f_best),
+                qKnowledgeGradient(
+                    self.model, num_fantasies=4, current_value=self.f_best
+                ),
+            ]
+            for acqf in acqfs:
+                ics = self.initial_conditions
+                if isinstance(acqf, qKnowledgeGradient):
+                    ics = ics.repeat(5, 1)
+
+                candidates, _ = gen_candidates(
+                    initial_conditions=ics,
+                    acquisition_function=acqf,
+                    lower_bounds=0,
+                    upper_bounds=1,
+                    fixed_features={1: None},
+                    options=options or {},
+                )
+                if isinstance(acqf, qKnowledgeGradient):
+                    candidates = acqf.extract_candidates(candidates)
+                candidates = candidates.squeeze(0)
+                self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
+                self.assertTrue(candidates[1].item() == 1.0)
 
     def test_gen_candidates_torch_with_none_fixed_features(self):
         self.test_gen_candidates_with_none_fixed_features(
@@ -121,18 +146,32 @@ def test_gen_candidates_with_fixed_features(
         options = {**options, "maxiter": 5}
         for double in (True, False):
             self._setUp(double=double, expand=True)
-            qEI = qExpectedImprovement(self.model, best_f=self.f_best)
-            candidates, _ = gen_candidates(
-                initial_conditions=self.initial_conditions,
-                acquisition_function=qEI,
-                lower_bounds=0,
-                upper_bounds=1,
-                fixed_features={1: 0.25},
-                options=options,
-            )
-            candidates = candidates.squeeze(0)
-            self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
-            self.assertTrue(candidates[1].item() == 0.25)
+            acqfs = [
+                qExpectedImprovement(self.model, best_f=self.f_best),
+                qKnowledgeGradient(
+                    self.model, num_fantasies=4, current_value=self.f_best
+                ),
+            ]
+            for acqf in acqfs:
+                ics = self.initial_conditions
+                if isinstance(acqf, qKnowledgeGradient):
+                    ics = ics.repeat(5, 1)
+
+                candidates, _ = gen_candidates(
+                    initial_conditions=ics,
+                    acquisition_function=acqf,
+                    lower_bounds=0,
+                    upper_bounds=1,
+                    fixed_features={1: 0.25},
+                    options=options,
+                )
+
+                if isinstance(acqf, qKnowledgeGradient):
+                    candidates = acqf.extract_candidates(candidates)
+
+                candidates = candidates.squeeze(0)
+                self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
+                self.assertTrue(candidates[1].item() == 0.25)
 
     def test_gen_candidates_scipy_with_fixed_features_inequality_constraints(self):
         options = {"maxiter": 5}