Re-index TorchOptConfig.objective_thresholds from (n_outcomes,) to (n_objectives,)

Summary:
Re-index `objective_thresholds` in `TorchOptConfig` from `(n_outcomes,)` with NaN
for non-objective outcomes to `(n_objectives,)` or `None`. This is a prerequisite
for supporting `ScalarizedObjective` as a sub-objective of `MultiObjective`.

Key changes:
- `extract_objective_thresholds` returns `(n_objectives,)` maximization-aligned
  array (sign-flipped for minimize objectives), or `None` for single-objective.
- `_untransform_objective_thresholds` indexes by objective index and un-flips
  the sign when converting back to raw `ObjectiveThreshold.bound`.
- Replaced `get_weighted_mc_objective_and_objective_thresholds` with
  `get_weighted_mc_objective`, which only returns the objective (thresholds
  no longer need transformation so callers use them directly).
- `infer_objective_thresholds` returns `(n_objectives,)` maximization-aligned.
- Removed `objective_thresholds` from `SubsetModelData` and `subset_model`
  (thresholds are per-objective, not per-outcome, so subsetting doesn't apply).
- Simplified `_objective_threshold_to_outcome_constraints` and pruning logic.
- Merged `_full_objective_thresholds` and `_objective_thresholds` into single
  `_objective_thresholds` in the `Acquisition` class.
- Pass thresholds directly as `ref_point` to BoTorch input constructors (using
  the new `ref_point` parameter from D96473523), avoiding the need to convert
  back to outcome space.

Differential Revision: D96391935

Author: saitcakmak

ax/adapter/adapter_utils.py

@@ -41,7 +41,7 @@
 from ax.core.types import TBounds, TCandidateMetadata, TNumeric
 from ax.exceptions.core import DataRequiredError, UserInputError
 from ax.generators.torch.botorch_moo_utils import (
-    get_weighted_mc_objective_and_objective_thresholds,
+    get_weighted_mc_objective,
     pareto_frontier_evaluator,
 )
 from ax.utils.common.logger import get_logger
@@ -205,23 +205,26 @@ def extract_objective_thresholds(
     objective: Objective,
     outcomes: list[str],
 ) -> npt.NDArray | None:
-    """Extracts objective thresholds' values, in the order of `outcomes`.
+    """Extracts objective thresholds' values, in the order of objectives.
 
-    Will return None if no objective thresholds, otherwise the extracted tensor
-    will be the same length as `outcomes`.
+    Will return None if no objective thresholds or if the objective is single-
+    objective. Otherwise the extracted array will have length ``n_objectives``
+    (matching the rows of the objective weight matrix).
 
-    Outcomes that are not part of an objective and the objectives that do no have
-    a corresponding objective threshold will be given a threshold of NaN. We will
-    later infer appropriate threshold values for the objectives that are given a
-    threshold of NaN.
+    Objectives that do not have a corresponding objective threshold will be
+    given a threshold of NaN. We will later infer appropriate threshold values
+    for those objectives.
+
+    The returned thresholds are maximization-aligned: for minimize objectives,
+    the threshold is negated.
 
     Args:
         objective_thresholds: Objective thresholds to extract values from.
         objective: The corresponding Objective, for validation purposes.
         outcomes: n-length list of names of metrics.
 
     Returns:
-        (n,) array of thresholds
+        ``(n_objectives,)`` array of maximization-aligned thresholds, or None.
     """
     if len(objective_thresholds) == 0:
         return None
@@ -242,11 +245,19 @@ def extract_objective_thresholds(
             f"Got {objective_thresholds=} and {objective=}."
         )
 
-    # Initialize these to be NaN to make sure that objective thresholds for
-    # non-objective metrics are never used.
-    obj_t = np.full(len(outcomes), float("nan"))
-    for metric, threshold in objective_threshold_dict.items():
-        obj_t[outcomes.index(metric)] = threshold
+    if not isinstance(objective, MultiObjective):
+        # Single objective — thresholds not applicable.
+        return None
+
+    n_objectives = len(objective.objectives)
+    obj_t = np.full(n_objectives, float("nan"))
+    for i, sub_obj in enumerate(objective.objectives):
+        if isinstance(sub_obj, ScalarizedObjective):
+            continue  # NaN — will be inferred later.
+        metric_name = sub_obj.metric.signature
+        if metric_name in objective_threshold_dict:
+            sign = -1.0 if sub_obj.minimize else 1.0
+            obj_t[i] = sign * objective_threshold_dict[metric_name]
     return obj_t
 
 
@@ -743,10 +754,8 @@ def pareto_frontier(
     if obj_t is None:
         return frontier_observations
 
-    # Apply appropriate weights and thresholds
-    obj, obj_t = get_weighted_mc_objective_and_objective_thresholds(
-        objective_weights=obj_w, objective_thresholds=obj_t
-    )
+    # Apply appropriate weights
+    obj = get_weighted_mc_objective(objective_weights=obj_w)
     f_t = obj(f)
 
     # Compute individual hypervolumes by taking the difference between the observation
@@ -911,10 +920,9 @@ def hypervolume(
         dtype=torch.bool,
         device=f.device,
     )
-    # Apply appropriate weights and thresholds
-    obj, obj_t = get_weighted_mc_objective_and_objective_thresholds(
-        objective_weights=obj_w, objective_thresholds=none_throws(obj_t)
-    )
+    # Apply appropriate weights
+    obj = get_weighted_mc_objective(objective_weights=obj_w)
+    obj_t = none_throws(obj_t)
     f_t = obj(f)
     obj_mask = (obj_w != 0).any(dim=0).nonzero().view(-1)
     selected_metrics_mask = selected_metrics_mask[obj_mask]

ax/adapter/tests/test_torch_moo_adapter.py

@@ -191,10 +191,10 @@ def helper_test_pareto_frontier(
             )
         )
         self.assertTrue(obj_t is not None)
+        # Thresholds are now (n_objectives,) and maximization-aligned.
+        # LEQ thresholds with bound=5.0 become -5.0 after sign flip.
         self.assertTrue(
-            torch.equal(
-                none_throws(obj_t)[:2], torch.full((2,), 5.0, dtype=torch.double)
-            )
+            torch.equal(none_throws(obj_t), torch.full((2,), -5.0, dtype=torch.double))
         )
         observed_frontier2 = pareto_frontier(
             adapter=adapter,

ax/adapter/tests/test_utils.py

@@ -122,32 +122,33 @@ def test_extract_objective_thresholds(self) -> None:
             for i, name in enumerate(outcomes[:3])
         ]
 
-        # None of no thresholds
+        # None if no thresholds
         self.assertIsNone(
             extract_objective_thresholds(
                 objective_thresholds=[], objective=objective, outcomes=outcomes
             )
         )
 
-        # Working case
+        # Working case: 3 objectives (all maximize), shape is (3,)
         obj_t = extract_objective_thresholds(
             objective_thresholds=objective_thresholds,
             objective=objective,
             outcomes=outcomes,
         )
-        expected_obj_t_not_nan = np.array([2.0, 3.0, 4.0])
-        self.assertTrue(np.array_equal(obj_t[:3], expected_obj_t_not_nan[:3]))
-        self.assertTrue(np.isnan(obj_t[-1]))
-        self.assertEqual(obj_t.shape[0], 4)
+        # All maximize, so thresholds are unchanged (sign = +1).
+        expected_obj_t = np.array([2.0, 3.0, 4.0])
+        self.assertTrue(np.array_equal(obj_t, expected_obj_t))
+        self.assertEqual(obj_t.shape[0], 3)
 
         # Returns NaN for objectives without a threshold.
         obj_t = extract_objective_thresholds(
             objective_thresholds=objective_thresholds[:2],
             objective=objective,
             outcomes=outcomes,
         )
-        self.assertTrue(np.array_equal(obj_t[:2], expected_obj_t_not_nan[:2]))
-        self.assertTrue(np.isnan(obj_t[-2:]).all())
+        self.assertTrue(np.array_equal(obj_t[:2], expected_obj_t[:2]))
+        self.assertTrue(np.isnan(obj_t[2]))
+        self.assertEqual(obj_t.shape[0], 3)
 
         # Fails if a threshold does not have a corresponding metric.
         objective2 = Objective(Metric("m1"), minimize=False)
@@ -158,15 +159,46 @@ def test_extract_objective_thresholds(self) -> None:
                 outcomes=outcomes,
             )
 
-        # Works with a single objective, single threshold
+        # Single objective returns None.
+        self.assertIsNone(
+            extract_objective_thresholds(
+                objective_thresholds=objective_thresholds[:1],
+                objective=objective2,
+                outcomes=outcomes,
+            )
+        )
+
+        # Maximize-alignment: minimize objectives get negated thresholds.
+        objective_with_min = MultiObjective(
+            objectives=[
+                Objective(metric=Metric("m1"), minimize=False),
+                Objective(metric=Metric("m2"), minimize=True),
+            ]
+        )
+        obj_thresholds_for_min = [
+            ObjectiveThreshold(
+                metric=Metric("m1"),
+                op=ComparisonOp.LEQ,
+                bound=2.0,
+                relative=False,
+            ),
+            ObjectiveThreshold(
+                metric=Metric("m2"),
+                op=ComparisonOp.LEQ,
+                bound=3.0,
+                relative=False,
+            ),
+        ]
         obj_t = extract_objective_thresholds(
-            objective_thresholds=objective_thresholds[:1],
-            objective=objective2,
+            objective_thresholds=obj_thresholds_for_min,
+            objective=objective_with_min,
             outcomes=outcomes,
         )
+        # m1 maximize: sign=+1, threshold=2.0 → 2.0
+        # m2 minimize: sign=-1, threshold=3.0 → -3.0
+        self.assertEqual(obj_t.shape[0], 2)
         self.assertEqual(obj_t[0], 2.0)
-        self.assertTrue(np.all(np.isnan(obj_t[1:])))
-        self.assertEqual(obj_t.shape[0], 4)
+        self.assertEqual(obj_t[1], -3.0)
 
         # Fails if relative
         objective_thresholds[2] = ObjectiveThreshold(

ax/adapter/torch.py

@@ -1111,12 +1111,14 @@ def _untransform_objective_thresholds(
         """
         obj_indices, obj_weights = extract_objectives(objective_weights)
         thresholds = []
-        for idx, w in zip(obj_indices, obj_weights):
+        for i, (idx, w) in enumerate(zip(obj_indices, obj_weights)):
             sign = torch.sign(w)
+            # Thresholds are maximization-aligned; undo sign flip to get raw bound.
+            raw_bound = float(sign * objective_thresholds[i].item())
             thresholds.append(
                 ObjectiveThreshold(
                     metric=opt_config_metrics[self.outcomes[idx]],
-                    bound=float(objective_thresholds[idx].item()),
+                    bound=raw_bound,
                     relative=False,
                     op=ComparisonOp.LEQ if sign < 0 else ComparisonOp.GEQ,
                 )

ax/generators/tests/test_botorch_moo_utils.py

@@ -11,12 +11,11 @@
 from unittest import mock
 from warnings import catch_warnings, simplefilter
 
-import numpy as np
 import torch
 from ax.core.search_space import SearchSpaceDigest
 from ax.generators.torch.botorch_modular.generator import BoTorchGenerator
 from ax.generators.torch.botorch_moo_utils import (
-    get_weighted_mc_objective_and_objective_thresholds,
+    get_weighted_mc_objective,
     infer_objective_thresholds,
     pareto_frontier_evaluator,
 )
@@ -213,19 +212,13 @@ def test_pareto_frontier_evaluator_with_nan(self) -> None:
 
 
 class BotorchMOOUtilsTest(TestCase):
-    def test_get_weighted_mc_objective_and_objective_thresholds(self) -> None:
+    def test_get_weighted_mc_objective(self) -> None:
         objective_weights = torch.tensor([[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]])
-        objective_thresholds = torch.arange(4, dtype=torch.float)
-        (
-            weighted_obj,
-            new_obj_thresholds,
-        ) = get_weighted_mc_objective_and_objective_thresholds(
+        weighted_obj = get_weighted_mc_objective(
             objective_weights=objective_weights,
-            objective_thresholds=objective_thresholds,
         )
         self.assertTrue(torch.equal(weighted_obj.weights, torch.tensor([1.0, 1.0])))
         self.assertEqual(weighted_obj.outcomes.tolist(), [1, 3])
-        self.assertTrue(torch.equal(new_obj_thresholds, objective_thresholds[[1, 3]]))
 
     # test infer objective thresholds alone
     @mock.patch(  # pyre-ignore
@@ -255,6 +248,10 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
             objective_weights = torch.tensor(
                 [[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]], **tkwargs
             )
+            # Expected: infer_reference_point returns (n_objectives,) in
+            # maximization-aligned space. With pareto_Y=[[-9, -3]] and
+            # scale=0.1, the result is [-9.9, -3.3].
+            expected_thresholds = torch.tensor([-9.9, -3.3], **tkwargs)
             with ExitStack() as es:
                 _mock_infer_reference_point = es.enter_context(
                     mock.patch(
@@ -282,10 +279,9 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
                     torch.tensor([[-9.0, -3.0]], **tkwargs),
                 )
             )
-            self.assertTrue(
-                torch.equal(obj_thresholds[:2], torch.tensor([9.9, 3.3], **tkwargs))
-            )
-            self.assertTrue(np.isnan(obj_thresholds[2].item()))
+            # Result is (n_objectives,) maximization-aligned.
+            self.assertEqual(obj_thresholds.shape[0], 2)
+            self.assertTrue(torch.equal(obj_thresholds, expected_thresholds))
 
             # test subset_model without subset_idcs
             with mock.patch.object(model, "posterior", return_value=posterior):
@@ -295,10 +291,8 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
                     outcome_constraints=outcome_constraints,
                     X_observed=Xs[0],
                 )
-            self.assertTrue(
-                torch.equal(obj_thresholds[:2], torch.tensor([9.9, 3.3], **tkwargs))
-            )
-            self.assertTrue(np.isnan(obj_thresholds[2].item()))
+            self.assertEqual(obj_thresholds.shape[0], 2)
+            self.assertTrue(torch.equal(obj_thresholds, expected_thresholds))
             # test passing subset_idcs
             subset_idcs = torch.tensor(
                 [0, 1], dtype=torch.long, device=tkwargs["device"]
@@ -312,10 +306,8 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
                     X_observed=Xs[0],
                     subset_idcs=subset_idcs,
                 )
-            self.assertTrue(
-                torch.equal(obj_thresholds[:2], torch.tensor([9.9, 3.3], **tkwargs))
-            )
-            self.assertTrue(np.isnan(obj_thresholds[2].item()))
+            self.assertEqual(obj_thresholds.shape[0], 2)
+            self.assertTrue(torch.equal(obj_thresholds, expected_thresholds))
             # test without subsetting (e.g. if there are
             # 3 metrics for 2 objectives + 1 outcome constraint)
             outcome_constraints = (
@@ -350,10 +342,8 @@ def test_infer_objective_thresholds(self, _, cuda: bool = False) -> None:
                     X_observed=Xs[0],
                     outcome_constraints=outcome_constraints,
                 )
-            self.assertTrue(
-                torch.equal(obj_thresholds[:2], torch.tensor([9.9, 3.3], **tkwargs))
-            )
-            self.assertTrue(np.isnan(obj_thresholds[2].item()))
+            self.assertEqual(obj_thresholds.shape[0], 2)
+            self.assertTrue(torch.equal(obj_thresholds, expected_thresholds))
 
     def test_infer_objective_thresholds_cuda(self) -> None:
         if torch.cuda.is_available():