Repository: microsoft/MLOS

Issue:
Specify prior distributions for the Tunable parameters
Currently we assume that all numeric tunables have uniform distribution of the values within the specified range.
We have to give the user an ability to specify other prior distributions and their parameters. We need to support at least Beta and Normal distributions.


Discussion:

Resolving patch:
diff --git a/mlos_bench/mlos_bench/optimizers/convert_configspace.py b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
--- a/mlos_bench/mlos_bench/optimizers/convert_configspace.py
+++ b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
@@ -8,15 +8,18 @@
 
 import logging
 
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from ConfigSpace import (
+    Beta,
     CategoricalHyperparameter,
     Configuration,
     ConfigurationSpace,
     EqualsCondition,
-    UniformFloatHyperparameter,
-    UniformIntegerHyperparameter,
+    Float,
+    Integer,
+    Normal,
+    Uniform,
 )
 from mlos_bench.tunables.tunable import Tunable, TunableValue
 from mlos_bench.tunables.tunable_groups import TunableGroups
@@ -76,24 +79,51 @@ def _tunable_to_configspace(
                 meta=meta)
         })
 
+    distribution: Union[Uniform, Normal, Beta, None] = None
+    if tunable.distribution == "uniform":
+        distribution = Uniform()
+    elif tunable.distribution == "normal":
+        distribution = Normal(
+            mu=tunable.distribution_params["mu"],
+            sigma=tunable.distribution_params["sigma"]
+        )
+    elif tunable.distribution == "beta":
+        distribution = Beta(
+            alpha=tunable.distribution_params["alpha"],
+            beta=tunable.distribution_params["beta"]
+        )
+    elif tunable.distribution is not None:
+        raise TypeError(f"Invalid Distribution Type: {tunable.distribution}")
+
     if tunable.type == "int":
-        hp_type = UniformIntegerHyperparameter
+        range_hp = Integer(
+            name=tunable.name,
+            bounds=(int(tunable.range[0]), int(tunable.range[1])),
+            log=bool(tunable.is_log),
+            q=int(tunable.quantization) if tunable.quantization is not None else None,
+            distribution=distribution,
+            default=(int(tunable.default)
+                     if tunable.in_range(tunable.default) and tunable.default is not None
+                     else None),
+            meta=meta
+        )
     elif tunable.type == "float":
-        hp_type = UniformFloatHyperparameter
+        range_hp = Float(
+            name=tunable.name,
+            bounds=tunable.range,
+            log=bool(tunable.is_log),
+            q=tunable.quantization,     # type: ignore[arg-type]
+            distribution=distribution,  # type: ignore[arg-type]
+            default=(float(tunable.default)
+                     if tunable.in_range(tunable.default) and tunable.default is not None
+                     else None),
+            meta=meta
+        )
     else:
-        raise TypeError(f"Undefined Parameter Type: {tunable.type}")
+        raise TypeError(f"Invalid Parameter Type: {tunable.type}")
 
     if not tunable.special:
-        return ConfigurationSpace({
-            tunable.name: hp_type(
-                name=tunable.name,
-                lower=tunable.range[0],
-                upper=tunable.range[1],
-                log=tunable.is_log,
-                q=tunable.quantization,
-                default_value=tunable.default if tunable.in_range(tunable.default) else None,
-                meta=meta)
-        })
+        return ConfigurationSpace({tunable.name: range_hp})
 
     # Compute the probabilities of switching between regular and special values.
     special_weights: Optional[List[float]] = None
@@ -106,15 +136,7 @@ def _tunable_to_configspace(
     # one for special values, and one to choose between the two.
     (special_name, type_name) = special_param_names(tunable.name)
     conf_space = ConfigurationSpace({
-        tunable.name: hp_type(
-            name=tunable.name,
-            lower=tunable.range[0],
-            upper=tunable.range[1],
-            log=tunable.is_log,
-            q=tunable.quantization,
-            default_value=tunable.default if tunable.in_range(tunable.default) else None,
-            meta=meta
-        ),
+        tunable.name: range_hp,
         special_name: CategoricalHyperparameter(
             name=special_name,
             choices=tunable.special,
diff --git a/mlos_bench/mlos_bench/tunables/tunable.py b/mlos_bench/mlos_bench/tunables/tunable.py
--- a/mlos_bench/mlos_bench/tunables/tunable.py
+++ b/mlos_bench/mlos_bench/tunables/tunable.py
@@ -9,7 +9,7 @@
 import collections
 import logging
 
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, TypedDict, Union
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Type, TypedDict, Union
 
 _LOG = logging.getLogger(__name__)
 
@@ -17,6 +17,17 @@
 """A tunable parameter value type alias."""
 TunableValue = Union[int, float, Optional[str]]
 
+DistributionName = Literal["uniform", "normal", "beta"]
+
+
+class DistributionDict(TypedDict, total=False):
+    """
+    A typed dict for tunable parameters' distributions.
+    """
+
+    type: DistributionName
+    params: Optional[Dict[str, float]]
+
 
 class TunableDict(TypedDict, total=False):
     """
@@ -34,6 +45,7 @@ class TunableDict(TypedDict, total=False):
     range: Optional[Union[Sequence[int], Sequence[float]]]
     quantization: Optional[Union[int, float]]
     log: Optional[bool]
+    distribution: Optional[DistributionDict]
     special: Optional[Union[List[int], List[float]]]
     values_weights: Optional[List[float]]
     special_weights: Optional[List[float]]
@@ -41,7 +53,7 @@ class TunableDict(TypedDict, total=False):
     meta: Dict[str, Any]
 
 
-class Tunable:  # pylint: disable=too-many-instance-attributes
+class Tunable:  # pylint: disable=too-many-instance-attributes,too-many-public-methods
     """
     A tunable parameter definition and its current value.
     """
@@ -64,7 +76,7 @@ def __init__(self, name: str, config: TunableDict):
         config : dict
             Python dict that represents a Tunable (e.g., deserialized from JSON)
         """
-        if '!' in name:  # TODO: Use a regex here and in JSON schema
+        if not isinstance(name, str) or '!' in name:  # TODO: Use a regex here and in JSON schema
             raise ValueError(f"Invalid name of the tunable: {name}")
         self._name = name
         self._type = config["type"]  # required
@@ -80,6 +92,12 @@ def __init__(self, name: str, config: TunableDict):
         self._range: Optional[Union[Tuple[int, int], Tuple[float, float]]] = None
         self._quantization: Optional[Union[int, float]] = config.get("quantization")
         self._log: Optional[bool] = config.get("log")
+        self._distribution: Optional[DistributionName] = None
+        self._distribution_params: Dict[str, float] = {}
+        distr = config.get("distribution")
+        if distr:
+            self._distribution = distr["type"]  # required
+            self._distribution_params = distr.get("params") or {}
         config_range = config.get("range")
         if config_range is not None:
             assert len(config_range) == 2, f"Invalid range: {config_range}"
@@ -95,60 +113,81 @@ def __init__(self, name: str, config: TunableDict):
         self.value = self._default
 
     def _sanity_check(self) -> None:
-        # pylint: disable=too-complex,too-many-branches
         """
         Check if the status of the Tunable is valid, and throw ValueError if it is not.
         """
         if self.is_categorical:
-            if not (self._values and isinstance(self._values, collections.abc.Iterable)):
-                raise ValueError(f"Must specify values for the categorical type tunable {self}")
-            if self._range is not None:
-                raise ValueError(f"Range must be None for the categorical type tunable {self}")
-            if len(set(self._values)) != len(self._values):
-                raise ValueError(f"Values must be unique for the categorical type tunable {self}")
-            if self._special:
-                raise ValueError(f"Categorical tunable cannot have special values: {self}")
-            if self._range_weight is not None:
-                raise ValueError(f"Categorical tunable cannot have range_weight: {self}")
-            if self._log is not None:
-                raise ValueError(f"Categorical tunable cannot have log parameter: {self}")
-            if self._quantization is not None:
-                raise ValueError(f"Categorical tunable cannot have quantization parameter: {self}")
-            if self._weights:
-                if len(self._weights) != len(self._values):
-                    raise ValueError(f"Must specify weights for all values: {self}")
-                if any(w < 0 for w in self._weights):
-                    raise ValueError(f"All weights must be non-negative: {self}")
+            self._sanity_check_categorical()
         elif self.is_numerical:
-            if self._values is not None:
-                raise ValueError(f"Values must be None for the numerical type tunable {self}")
-            if not self._range or len(self._range) != 2 or self._range[0] >= self._range[1]:
-                raise ValueError(f"Invalid range for tunable {self}: {self._range}")
-            if self._quantization is not None:
-                if self.dtype == int:
-                    if not isinstance(self._quantization, int):
-                        raise ValueError(f"Quantization of a int param should be an int: {self}")
-                    if self._quantization <= 1:
-                        raise ValueError(f"Number of quantization points is <= 1: {self}")
-                if self.dtype == float:
-                    if not isinstance(self._quantization, (float, int)):
-                        raise ValueError(f"Quantization of a float param should be a float or int: {self}")
-                    if self._quantization <= 0:
-                        raise ValueError(f"Number of quantization points is <= 0: {self}")
-            if self._weights:
-                if self._range_weight is None:
-                    raise ValueError(f"Must specify weight for the range: {self}")
-                if len(self._weights) != len(self._special):
-                    raise ValueError("Must specify weights for all special values {self}")
-                if any(w < 0 for w in self._weights + [self._range_weight]):
-                    raise ValueError(f"All weights must be non-negative: {self}")
-            elif self._range_weight is not None:
-                raise ValueError(f"Must specify both weights and range_weight or none: {self}")
+            self._sanity_check_numerical()
         else:
             raise ValueError(f"Invalid parameter type for tunable {self}: {self._type}")
         if not self.is_valid(self.default):
             raise ValueError(f"Invalid default value for tunable {self}: {self.default}")
 
+    def _sanity_check_categorical(self) -> None:
+        """
+        Check if the status of the categorical Tunable is valid, and throw ValueError if it is not.
+        """
+        # pylint: disable=too-complex
+        assert self.is_categorical
+        if not (self._values and isinstance(self._values, collections.abc.Iterable)):
+            raise ValueError(f"Must specify values for the categorical type tunable {self}")
+        if self._range is not None:
+            raise ValueError(f"Range must be None for the categorical type tunable {self}")
+        if len(set(self._values)) != len(self._values):
+            raise ValueError(f"Values must be unique for the categorical type tunable {self}")
+        if self._special:
+            raise ValueError(f"Categorical tunable cannot have special values: {self}")
+        if self._range_weight is not None:
+            raise ValueError(f"Categorical tunable cannot have range_weight: {self}")
+        if self._log is not None:
+            raise ValueError(f"Categorical tunable cannot have log parameter: {self}")
+        if self._quantization is not None:
+            raise ValueError(f"Categorical tunable cannot have quantization parameter: {self}")
+        if self._distribution is not None:
+            raise ValueError(f"Categorical parameters do not support `distribution`: {self}")
+        if self._weights:
+            if len(self._weights) != len(self._values):
+                raise ValueError(f"Must specify weights for all values: {self}")
+            if any(w < 0 for w in self._weights):
+                raise ValueError(f"All weights must be non-negative: {self}")
+
+    def _sanity_check_numerical(self) -> None:
+        """
+        Check if the status of the numerical Tunable is valid, and throw ValueError if it is not.
+        """
+        # pylint: disable=too-complex,too-many-branches
+        assert self.is_numerical
+        if self._values is not None:
+            raise ValueError(f"Values must be None for the numerical type tunable {self}")
+        if not self._range or len(self._range) != 2 or self._range[0] >= self._range[1]:
+            raise ValueError(f"Invalid range for tunable {self}: {self._range}")
+        if self._quantization is not None:
+            if self.dtype == int:
+                if not isinstance(self._quantization, int):
+                    raise ValueError(f"Quantization of a int param should be an int: {self}")
+                if self._quantization <= 1:
+                    raise ValueError(f"Number of quantization points is <= 1: {self}")
+            if self.dtype == float:
+                if not isinstance(self._quantization, (float, int)):
+                    raise ValueError(f"Quantization of a float param should be a float or int: {self}")
+                if self._quantization <= 0:
+                    raise ValueError(f"Number of quantization points is <= 0: {self}")
+        if self._distribution is not None and self._distribution not in {"uniform", "normal", "beta"}:
+            raise ValueError(f"Invalid distribution: {self}")
+        if self._distribution_params and self._distribution is None:
+            raise ValueError(f"Must specify the distribution: {self}")
+        if self._weights:
+            if self._range_weight is None:
+                raise ValueError(f"Must specify weight for the range: {self}")
+            if len(self._weights) != len(self._special):
+                raise ValueError("Must specify weights for all special values {self}")
+            if any(w < 0 for w in self._weights + [self._range_weight]):
+                raise ValueError(f"All weights must be non-negative: {self}")
+        elif self._range_weight is not None:
+            raise ValueError(f"Must specify both weights and range_weight or none: {self}")
+
     def __repr__(self) -> str:
         """
         Produce a human-readable version of the Tunable (mostly for logging).
@@ -533,6 +572,31 @@ def is_log(self) -> Optional[bool]:
         assert self.is_numerical
         return self._log
 
+    @property
+    def distribution(self) -> Optional[DistributionName]:
+        """
+        Get the name of the distribution (uniform, normal, or beta) if specified.
+
+        Returns
+        -------
+        distribution : str
+            Name of the distribution (uniform, normal, or beta) or None.
+        """
+        return self._distribution
+
+    @property
+    def distribution_params(self) -> Dict[str, float]:
+        """
+        Get the parameters of the distribution, if specified.
+
+        Returns
+        -------
+        distribution_params : Dict[str, float]
+            Parameters of the distribution or None.
+        """
+        assert self._distribution is not None
+        return self._distribution_params
+
     @property
     def categories(self) -> List[Optional[str]]:
         """


Test patch:
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc
@@ -3,7 +3,7 @@
         "cost": 1,
         "description": "Covariant group description",
         "params": {
-            "int": {
+            "beta-int": {
                 "description": "Int",
                 "type": "int",
                 "default": 10,
@@ -13,9 +13,35 @@
                 "special_weights": [0.1],
                 "range_weight": 0.9,
                 "quantization": 50,
+                "distribution": {
+                    "type": "beta",
+                    "params": {
+                        "alpha": 0.1,
+                        "beta": 0.1
+                    }
+                },
                 "log": true
             },
-            "float": {
+            "normal-int": {
+                "description": "Int",
+                "type": "int",
+                "default": 10,
+                "range": [1, 500],
+                "meta": {"suffix": "MB"},
+                "special": [-1],
+                "special_weights": [0.1],
+                "range_weight": 0.9,
+                "quantization": 50,
+                "distribution": {
+                    "type": "normal",
+                    "params": {
+                        "mu": 0,
+                        "sigma": 0.1
+                    }
+                },
+                "log": true
+            },
+            "uniform-float": {
                 "description": "Float",
                 "type": "float",
                 "default": 10.1,
@@ -23,6 +49,9 @@
                 "range": [1.1, 111.1],
                 "special": [-1.1],
                 "quantization": 10,
+                "distribution": {
+                    "type": "uniform"
+                },
                 "log": false
             },
             "cat": {
diff --git a/mlos_bench/mlos_bench/tests/tunables/tunable_distributions_test.py b/mlos_bench/mlos_bench/tests/tunables/tunable_distributions_test.py
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_distributions_test.py
@@ -0,0 +1,110 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests for checking tunable parameters' distributions.
+"""
+
+import json5 as json
+import pytest
+
+from mlos_bench.tunables.tunable import Tunable
+
+
+def test_categorical_distribution() -> None:
+    """
+    Try to instantiate a categorical tunable with distribution specified.
+    """
+    with pytest.raises(ValueError):
+        Tunable(name='test', config={
+            "type": "categorical",
+            "values": ["foo", "bar", "baz"],
+            "distribution": {
+                "type": "uniform"
+            },
+            "default": "foo"
+        })
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_distribution_uniform(tunable_type: str) -> None:
+    """
+    Create a numeric Tunable with explicit uniform distribution.
+    """
+    tunable = Tunable(name="test", config={
+        "type": tunable_type,
+        "range": [0, 10],
+        "distribution": {
+            "type": "uniform"
+        },
+        "default": 0
+    })
+    assert tunable.is_numerical
+    assert tunable.distribution == "uniform"
+    assert not tunable.distribution_params
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_distribution_normal(tunable_type: str) -> None:
+    """
+    Create a numeric Tunable with explicit Gaussian distribution specified.
+    """
+    tunable = Tunable(name="test", config={
+        "type": tunable_type,
+        "range": [0, 10],
+        "distribution": {
+            "type": "normal",
+            "params": {
+                "mu": 0,
+                "sigma": 1.0
+            }
+        },
+        "default": 0
+    })
+    assert tunable.distribution == "normal"
+    assert tunable.distribution_params == {"mu": 0, "sigma": 1.0}
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_distribution_beta(tunable_type: str) -> None:
+    """
+    Create a numeric Tunable with explicit Beta distribution specified.
+    """
+    tunable = Tunable(name="test", config={
+        "type": tunable_type,
+        "range": [0, 10],
+        "distribution": {
+            "type": "beta",
+            "params": {
+                "alpha": 2,
+                "beta": 5
+            }
+        },
+        "default": 0
+    })
+    assert tunable.distribution == "beta"
+    assert tunable.distribution_params == {"alpha": 2, "beta": 5}
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_distribution_unsupported(tunable_type: str) -> None:
+    """
+    Create a numeric Tunable with unsupported distribution.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 10],
+        "distribution": {{
+            "type": "poisson",
+            "params": {{
+                "lambda": 1.0
+            }}
+        }},
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name="test", config=config)
diff --git a/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_distr_test.py b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_distr_test.py
new file mode 100644
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_distr_test.py
@@ -0,0 +1,88 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Unit tests for converting tunable parameters with explicitly
+specified distributions to ConfigSpace.
+"""
+
+import pytest
+
+from ConfigSpace import (
+    CategoricalHyperparameter,
+    BetaFloatHyperparameter,
+    BetaIntegerHyperparameter,
+    NormalFloatHyperparameter,
+    NormalIntegerHyperparameter,
+    UniformFloatHyperparameter,
+    UniformIntegerHyperparameter,
+)
+
+from mlos_bench.tunables.tunable import DistributionName
+from mlos_bench.tunables.tunable_groups import TunableGroups
+from mlos_bench.optimizers.convert_configspace import (
+    special_param_names,
+    tunable_groups_to_configspace,
+)
+
+
+_CS_HYPERPARAMETER = {
+    ("float", "beta"): BetaFloatHyperparameter,
+    ("int", "beta"): BetaIntegerHyperparameter,
+    ("float", "normal"): NormalFloatHyperparameter,
+    ("int", "normal"): NormalIntegerHyperparameter,
+    ("float", "uniform"): UniformFloatHyperparameter,
+    ("int", "uniform"): UniformIntegerHyperparameter,
+}
+
+
+@pytest.mark.parametrize("param_type", ["int", "float"])
+@pytest.mark.parametrize("distr_name,distr_params", [
+    ("normal", {"mu": 0.0, "sigma": 1.0}),
+    ("beta", {"alpha": 2, "beta": 5}),
+    ("uniform", {}),
+])
+def test_convert_numerical_distributions(param_type: str,
+                                         distr_name: DistributionName,
+                                         distr_params: dict) -> None:
+    """
+    Convert a numerical Tunable with explicit distribution to ConfigSpace.
+    """
+    tunable_name = "x"
+    tunable_groups = TunableGroups({
+        "tunable_group": {
+            "cost": 1,
+            "params": {
+                tunable_name: {
+                    "type": param_type,
+                    "range": [0, 100],
+                    "special": [-1, 0],
+                    "special_weights": [0.1, 0.2],
+                    "range_weight": 0.7,
+                    "distribution": {
+                        "type": distr_name,
+                        "params": distr_params
+                    },
+                    "default": 0
+                }
+            }
+        }
+    })
+
+    (tunable, _group) = tunable_groups.get_tunable(tunable_name)
+    assert tunable.distribution == distr_name
+    assert tunable.distribution_params == distr_params
+
+    space = tunable_groups_to_configspace(tunable_groups)
+
+    (tunable_special, tunable_type) = special_param_names(tunable_name)
+    assert set(space.keys()) == {tunable_name, tunable_type, tunable_special}
+
+    assert isinstance(space[tunable_special], CategoricalHyperparameter)
+    assert isinstance(space[tunable_type], CategoricalHyperparameter)
+
+    cs_param = space[tunable_name]
+    assert isinstance(cs_param, _CS_HYPERPARAMETER[param_type, distr_name])
+    for (key, val) in distr_params.items():
+        assert getattr(cs_param, key) == val