Refactor Objective, OutcomeConstraint to be expression-based; remove Metrics from OptimizationConfig (#5000)

ax/adapter/adapter_utils.py

@@ -23,7 +23,7 @@
 )
 from ax.core.arm import Arm
 from ax.core.experiment import Experiment
-from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
+from ax.core.objective import Objective
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
@@ -206,6 +206,7 @@ def extract_objective_thresholds(
     objective_thresholds: TRefPoint,
     objective: Objective,
     outcomes: list[str],
+    experiment: Experiment,
 ) -> npt.NDArray | None:
     """Extracts objective thresholds' values, in the order of `outcomes`.
 
@@ -221,6 +222,7 @@ def extract_objective_thresholds(
         objective_thresholds: Objective thresholds to extract values from.
         objective: The corresponding Objective, for validation purposes.
         outcomes: n-length list of names of metrics.
+        experiment: The experiment, used to map metric names to signatures.
 
     Returns:
         (n,) array of thresholds
@@ -230,15 +232,19 @@ def extract_objective_thresholds(
 
     objective_threshold_dict = {}
     for ot in objective_thresholds:
+        ot_signature = experiment.get_metric(ot.metric_names[0]).signature
         if ot.relative:
             raise ValueError(
-                f"Objective {ot.metric.signature} has a relative threshold that "
+                f"Objective {ot_signature} has a relative threshold that "
                 f"is not supported here."
             )
-        objective_threshold_dict[ot.metric.signature] = ot.bound
+        objective_threshold_dict[ot_signature] = ot.bound
 
     # Check that all thresholds correspond to a metric.
-    if set(objective_threshold_dict.keys()).difference(set(objective.metric_names)):
+    obj_metric_signatures = [
+        experiment.get_metric(name).signature for name in objective.metric_names
+    ]
+    if set(objective_threshold_dict.keys()).difference(set(obj_metric_signatures)):
         raise ValueError(
             "Some objective thresholds do not have corresponding metrics. "
             f"Got {objective_thresholds=} and {objective=}."
@@ -252,7 +258,9 @@ def extract_objective_thresholds(
     return obj_t
 
 
-def extract_objective_weights(objective: Objective, outcomes: list[str]) -> npt.NDArray:
+def extract_objective_weights(
+    objective: Objective, outcomes: list[str], experiment: Experiment
+) -> npt.NDArray:
     """Extract a weights for objectives.
 
     Weights are for a maximization problem.
@@ -268,29 +276,24 @@ def extract_objective_weights(objective: Objective, outcomes: list[str]) -> npt.
 
     Args:
         objective: Objective to extract weights from.
-        outcomes: n-length list of names of metrics.
+        outcomes: n-length list of metric signatures.
+        experiment: The experiment, used to map metric names to signatures.
 
     Returns:
         n-length array of weights.
 
     """
     objective_weights = np.zeros(len(outcomes))
-    if isinstance(objective, ScalarizedObjective):
-        s = -1.0 if objective.minimize else 1.0
-        for obj_metric, obj_weight in objective.metric_weights:
-            objective_weights[outcomes.index(obj_metric.signature)] = obj_weight * s
-    elif isinstance(objective, MultiObjective):
-        for obj in objective.objectives:
-            s = -1.0 if obj.minimize else 1.0
-            objective_weights[outcomes.index(obj.metric.signature)] = s
-    else:
-        s = -1.0 if objective.minimize else 1.0
-        objective_weights[outcomes.index(objective.metric.signature)] = s
+    # metric_weights returns sign-encoded (name, weight) tuples for all
+    # objective types (single, scalarized, multi).
+    for obj_metric_name, obj_weight in objective.metric_weights:
+        sig = experiment.get_metric(obj_metric_name).signature
+        objective_weights[outcomes.index(sig)] = obj_weight
     return objective_weights
 
 
 def extract_objective_weight_matrix(
-    objective: Objective, outcomes: list[str]
+    objective: Objective, outcomes: list[str], experiment: Experiment
 ) -> npt.NDArray:
     """Extract a 2D weight matrix for objectives.
 
@@ -304,23 +307,31 @@ def extract_objective_weight_matrix(
 
     Args:
         objective: Objective to extract weights from.
-        outcomes: n-length list of names of metrics.
+        outcomes: n-length list of signatures of metrics.
 
     Returns:
         ``(n_objectives, n)`` array of weights.
     """
-    if isinstance(objective, MultiObjective):
+    if objective.is_multi_objective:
         rows: list[npt.NDArray] = []
-        for obj in objective.objectives:
-            rows.append(extract_objective_weights(obj, outcomes))
+        for name, weight in objective.metric_weights:
+            rows.append(
+                extract_objective_weights(
+                    objective=Objective(expression=f"{weight} * {name}"),
+                    outcomes=outcomes,
+                    experiment=experiment,
+                )
+            )
         return np.stack(rows, axis=0)
     else:
         # Single row – covers Objective and ScalarizedObjective
-        return extract_objective_weights(objective, outcomes).reshape(1, -1)
+        return extract_objective_weights(objective, outcomes, experiment).reshape(1, -1)
 
 
 def extract_outcome_constraints(
-    outcome_constraints: list[OutcomeConstraint], outcomes: list[str]
+    outcome_constraints: list[OutcomeConstraint],
+    outcomes: list[str],
+    experiment: Experiment,
 ) -> TBounds:
     if len(outcome_constraints) == 0:
         return None
@@ -330,11 +341,11 @@ def extract_outcome_constraints(
     for i, c in enumerate(outcome_constraints):
         s = 1 if c.op == ComparisonOp.LEQ else -1
         if isinstance(c, ScalarizedOutcomeConstraint):
-            for c_metric, c_weight in c.metric_weights:
-                j = outcomes.index(c_metric.signature)
+            for c_metric_name, c_weight in c.metric_weights:
+                j = outcomes.index(experiment.get_metric(c_metric_name).signature)
                 A[i, j] = s * c_weight
         else:
-            j = outcomes.index(c.metric.signature)
+            j = outcomes.index(experiment.get_metric(c.metric_names[0]).signature)
             A[i, j] = s
         b[i, 0] = s * c.bound
     return (A, b)
@@ -645,16 +656,20 @@ def get_pareto_frontier_and_configs(
     )
     # Extract weights, constraints, and objective_thresholds
     objective_weights = extract_objective_weight_matrix(
-        objective=optimization_config.objective, outcomes=adapter.outcomes
+        objective=optimization_config.objective,
+        outcomes=adapter.outcomes,
+        experiment=adapter._experiment,
     )
     outcome_constraints = extract_outcome_constraints(
         outcome_constraints=optimization_config.outcome_constraints,
         outcomes=adapter.outcomes,
+        experiment=adapter._experiment,
     )
     obj_t = extract_objective_thresholds(
         objective_thresholds=optimization_config.objective_thresholds,
         objective=optimization_config.objective,
         outcomes=adapter.outcomes,
+        experiment=adapter._experiment,
     )
     if obj_t is not None:
         obj_t = array_to_tensor(obj_t)
@@ -1113,6 +1128,7 @@ def observation_features_to_array(
 def feasible_hypervolume(
     optimization_config: MultiObjectiveOptimizationConfig,
     values: dict[str, npt.NDArray],
+    experiment: Experiment,
 ) -> npt.NDArray:
     """Compute the feasible hypervolume each iteration.
 
@@ -1121,34 +1137,35 @@ def feasible_hypervolume(
         values: Dictionary from metric name to array of value at each
             iteration (each array is `n`-dim). If optimization config contains
             outcome constraints, values for them must be present in `values`.
+        experiment: The experiment, used to map metric names to signatures.
 
     Returns: Array of feasible hypervolumes.
     """
     # Get objective at each iteration
     obj_threshold_dict = {
-        ot.metric.signature: ot.bound for ot in optimization_config.objective_thresholds
+        experiment.get_metric(ot.metric_names[0]).signature: ot.bound
+        for ot in optimization_config.objective_thresholds
     }
-    f_vals = np.hstack(
-        [
-            values[m.signature].reshape(-1, 1)
-            for m in optimization_config.objective.metrics
-        ]
-    )
-    obj_thresholds = np.array(
-        [obj_threshold_dict[m.signature] for m in optimization_config.objective.metrics]
-    )
+    obj_metric_names = optimization_config.objective.metric_names
+    obj_metrics = [experiment.get_metric(name) for name in obj_metric_names]
+    f_vals = np.hstack([values[m.signature].reshape(-1, 1) for m in obj_metrics])
+    obj_thresholds = np.array([obj_threshold_dict[m.signature] for m in obj_metrics])
     # Set infeasible points to be the objective threshold
     for oc in optimization_config.outcome_constraints:
         if oc.relative:
             raise ValueError(
                 "Benchmark aggregation does not support relative constraints"
             )
-        g = values[oc.metric.signature]
+        oc_sig = experiment.get_metric(oc.metric_names[0]).signature
+        g = values[oc_sig]
         feas = g <= oc.bound if oc.op == ComparisonOp.LEQ else g >= oc.bound
         f_vals[~feas] = obj_thresholds
 
+    # Derive objective directions from the objective's metric_weights.
+    # Positive weight = maximize, negative weight = minimize.
+    obj_weight_dict = dict(optimization_config.objective.metric_weights)
     obj_weights = np.array(
-        [-1 if m.lower_is_better else 1 for m in optimization_config.objective.metrics]
+        [1 if obj_weight_dict[name] > 0 else -1 for name in obj_metric_names]
     )
     obj_thresholds = obj_thresholds * obj_weights
     f_vals = f_vals * obj_weights

ax/adapter/base.py

@@ -201,7 +201,7 @@ def __init__(
                     "Optimization config is required when "
                     "`fit_tracking_metrics` is False."
                 )
-            self.outcomes = sorted(self._optimization_config.metrics.keys())
+            self.outcomes = sorted(self._optimization_config.metric_names)
 
         # Set training data (in the raw / untransformed space). This also omits
         # out-of-design and abandoned observations depending on the corresponding flags.
@@ -466,10 +466,14 @@ def _set_status_quo(self, experiment: Experiment) -> None:
                 )
                 return
 
-            if has_map_metrics(optimization_config=self._optimization_config):
+            if has_map_metrics(
+                metrics=experiment.get_metrics(
+                    metric_names=[*self._optimization_config.metric_names]
+                )
+            ):
                 self._status_quo = _combine_multiple_status_quo_observations(
                     status_quo_observations=status_quo_observations,
-                    metrics=set(none_throws(self._optimization_config).metrics),
+                    metrics=none_throws(self._optimization_config).metric_names,
                 )
             else:
                 logger.warning(
@@ -689,7 +693,7 @@ def _get_transformed_gen_args(
                 # Check that the optimization config has the same metrics as
                 # the original one. Otherwise, we may attempt to optimize over
                 # metrics that do not have a fitted model.
-                outcomes = set(optimization_config.metrics.keys())
+                outcomes = optimization_config.metric_names
                 if not outcomes.issubset(self.outcomes):
                     raise UnsupportedError(
                         "When fit_tracking_metrics is False, the optimization config "

ax/adapter/cross_validation.py

@@ -21,6 +21,7 @@
 from ax.adapter.data_utils import ExperimentData
 from ax.adapter.observation_utils import unwrap_observation_data
 from ax.adapter.torch import TorchAdapter
+from ax.core.experiment import Experiment
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
 from ax.core.optimization_config import OptimizationConfig
 from ax.exceptions.core import UnsupportedError
@@ -573,6 +574,7 @@ def assess_model_fit(
 def has_good_opt_config_model_fit(
     optimization_config: OptimizationConfig,
     assess_model_fit_result: AssessModelFitResult,
+    experiment: Experiment,
 ) -> bool:
     """Assess model fit for given diagnostics results across the optimization
     config metrics
@@ -584,7 +586,8 @@ def has_good_opt_config_model_fit(
 
     Args:
         optimization_config: Objective/Outcome constraint metrics to assess
-        diagnostics: Output of compute_diagnostics
+        assess_model_fit_result: Output of assess_model_fit
+        experiment: The experiment, used to map metric names to signatures.
 
     Returns:
         Two dictionaries, one for good metrics, one for bad metrics, each
@@ -594,8 +597,9 @@ def has_good_opt_config_model_fit(
     # Bad fit criteria: Any objective metrics are poorly fit
     # TODO[]: Incl. outcome constraints in assessment
     has_good_opt_config_fit = all(
-        (m.signature in assess_model_fit_result.good_fit_metrics_to_fisher_score)
-        for m in optimization_config.objective.metrics
+        experiment.get_metric(name).signature
+        in assess_model_fit_result.good_fit_metrics_to_fisher_score
+        for name in optimization_config.objective.metric_names
     )
     return has_good_opt_config_fit
 

ax/adapter/discrete.py

@@ -163,13 +163,18 @@ def _gen(
             objective_weights = None
             outcome_constraints = None
         else:
-            validate_transformed_optimization_config(optimization_config, self.outcomes)
+            validate_transformed_optimization_config(
+                optimization_config, self.outcomes, experiment=self._experiment
+            )
             objective_weights = extract_objective_weights(
-                objective=optimization_config.objective, outcomes=self.outcomes
+                objective=optimization_config.objective,
+                outcomes=self.outcomes,
+                experiment=self._experiment,
             )
             outcome_constraints = extract_outcome_constraints(
                 outcome_constraints=optimization_config.outcome_constraints,
                 outcomes=self.outcomes,
+                experiment=self._experiment,
             )
 
         # Get fixed features

ax/adapter/tests/test_adapter_utils.py

@@ -43,6 +43,7 @@
 from ax.utils.common.hash_utils import compute_lilo_input_hash
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
+    get_branin_experiment,
     get_experiment_with_observations,
     get_hierarchical_search_space,
     get_search_space_for_range_values,
@@ -61,7 +62,7 @@ def test_feasible_hypervolume(self) -> None:
             ),
             outcome_constraints=[
                 OutcomeConstraint(
-                    mc,
+                    metric=mc,
                     op=ComparisonOp.GEQ,
                     bound=0,
                     relative=False,
@@ -78,6 +79,11 @@ def test_feasible_hypervolume(self) -> None:
                 ),
             ],
         )
+        experiment = Experiment(
+            search_space=SearchSpace(parameters=[]),
+            optimization_config=optimization_config,
+            tracking_metrics=[mc],
+        )
         feas_hv = feasible_hypervolume(
             optimization_config,
             values={
@@ -106,6 +112,7 @@ def test_feasible_hypervolume(self) -> None:
                     ]
                 ),
             },
+            experiment=experiment,
         )
         self.assertEqual(list(feas_hv), [0.0, 0.0, 1.0, 1.0])
 
@@ -537,20 +544,24 @@ def test_can_map_to_binary(self) -> None:
     def test_extract_objective_weight_matrix(self) -> None:
         m1, m2, m3 = Metric(name="m1"), Metric(name="m2"), Metric(name="m3")
         outcomes = ["m1", "m2", "m3"]
+        experiment = get_branin_experiment()
+        experiment.add_metric(m1)
+        experiment.add_metric(m2)
+        experiment.add_metric(m3)
 
         # Single Objective: one row, nonzero only in matching column.
         obj = Objective(metric=m1, minimize=False)
-        result = extract_objective_weight_matrix(obj, outcomes)
+        result = extract_objective_weight_matrix(obj, outcomes, experiment)
         np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0]])
 
         # Minimization flips the sign.
         obj_min = Objective(metric=m2, minimize=True)
-        result = extract_objective_weight_matrix(obj_min, outcomes)
+        result = extract_objective_weight_matrix(obj_min, outcomes, experiment)
         np.testing.assert_array_equal(result, [[0.0, -1.0, 0.0]])
 
         # ScalarizedObjective: single row with multiple nonzero entries.
         scal = ScalarizedObjective(metrics=[m1, m3], weights=[0.3, 0.7], minimize=False)
-        result = extract_objective_weight_matrix(scal, outcomes)
+        result = extract_objective_weight_matrix(scal, outcomes, experiment)
         np.testing.assert_array_almost_equal(result, [[0.3, 0.0, 0.7]])
 
         # MultiObjective: one row per sub-objective.
@@ -560,7 +571,7 @@ def test_extract_objective_weight_matrix(self) -> None:
                 Objective(metric=m3, minimize=True),
             ]
         )
-        result = extract_objective_weight_matrix(multi, outcomes)
+        result = extract_objective_weight_matrix(multi, outcomes, experiment)
         np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0]])
 
     def test_get_fresh_pairwise_trial_indices(self) -> None: