Support BOPE experiments in get_trace with preference model utilities (#4792)

Summary:

This diff updates the `get_trace` function in `ax/service/utils/best_point.py` to support preference learning (BOPE) experiments with `PreferenceOptimizationConfig`.

When a BOPE experiment has an associated PE_EXPERIMENT auxiliary experiment with preference data, `get_trace` now:
1. Fits a PairwiseGP preference model to the PE_EXPERIMENT data
2. Uses the learned preference model to predict utility values for each arm's metric values
3. Returns a trace based on predicted utilities

Adds `_compute_utility_from_preference_model()` helper function and corresponding unit tests.

Differential Revision: D91073267

Author: shrutipatel31

ax/service/tests/test_best_point.py

@@ -10,13 +10,23 @@
 import pandas as pd
 from ax.adapter.registry import Generators
 from ax.core.arm import Arm
+from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
 from ax.core.batch_trial import BatchTrial
 from ax.core.data import Data
-from ax.core.optimization_config import MultiObjectiveOptimizationConfig
+from ax.core.experiment import Experiment
+from ax.core.metric import Metric
+from ax.core.objective import MultiObjective, Objective
+from ax.core.optimization_config import (
+    MultiObjectiveOptimizationConfig,
+    PreferenceOptimizationConfig,
+)
+from ax.core.parameter import ParameterType, RangeParameter
+from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
-from ax.exceptions.core import DataRequiredError
+from ax.exceptions.core import DataRequiredError, UserInputError
 from ax.service.utils.best_point import get_trace
 from ax.service.utils.best_point_mixin import BestPointMixin
+from ax.utils.common.constants import Keys
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
     get_experiment_with_batch_trial,
@@ -367,3 +377,165 @@ def test_get_best_observed_value(self) -> None:
             minimize=True,
         )
         self.assertEqual(get_best(exp), 10)  # 5 and 9 are out of design
+
+    def _get_pe_search_space(self) -> SearchSpace:
+        """Create a standard PE_EXPERIMENT search space with m1 and m2 parameters."""
+        return SearchSpace(
+            parameters=[
+                RangeParameter(
+                    name="m1",
+                    parameter_type=ParameterType.FLOAT,
+                    lower=0.0,
+                    upper=10.0,
+                ),
+                RangeParameter(
+                    name="m2",
+                    parameter_type=ParameterType.FLOAT,
+                    lower=0.0,
+                    upper=10.0,
+                ),
+            ]
+        )
+
+    def _make_pref_opt_config(self, profile_name: str) -> PreferenceOptimizationConfig:
+        """Create a PreferenceOptimizationConfig with m1 and m2 objectives."""
+        return PreferenceOptimizationConfig(
+            objective=MultiObjective(
+                objectives=[
+                    Objective(metric=Metric(name="m1"), minimize=False),
+                    Objective(metric=Metric(name="m2"), minimize=False),
+                ]
+            ),
+            preference_profile_name=profile_name,
+        )
+
+    def _assert_valid_trace(
+        self,
+        trace: list[float],
+        expected_len: int,
+        check_monotonic: bool = True,
+    ) -> None:
+        """Assert trace has expected length, contains floats, and is monotonic."""
+        self.assertEqual(len(trace), expected_len)
+        for value in trace:
+            self.assertIsInstance(value, float)
+        if check_monotonic:
+            for i in range(1, len(trace)):
+                self.assertGreaterEqual(
+                    trace[i],
+                    trace[i - 1],
+                    msg=f"Trace not monotonic at index {i}: {trace}",
+                )
+
+    def test_get_trace_preference_learning_config(self) -> None:
+        """Test that get_trace works correctly with PreferenceOptimizationConfig.
+
+        This test verifies various scenarios for BOPE experiments,
+        including cases with and without PE_EXPERIMENT data.
+        """
+        with self.subTest("without_pe_experiment_raises_error"):
+            # Setup: Create a multi-objective experiment WITHOUT PE_EXPERIMENT
+            exp = get_experiment_with_observations(
+                observations=[[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]],
+            )
+            exp.name = "main_experiment"
+            pref_opt_config = self._make_pref_opt_config(
+                profile_name="nonexistent_profile"
+            )
+
+            # Execute & Assert: Should raise UserInputError without PE_EXPERIMENT
+            with self.assertRaisesRegex(
+                UserInputError,
+                "Preference profile 'nonexistent_profile' not found",
+            ):
+                get_trace(exp, pref_opt_config)
+
+        with self.subTest("with_pe_experiment_empty_data_raises_error"):
+            # Setup: Create main experiment
+            exp = get_experiment_with_observations(
+                observations=[[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]],
+            )
+            exp.name = "main_experiment_empty"
+
+            # Setup: Create PE_EXPERIMENT with no preference comparisons
+            pe_experiment = Experiment(
+                name="test_profile_empty",
+                search_space=self._get_pe_search_space(),
+            )
+
+            # Setup: Attach PE_EXPERIMENT without any data
+            aux_exp = AuxiliaryExperiment(experiment=pe_experiment, data=None)
+            exp.add_auxiliary_experiment(
+                auxiliary_experiment=aux_exp,
+                purpose=AuxiliaryExperimentPurpose.PE_EXPERIMENT,
+            )
+            pref_opt_config = self._make_pref_opt_config(
+                profile_name="test_profile_empty"
+            )
+
+            # Execute & Assert: Should raise DataRequiredError due to empty data
+            with self.assertRaisesRegex(
+                DataRequiredError,
+                "No preference data found in PE_EXPERIMENT",
+            ):
+                get_trace(exp, pref_opt_config)
+
+        with self.subTest("with_pe_experiment_valid_data_computes_utility"):
+            # This subtest verifies that when PE_EXPERIMENT exists with data,
+            # the code attempts to use the preference model (may fall back).
+
+            # Setup: Create main experiment with tracking data
+            exp = get_experiment_with_observations(
+                observations=[[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]],
+            )
+            exp.name = "main_experiment_with_pe"
+
+            # Setup: Create PE_EXPERIMENT with minimal preference data
+            pe_experiment = Experiment(
+                name="test_profile_with_minimal_data",
+                search_space=self._get_pe_search_space(),
+            )
+
+            # Setup: Add one pairwise preference comparison (minimal data)
+            trial1 = pe_experiment.new_batch_trial()
+            trial1.add_arm(Arm(name="0_0", parameters={"m1": 0.5, "m2": 1.0}))
+            trial1.add_arm(Arm(name="0_1", parameters={"m1": 1.0, "m2": 1.5}))
+            trial1.mark_running(no_runner_required=True).mark_completed()
+
+            # Setup: Create minimal preference data
+            pe_data_records = [
+                {
+                    "trial_index": 0,
+                    "arm_name": "0_0",
+                    "metric_name": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                    "mean": 0.0,
+                    "sem": 0.0,
+                    "metric_signature": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                },
+                {
+                    "trial_index": 0,
+                    "arm_name": "0_1",
+                    "metric_name": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                    "mean": 1.0,
+                    "sem": 0.0,
+                    "metric_signature": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                },
+            ]
+            pe_data = Data(df=pd.DataFrame.from_records(pe_data_records))
+            pe_experiment.attach_data(pe_data)
+
+            # Setup: Attach PE_EXPERIMENT to main experiment
+            aux_exp = AuxiliaryExperiment(experiment=pe_experiment, data=pe_data)
+            exp.add_auxiliary_experiment(
+                auxiliary_experiment=aux_exp,
+                purpose=AuxiliaryExperimentPurpose.PE_EXPERIMENT,
+            )
+            pref_opt_config = self._make_pref_opt_config(
+                profile_name="test_profile_with_minimal_data"
+            )
+
+            # Execute: With minimal data, model may fail and fall back gracefully
+            trace = get_trace(exp, pref_opt_config)
+
+            # Assert: Verify trace is valid and monotonically increasing
+            self._assert_valid_trace(trace, expected_len=3, check_monotonic=True)