From fdb718e300025fe2864c47dd8927a33b67fb1278 Mon Sep 17 00:00:00 2001
From: Louis Tiao <tiao@meta.com>
Date: Tue, 10 Sep 2024 11:40:35 -0700
Subject: [PATCH] Remove extraneous `current_value` computation from input
 constructor for `construct_inputs_qKG` (#2520)

Summary:
Pull Request resolved: https://github.com/pytorch/botorch/pull/2520

The input constructor `construct_inputs_qKG` previously computed `current_value` which is extraneous (it is merely an additive constant in the acquisition function).

However, `construct_inputs_qMFKG`, which *does* require an explicit `current_value` (in the cost-aware setting), was piggy-backing off `construct_inputs_qKG`, which I suspect is why this persisted.

Having said that, `construct_inputs_qKG` computes `current_value` incorrectly for the purposes of  `construct_inputs_qMFKG` anyway (as it failed to take fidelity input dimension into consideration), so it wasn't actually useful or beneficial under any circumstances.

With D62391106, which calculates `current_value` correctly for `construct_inputs_qMFKG`, we can safely remove this computation altogether from `construct_inputs_qKG`

Reviewed By: Balandat

Differential Revision: D62415444
---
 botorch/acquisition/input_constructors.py   | 33 ++++++++++---------
 test/acquisition/test_input_constructors.py | 35 +++++++++++++++------
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/botorch/acquisition/input_constructors.py b/botorch/acquisition/input_constructors.py
index 06c069fd71..0714f79fec 100644
--- a/botorch/acquisition/input_constructors.py
+++ b/botorch/acquisition/input_constructors.py
@@ -1240,30 +1240,35 @@ def construct_inputs_qKG(
     objective: Optional[MCAcquisitionObjective] = None,
     posterior_transform: Optional[PosteriorTransform] = None,
     num_fantasies: int = 64,
+    with_current_value: bool = False,
     **optimize_objective_kwargs: TOptimizeObjectiveKwargs,
 ) -> dict[str, Any]:
     r"""Construct kwargs for `qKnowledgeGradient` constructor."""
 
-    X = _get_dataset_field(training_data, "X", first_only=True)
-    _bounds = torch.as_tensor(bounds, dtype=X.dtype, device=X.device)
-
-    _, current_value = optimize_objective(
-        model=model,
-        bounds=_bounds.t(),
-        q=1,
-        objective=objective,
-        posterior_transform=posterior_transform,
-        **optimize_objective_kwargs,
-    )
-
-    return {
+    inputs_qkg = {
         "model": model,
         "objective": objective,
         "posterior_transform": posterior_transform,
         "num_fantasies": num_fantasies,
-        "current_value": current_value.detach().cpu().max(),
     }
 
+    if with_current_value:
+
+        X = _get_dataset_field(training_data, "X", first_only=True)
+        _bounds = torch.as_tensor(bounds, dtype=X.dtype, device=X.device)
+
+        _, current_value = optimize_objective(
+            model=model,
+            bounds=_bounds.t(),
+            q=1,
+            objective=objective,
+            posterior_transform=posterior_transform,
+            **optimize_objective_kwargs,
+        )
+        inputs_qkg["current_value"] = current_value.detach().cpu().max()
+
+    return inputs_qkg
+
 
 @acqf_input_constructor(qMultiFidelityKnowledgeGradient)
 def construct_inputs_qMFKG(
diff --git a/test/acquisition/test_input_constructors.py b/test/acquisition/test_input_constructors.py
index f07bf002ef..3d2abf13df 100644
--- a/test/acquisition/test_input_constructors.py
+++ b/test/acquisition/test_input_constructors.py
@@ -1257,24 +1257,39 @@ def test_construct_inputs_qLogNParEGO(self) -> None:
 
 class TestKGandESAcquisitionFunctionInputConstructors(InputConstructorBaseTestCase):
     def test_construct_inputs_kg(self) -> None:
-        current_value = torch.tensor(1.23)
-        with mock.patch(
-            target="botorch.acquisition.input_constructors.optimize_objective",
-            return_value=(None, current_value),
-        ):
-            from botorch.acquisition import input_constructors
+        func = get_acqf_input_constructor(qKnowledgeGradient)
 
-            func = input_constructors.get_acqf_input_constructor(qKnowledgeGradient)
+        with self.subTest("test_with_current_value"):
+
+            current_value = torch.tensor(1.23)
+
+            with mock.patch(
+                target="botorch.acquisition.input_constructors.optimize_objective",
+                return_value=(None, current_value),
+            ):
+
+                kwargs = func(
+                    model=mock.Mock(),
+                    training_data=self.blockX_blockY,
+                    objective=LinearMCObjective(torch.rand(2)),
+                    bounds=self.bounds,
+                    num_fantasies=33,
+                    with_current_value=True,
+                )
+
+                self.assertEqual(kwargs["num_fantasies"], 33)
+                self.assertEqual(kwargs["current_value"], current_value)
+
+        with self.subTest("test_without_current_value"):
             kwargs = func(
                 model=mock.Mock(),
                 training_data=self.blockX_blockY,
                 objective=LinearMCObjective(torch.rand(2)),
                 bounds=self.bounds,
                 num_fantasies=33,
+                with_current_value=False,
             )
-
-            self.assertEqual(kwargs["num_fantasies"], 33)
-            self.assertEqual(kwargs["current_value"], current_value)
+            self.assertNotIn("current_value", kwargs)
 
     def test_construct_inputs_mes(self) -> None:
         func = get_acqf_input_constructor(qMaxValueEntropy)