Add objectives: list[Objective] to OptimizationConfig (#5150)

Summary:

Part of the Restrict Objective to Single/Scalarized & Simplify
OptimizationConfig design (see design doc:
https://docs.google.com/document/d/1EGQYmBjiNGtYapXu1RLHEBdA5Yz2c7q17acX3es0yV8/edit).

This is Diff 1 of the stack: enables the new
`OptimizationConfig(objectives=[...])` construction path without
breaking any existing code.

Changes:
- `OptimizationConfig.__init__` and `clone_with_args` are now
  keyword-only across `OptimizationConfig`, `MOOC`, and
  `PreferenceOptimizationConfig`. All positional callers updated.
- New kwarg `objectives: list[Objective] | None = None`, mutually
  exclusive with `objective`, on both `__init__` and `clone_with_args`.
- Internally stores `self._objectives: list[Objective]` (both paths).
- New `objectives` property returns the list.
- `objective` property raises `UnsupportedError` if `len > 1`.
- `is_moo_problem` property: True when multiple objectives or legacy
  multi-objective expression.
- `metric_names`, `metric_name_to_signature`, `metric_signatures`
  aggregate across all objectives + constraints.
- `__repr__` always uses `objectives=`.
- JSON storage: encoder uses `objectives` key; decoder has backward
  compat to convert old `objective` key to `objectives` list.
- SQA storage: encoder iterates `objectives` to encode each one;
  decoder collects multiple OBJECTIVE rows and reconstructs
  `OptimizationConfig(objectives=...)` when `len > 1`.
- Validation ensures no duplicate metrics across objectives and no
  multi-objective expressions in individual list elements.

Differential Revision: D99387020

Author: saitcakmak

ax/adapter/tests/test_torch_adapter.py

@@ -1205,7 +1205,7 @@ def test_pairwise_preference_generator(self) -> None:
                     surrogate=surrogate,
                 ),
                 optimization_config=OptimizationConfig(
-                    Objective(
+                    objective=Objective(
                         metric=Metric(Keys.PAIRWISE_PREFERENCE_QUERY.value),
                         minimize=False,
                     )

ax/core/optimization_config.py

@@ -154,7 +154,7 @@ def test_setting_opt_config(self) -> None:
         m3 = BraninMetric("m3", ["x1", "x2"])
         self.experiment.add_tracking_metric(m3)
         self.experiment.optimization_config = OptimizationConfig(
-            Objective(metric=m3, minimize=True)
+            objective=Objective(metric=m3, minimize=True)
         )
         self.assertDictEqual(
             self.experiment._metric_to_trial_type,

ax/core/tests/test_multi_type_experiment.py

@@ -21,14 +21,14 @@
     ScalarizedOutcomeConstraint,
 )
 from ax.core.types import ComparisonOp
-from ax.exceptions.core import UserInputError
+from ax.exceptions.core import UnsupportedError, UserInputError
 from ax.utils.common.testutils import TestCase
 from pyre_extensions import assert_is_instance
 
 
 OC_STR = (
     "OptimizationConfig("
-    'objective=Objective(expression="m1"), '
+    'objectives=[Objective(expression="m1")], '
     "outcome_constraints=[OutcomeConstraint(m3 >= -0.25), "
     "OutcomeConstraint(m4 <= 0.25), "
     "ScalarizedOutcomeConstraint(0.5*m3 + 0.5*m4 >= 0.9975 * baseline)])"
@@ -271,6 +271,111 @@ def test_CloneWithArgs(self) -> None:
         )
 
 
+class OptimizationConfigObjectivesListTest(TestCase):
+    """Tests for the new `OptimizationConfig(objectives=[...])` construction path."""
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.metrics = {
+            "m1": Metric(name="m1"),
+            "m2": Metric(name="m2"),
+            "m3": Metric(name="m3"),
+        }
+        self.sig = {m: m for m in self.metrics}
+        self.obj1 = Objective(expression="m1", metric_name_to_signature=self.sig)
+        self.obj2 = Objective(expression="-m2", metric_name_to_signature=self.sig)
+        self.scalarized_obj = Objective(
+            expression="2*m1 + m2", metric_name_to_signature=self.sig
+        )
+
+    def test_objectives_kwarg_construction(self) -> None:
+        """Test single and multi-objective construction via objectives kwarg."""
+        # Single objective
+        config = OptimizationConfig(objectives=[self.obj1])
+        self.assertEqual(config.objectives, [self.obj1])
+        self.assertEqual(config.objective, self.obj1)
+        self.assertFalse(config.is_moo_problem)
+
+        # Multi-objective
+        config = OptimizationConfig(objectives=[self.obj1, self.obj2])
+        self.assertEqual(config.objectives, [self.obj1, self.obj2])
+        self.assertTrue(config.is_moo_problem)
+        with self.assertRaisesRegex(UnsupportedError, "multiple objectives"):
+            config.objective
+
+    def test_objectives_kwarg_metric_aggregation(self) -> None:
+        """Test metric_names, metric_name_to_signature, metric_signatures."""
+        constraint = OutcomeConstraint(
+            expression="m3 >= 0.5", metric_name_to_signature=self.sig
+        )
+        config = OptimizationConfig(
+            objectives=[self.obj1, self.obj2],
+            outcome_constraints=[constraint],
+        )
+        self.assertEqual(config.metric_names, {"m1", "m2", "m3"})
+        self.assertEqual(
+            config.metric_name_to_signature, {"m1": "m1", "m2": "m2", "m3": "m3"}
+        )
+        self.assertEqual(config.metric_signatures, {"m1", "m2", "m3"})
+
+    def test_objectives_kwarg_validation(self) -> None:
+        """Test validation errors for objectives kwarg."""
+        with self.subTest("mutual_exclusivity"):
+            with self.assertRaisesRegex(UserInputError, "Cannot specify both"):
+                OptimizationConfig(objective=self.obj1, objectives=[self.obj1])
+
+        with self.subTest("neither_specified"):
+            with self.assertRaisesRegex(UserInputError, "Must specify either"):
+                OptimizationConfig()
+
+        with self.subTest("empty_list"):
+            with self.assertRaisesRegex(UserInputError, "must not be empty"):
+                OptimizationConfig(objectives=[])
+
+        with self.subTest("multi_objective_expression"):
+            multi_obj = Objective(
+                expression="m1, -m2", metric_name_to_signature=self.sig
+            )
+            with self.assertRaisesRegex(ValueError, "single or scalarized"):
+                OptimizationConfig(objectives=[multi_obj])
+
+        with self.subTest("duplicate_metric_names"):
+            obj_dup = Objective(expression="m1", metric_name_to_signature=self.sig)
+            with self.assertRaisesRegex(UserInputError, "appears in multiple"):
+                OptimizationConfig(objectives=[self.obj1, obj_dup])
+
+    def test_objectives_kwarg_clone_and_repr(self) -> None:
+        """Test clone, clone_with_args, and repr for objectives-list configs."""
+        config = OptimizationConfig(objectives=[self.obj1, self.obj2])
+
+        # clone preserves objectives
+        cloned = config.clone()
+        self.assertEqual(len(cloned.objectives), 2)
+        self.assertEqual(cloned.objectives[0].expression, "m1")
+        self.assertEqual(cloned.objectives[1].expression, "-m2")
+        self.assertTrue(cloned.is_moo_problem)
+
+        # clone_with_args(objective=) replaces the list with a single objective
+        cloned = config.clone_with_args(objective=self.obj1)
+        self.assertEqual(len(cloned.objectives), 1)
+        self.assertFalse(cloned.is_moo_problem)
+
+        # clone_with_args(objectives=) replaces the list
+        obj3 = Objective(expression="m3", metric_name_to_signature=self.sig)
+        cloned = config.clone_with_args(objectives=[self.obj1, obj3])
+        self.assertEqual(len(cloned.objectives), 2)
+        self.assertEqual(cloned.objectives[1].expression, "m3")
+
+        # objective= and objectives= are mutually exclusive in clone_with_args
+        with self.assertRaisesRegex(UserInputError, "Cannot specify both"):
+            config.clone_with_args(objective=self.obj1, objectives=[self.obj1])
+
+        # repr always uses "objectives="
+        self.assertIn("objectives=", repr(config))
+        single_config = OptimizationConfig(objectives=[self.obj1])
+        self.assertIn("objectives=", repr(single_config))
+
+
 class MultiObjectiveOptimizationConfigTest(TestCase):
     def setUp(self) -> None:
         super().setUp()

ax/core/tests/test_optimization_config.py

@@ -2727,7 +2727,7 @@ def test_generate_candidates_does_not_generate_if_missing_data(self) -> None:
         )
         self.branin_experiment.add_tracking_metric(custom_metric)
         self.branin_experiment.optimization_config = OptimizationConfig(
-            Objective(
+            objective=Objective(
                 metric=CustomTestMetric(
                     name="custom_test_metric", test_attribute="test"
                 ),
@@ -2974,7 +2974,7 @@ def setUp(self) -> None:
         self.branin_experiment_no_impl_runner_or_metrics = MultiTypeExperiment(
             search_space=get_branin_search_space(),
             optimization_config=OptimizationConfig(
-                Objective(metric=Metric(name="branin"), minimize=True)
+                objective=Objective(metric=Metric(name="branin"), minimize=True)
             ),
             default_trial_type="type1",
             default_runner=None,

ax/orchestration/tests/test_orchestrator.py

@@ -615,7 +615,7 @@ def test_best_raw_objective_point_scalarized(self) -> None:
         exp = get_branin_experiment()
         gs = choose_generation_strategy_legacy(search_space=exp.search_space)
         exp.optimization_config = OptimizationConfig(
-            ScalarizedObjective(metrics=[get_branin_metric()], minimize=True)
+            objective=ScalarizedObjective(metrics=[get_branin_metric()], minimize=True)
         )
         with self.assertRaisesRegex(ValueError, "Cannot identify best "):
             get_best_raw_objective_point_with_trial_index(exp)
@@ -637,7 +637,7 @@ def test_best_raw_objective_point_scalarized_multi(self) -> None:
         exp = get_branin_experiment()
         gs = choose_generation_strategy_legacy(search_space=exp.search_space)
         exp.optimization_config = OptimizationConfig(
-            ScalarizedObjective(
+            objective=ScalarizedObjective(
                 metrics=[get_branin_metric(), get_branin_metric(lower_is_better=False)],
                 weights=[0.1, -0.9],
                 minimize=True,
@@ -1037,7 +1037,7 @@ def test_best_parameters_from_model_predictions_scalarized(self) -> None:
         )
         exp.add_tracking_metric(metric2)
         exp.optimization_config = OptimizationConfig(
-            ScalarizedObjective(
+            objective=ScalarizedObjective(
                 metrics=[metric1, metric2],
                 weights=[0.5, 0.5],
                 minimize=True,

ax/service/tests/test_best_point_utils.py

@@ -449,10 +449,12 @@ def _test_get_standard_plots_moo_relative_constraints(
         names = obj.metric_names
         # Create a new Objective rather than mutating _expression_str to
         # avoid stale _parsed cached_property.
-        none_throws(exp.optimization_config)._objective = Objective(
-            expression=f"{names[0]}, -{names[1]}",
-            metric_name_to_signature={n: n for n in names},
-        )
+        none_throws(exp.optimization_config)._objectives = [
+            Objective(
+                expression=f"{names[0]}, -{names[1]}",
+                metric_name_to_signature={n: n for n in names},
+            )
+        ]
         exp.get_metric(names[0]).lower_is_better = False
         assert_is_instance(
             exp.optimization_config, MultiObjectiveOptimizationConfig
@@ -494,10 +496,12 @@ def test_get_standard_plots_moo_no_objective_thresholds(self) -> None:
         # first objective to maximize, second to minimize
         obj = none_throws(exp.optimization_config).objective
         names = obj.metric_names
-        none_throws(exp.optimization_config)._objective = Objective(
-            expression=f"{names[0]}, -{names[1]}",
-            metric_name_to_signature={n: n for n in names},
-        )
+        none_throws(exp.optimization_config)._objectives = [
+            Objective(
+                expression=f"{names[0]}, -{names[1]}",
+                metric_name_to_signature={n: n for n in names},
+            )
+        ]
         exp.trials[0].run()
         plots = get_standard_plots(
             experiment=exp,

ax/service/tests/test_report_utils.py

@@ -351,6 +351,13 @@ def object_from_json(
             object_json = _sanitize_inputs_to_surrogate_spec(object_json=object_json)
         if isclass(_class) and issubclass(_class, OptimizationConfig):
             object_json.pop("risk_measure", None)  # Deprecated.
+            # Backward compat: old JSON uses "objective", new uses "objectives".
+            if (
+                _class is OptimizationConfig
+                and "objective" in object_json
+                and "objectives" not in object_json
+            ):
+                object_json["objectives"] = [object_json.pop("objective")]
         return ax_class_from_json_dict(
             _class=_class, object_json=object_json, **vars(registry_kwargs)
         )

ax/storage/json_store/decoder.py

@@ -383,7 +383,7 @@ def optimization_config_to_dict(
     """Convert Ax optimization config to a dictionary."""
     return {
         "__type": optimization_config.__class__.__name__,
-        "objective": optimization_config.objective,
+        "objectives": optimization_config.objectives,
         "outcome_constraints": optimization_config.outcome_constraints,
         "pruning_target_parameterization": (
             optimization_config.pruning_target_parameterization
@@ -782,16 +782,17 @@ def _build_opt_config_dict(
     will then recursively encode them via ``metric_to_dict``, capturing the
     full metric type.
     """
-    objective_dict = _build_objective_dict(
-        objective=opt_config.objective, experiment_metrics=experiment_metrics
-    )
+    objective_dicts = [
+        _build_objective_dict(objective=obj, experiment_metrics=experiment_metrics)
+        for obj in opt_config.objectives
+    ]
     constraint_dicts = [
         _build_constraint_dict(constraint=c, experiment_metrics=experiment_metrics)
         for c in opt_config.outcome_constraints
     ]
     result: dict[str, Any] = {
         "__type": opt_config.__class__.__name__,
-        "objective": objective_dict,
+        "objectives": objective_dicts,
         "outcome_constraints": constraint_dicts,
         "pruning_target_parameterization": opt_config.pruning_target_parameterization,
     }

ax/storage/json_store/encoders.py

@@ -137,6 +137,7 @@
     get_metric,
     get_mll_type,
     get_model_type,
+    get_moo_optimization_config,
     get_multi_objective,
     get_multi_objective_optimization_config,
     get_multi_type_experiment,
@@ -380,6 +381,7 @@
     ("Objective", get_objective),
     ("ObjectiveThreshold", get_objective_threshold),
     ("OptimizationConfig", get_optimization_config),
+    ("OptimizationConfig", get_moo_optimization_config),
     ("OrEarlyStoppingStrategy", get_or_early_stopping_strategy),
     ("OrderConstraint", get_order_constraint),
     ("OutcomeConstraint", get_outcome_constraint),