Fix bug where StandardizeY doesn't update the weights in ScalarizedObjective

ax/adapter/tests/test_base_adapter.py

@@ -36,7 +36,7 @@
 from ax.core.experiment import Experiment
 from ax.core.map_data import MapData
 from ax.core.metric import Metric
-from ax.core.objective import Objective, ScalarizedObjective
+from ax.core.objective import Objective
 from ax.core.observation import ObservationData, ObservationFeatures
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.outcome_constraint import ComparisonOp, OutcomeConstraint
@@ -297,11 +297,9 @@ def test_gen_base(self, mock_fit: Mock, mock_gen_arms: Mock) -> None:
             model_gen_options=None,
         )
 
-        # Gen with multi-objective optimization config.
+        # Gen with a different optimization config.
         oc2 = OptimizationConfig(
-            objective=ScalarizedObjective(
-                metrics=[Metric(name="test_metric"), Metric(name="test_metric_2")]
-            )
+            objective=Objective(metric=Metric(name="branin"), minimize=True)
         )
         with mock.patch(ADAPTER__GEN_PATH, return_value=mock_return_value) as mock_gen:
             adapter.gen(n=1, search_space=search_space, optimization_config=oc2)

ax/adapter/transforms/standardize_y.py

@@ -13,6 +13,8 @@
 import numpy as np
 from ax.adapter.data_utils import ExperimentData
 from ax.adapter.transforms.base import Transform
+from ax.core.metric import Metric
+from ax.core.objective import ScalarizedObjective
 from ax.core.observation import ObservationData, ObservationFeatures
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint
@@ -74,53 +76,72 @@ def _transform_observation_data(
             obsd.covariance /= np.dot(stds[:, None], stds[:, None].transpose())
         return observation_data
 
+    def _check_metrics_available(self, metrics: list[Metric], context: str) -> set[str]:
+        """Check that all metrics are available and return the set of metrics."""
+        available_metrics = set(self.Ymean).intersection(set(self.Ystd))
+        required_metrics = {metric.signature for metric in metrics}
+        if len(required_metrics - available_metrics) > 0:
+            raise DataRequiredError(
+                f"`StandardizeY` transform requires {context} metric(s) "
+                f"{required_metrics} but received only {available_metrics}."
+            )
+        return available_metrics
+
+    def _transform_scalarized_weights(
+        self, metrics: list[Metric], weights: list[float]
+    ) -> list[float]:
+        """Transform weights for scalarized objectives/constraints.
+
+        When standardizing yi to zi = (yi - μi) / σi, the scalarized term
+        Σ(wi * yi) transforms to Σ(wi * σi * zi) + Σ(wi * μi).
+        This method returns the new weights: new_wi = wi * σi.
+
+        Args:
+            metrics: List of metrics in the scalarized term.
+            weights: Original weights for each metric.
+
+        Returns:
+            Transformed weights scaled by standard deviations.
+        """
+        return [
+            weights[i] * float(self.Ystd[metric.signature])
+            for i, metric in enumerate(metrics)
+        ]
+
     def transform_optimization_config(
         self,
         optimization_config: OptimizationConfig,
         adapter: Optional["base_adapter.Adapter"] = None,
         fixed_features: ObservationFeatures | None = None,
     ) -> OptimizationConfig:
+        # Handle ScalarizedObjective
+        if isinstance(optimization_config.objective, ScalarizedObjective):
+            objective = optimization_config.objective
+            self._check_metrics_available(objective.metrics, context="objective")
+            objective.weights = self._transform_scalarized_weights(
+                objective.metrics, objective.weights
+            )
+
         for c in optimization_config.all_constraints:
             if c.relative:
                 raise ValueError(
                     f"StandardizeY transform does not support relative constraint {c}"
                 )
-            # For required data checks, metrics must be available in Ymean and Ystd.
-            available_metrics = set(self.Ymean).intersection(set(self.Ystd))
             if isinstance(c, ScalarizedOutcomeConstraint):
-                # check metrics are present.
-                constraint_metrics = {metric.signature for metric in c.metrics}
-                if len(constraint_metrics - available_metrics) > 0:
-                    raise DataRequiredError(
-                        "`StandardizeY` transform requires constraint metric(s) "
-                        f"{constraint_metrics} but received only {available_metrics}."
-                    )
-
-                # transform \sum (wi * yi) <= C to
-                # \sum (wi * si * zi) <= C - \sum (wi * mu_i) that zi = (yi - mu_i) / si
-
-                # update bound C to new c = C.bound - sum_i (wi * mu_i)
+                self._check_metrics_available(c.metrics, context="constraint")
+
+                # Transform Σ(wi * yi) <= C to Σ(wi * σi * zi) <= C - Σ(wi * μi)
+                # where zi = (yi - μi) / σi
                 agg_mean = np.sum(
                     [
-                        c.weights[i] * self.Ymean[metric.signature]
+                        c.weights[i] * float(self.Ymean[metric.signature])
                         for i, metric in enumerate(c.metrics)
                     ]
                 )
                 c.bound = float(c.bound - agg_mean)
-
-                # update the weights in the scalarized constraint
-                # new wi = wi * si
-                new_weight = [
-                    c.weights[i] * self.Ystd[metric.signature]
-                    for i, metric in enumerate(c.metrics)
-                ]
-                c.weights = new_weight
+                c.weights = self._transform_scalarized_weights(c.metrics, c.weights)
             else:
-                if c.metric.signature not in available_metrics:
-                    raise DataRequiredError(
-                        "`StandardizeY` transform requires constraint metric(s) "
-                        f"{c.metric.signature} but got {available_metrics}"
-                    )
+                self._check_metrics_available([c.metric], context="constraint")
                 c.bound = float(
                     (c.bound - self.Ymean[c.metric.signature])
                     / self.Ystd[c.metric.signature]

ax/adapter/transforms/tests/test_standardize_y_transform.py

@@ -14,7 +14,7 @@
 from ax.adapter.data_utils import extract_experiment_data
 from ax.adapter.transforms.standardize_y import StandardizeY
 from ax.core.metric import Metric
-from ax.core.objective import Objective
+from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
 from ax.core.observation import ObservationData
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint
@@ -24,6 +24,7 @@
 from ax.utils.testing.core_stubs import get_experiment_with_observations
 from pandas import DataFrame
 from pandas.testing import assert_frame_equal
+from pyre_extensions import assert_is_instance
 
 
 class StandardizeYTransformTest(TestCase):
@@ -213,6 +214,67 @@ def test_transform_experiment_data(self) -> None:
         )
         assert_frame_equal(observation_data["sem"], expected_sems)
 
+    def test_TransformOptimizationConfigWithScalarizedObjective(self) -> None:
+        # Test with ScalarizedObjective
+        # Given: objective = w1*m1 + w2*m2 with w1=0.5, w2=0.5
+        # After standardization: zi = (yi - mu_i) / si
+        # The objective becomes: w1*s1*z1 + w2*s2*z2 (constant term doesn't matter)
+        # Expected weights: [0.5 * 1.0, 0.5 * sqrt(1/3)]
+        m1 = Metric(name="m1")
+        m2 = Metric(name="m2")
+        m3 = Metric(name="m3")
+
+        # Test with ScalarizedObjective that has all required metrics
+        objective = ScalarizedObjective(
+            metrics=[m1, m2], weights=[0.5, 0.5], minimize=False
+        )
+        oc = OptimizationConfig(objective=objective)
+        oc_transformed = self.t.transform_optimization_config(oc, None, None)
+
+        # Check that weights are scaled by standard deviations
+        expected_weights = [0.5 * 1.0, 0.5 * sqrt(1 / 3)]
+        transformed_objective = assert_is_instance(
+            oc_transformed.objective, ScalarizedObjective
+        )
+        self.assertTrue(np.allclose(transformed_objective.weights, expected_weights))
+
+        # Test with ScalarizedObjective missing a metric
+        objective_missing = ScalarizedObjective(
+            metrics=[m1, m3], weights=[0.5, 0.5], minimize=False
+        )
+        oc_missing = OptimizationConfig(objective=objective_missing)
+        with self.assertRaisesRegex(
+            DataRequiredError, "`StandardizeY` transform requires objective metric"
+        ):
+            self.t.transform_optimization_config(oc_missing, None, None)
+
+        # Test with different weights and minimize=True
+        objective_minimize = ScalarizedObjective(
+            metrics=[m1, m2], weights=[1.0, -2.0], minimize=True
+        )
+        oc_minimize = OptimizationConfig(objective=objective_minimize)
+        oc_minimize_transformed = self.t.transform_optimization_config(
+            oc_minimize, None, None
+        )
+
+        # Check that weights are scaled by standard deviations
+        expected_weights_minimize = [1.0 * 1.0, -2.0 * sqrt(1 / 3)]
+        transformed_objective_minimize = assert_is_instance(
+            oc_minimize_transformed.objective, ScalarizedObjective
+        )
+        self.assertTrue(
+            np.allclose(
+                transformed_objective_minimize.weights, expected_weights_minimize
+            )
+        )
+
+        # Multi-objective with scalarized objective should error out
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            "Scalarized objectives are not supported for a `MultiObjective`.",
+        ):
+            MultiObjective([objective_minimize, Objective(metric=m3, minimize=False)])
+
 
 def osd_allclose(osd1: ObservationData, osd2: ObservationData) -> bool:
     if osd1.metric_signatures != osd2.metric_signatures: