"""
Hartmann objective function, relying on the botorch implementation.
See:
https://botorch.org/api/test_functions.html#botorch.test_functions.synthetic.Hartmann
This code is based on the implementation by Nikita Saxena (nikita-0209) in
https://github.com/alexhernandezgarcia/activelearning
The implementation assumes that the inputs will be on [-1, 1]^6 as is typical in the
uses of the Hartmann function. The original range is negative and is a minimisation
problem. By default, the proxy values remain in this range and the absolute value of
the proxy values is used as the reward function.
"""
from typing import Callable, Optional, Union
import numpy as np
import torch
from botorch.test_functions.multi_fidelity import AugmentedHartmann
from torchtyping import TensorType
from gflownet.proxy.base import Proxy
# Global optimum, according to BoTorch
# A rough estimate of modes
[docs]
MODES = [
[0.2, 0.2, 0.5, 0.3, 0.3, 0.7],
[0.4, 0.9, 0.9, 0.6, 0.1, 0.0],
[0.3, 0.1, 0.4, 0.3, 0.3, 0.7],
[0.4, 0.9, 0.4, 0.6, 0.0, 0.0],
[0.4, 0.9, 0.6, 0.6, 0.3, 0.0],
]
[docs]
class Hartmann(Proxy):
def __init__(
self,
fidelity=1.0,
do_domain_map: bool = True,
negate: bool = False,
reward_function: Optional[Union[Callable, str]] = "absolute",
**kwargs,
):
"""
Parameters
----------
fidelity : float
Fidelity of the Hartmann oracle. 1.0 corresponds to the original Hartmann.
Smaller values (up to 0.0) reduce the fidelity of the oracle.
do_domain_map : bool
If True, the states are assumed to be in [-1, 1]^6 and are re-mapped to the
standard domain in [0, 1]^6 before calling the botorch method. If False,
the botorch method is called directly on the states values.
negate : bool
If True, proxy values are multiplied by -1.
reward_function : str or Callable
The transformation applied to the proxy outputs to obtain a GFlowNet
reward. By default, the reward function is the absolute value of proxy
outputs.
See: https://botorch.org/api/test_functions.html
"""
# Replace the value of reward_function in kwargs by the one passed explicitly
# as a parameter
kwargs["reward_function"] = reward_function
super().__init__(**kwargs)
[docs]
self.fidelity = fidelity
[docs]
self.do_domain_map = do_domain_map
[docs]
self.function_mf_botorch = AugmentedHartmann(negate=negate)
# Optimum
self._optimum = torch.tensor(OPTIMUM, device=self.device, dtype=self.float)
if negate:
self._optimum *= -1.0
[docs]
def __call__(self, states: TensorType["batch", "state_dim"]) -> TensorType["batch"]:
if states.shape[1] != 6:
raise ValueError(
"Inputs to the Hartmann function must be 6-dimensional, "
f"but inputs with {states.shape[1]} dimensions were passed."
)
# Map states to Hartmann domain
if self.do_domain_map:
states = self.map_to_standard_domain(states)
# Append fidelity as a new dimension of states
states = torch.cat(
[
states,
self.fidelity
* torch.ones(
states.shape[0], device=self.device, dtype=self.float
).unsqueeze(-1),
],
dim=1,
)
return self.function_mf_botorch(states)
[docs]
def map_to_standard_domain(
self,
states: TensorType["batch", "6"],
) -> TensorType["batch", "6"]:
"""
Maps a batch of input states onto the domain typically used to evaluate the
Hartmann function, that is [0, 1]^6. See DOMAIN and LENGTH. It assumes that the
inputs are on [-1, 1]^6
"""
return (states + 1.0) / 2.0