Specify multitask trial_type as discrete variable

Author: shuds13

optimas/generators/ax/developer/multitask.py

@@ -67,7 +67,7 @@
     TrialStatus,
 )
 from .ax_metric import AxMetric
-from generator_standard.vocs import VOCS
+from generator_standard.vocs import VOCS, DiscreteVariable
 
 # Define generator states.
 NOT_STARTED = "not_started"
@@ -193,13 +193,11 @@ def __init__(
         model_save_period: Optional[int] = 5,
         model_history_dir: Optional[str] = "model_history",
     ) -> None:
-        # As trial parameters these get written to history array
-        # Ax trial_index and arm toegther locate a point
-        # Multiple points (Optimas trials) can share the same Ax trial_index
-        custom_trial_parameters = [
-            TrialParameter("trial_type", "ax_trial_type", dtype="U32"),
-        ]
         self._check_inputs(vocs, lofi_task, hifi_task)
+        
+        # Convert discrete variables to trial parameters before calling super().__init__
+        custom_trial_parameters = self._convert_discrete_variables_to_trial_parameters(vocs)
+        
         super().__init__(
             vocs=vocs,
             use_cuda=use_cuda,
@@ -223,11 +221,22 @@ def __init__(
         self.current_trial = None
         self.gr_lofi = None
         self._experiment = self._create_experiment()
-
+        
         # Internal mapping: _id -> (arm_name, ax_trial_id, trial_type)
         self._id_mapping = {}
         self._next_id = 0
 
+    def _convert_discrete_variables_to_trial_parameters(self, vocs: VOCS) -> List[TrialParameter]:
+        """Convert discrete variables from VOCS to TrialParameter objects."""
+        trial_parameters = []
+        for var_name, var_spec in vocs.variables.items():
+            if isinstance(var_spec, DiscreteVariable):
+                # Convert discrete variable to trial parameter
+                max_len = max(len(str(val)) for val in var_spec.values)
+                trial_param = TrialParameter(var_name, var_name, dtype=f"U{max_len}")
+                trial_parameters.append(trial_param)
+        return trial_parameters
+
     def get_gen_specs(
         self, sim_workers: int, run_params: Dict, sim_max: int
     ) -> Dict:
@@ -239,19 +248,30 @@ def get_gen_specs(
         gen_specs["out"].append(("task", str, max_length))
         return gen_specs
 
+    def _validate_vocs(self, vocs: VOCS) -> None:
+        """Validate VOCS for multitask generator."""
+        super()._validate_vocs(vocs)
+        # Check that only one objective has been given.
+        n_objectives = len(vocs.objectives)
+        assert n_objectives == 1, (
+            "Multitask generator supports only a single objective. "
+            "Objectives given: {}.".format(n_objectives)
+        )
+        # Check that there is a discrete variable called 'trial_type'
+        assert "trial_type" in vocs.variables, (
+            "Multitask generator requires a discrete variable named 'trial_type'"
+        )
+        assert isinstance(vocs.variables["trial_type"], DiscreteVariable), (
+            "Variable 'trial_type' must be a discrete variable"
+        )
+
     def _check_inputs(
         self,
         vocs: VOCS,
         lofi_task: Task,
         hifi_task: Task,
     ) -> None:
         """Check that the generator inputs are valid."""
-        # Check that only one objective has been given.
-        n_objectives = len(vocs.objectives)
-        assert n_objectives == 1, (
-            "Multitask generator supports only a single objective. "
-            "Objectives given: {}.".format(n_objectives)
-        )
         # Check that the number of low-fidelity trials per iteration is larger
         # than that of high-fidelity trials.
         assert lofi_task.n_opt >= hifi_task.n_opt, (
@@ -270,6 +290,12 @@ def suggest(self, num_points: Optional[int]) -> List[dict]:
                     var.name: arm.parameters.get(var.name)
                     for var in self._varying_parameters
                 }
+                # SH We can use a discrete var here in vocs (converted for now to trial parameters)
+                # But unlike varying parameters the name refers to a fixed generator concept.
+                for trial_param in self._custom_trial_parameters:
+                    if trial_param.name == "trial_type":
+                        point[trial_param.name] = trial_type
+                
                 # Generate unique _id and store mapping
                 current_id = self._next_id
                 self._id_mapping[current_id] = {
@@ -278,7 +304,6 @@ def suggest(self, num_points: Optional[int]) -> List[dict]:
                     "trial_type": trial_type,
                 }
                 point["_id"] = current_id
-                point["trial_type"] = trial_type  # Keep trial_type for now
                 self._next_id += 1
                 points.append(point)
         return points

optimas/generators/base.py

@@ -20,7 +20,7 @@
     TrialParameter,
     TrialStatus,
 )
-from generator_standard.vocs import VOCS
+from generator_standard.vocs import VOCS, ContinuousVariable
 from generator_standard.generator import Generator as StandardGenerator
 
 logger = get_logger(__name__)
@@ -127,13 +127,14 @@ def _convert_vocs_variables_to_varying_parameters(
         varying_parameters = []
         for var_name, var_spec in self._vocs.variables.items():
             # Only handle ContinuousVariable for now
-            vp = VaryingParameter(
-                name=var_name,
-                lower_bound=var_spec.domain[0],
-                upper_bound=var_spec.domain[1],
-                default_value=var_spec.default_value,
-            )
-            varying_parameters.append(vp)
+            if isinstance(var_spec, ContinuousVariable):
+                vp = VaryingParameter(
+                    name=var_name,
+                    lower_bound=var_spec.domain[0],
+                    upper_bound=var_spec.domain[1],
+                    default_value=var_spec.default_value,
+                )
+                varying_parameters.append(vp)
         return varying_parameters
 
     def _convert_vocs_objectives_to_objectives(self) -> List[Objective]:

tests/test_ax_generators.py

@@ -534,7 +534,7 @@ def test_ax_multitask():
     """Test that an exploration with a multitask generator runs"""
 
     vocs = VOCS(
-        variables={"x0": [-50.0, 5.0], "x1": [-5.0, 15.0]},
+        variables={"x0": [-50.0, 5.0], "x1": [-5.0, 15.0], "trial_type": {"task_1", "task_2"}},
         objectives={"f": "MAXIMIZE"},
     )
 
@@ -721,7 +721,7 @@ def test_ax_multitask_with_history():
     """
 
     vocs = VOCS(
-        variables={"x0": [-50.0, 5.0], "x1": [-5.0, 15.0]},
+        variables={"x0": [-50.0, 5.0], "x1": [-5.0, 15.0], "trial_type": {"task_1", "task_2"}},
         objectives={"f": "MAXIMIZE"},
     )