<reponame>facebookresearch/robust_mobo<filename>robust_mobo/ch_var_ucb.py
#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import math
from typing import Union

import torch
from botorch.acquisition.multi_objective.analytic import (
    MultiObjectiveAnalyticAcquisitionFunction,
)
from botorch.acquisition.risk_measures import VaR
from botorch.models import SingleTaskGP, ModelListGP
from botorch.utils import t_batch_mode_transform
from botorch.utils.multi_objective import get_chebyshev_scalarization
from torch import Tensor


class ChVUCB(MultiObjectiveAnalyticAcquisitionFunction):
    r"""The upper confidence bound acquisition function for
    VaR of Chebyshev scalarizations.
    """

    def __init__(
        self,
        model: Union[SingleTaskGP, ModelListGP],
        weights: Tensor,
        var: VaR,
        t: int,
    ) -> None:
        r"""VaR UCB acquisition function supporting `m >= 2` objectives
        via Chebyshev scalarizations.

        Args:
            model: A SingleTaskGP model, modeling `m` outcomes independently.
                `model` must have an `InputPerturbation`.
            weights: An `m`-dim tensor denoting the weights to use for
                Chebyshev scalarization.
            var: A VaR module used to calculate the VaR of UCB of scalarizations.
                Must be compatible with `model`'s input perturbation.
            t: An integer denoting the iteration count.
        """
        super().__init__(model=model)
        if isinstance(model, ModelListGP):
            train_targets = torch.stack(model.train_targets, dim=-1)
        else:
            train_targets = model.train_targets.t()
        self.m = self.model.num_outputs
        self.betas = self.get_betas(torch.as_tensor(t).to(device=train_targets.device))
        self.var = var
        self.scalarization = get_chebyshev_scalarization(
            weights=weights, Y=train_targets, alpha=0
        )
        self.weights_sign = torch.sign(weights)

    def get_betas(self, t: Tensor) -> Tensor:
        r"""Get the betas corresponding to each of `m` outcomes.

        NOTE: Following Nguyen et al., 2021, we use beta = 2 log(t^2 \pi^2 / 0.6),
        which has nothing to do with the theoretical values.

        Args:
            t: A scalar tensor denoting the iteration number.

        Returns:
            An `m`-dim tensor of betas.
        """
        return (2 * torch.log(t.pow(2) * math.pi ** 2 / 0.6)).sqrt().expand(self.m)

    @t_batch_mode_transform(expected_q=1)
    def forward(self, X: Tensor) -> Tensor:
        r"""Calculate VaR(u(s[f(X + \xi), w])).

        Args:
            X: A `batch_shape x 1 x d`-dim tensor of candidate solutions.

        Returns:
            A `batch_shape`-dim tensor of UCB values.
        """
        posterior = self.model.posterior(X)
        mean = posterior.mean
        sigma = posterior.variance.sqrt()
        # If weights are negative, we want the LCB.
        individual_ucb = mean + self.weights_sign * self.betas * sigma
        scalarized_ucb = self.scalarization(individual_ucb, X)
        return self.var(scalarized_ucb.unsqueeze(-1)).squeeze(-1)

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import math
from typing import Union

import torch
from botorch.acquisition.multi_objective.analytic import (
    MultiObjectiveAnalyticAcquisitionFunction,
)
from botorch.acquisition.risk_measures import VaR
from botorch.models import SingleTaskGP, ModelListGP
from botorch.utils import t_batch_mode_transform
from botorch.utils.multi_objective import get_chebyshev_scalarization
from torch import Tensor


class ChVUCB(MultiObjectiveAnalyticAcquisitionFunction):
    r"""The upper confidence bound acquisition function for
    VaR of Chebyshev scalarizations.
    """

    def __init__(
        self,
        model: Union[SingleTaskGP, ModelListGP],
        weights: Tensor,
        var: VaR,
        t: int,
    ) -> None:
        r"""VaR UCB acquisition function supporting `m >= 2` objectives
        via Chebyshev scalarizations.

        Args:
            model: A SingleTaskGP model, modeling `m` outcomes independently.
                `model` must have an `InputPerturbation`.
            weights: An `m`-dim tensor denoting the weights to use for
                Chebyshev scalarization.
            var: A VaR module used to calculate the VaR of UCB of scalarizations.
                Must be compatible with `model`'s input perturbation.
            t: An integer denoting the iteration count.
        """
        super().__init__(model=model)
        if isinstance(model, ModelListGP):
            train_targets = torch.stack(model.train_targets, dim=-1)
        else:
            train_targets = model.train_targets.t()
        self.m = self.model.num_outputs
        self.betas = self.get_betas(torch.as_tensor(t).to(device=train_targets.device))
        self.var = var
        self.scalarization = get_chebyshev_scalarization(
            weights=weights, Y=train_targets, alpha=0
        )
        self.weights_sign = torch.sign(weights)

    def get_betas(self, t: Tensor) -> Tensor:
        r"""Get the betas corresponding to each of `m` outcomes.

        NOTE: Following Nguyen et al., 2021, we use beta = 2 log(t^2 \pi^2 / 0.6),
        which has nothing to do with the theoretical values.

        Args:
            t: A scalar tensor denoting the iteration number.

        Returns:
            An `m`-dim tensor of betas.
        """
        return (2 * torch.log(t.pow(2) * math.pi ** 2 / 0.6)).sqrt().expand(self.m)

    @t_batch_mode_transform(expected_q=1)
    def forward(self, X: Tensor) -> Tensor:
        r"""Calculate VaR(u(s[f(X + \xi), w])).

        Args:
            X: A `batch_shape x 1 x d`-dim tensor of candidate solutions.

        Returns:
            A `batch_shape`-dim tensor of UCB values.
        """
        posterior = self.model.posterior(X)
        mean = posterior.mean
        sigma = posterior.variance.sqrt()
        # If weights are negative, we want the LCB.
        individual_ucb = mean + self.weights_sign * self.betas * sigma
        scalarized_ucb = self.scalarization(individual_ucb, X)
        return self.var(scalarized_ucb.unsqueeze(-1)).squeeze(-1)
#!/usr/bin/env python3 # Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. The upper confidence bound acquisition function for     VaR of Chebyshev scalarizations. VaR UCB acquisition function supporting `m >= 2` objectives         via Chebyshev scalarizations.          Args:             model: A SingleTaskGP model, modeling `m` outcomes independently.                 `model` must have an `InputPerturbation`.             weights: An `m`-dim tensor denoting the weights to use for                 Chebyshev scalarization.             var: A VaR module used to calculate the VaR of UCB of scalarizations.                 Must be compatible with `model`'s input perturbation.             t: An integer denoting the iteration count. Get the betas corresponding to each of `m` outcomes.          NOTE: Following Nguyen et al., 2021, we use beta = 2 log(t^2 \pi^2 / 0.6),         which has nothing to do with the theoretical values.          Args:             t: A scalar tensor denoting the iteration number.          Returns:             An `m`-dim tensor of betas. Calculate VaR(u(s[f(X + \xi), w])).          Args:             X: A `batch_shape x 1 x d`-dim tensor of candidate solutions.          Returns:             A `batch_shape`-dim tensor of UCB values. # If weights are negative, we want the LCB.