Refactor Objective, OutcomeConstraint to be expression-based; remove Metrics from OptimizationConfig

Summary:
This diff is a major refactor.

In this diff we remove Metric instances from both Objective and OutcomeConstraint (and transitively, also off OptimizationConfig) and instead add them directly to the Experiment. In their place, we use the "expression" syntax and Sympy parsing from the Ax API directly on these classes, allowing users to specify `Objective(expression="ne1 + 2*ne2")` or `OutcomeConstraint(expression="qps >= 7000")` directly. This has a number of benefits:
* Dependency spread is contained. OptimizationConfig is now able to be loaded without loading the any specific Metric implementation
* Major separation of concerns: OptimizationConfig is now solely concerned with the goals of the optimization and not also implicitly concerned with fetching logic
* Bringing Expressions directly onto Objective and OutcomeConstraint allows us to deprecated ScalarizedObjective, MultiObjective, ScalarizedOutcomeConstraint, and ObjectiveThreshold
* Massive simplification for eventual upstreaming of MultiTypeExperiment functionality into base Experiment

This is also a necessary step in our storage redesign in order to make metrics "live" only in one place. See this doc for details: https://docs.google.com/document/d/1I8FSQJ05_WHXFHBtNaL3dTAdAbvKfADwMRlIdq6BEV8/edit?tab=t.0#heading=h.s4r4rlrnva07

The refactor touches several core modules, please focus review on these:
* ax/core/objective.py
* ax/core/outcome_constraint.py
* ax/core/optimization_config.py
* ax/core/experiment.py

These changes mandate many many updates to callsites which is ballooning the size of the diff. Changes are summarized here.

# Core changes:
* Objective, OutcomeConstraint classes now takes an expression string (e.g., "accuracy", "-loss", "2*acc + recall", "acc, -loss")
  * Old-style constructors still work, but they discard the Metric after initialization
* Objective.metric property no longer exists. If a user wants to get the specific metric they retrieve its name via get_names and pass that string into Experiment.get_metric(name: str) -> Metric
* isinstance checks for MultiObective, ScalarizedObjective, etc are replaced with is_multi_objective and is_scalarized_objective properties
* Experiment.tracking_metrics is a property not an attr, now experiment simply has a collection of Metrics called "metrics" and tracking_metrics is a property which returns the metrics not on the optimization config
* **Parsing is directly lifted and shifted from ax/api**
* **Both JSON and SQL encoders/decoders carefully reconstruct the pre-refactor structure, allowing existing storage to be unaffected**

# Other changes
## **ax/adapter/**
- **adapter_utils.py**
  - Functions updated to accept experiment metrics for lookup.
  - All extraction functions use metric names.
- **base.py, torch.py, transforms/**
  - Updated to thread experiment metrics through extraction functions.
  - Metric access patterns updated.

## **ax/analysis/**
- Updated all metric access patterns to use metric names, not Metric objects.

## **ax/api/**
- **client.py**
  - Simplified metric overwriting logic: metrics only live on Experiment.
- **utils/instantiation/from_string.py**
  - Simplified parsing: constructs Objective directly from expression string.

## **ax/benchmark/**
- Updated objective construction to use expression-based interface.

## **ax/core/**
- **experiment.py**
  - Experiment now maintains a flat list of metrics (`_metrics: dict[str, Metric]`).
  - No distinction between tracking and optimization metrics.
  - Objective and constraints reference metrics by name, not object.
- **objectives.py**
  - Unified Objective class: now takes an expression string, supporting single, scalarized, and multi-objective optimization.
  - Deprecated MultiObjective and ScalarizedObjective classes; their functionality is merged into Objective.
  - All metric access is via metric names/weights parsed from the expression, not direct Metric objects.
- **optimization_config.py**
  - Unified OptimizationConfig class: handles both single and multi-objective cases.
  - Removed MultiObjectiveOptimizationConfig.
  - All config logic uses expression-based Objective and OutcomeConstraint.
- **outcome_constraint.py**
  - Unified OutcomeConstraint class: now takes an expression string, supporting scalarized constraints.
  - Removed ObjectiveThreshold and ScalarizedOutcomeConstraint classes.
  - All validation and access is via metric names/weights parsed from the expression.
- **trial.py**
  - Updated to reference metrics via `objective.metric_names[0]` instead of `objective.metric.name`.
- **utils.py**
  - Updated metric access patterns to use metric names.

## **ax/early_stopping/**
- Updated to reference metric signatures via experiment metrics, not via Objective.

## **ax/plot/**
- Updated metric access patterns to use metric names.

## **ax/service/**
- **ax_client.py**
  - Updated to reference metrics via `objective.metric_names[0]`.
  - Simplified metric overwriting logic.
- **utils/best_point.py**
  - Updated metric access patterns to use metric names.

## **ax/storage/json_store/**
- Encoder: serializes objectives as expression strings.
- Decoder: converts legacy formats (with embedded Metric dicts) to expression-based objects.

## **ax/storage/sqa_store/**
- Encoder: stores expression string in JSON column.
- Decoder: reconstructs expression strings from MetricIntent-based SQA rows.

## **Tests**
- All test constructors rewritten to use expression strings.
- Assertions updated for new property names.
- Added tests for expression parsing, multi-objective, scalarized, and minimize detection.

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