|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +BayesianOptimization class for performing safe multi-task Bayesian optimization. |
| 4 | +
|
| 5 | +Attributes: |
| 6 | + obj: The objective function to be optimized. |
| 7 | + tasks: List of tasks for multi-task optimization. |
| 8 | + bounds: Tensor specifying the bounds for the optimization variables. |
| 9 | + threshold: The safety threshold for the optimization. |
| 10 | + targets_mean_std: Tuple containing the mean and standard deviation of the targets. |
| 11 | + num_acq_samps: List specifying the number of acquisition samples for each task. |
| 12 | + boundary_T: Boundary threshold for the optimization. |
| 13 | + run: Counter for the number of optimization steps. |
| 14 | + best_y: List of best observed values. |
| 15 | + best_x: List of best observed inputs. |
| 16 | + dim: Dimensionality of the optimization problem. |
| 17 | + gp: Gaussian process model used for optimization. |
| 18 | +
|
| 19 | +Methods: |
| 20 | + __init__(self, obj, tasks, bounds, threshold, targets_mean_std, num_acq_samps=[1, 1], boundary_T=-15.0): |
| 21 | + Initializes the BayesianOptimization class with the given parameters. |
| 22 | +
|
| 23 | + step(self): |
| 24 | + Performs one step of the Bayesian optimization loop. |
| 25 | +
|
| 26 | + inequality_consts(self, input: Tensor): |
| 27 | + Computes the inequality constraints for the given input. |
| 28 | +
|
| 29 | + update_gp(self, gp, sqrtbeta): |
| 30 | + Updates the Gaussian process model and related attributes. |
| 31 | +
|
| 32 | + _line_search(self, initial_condition, maxiter=20, step_size=2.0): |
| 33 | + Performs a line search to find a feasible initial condition. |
| 34 | +
|
| 35 | + _get_max_observed(self): |
| 36 | + Returns the maximum observed values for each task. |
| 37 | +
|
| 38 | + _get_min_observed(self): |
| 39 | + Returns the minimum observed values for each task. |
| 40 | +
|
| 41 | + _get_best_input(self): |
| 42 | + Returns the best input values for each task. |
| 43 | +
|
| 44 | + _get_initial_cond(self): |
| 45 | + Returns the initial conditions for the optimization. |
| 46 | +
|
| 47 | + get_next_point(self, task, posterior_transform): |
| 48 | + Returns the next point to evaluate for the given task. |
| 49 | +""" |
| 50 | + |
| 51 | + |
| 52 | +import torch |
| 53 | +from torch import Tensor |
| 54 | +from botorch.optim.optimize import optimize_acqf |
| 55 | +from botorch.acquisition import qUpperConfidenceBound |
| 56 | +from utils.utils import concat_data, unstandardize, standardize |
| 57 | +from botorch.acquisition.objective import ScalarizedPosteriorTransform |
| 58 | + |
| 59 | +N_TOL = -1e-6 |
| 60 | + |
| 61 | + |
| 62 | +class BayesianOptimization: |
| 63 | + def __init__( |
| 64 | + self, |
| 65 | + obj, |
| 66 | + tasks, |
| 67 | + bounds, |
| 68 | + threshold, |
| 69 | + targets_mean_std, |
| 70 | + num_acq_samps: list = [1, 1], |
| 71 | + boundary_T=-15.0, |
| 72 | + ): |
| 73 | + self.obj = obj |
| 74 | + self.bounds = bounds |
| 75 | + self.threshold = threshold |
| 76 | + self.boundary_T = boundary_T |
| 77 | + self.mu, self.std = targets_mean_std |
| 78 | + self.num_acq_samps = num_acq_samps |
| 79 | + self.tasks = tasks |
| 80 | + if len(self.num_acq_samps) != len(self.tasks): |
| 81 | + raise ValueError("Number of tasks and number of samples must match") |
| 82 | + self.run = 0 |
| 83 | + self.best_y = [] |
| 84 | + self.best_x = [] |
| 85 | + self.dim = bounds.size(-1) |
| 86 | + self.gp = None |
| 87 | + |
| 88 | + def step(self): |
| 89 | + self.run += 1 |
| 90 | + print("Run : ", self.run) |
| 91 | + print(f"Best value found: {self.observed_max[0]: .3f}") |
| 92 | + print(f"Worst value: {self._get_min_observed()[0]}") |
| 93 | + W = torch.eye(len(self.tasks)) |
| 94 | + for i in self.tasks: |
| 95 | + posterior_transform = ScalarizedPosteriorTransform(W[:, i].squeeze()) |
| 96 | + new_point = self.get_next_point(i, posterior_transform) |
| 97 | + if i == 0: |
| 98 | + print(f"New Point: {new_point}") |
| 99 | + new_point_task0 = new_point |
| 100 | + if i != 0: |
| 101 | + new_point = torch.vstack((new_point, new_point_task0)) |
| 102 | + new_result = self.obj.f(new_point, i) |
| 103 | + self.train_inputs, self.train_tasks, self.unstd_train_targets = concat_data( |
| 104 | + (new_point, i * torch.ones(new_point.shape[0], 1), new_result), |
| 105 | + (self.train_inputs, self.train_tasks, self.unstd_train_targets), |
| 106 | + ) |
| 107 | + threshold = unstandardize(self.threshold, self.mu, self.std) |
| 108 | + self.train_targets, self.mu, self.std = standardize( |
| 109 | + self.unstd_train_targets, train_task=self.train_tasks |
| 110 | + ) |
| 111 | + self.threshold, _, _ = standardize(threshold, self.mu, self.std) |
| 112 | + self.observed_max = self._get_max_observed() |
| 113 | + self.best_y.append(self.observed_max[0]) |
| 114 | + self.best_x.append(self._get_best_input()[0]) |
| 115 | + return self.train_inputs, self.train_tasks, self.train_targets |
| 116 | + |
| 117 | + def inequality_consts(self, input: Tensor): |
| 118 | + self.gp.eval() |
| 119 | + inputx = input.view(int(input.numel() / self.dim), self.dim) |
| 120 | + output = self.gp(torch.hstack((inputx, torch.zeros(inputx.size(0), 1)))) |
| 121 | + val = ( |
| 122 | + output.mean |
| 123 | + - output.covariance_matrix.diag().sqrt() * self.sqrtbeta |
| 124 | + - self.threshold |
| 125 | + ) |
| 126 | + return val.view(inputx.shape[0], 1) |
| 127 | + |
| 128 | + def update_gp(self, gp, sqrtbeta): |
| 129 | + with torch.no_grad(): |
| 130 | + self.train_inputs = gp.train_inputs[0][..., :-1] |
| 131 | + self.train_tasks = gp.train_inputs[0][..., -1:].to(dtype=torch.int32) |
| 132 | + self.train_targets = gp.train_targets.unsqueeze(-1) |
| 133 | + self.unstd_train_targets = unstandardize( |
| 134 | + self.train_targets, self.mu, self.std, self.train_tasks |
| 135 | + ) |
| 136 | + self.sqrtbeta = sqrtbeta.detach() |
| 137 | + if self.gp is None: |
| 138 | + self.observed_max = self._get_max_observed() |
| 139 | + self.best_y.append(self.observed_max[0]) |
| 140 | + self.best_x.append(self._get_best_input()[0]) |
| 141 | + self.gp = gp |
| 142 | + pass |
| 143 | + |
| 144 | + def _line_search(self, initial_condition, maxiter=20, step_size=2.0): |
| 145 | + k = 1000 |
| 146 | + direction = torch.randn(initial_condition.size()) |
| 147 | + direction /= ( |
| 148 | + torch.linalg.norm(direction, dim=-1, ord=2) |
| 149 | + .unsqueeze(-1) |
| 150 | + .repeat(1, 1, self.dim) |
| 151 | + ) |
| 152 | + steps = torch.linspace(0, step_size, k).view(1, k, 1) - step_size / 2 |
| 153 | + line_search = initial_condition + steps * direction |
| 154 | + inds = ( |
| 155 | + (self.inequality_consts(line_search) >= 0).view( |
| 156 | + initial_condition.size(0), -1 |
| 157 | + ) |
| 158 | + & torch.all(line_search <= self.bounds[1, :].view(1, 1, self.dim), dim=-1) |
| 159 | + & torch.all(line_search >= self.bounds[0, :].view(1, 1, self.dim), dim=-1) |
| 160 | + ) |
| 161 | + for id in range(inds.size(0)): |
| 162 | + possible_steps = steps[:, inds[id, :].squeeze(), :].squeeze() |
| 163 | + if possible_steps.numel() <= 1: |
| 164 | + return initial_condition |
| 165 | + max_step_ind = possible_steps.abs().argmax() |
| 166 | + initial_condition[id] = ( |
| 167 | + initial_condition[id] + possible_steps[max_step_ind] * direction[id] |
| 168 | + ) |
| 169 | + return initial_condition |
| 170 | + |
| 171 | + def _get_max_observed(self): |
| 172 | + return [ |
| 173 | + torch.max(self.unstd_train_targets[self.train_tasks == i]).item() |
| 174 | + for i in self.tasks |
| 175 | + ] |
| 176 | + |
| 177 | + def _get_min_observed(self): |
| 178 | + return [ |
| 179 | + torch.min(self.unstd_train_targets[self.train_tasks == i]).item() |
| 180 | + for i in self.tasks |
| 181 | + ] |
| 182 | + |
| 183 | + def _get_best_input(self): |
| 184 | + return [ |
| 185 | + self.train_inputs[self.train_tasks.squeeze() == i, ...][ |
| 186 | + torch.argmax(self.train_targets[self.train_tasks == i]) |
| 187 | + ] |
| 188 | + for i in self.tasks |
| 189 | + ] |
| 190 | + |
| 191 | + def _get_initial_cond(self): |
| 192 | + _, ind = self.train_targets.sort(dim=0, descending=True) |
| 193 | + sorted_train_inp = self.train_inputs[ind.squeeze(), ...] |
| 194 | + eqfull = self.inequality_consts(sorted_train_inp).squeeze() |
| 195 | + pot_cond = sorted_train_inp.view(self.train_inputs.size())[eqfull >= 0, ...][ |
| 196 | + :5, ... |
| 197 | + ] |
| 198 | + return pot_cond.view(pot_cond.size(0), 1, self.dim) |
| 199 | + |
| 200 | + def get_next_point(self, task, posterior_transform): |
| 201 | + if task == 0: |
| 202 | + init_cond = self._get_initial_cond() |
| 203 | + if init_cond.numel() == 0: |
| 204 | + print( |
| 205 | + "No feasible initial condition found. Randomly sampling a new one." |
| 206 | + ) |
| 207 | + x_new = self.train_inputs[ |
| 208 | + self.train_targets[self.train_tasks == 0].argmax(), : |
| 209 | + ].view(1, self.dim) |
| 210 | + offset = torch.randn(1, self.dim) * 0.005 |
| 211 | + ind = (x_new + offset <= self.bounds[1, :].view(1, self.dim)) & ( |
| 212 | + x_new + offset >= self.bounds[0, :].view(1, self.dim) |
| 213 | + ) |
| 214 | + x_new[ind] = x_new[ind] + offset[ind] |
| 215 | + x_new[~ind] = x_new[~ind] - offset[~ind] |
| 216 | + return x_new |
| 217 | + else: |
| 218 | + init_cond = self._line_search(init_cond) |
| 219 | + acq = qUpperConfidenceBound( |
| 220 | + self.gp, |
| 221 | + self.sqrtbeta, |
| 222 | + posterior_transform=posterior_transform, |
| 223 | + ) |
| 224 | + # if different acquisitions should be used |
| 225 | + else: |
| 226 | + acq = qUpperConfidenceBound( |
| 227 | + self.gp, |
| 228 | + beta=self.sqrtbeta, |
| 229 | + posterior_transform=posterior_transform, |
| 230 | + ) |
| 231 | + candidate, tt = optimize_acqf( |
| 232 | + acq_function=acq, |
| 233 | + bounds=( |
| 234 | + self.bounds |
| 235 | + if task == 0 |
| 236 | + else self.bounds |
| 237 | + + torch.tensor( |
| 238 | + [[self.obj.max_disturbance], [-self.obj.max_disturbance]] # max_disturbance is zero for LbSync (only shifts) |
| 239 | + ) |
| 240 | + ), |
| 241 | + q=self.num_acq_samps[task], |
| 242 | + num_restarts=init_cond.size(0) if task == 0 else 1, |
| 243 | + raw_samples=512 if task != 0 else None, |
| 244 | + nonlinear_inequality_constraints=( |
| 245 | + [self.inequality_consts] if task == 0 else None |
| 246 | + ), |
| 247 | + batch_initial_conditions=init_cond if task == 0 else None, |
| 248 | + options={"maxiter": 200}, |
| 249 | + ) |
| 250 | + return candidate |
0 commit comments