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,219 @@ 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) -> None:
+        """Assert trace has expected length, contains floats, is non-decreasing and has
+        more than one unique value."""
+        self.assertEqual(len(trace), expected_len)
+        for value in trace:
+            self.assertIsInstance(value, float)
+        for i in range(1, len(trace)):
+            self.assertGreaterEqual(
+                trace[i],
+                trace[i - 1],
+                msg=f"Trace not monotonically increasing at index {i}: {trace}",
+            )
+        unique_values = set(trace)
+        self.assertGreater(
+            len(unique_values),
+            1,
+            msg=f"Trace has only trivial values (all same): {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 preference profile",
+            ):
+                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 valid data,
+            # the code uses the preference model to compute utility-based traces.
+
+            # Setup: Create main experiment with tracking data
+            # Values are spread across the parameter space to get different utilities
+            exp = get_experiment_with_observations(
+                observations=[[1.0, 1.0], [5.0, 5.0], [9.0, 9.0]],
+            )
+            exp.name = "main_experiment_with_pe"
+
+            # Setup: Create PE_EXPERIMENT with sufficient preference data
+            pe_experiment = Experiment(
+                name="test_profile_with_data",
+                search_space=self._get_pe_search_space(),
+            )
+
+            # Setup: Add multiple pairwise preference comparisons
+            # Trial 1: User prefers higher values (0_1 over 0_0)
+            trial1 = pe_experiment.new_batch_trial()
+            trial1.add_arm(Arm(name="0_0", parameters={"m1": 1.0, "m2": 1.0}))
+            trial1.add_arm(Arm(name="0_1", parameters={"m1": 5.0, "m2": 5.0}))
+            trial1.mark_running(no_runner_required=True).mark_completed()
+
+            # Trial 2: Another comparison reinforcing preference for higher values
+            trial2 = pe_experiment.new_batch_trial()
+            trial2.add_arm(Arm(name="1_0", parameters={"m1": 3.0, "m2": 3.0}))
+            trial2.add_arm(Arm(name="1_1", parameters={"m1": 8.0, "m2": 8.0}))
+            trial2.mark_running(no_runner_required=True).mark_completed()
+
+            # Trial 3: Another comparison to strengthen the preference model
+            trial3 = pe_experiment.new_batch_trial()
+            trial3.add_arm(Arm(name="2_0", parameters={"m1": 2.0, "m2": 2.0}))
+            trial3.add_arm(Arm(name="2_1", parameters={"m1": 7.0, "m2": 7.0}))
+            trial3.mark_running(no_runner_required=True).mark_completed()
+
+            # Setup: Create preference data indicating user prefers higher metric values
+            # In each trial, the second arm (higher values) is preferred (mean=1.0)
+            pe_data_records = [
+                # Trial 1
+                {
+                    "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,
+                },
+                # Trial 2
+                {
+                    "trial_index": 1,
+                    "arm_name": "1_0",
+                    "metric_name": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                    "mean": 0.0,
+                    "sem": 0.0,
+                    "metric_signature": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                },
+                {
+                    "trial_index": 1,
+                    "arm_name": "1_1",
+                    "metric_name": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                    "mean": 1.0,
+                    "sem": 0.0,
+                    "metric_signature": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                },
+                # Trial 3
+                {
+                    "trial_index": 2,
+                    "arm_name": "2_0",
+                    "metric_name": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                    "mean": 0.0,
+                    "sem": 0.0,
+                    "metric_signature": Keys.PAIRWISE_PREFERENCE_QUERY.value,
+                },
+                {
+                    "trial_index": 2,
+                    "arm_name": "2_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_data"
+            )
+
+            # Execute: With valid data, model computes utility-based trace
+            trace = get_trace(exp, pref_opt_config)
+
+            # Assert: Verify trace is valid, monotonically increasing, and non-trivial
+            self._assert_valid_trace(
+                trace,
+                expected_len=3,
+            )