Enable support for multi-task BO

README.md

@@ -5,8 +5,6 @@
 Install PyTorch first, then:
 
 ```
-pip install --upgrade pip wheel packaging
-pip install git+https://github.com/aleximmer/[email protected]
 pip install laplace-bayesopt
 ```
 
@@ -28,12 +26,12 @@ train_X, train_Y = ..., ...
 
 model = LaplaceBoTorch(get_net, train_X, train_Y)
 
-# Use this model in your existing BoTorch loop, e.g. to replace BoTorch's MultiTaskGP model.
+# Use this model in your existing BoTorch loop, e.g. to replace BoTorch's SingleTaskGP model.
 ```
 
 The full arguments of `LaplaceBoTorch` can be found in the class documentation.
 
-Check out a full BoTorch example in `examples/botorch/experiments.py`.
+Check out examples in `examples/`.
 
 ## Useful References
 

examples/multitask.py

@@ -0,0 +1,121 @@
+"""
+Following https://botorch.org/tutorials/composite_mtbo
+"""
+
+from __future__ import annotations
+import warnings
+
+warnings.filterwarnings("ignore")
+
+from botorch.acquisition.logei import qLogExpectedImprovement
+from botorch.acquisition import GenericMCObjective
+from botorch.sampling import IIDNormalSampler
+import numpy as np
+import torch
+from torch import nn
+import tqdm
+
+
+from botorch.test_functions import Hartmann
+from botorch.optim.optimize import optimize_acqf
+
+from laplace_bayesopt.botorch import LaplaceBoTorch
+
+
+np.random.seed(10)
+torch.set_default_dtype(torch.float64)
+torch.manual_seed(10)
+
+
+class ContextualHartmann6(Hartmann):
+    def __init__(self, num_tasks: int = 3, noise_std=None, negate=False):
+        super().__init__(dim=6, noise_std=noise_std, negate=negate)
+        self.task_range = torch.linspace(0, 1, num_tasks).unsqueeze(-1)
+        self._bounds = [(0.0, 1.0) for _ in range(self.dim - 1)]
+        self.bounds = torch.tensor(self._bounds).t()
+
+    def evaluate_true(self, X: torch.Tensor) -> torch.Tensor:
+        batch_X = X.unsqueeze(-2)
+        batch_dims = X.ndim - 1
+
+        expanded_task_range = self.task_range
+        for _ in range(batch_dims):
+            expanded_task_range = expanded_task_range.unsqueeze(0)
+        task_range = expanded_task_range.repeat(*X.shape[:-1], 1, 1).to(X)
+        concatenated_X = torch.cat(
+            (
+                batch_X.repeat(*[1] * batch_dims, self.task_range.shape[0], 1),
+                task_range,
+            ),
+            dim=-1,
+        )
+        return super().evaluate_true(concatenated_X)
+
+
+NUM_TASKS = 3
+problem = ContextualHartmann6(num_tasks=NUM_TASKS, noise_std=0.001, negate=True)
+weights = torch.randn(NUM_TASKS)
+
+
+def callable_func(samples, X=None):
+    res = -torch.cos((samples**2) + samples * weights)
+    return res.squeeze().sum(dim=-1)
+
+
+objective = GenericMCObjective(callable_func)
+bounds = problem.bounds
+
+n_init = 5
+train_x = (bounds[1] - bounds[0]) * torch.rand(n_init, bounds.shape[1]) + bounds[0]
+train_y = problem(train_x)
+
+
+def get_net():
+    return torch.nn.Sequential(
+        nn.Linear(train_x.shape[-1], 20),
+        nn.ReLU(),
+        nn.Linear(20, 20),
+        nn.ReLU(),
+        nn.Linear(20, NUM_TASKS),
+    )
+
+
+model = LaplaceBoTorch(get_net, train_x, train_y)
+
+best_objective = objective(train_y).max()
+pbar = tqdm.trange(100)
+pbar.set_description(f"[Best objective = {best_objective:.3f}]")
+
+# For qEI
+NUM_SAMPLES = 4
+
+for i in pbar:
+    sampler = IIDNormalSampler(sample_shape=torch.Size([NUM_SAMPLES]))
+    acq_f = qLogExpectedImprovement(
+        model, best_f=best_objective, sampler=sampler, objective=objective
+    )
+
+    # Get a proposal for new x
+    new_x, val = optimize_acqf(
+        acq_f,
+        bounds=bounds,
+        q=4,
+        num_restarts=11,
+        raw_samples=22,
+    )
+
+    if len(new_x.shape) == 1:
+        new_x = new_x.unsqueeze(0)
+
+    # Evaluate the objectives of all tasks on the proposed x
+    new_y = problem(new_x)  # (q, NUM_TASKS)
+
+    # Update posterior
+    model = model.condition_on_observations(new_x, new_y)
+
+    # Evaluate the summarized objective (a scalar)
+    curr_objective = objective(new_y).max()
+    best_objective = objective(model.train_Y).max()
+    pbar.set_description(
+        f"[Best objective = {best_objective:.3f}, curr objective = {curr_objective:.3f}]"
+    )

examples/botorch/experiment.py → examples/singletask.py

@@ -1,4 +1,10 @@
 from __future__ import annotations
+
+import warnings
+
+warnings.filterwarnings("ignore")
+
+
 import numpy as np
 import torch
 from torch import nn
@@ -52,7 +58,7 @@ def get_net():
 
 
 def evaluate_model(model):
-    pred, _ = model.get_prediction(test_x, use_test_loader=True, joint=False)
+    pred, _ = model._get_prediction(test_x, use_test_loader=True, joint=False)
     return F.mse_loss(pred, test_y).squeeze().item()
 
 

laplace_bayesopt/botorch.py

@@ -208,7 +208,13 @@ def posterior(
         cov_y = torch.einsum("bqkbrl->bqkrl", cov_y)  # (B, Q, K, Q, K)
         cov_y = cov_y.reshape(B, Q * K, Q * K)
 
-        dist = gdists.MultivariateNormal(mean_y, covariance_matrix=cov_y)
+        if K > 1:
+            dist = gdists.MultitaskMultivariateNormal(
+                mean_y.reshape(B, Q, K), covariance_matrix=cov_y
+            )
+        else:
+            dist = gdists.MultivariateNormal(mean_y, covariance_matrix=cov_y)
+
         post_pred = GPyTorchPosterior(dist)
 
         if hasattr(self, "outcome_transform"):

pyproject.toml

@@ -1,23 +1,22 @@
 [project]
 name = "laplace-bayesopt"
-version = "0.1.4"
+version = "0.1.5"
 description = "Bayesian optimization interface for the laplace-torch library"
-authors = [
-    {name = "Agustinus Kristiadi", email = "[email protected]"},
-]
+authors = [{ name = "Agustinus Kristiadi", email = "[email protected]" }]
 dependencies = [
     "asdfghjkl",
     "backpack-for-pytorch",
     "botorch",
     "gpytorch",
+    "laplace-torch",
     "opt_einsum",
     "torch",
     "torchaudio",
     "torchvision",
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-license = {text = "MIT"}
+license = { text = "MIT" }
 classifiers = [
     "Development Status :: 3 - Alpha",
     "License :: OSI Approved :: MIT License",
@@ -30,12 +29,7 @@ Homepage = "https://github.com/wiseodd/laplace-bayesopt"
 "Bug Tracker" = "https://github.com/wiseodd/laplace-bayesopt/issues"
 
 [project.optional-dependencies]
-tests = [
-    "coveralls",
-    "pytest",
-    "pytest-cov",
-    "scipy",
-]
+tests = ["coveralls", "pytest", "pytest-cov", "scipy"]
 [build-system]
 requires = ["pdm-backend", "packaging"]
 build-backend = "pdm.backend"

tests/botorch/test_laplace_botorch.py

@@ -1,3 +1,4 @@
+from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
 import pytest
 import torch
 from torch import nn
@@ -47,6 +48,7 @@ def test_posterior_singletask(net_singletask, reg_data_singletask):
     B, Q, D, K = 15, 4, train_X.shape[-1], 1
     post = model.posterior(torch.randn(B, Q, D))
 
+    assert isinstance(post.mvn, MultivariateNormal)
     assert post.mean.shape == (B, Q * K, 1)  # Quirk of GPyTorch
     assert post.covariance_matrix.shape == (B, Q * K, Q * K)
 
@@ -60,7 +62,8 @@ def test_posterior_multitask(net_multitask, reg_data_multitask):
     B, Q, D, K = 15, 4, train_X.shape[-1], train_Y.shape[-1]
     post = model.posterior(torch.randn(B, Q, D))
 
-    assert post.mean.shape == (B, Q * K, 1)  # Quirk of GPyTorch
+    assert isinstance(post.mvn, MultitaskMultivariateNormal)
+    assert post.mean.shape == (B, Q, K)
     assert post.covariance_matrix.shape == (B, Q * K, Q * K)