Use discrete_values and post_processing_func in optimize_with_nsgaii (#4971)

Summary:
Pull Request resolved: #4971

`optimize_with_nsgaii` may produce invalid candidates as it currently doesn't take into account whether a parameter is float or integer-valued. This diff addresses this by passing in `rounding_func` by default when `optimize_with_nsgaii` is called through MBM.

Reviewed By: sdaulton, saitcakmak

Differential Revision: D95243318

Author: David Eriksson

ax/generators/torch/botorch_modular/acquisition.py

@@ -47,7 +47,10 @@
 from botorch.acquisition.input_constructors import get_acqf_input_constructor
 from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
 from botorch.acquisition.logei import qLogProbabilityOfFeasibility
-from botorch.acquisition.multioutput_acquisition import MultiOutputAcquisitionFunction
+from botorch.acquisition.multioutput_acquisition import (
+    MultiOutputAcquisitionFunction,
+    MultiOutputAcquisitionFunctionWrapper,
+)
 from botorch.acquisition.objective import MCAcquisitionObjective, PosteriorTransform
 from botorch.exceptions.errors import BotorchError, InputDataError
 from botorch.generation.sampling import SamplingStrategy
@@ -66,7 +69,7 @@
 )
 from botorch.optim.parameter_constraints import evaluate_feasibility
 from botorch.utils.constraints import get_outcome_constraint_transforms
-from pyre_extensions import none_throws
+from pyre_extensions import assert_is_instance, none_throws
 from torch import Tensor
 
 try:
@@ -816,18 +819,17 @@ def optimize(
             )
         elif optimizer == "optimize_with_nsgaii":
             if optimize_with_nsgaii is not None:
-                # TODO: support post_processing_func
+                acqf = assert_is_instance(
+                    self.acqf, MultiOutputAcquisitionFunctionWrapper
+                )
                 candidates, acqf_values = optimize_with_nsgaii(
                     acq_function=self.acqf,
                     bounds=bounds,
                     q=n,
                     fixed_features=fixed_features,
-                    # We use pyre-ignore here to avoid a circular import.
-                    # pyre-ignore [6]: Incompatible parameter type [6]: In call `len`,
-                    # for 1st positional argument, expected
-                    # `pyre_extensions.PyreReadOnly[Sized]` but got `Union[Tensor,
-                    # Module]`.
-                    num_objectives=len(self.acqf.acqfs),
+                    num_objectives=len(acqf.acqfs),
+                    discrete_choices=discrete_choices if discrete_choices else None,
+                    post_processing_func=rounding_func,
                     **optimizer_options_with_defaults,
                 )
             else:

ax/generators/torch/tests/test_acquisition.py

@@ -2228,6 +2228,7 @@ def test_optimize(self) -> None:
             q=n,
             num_objectives=2,
             fixed_features=self.fixed_features,
+            post_processing_func=self.rounding_func,
             **optimizer_options,
         )
         # can't use assert_called_with on bounds due to ambiguous bool comparison
@@ -2242,6 +2243,44 @@ def test_optimize(self) -> None:
             )
         )
 
+    @skip_if_import_error
+    def test_optimize_with_nsgaii_applies_rounding(self) -> None:
+        """Test that optimize_with_nsgaii applies the rounding function correctly.
+
+        This tests a mixed setting where only some parameters are rounded (e.g.,
+        integer-valued), while others remain continuous.
+        """
+
+        # Create a rounding function that only rounds the first parameter (index 0)
+        # This simulates a mixed search space with one integer parameter
+        def rounding_func(X: Tensor) -> Tensor:
+            X_rounded = X.clone()
+            X_rounded[..., 0] = X_rounded[..., 0].round()
+            return X_rounded
+
+        acquisition = self.get_acquisition_function(fixed_features=self.fixed_features)
+        n = 3
+        optimizer_options = {"max_gen": 5, "population_size": 20, "seed": 0}
+
+        candidates, _, _ = acquisition.optimize(
+            n=n,
+            search_space_digest=self.search_space_digest,
+            inequality_constraints=self.inequality_constraints,
+            fixed_features=self.fixed_features,
+            rounding_func=rounding_func,
+            optimizer_options=optimizer_options,
+        )
+
+        # Assert that the first parameter (index 0) is rounded to integers
+        self.assertTrue(
+            torch.equal(candidates[:, 0], candidates[:, 0].round()),
+            f"First parameter should be rounded but got: {candidates[:, 0]}",
+        )
+        # Assert that other parameters (indices 1, 2) may have non-integer values
+        # (they won't necessarily be non-integers, but they weren't explicitly rounded)
+        # The key assertion is that the rounding function was applied to produce
+        # integer values for the first parameter
+
     def test_evaluate(self) -> None:
         acquisition = self.get_acquisition_function()
         with mock.patch.object(acquisition.acqf, "forward") as mock_forward: