Skip to content

Noisy constant value handling improvement for RobustGPSampler #330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+95 −11
Open

Noisy constant value handling improvement for RobustGPSampler #330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions package/samplers/value_at_risk/_gp/acqf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,8 @@ def __init__(
threshold_list: list[float],
n_input_noise_samples: int,
qmc_seed: int | None,
fixed_indices: torch.Tensor,
fixed_values: torch.Tensor,
uniform_input_noise_rads: torch.Tensor | None = None,
normal_input_noise_stdevs: torch.Tensor | None = None,
stabilizing_noise: float = 1e-12,
Expand All @@ -133,12 +135,22 @@ def __init__(
)
self._stabilizing_noise = stabilizing_noise
self._confidence_level = confidence_level
self._fixed_indices = fixed_indices
self._fixed_values = fixed_values
super().__init__(
length_scales=np.mean([gpr.length_scales for gpr in gpr_list], axis=0),
search_space=search_space,
)

def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
# The search space of constant noisy parameters is internally replaced with IntDistribution(0, 0),
# so that their normalized values passed as x is always 0.5. However, values passed to
# const_noisy_param_values argument of RobustGPSampler._get_value_at_risk may be normalized to
# different values under the original search space used by GPRegressor. So we carry around the
# normalized version of const_noisy_param_values explicity and use them instead.
x = x.clone()
x[:, self._fixed_indices] = self._fixed_values

x_noisy = x.unsqueeze(-2) + self._input_noise
log_feas_probs = torch.zeros(x_noisy.shape[:-1], dtype=torch.float64)
for gpr, threshold in zip(self._gpr_list, self._threshold_list):
Expand Down Expand Up @@ -168,6 +180,8 @@ def __init__(
n_qmc_samples: int,
qmc_seed: int | None,
acqf_type: str,
fixed_indices: torch.Tensor,
fixed_values: torch.Tensor,
uniform_input_noise_rads: torch.Tensor | None = None,
normal_input_noise_stdevs: torch.Tensor | None = None,
) -> None:
Expand All @@ -187,6 +201,8 @@ def __init__(
seed=rng.random_integers(0, 2**31 - 1, size=1).item(),
)
self._acqf_type = acqf_type
self._fixed_indices = fixed_indices
self._fixed_values = fixed_values
super().__init__(length_scales=gpr.length_scales, search_space=search_space)

def _value_at_risk(self, x: torch.Tensor) -> torch.Tensor:
Expand All @@ -206,6 +222,15 @@ def eval_acqf(self, x: torch.Tensor) -> torch.Tensor:
4. Then compute (mc_value_at_risk - f0).clamp_min(0).mean()
Appendix B.2 of https://www.robots.ox.ac.uk/~mosb/public/pdf/136/full_thesis.pdf
"""

# The search space of constant noisy parameters is internally replaced with IntDistribution(0, 0),
# so that their normalized values passed as x is always 0.5. However, values passed to
# const_noisy_param_values argument of RobustGPSampler._get_value_at_risk may be normalized to
# different values under the original search space used by GPRegressor. So we carry around the
# normalized version of const_noisy_param_values explicity and use them instead.
x = x.clone()
x[:, self._fixed_indices] = self._fixed_values

if self._acqf_type == "mean":
return self._value_at_risk(x).mean(dim=-1)
elif self._acqf_type == "nei":
Expand All @@ -227,6 +252,8 @@ def __init__(
n_qmc_samples: int,
qmc_seed: int | None,
acqf_type: str,
fixed_indices: torch.Tensor,
fixed_values: torch.Tensor,
uniform_input_noise_rads: torch.Tensor | None = None,
normal_input_noise_stdevs: torch.Tensor | None = None,
stabilizing_noise: float = 1e-12,
Expand All @@ -241,6 +268,8 @@ def __init__(
acqf_type=acqf_type,
uniform_input_noise_rads=uniform_input_noise_rads,
normal_input_noise_stdevs=normal_input_noise_stdevs,
fixed_indices=fixed_indices,
fixed_values=fixed_values,
)
self._log_prob_at_risk = LogCumulativeProbabilityAtRisk(
gpr_list=constraints_gpr_list,
Expand All @@ -252,6 +281,8 @@ def __init__(
uniform_input_noise_rads=uniform_input_noise_rads,
normal_input_noise_stdevs=normal_input_noise_stdevs,
stabilizing_noise=stabilizing_noise,
fixed_indices=fixed_indices,
fixed_values=fixed_values,
)
assert torch.allclose(
self._log_prob_at_risk._input_noise, self._value_at_risk._input_noise
Expand Down
75 changes: 64 additions & 11 deletions package/samplers/value_at_risk/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,8 @@ def _get_constraints_acqf_args(
return constraints_gprs, constraints_threshold_list

def _get_internal_search_space_with_fixed_params(
self, search_space: dict[str, BaseDistribution]
self,
search_space: dict[str, BaseDistribution],
) -> gp_search_space.SearchSpace:
search_space_with_fixed_params = search_space.copy()
for param_name in self._const_noisy_param_names:
Expand All @@ -275,6 +276,7 @@ def _get_value_at_risk(
internal_search_space: gp_search_space.SearchSpace,
search_space: dict[str, BaseDistribution],
acqf_type: str,
const_noisy_param_values: dict[str, float],
constraints_gpr_list: list[gp.GPRegressor] | None = None,
constraints_threshold_list: list[float] | None = None,
) -> acqf_module.ValueAtRisk | acqf_module.ConstrainedLogValueAtRisk:
Expand Down Expand Up @@ -313,17 +315,37 @@ def _get_scaled_input_noise_params(
for i, param_name in enumerate(search_space)
if param_name in self._const_noisy_param_names
]

def normalize(dist: BaseDistribution, x: float) -> float:
assert isinstance(
dist,
(optuna.distributions.IntDistribution, optuna.distributions.FloatDistribution),
)
return (x - dist.low) / (dist.high - dist.low)

const_noisy_param_normalized_values = [
normalize(dist, const_noisy_param_values[param_name])
if param_name in const_noisy_param_values
else 0.5
for i, (param_name, dist) in enumerate(search_space.items())
if param_name in self._const_noisy_param_names
]

if self._uniform_input_noise_rads is not None:
scaled_input_noise_params = _get_scaled_input_noise_params(
self._uniform_input_noise_rads, "uniform_input_noise_rads"
)
# FIXME(sakai): If the fixed value is not at the center of the range,
# \pm 0.5 may not cover the domain.
scaled_input_noise_params[const_noise_param_inds] = 0.5
noise_kwargs["uniform_input_noise_rads"] = scaled_input_noise_params
elif self._normal_input_noise_stdevs is not None:
scaled_input_noise_params = _get_scaled_input_noise_params(
self._normal_input_noise_stdevs, "normal_input_noise_stdevs"
)
# NOTE(nabenabe): \pm 2 sigma will cover the domain.
# FIXME(sakai): If the fixed value is not at the center of the range,
# \pm 2 sigma may not cover the domain.
scaled_input_noise_params[const_noise_param_inds] = 0.25
noise_kwargs["normal_input_noise_stdevs"] = scaled_input_noise_params
else:
Expand All @@ -341,6 +363,10 @@ def _get_scaled_input_noise_params(
n_qmc_samples=self._n_qmc_samples,
qmc_seed=self._rng.rng.randint(1 << 30),
acqf_type=acqf_type,
fixed_indices=torch.tensor(const_noise_param_inds, dtype=torch.int64),
fixed_values=torch.tensor(
const_noisy_param_normalized_values, dtype=torch.float64
),
**noise_kwargs,
)
else:
Expand All @@ -355,6 +381,10 @@ def _get_scaled_input_noise_params(
n_qmc_samples=self._n_qmc_samples,
qmc_seed=self._rng.rng.randint(1 << 30),
acqf_type=acqf_type,
fixed_indices=torch.tensor(const_noise_param_inds, dtype=torch.int64),
fixed_values=torch.tensor(
const_noisy_param_normalized_values, dtype=torch.float64
),
**noise_kwargs,
)

Expand Down Expand Up @@ -419,7 +449,11 @@ def _get_gpr_list(
return gprs_list

def _optimize_params(
self, study: Study, trials: list[FrozenTrial], search_space: dict[str, BaseDistribution]
self,
study: Study,
trials: list[FrozenTrial],
search_space: dict[str, BaseDistribution],
const_noisy_param_values: dict[str, float],
) -> dict[str, Any]:
if search_space == {}:
return {}
Expand All @@ -437,6 +471,7 @@ def _optimize_params(
internal_search_space,
search_space,
acqf_type="mean",
const_noisy_param_values=const_noisy_param_values,
)
else:
constraint_vals, _ = _get_constraint_vals_and_feasibility(study, trials)
Expand All @@ -453,10 +488,18 @@ def _optimize_params(
acqf_type="mean",
constraints_gpr_list=constr_gpr_list,
constraints_threshold_list=constr_threshold_list,
const_noisy_param_values=const_noisy_param_values,
)

normalized_param = self._optimize_acqf(acqf)
return internal_search_space.get_unnormalized_param(normalized_param)
# The normalized values of constant noise parameters are fixed at 0.5 during search
# regardless of their original values given as const_noisy_param_values, so
# `internal_search_space.get_unnormalized_param` cannot decode them correctly.
# Therefore, we overwrite those values with their original values.
return (
internal_search_space.get_unnormalized_param(normalized_param)
| const_noisy_param_values
)

def sample_relative(
self, study: Study, trial: FrozenTrial, search_space: dict[str, BaseDistribution]
Expand All @@ -465,17 +508,18 @@ def sample_relative(
if len(trials) < self._n_startup_trials:
return {}

params = self._optimize_params(study, trials, search_space)

# Perturb constant noisy parameter uniformly
const_noisy_param_values = {}
for name in self._const_noisy_param_names:
dist = search_space[name]
assert isinstance(dist, optuna.distributions.FloatDistribution)
params[name] = self._rng.rng.uniform(dist.low, dist.high)
const_noisy_param_values[name] = self._rng.rng.uniform(dist.low, dist.high)

return params
return self._optimize_params(study, trials, search_space, const_noisy_param_values)

def get_robust_trial(self, study: Study) -> FrozenTrial:
def get_robust_trial(
self, study: Study, const_noisy_param_nominal_values: dict[str, float] | None = None
) -> FrozenTrial:
states = (TrialState.COMPLETE,)
trials = study._get_trials(deepcopy=False, states=states, use_cache=True)
search_space = self.infer_relative_search_space(study, trials[0])
Expand All @@ -485,7 +529,11 @@ def get_robust_trial(self, study: Study) -> FrozenTrial:
acqf: acqf_module.BaseAcquisitionFunc
if self._constraints_func is None:
acqf = self._get_value_at_risk(
gpr, internal_search_space, search_space, acqf_type="mean"
gpr,
internal_search_space,
search_space,
acqf_type="mean",
const_noisy_param_values=const_noisy_param_nominal_values or {},
)
else:
constraint_vals, _ = _get_constraint_vals_and_feasibility(study, trials)
Expand All @@ -502,16 +550,21 @@ def get_robust_trial(self, study: Study) -> FrozenTrial:
acqf_type="mean",
constraints_gpr_list=constr_gpr_list,
constraints_threshold_list=constr_threshold_list,
const_noisy_param_values=const_noisy_param_nominal_values or {},
)

best_idx = np.argmax(acqf.eval_acqf_no_grad(X_train)).item()
return trials[best_idx]

def get_robust_params(self, study: Study) -> dict[str, Any]:
def get_robust_params(
self, study: Study, const_noisy_param_nominal_values: dict[str, float] | None = None
) -> dict[str, Any]:
states = (TrialState.COMPLETE,)
trials = study._get_trials(deepcopy=False, states=states, use_cache=True)
search_space = self.infer_relative_search_space(study, trials[0])
return self._optimize_params(study, trials, search_space)
return self._optimize_params(
study, trials, search_space, const_noisy_param_nominal_values or {}
)

def sample_independent(
self,
Expand Down