allow TrialAsTask with less than two trials

Summary:
The transform is a no-op in that case.

This is necessary to enable STGP vs MTGP model selection without adding additional nodes to the GS, thus creating further complication of the GS DAG.

Reviewed By: mgarrard

Differential Revision: D69496509

ax/modelbridge/transforms/stratified_standardize_y.py

@@ -99,7 +99,9 @@ def __init__(
             ]
             if len(task_parameters) == 0:
                 raise ValueError(
-                    "Must specify parameter for stratified standardization"
+                    "Must specify parameter for stratified standardization. This can "
+                    "happen if TrialAsTask is a no-op, due to there only being a single"
+                    " task level."
                 )
             elif len(task_parameters) != 1:
                 raise ValueError(

ax/modelbridge/transforms/tests/test_trial_as_task_transform.py

@@ -248,3 +248,59 @@ def test_w_robust_search_space(self) -> None:
                 observations=[],
                 modelbridge=self.modelbridge,
             )
+
+    def test_less_than_two_trials(self) -> None:
+        # test transform is a no-op with less than two trials
+        exp = get_branin_experiment()
+        exp.new_trial().add_arm(Arm(parameters={"x": 1}))
+        modelbridge = Adapter(
+            search_space=exp.search_space,
+            model=Generator(),
+            experiment=exp,
+        )
+        training_obs = self.training_obs[:1]
+        t = TrialAsTask(
+            search_space=exp.search_space,
+            observations=training_obs,
+            modelbridge=modelbridge,
+        )
+        self.assertEqual(t.trial_level_map, {})
+        training_feats = [training_obs[0].features]
+        training_feats_clone = deepcopy(training_feats)
+        self.assertEqual(
+            t.transform_observation_features(training_feats_clone), training_feats
+        )
+        self.assertEqual(
+            t.untransform_observation_features(training_feats), training_feats_clone
+        )
+        ss2 = exp.search_space.clone()
+        self.assertEqual(t.transform_search_space(ss2), exp.search_space)
+
+    def test_less_than_two_levels(self) -> None:
+        # test transform is a no-op with less than two trials
+        exp = get_branin_experiment()
+        exp.new_trial().add_arm(Arm(parameters={"x": 1}))
+        exp.new_trial().add_arm(Arm(parameters={"x": 2}))
+        modelbridge = Adapter(
+            search_space=exp.search_space,
+            model=Generator(),
+            experiment=exp,
+        )
+        training_obs = self.training_obs[:1]
+        t = TrialAsTask(
+            search_space=exp.search_space,
+            observations=training_obs,
+            modelbridge=modelbridge,
+            config={"trial_level_map": {"t": {0: "v1", 1: "v1"}}},
+        )
+        self.assertEqual(t.trial_level_map, {})
+        training_feats = [training_obs[0].features]
+        training_feats_clone = deepcopy(training_feats)
+        self.assertEqual(
+            t.transform_observation_features(training_feats_clone), training_feats
+        )
+        self.assertEqual(
+            t.untransform_observation_features(training_feats), training_feats_clone
+        )
+        ss2 = exp.search_space.clone()
+        self.assertEqual(t.transform_search_space(ss2), exp.search_space)

ax/modelbridge/transforms/trial_as_task.py

@@ -51,6 +51,8 @@ class TrialAsTask(Transform):
 
     Will raise if trial not specified for every point in the training data.
 
+    If there are fewer than 2 trials or levels, the transform is a no-op.
+
     Transform is done in-place.
     """
 
@@ -109,7 +111,13 @@ def __init__(
             self.inverse_map = None
         # Compute target values
         self.target_values: dict[str, int | str] = {}
-        for p_name, trial_map in self.trial_level_map.items():
+
+        for p_name, trial_map in list(self.trial_level_map.items()):
+            if len(set(trial_map.values())) < 2:
+                # If there are less than two distinct levels, then we don't need to
+                # create a task parameter and the transform is a no-op.
+                del self.trial_level_map[p_name]
+                continue
             if config is not None and "target_trial" in config:
                 target_trial = int(config["target_trial"])  # pyre-ignore [6]
             else:
@@ -133,6 +141,9 @@ def transform_observation_features(
         trials. Trial indices set to None are probably pending points passed in by the
         user.
         """
+        if len(self.trial_level_map) == 0:
+            # no-op
+            return observation_features
         for obsf in observation_features:
             for p_name, level_dict in self.trial_level_map.items():
                 if obsf.trial_index is not None and int(obsf.trial_index) in level_dict:
@@ -147,6 +158,9 @@ def transform_observation_features(
         return observation_features
 
     def _transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
+        if len(self.trial_level_map) == 0:
+            # no-op
+            return search_space
         for p_name, level_dict in self.trial_level_map.items():
             level_values = sorted(set(level_dict.values()))
             if len(level_values) < 2:
@@ -174,6 +188,20 @@ def _transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
     def untransform_observation_features(
         self, observation_features: list[ObservationFeatures]
     ) -> list[ObservationFeatures]:
+        """If task parameters have been added to observation features by
+        this parameter, then remove those task parameters and restore
+        the trial index/
+
+        Args:
+            observation_features: List of observation features to untransform.
+
+        Returns:
+            List of observation features with task parameters removed and trial
+                index restored.
+        """
+        if len(self.trial_level_map) == 0:
+            # no-op
+            return observation_features
         for obsf in observation_features:
             for p_name in self.trial_level_map:
                 pval = obsf.parameters.pop(p_name)