Add BoTorch_Modular to config, giving custom GP and acq func ability

Now users can customize their GP (kernel, MLL, GP class) as well as the acq func in the config without writing any code. Giving users a lot more control and power from any language with little setup

Author: madeline-scyphers

boa/config/converters.py

@@ -5,8 +5,13 @@
 
 import ax.early_stopping.strategies as early_stopping_strats
 import ax.global_stopping.strategies as global_stopping_strats
+import botorch.acquisition
+import botorch.models
+import gpytorch.kernels
+import gpytorch.mlls
 from ax.modelbridge.generation_node import GenerationStep
 from ax.modelbridge.registry import Models
+from ax.models.torch.botorch_modular.surrogate import Surrogate
 from ax.service.utils.instantiation import TParameterRepresentation
 from ax.service.utils.scheduler_options import SchedulerOptions
 
@@ -49,6 +54,36 @@ def _gen_strat_converter(gs: Optional[dict] = None) -> dict:
                 gs["steps"][i] = step
                 steps.append(step)
                 continue
+            if "model_kwargs" in step:
+                if "botorch_acqf_class" in step["model_kwargs"] and not isinstance(
+                    step["model_kwargs"]["botorch_acqf_class"], botorch.acquisition.AcquisitionFunction
+                ):
+                    step["model_kwargs"]["botorch_acqf_class"] = getattr(
+                        botorch.acquisition, step["model_kwargs"]["botorch_acqf_class"]
+                    )
+
+                if "surrogate" in step["model_kwargs"]:
+                    if "mll_class" in step["model_kwargs"]["surrogate"] and not isinstance(
+                        step["model_kwargs"]["surrogate"]["mll_class"], gpytorch.mlls.MarginalLogLikelihood
+                    ):
+                        step["model_kwargs"]["surrogate"]["mll_class"] = getattr(
+                            gpytorch.mlls, step["model_kwargs"]["surrogate"]["mll_class"]
+                        )
+                    if "botorch_model_class" in step["model_kwargs"]["surrogate"] and not isinstance(
+                        step["model_kwargs"]["surrogate"]["botorch_model_class"], botorch.models.model.Model
+                    ):
+                        step["model_kwargs"]["surrogate"]["botorch_model_class"] = getattr(
+                            botorch.models, step["model_kwargs"]["surrogate"]["botorch_model_class"]
+                        )
+                    if "covar_module_class" in step["model_kwargs"]["surrogate"] and not isinstance(
+                        step["model_kwargs"]["surrogate"]["covar_module_class"], gpytorch.kernels.Kernel
+                    ):
+                        step["model_kwargs"]["surrogate"]["covar_module_class"] = getattr(
+                            gpytorch.kernels, step["model_kwargs"]["surrogate"]["covar_module_class"]
+                        )
+
+                    step["model_kwargs"]["surrogate"] = Surrogate(**step["model_kwargs"]["surrogate"])
+
             try:
                 step["model"] = Models[step["model"]]
             except KeyError:

boa/registry.py

@@ -1,4 +1,22 @@
-from ax.storage.json_store.registry import CORE_DECODER_REGISTRY, CORE_ENCODER_REGISTRY
+from __future__ import annotations
+
+from typing import Type
+
+import botorch.acquisition
+import gpytorch.kernels
+from ax.storage.botorch_modular_registry import (
+    ACQUISITION_FUNCTION_REGISTRY,
+    CLASS_TO_REGISTRY,
+    CLASS_TO_REVERSE_REGISTRY,
+)
+from ax.storage.json_store.registry import (
+    CORE_CLASS_DECODER_REGISTRY,
+    CORE_CLASS_ENCODER_REGISTRY,
+    CORE_DECODER_REGISTRY,
+    CORE_ENCODER_REGISTRY,
+    botorch_modular_to_dict,
+    class_from_json,
+)
 
 
 def config_to_dict(inst):
@@ -15,3 +33,25 @@ def _add_common_encodes_and_decodes():
     CORE_ENCODER_REGISTRY[BOAConfig] = config_to_dict
     # CORE_DECODER_REGISTRY[BOAConfig.__name__] = BOAConfig
     CORE_DECODER_REGISTRY[MetricType.__name__] = MetricType
+
+    CORE_CLASS_DECODER_REGISTRY["Type[Kernel]"] = class_from_json
+    CORE_CLASS_ENCODER_REGISTRY[gpytorch.kernels.Kernel] = botorch_modular_to_dict
+
+    KERNEL_REGISTRY = {getattr(gpytorch.kernels, kernel): kernel for kernel in gpytorch.kernels.__all__}
+
+    REVERSE_KERNEL_REGISTRY: dict[str, Type[gpytorch.kernels.Kernel]] = {v: k for k, v in KERNEL_REGISTRY.items()}
+
+    CLASS_TO_REGISTRY[gpytorch.kernels.Kernel] = KERNEL_REGISTRY
+    CLASS_TO_REVERSE_REGISTRY[gpytorch.kernels.Kernel] = REVERSE_KERNEL_REGISTRY
+
+    for acq_func_name in botorch.acquisition.__all__:
+        acq_func = getattr(botorch.acquisition, acq_func_name)
+        if acq_func not in ACQUISITION_FUNCTION_REGISTRY:
+            ACQUISITION_FUNCTION_REGISTRY[acq_func] = acq_func_name
+
+    REVERSE_ACQUISITION_FUNCTION_REGISTRY: dict[str, Type[botorch.acquisition.AcquisitionFunction]] = {
+        v: k for k, v in ACQUISITION_FUNCTION_REGISTRY.items()
+    }
+
+    CLASS_TO_REGISTRY[botorch.acquisition.AcquisitionFunction] = ACQUISITION_FUNCTION_REGISTRY
+    CLASS_TO_REVERSE_REGISTRY[botorch.acquisition.AcquisitionFunction] = REVERSE_ACQUISITION_FUNCTION_REGISTRY

tests/1unit_tests/test_generation_strategy.py

@@ -1,3 +1,7 @@
+import botorch.acquisition
+import botorch.models
+import gpytorch.kernels
+import gpytorch.mlls
 from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
 from ax.modelbridge.registry import Models
 
@@ -41,3 +45,32 @@ def test_auto_gen_use_saasbo(saasbo_config, tmp_path):
         assert "SAASBO" in gs.name
     else:
         assert "FullyBayesian" in gs.name
+
+
+def test_modular_botorch(gen_strat_modular_botorch_config, tmp_path):
+    controller = Controller(
+        config=gen_strat_modular_botorch_config,
+        wrapper=ScriptWrapper(config=gen_strat_modular_botorch_config, experiment_dir=tmp_path),
+    )
+    exp = get_experiment(
+        config=controller.config, runner=WrappedJobRunner(wrapper=controller.wrapper), wrapper=controller.wrapper
+    )
+    gs = get_generation_strategy(config=controller.config, experiment=exp)
+    cfg_botorch_modular = gen_strat_modular_botorch_config.orig_config["generation_strategy"]["steps"][-1]
+    step = gs._steps[-1]
+    assert step.model == Models.BOTORCH_MODULAR
+    mdl_kw = step.model_kwargs
+    assert mdl_kw["botorch_acqf_class"] == getattr(
+        botorch.acquisition, cfg_botorch_modular["model_kwargs"]["botorch_acqf_class"]
+    )
+    assert mdl_kw["acquisition_options"] == cfg_botorch_modular["model_kwargs"]["acquisition_options"]
+
+    assert mdl_kw["surrogate"].mll_class == getattr(
+        gpytorch.mlls, cfg_botorch_modular["model_kwargs"]["surrogate"]["mll_class"]
+    )
+    assert mdl_kw["surrogate"].botorch_model_class == getattr(
+        botorch.models, cfg_botorch_modular["model_kwargs"]["surrogate"]["botorch_model_class"]
+    )
+    assert mdl_kw["surrogate"].covar_module_class == getattr(
+        gpytorch.kernels, cfg_botorch_modular["model_kwargs"]["surrogate"]["covar_module_class"]
+    )

tests/conftest.py

@@ -85,6 +85,12 @@ def gen_strat1_config():
     return BOAConfig.from_jsonlike(file=config_path)
 
 
+@pytest.fixture
+def gen_strat_modular_botorch_config():
+    config_path = TEST_DIR / f"scripts/other_langs/r_package_streamlined/config_modular_botorch.yaml"
+    return BOAConfig.from_jsonlike(file=config_path)
+
+
 @pytest.fixture
 def synth_config():
     config_path = TEST_CONFIG_DIR / "test_config_synth.yaml"
@@ -233,3 +239,9 @@ def r_streamlined(tmp_path_factory, cd_to_root_and_back_session):
     config_path = TEST_DIR / f"scripts/other_langs/r_package_streamlined/config.yaml"
 
     yield cli_main(split_shell_command(f"--config-path {config_path} -td"), standalone_mode=False)
+
+
+@pytest.fixture(scope="session")
+def r_streamlined_botorch_modular(tmp_path_factory, cd_to_root_and_back_session):
+    config_path = TEST_DIR / f"scripts/other_langs/r_package_streamlined/config_modular_botorch.yaml"
+    return cli_main(split_shell_command(f"--config-path {config_path} -td"), standalone_mode=False)

tests/integration_tests/test_cli.py

@@ -75,9 +75,12 @@ def test_calling_command_line_test_script_doesnt_error_out_and_produces_correct_
 
 # parametrize the test to use the full version (all scripts) or the light version (only run_model.R)
 # or parametrize the test to use the streamlined version (doesn't use trial_status.json, only use output.json)
+# the botorch modular version is the same as the streamlined version, but also uses botorch modular
+# which uses a custom kernel, acquisition function, mll and botorch model class
+# (which can customize the GP process even more)
 @pytest.mark.parametrize(
     "r_scripts_run",
-    ["r_full", "r_light", "r_streamlined"],
+    ["r_full", "r_light", "r_streamlined", "r_streamlined_botorch_modular"],
 )
 @pytest.mark.skipif(not R_INSTALLED, reason="requires R to be installed")
 def test_calling_command_line_r_test_scripts(r_scripts_run, request):
@@ -92,7 +95,7 @@ def test_calling_command_line_r_test_scripts(r_scripts_run, request):
         assert "param_names" in data
         assert "metric_properties" in data
 
-    if "r_streamlined" == r_scripts_run:
+    if r_scripts_run in ("r_streamlined", "r_streamlined_botorch_modular"):
         with cd_and_cd_back(scheduler.wrapper.config_path.parent):
 
             pre_num_trials = len(scheduler.experiment.trials)

tests/scripts/other_langs/r_package_streamlined/config_modular_botorch.yaml

@@ -0,0 +1,60 @@
+objective:
+    metrics:
+        - name: metric
+scheduler:
+    n_trials: 15
+
+parameters:
+    x0:
+        'bounds': [ 0, 1 ]
+        'type': 'range'
+        'value_type': 'float'
+    x1:
+        'bounds': [ 0, 1]
+        'type': 'range'
+        'value_type': 'float'
+    x2:
+        'bounds': [ 0, 1 ]
+        'type': 'range'
+        'value_type': 'float'
+    x3:
+        'bounds': [ 0, 1]
+        'type': 'range'
+        'value_type': 'float'
+    x4:
+        'bounds': [ 0, 1 ]
+        'type': 'range'
+        'value_type': 'float'
+    x5:
+        'bounds': [ 0, 1]
+        'type': 'range'
+        'value_type': 'float'
+
+script_options:
+    # notice here that this is a shell command
+    # this is what BOA will do to launch your script
+    # it will also pass as a command line argument the current trial directory
+    # that is being parameterized
+
+    # This can either be a relative path or absolute path
+    # (by default when BOA launches from a config file
+    # it uses the config file directory as your working directory)
+    # here config.yaml and run_model.R are in the same directory
+    run_model: Rscript run_model.R
+    exp_name: "r_streamlined_botorch_modular"
+
+generation_strategy:
+    steps:
+        - model: SOBOL
+          num_trials: 5
+        - model: BOTORCH_MODULAR
+          num_trials: -1  # No limitation on how many trials should be produced from this step
+          model_kwargs:
+              surrogate:
+                  botorch_model_class: SingleTaskGP  # BoTorch model class name
+
+                  covar_module_class: RBFKernel  # GPyTorch kernel class name
+                  mll_class: LeaveOneOutPseudoLikelihood
+              botorch_acqf_class: qUpperConfidenceBound  # BoTorch acquisition function class name
+              acquisition_options:
+                  beta: 0.5