Refactor Objective and OutcomeConstraint to expression-based

Summary:
Refactors the Objective class to hold an expression string instead of Metric objects. This is Step 1 of decoupling Objective from Metric, enabling objectives to be defined purely by string expressions like "accuracy", "-loss", "2*acc + recall", or "acc, -loss" (multi-objective).

Key changes:
- Objective.__init__ now takes an `expression: str` as its primary argument, eagerly parsed via SymPy into cached metric names, weights, and objective type flags.
- The deprecated `metric`/`minimize` kwargs are preserved for backward compatibility and emit DeprecationWarning.
- New properties: `expression`, `metric_weights`, `metric_weights_by_objective`, `is_single_objective`, `is_scalarized_objective`, `is_multi_objective`.
- Old properties that returned Metric objects (`metric`, `metrics`, `metric_signatures`, `get_unconstrainable_metrics`) now emit DeprecationWarning and raise NotImplementedError.
- New `get_unconstrainable_metric_names()` replaces `get_unconstrainable_metrics()`.
- MultiObjective and ScalarizedObjective are preserved as deprecated subclasses that build expression strings internally, keeping isinstance checks working during transition.
- Added three expression-parsing helpers to ax/utils/common/sympy.py: `parse_objective_expression`, `extract_metric_names_from_objective_expr`, `extract_metric_weights_from_objective_expr`.
- All private cached attributes removed from Objective.__init__ except _expression_str. All derived values become computed properties.
- Refactored OutcomeConstraint, ObjectiveThreshold, and ScalarizedOutcomeConstraint to hold expression strings instead of Metric objects.
- OutcomeConstraint now takes an expression string (e.g. "qps >= 700", "loss <= 0.5%") as its primary constructor argument.

Differential Revision: D93520819

Author: mpolson64

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,
@@ -204,6 +204,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`.
 
@@ -219,6 +220,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
@@ -228,15 +230,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=}."
@@ -250,7 +256,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.
@@ -266,29 +274,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.
 
@@ -302,23 +305,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
@@ -328,11 +339,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)
@@ -643,16 +654,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)
@@ -1111,6 +1126,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.
 
@@ -1119,34 +1135,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