Change the implementation of the Trust Region in the Vizier Acquisition

function to reduce rounding errors.

PiperOrigin-RevId: 796448366

Author: vizier-team

vizier/_src/algorithms/designers/gp/acquisitions.py

@@ -162,7 +162,7 @@ def _apply_trust_region(
   acquisition = jnp.where(
       ((distance <= region.trust_radius) | (region.trust_radius > 0.5)),
       acquisition,
-      -1e12 - distance,
+      -1e4 - distance,
   )
   aux = aux | {
       'mean': pred.mean(),
@@ -420,7 +420,7 @@ def default_ucb_pi(cls, data: types.ModelData) -> 'AcquisitionTrustRegion':
     return cls(
         UCB(1.8),
         PI(best_labels),
-        bad_acq_value=-1e12,
+        bad_acq_value=-1e4,
         labels=data.labels,
         threshold=0.3,
         apply_tr_after=0,
@@ -432,7 +432,7 @@ def default_ucb_lcb(cls, data: types.ModelData) -> 'AcquisitionTrustRegion':
         UCB(1.8),
         LCB(1.8),
         labels=data.labels,
-        bad_acq_value=-1e12,
+        bad_acq_value=-1e4,
         threshold=None,
         apply_tr_after=0,
     )
@@ -445,7 +445,7 @@ def default_ucb_lcb_wide(
         UCB(1.8),
         LCB(2.5),
         labels=data.labels,
-        bad_acq_value=-1e12,
+        bad_acq_value=-1e4,
         threshold=None,
         apply_tr_after=0,
     )
@@ -458,7 +458,7 @@ def default_ucb_lcb_delay_tr(
         UCB(1.8),
         LCB(1.8),
         labels=data.labels,
-        bad_acq_value=-1e12,
+        bad_acq_value=-1e4,
         threshold=None,
         apply_tr_after=5,
     )

vizier/_src/algorithms/designers/gp/acquisitions_test.py

@@ -238,7 +238,7 @@ def test_acq_pi_tr_good_point(self):
     acq = acquisitions.AcquisitionTrustRegion.default_ucb_pi(data=data)
     self.assertAlmostEqual(
         acq(tfd.Normal(jnp.float64(0.1), 1)),
-        -1.0e12,
+        -1e4,
     )
 
   def test_acq_pi_tr_bad_point(self):
@@ -254,7 +254,7 @@ def test_acq_lcb_tr_bad_point(self):
     acq = acquisitions.AcquisitionTrustRegion.default_ucb_lcb(data=data)
     self.assertAlmostEqual(
         acq(tfd.Normal(jnp.float64(0.1), 1)),
-        jnp.array([-1e12]),
+        jnp.array([-1e4]),
         delta=2.0,
     )
 

vizier/_src/algorithms/designers/gp_ucb_pe.py

@@ -226,13 +226,13 @@ def _apply_trust_region(
     Acquisition function values with trust region applied.
   """
   distance = tr.min_linf_distance(xs)
-  # Due to output normalization, acquisition values can't be as low as -1e12.
+  # Due to output normalization, acquisition values can't be as low as -1e4.
   # We use a bad value that decreases in the distance to trust region so that
   # acquisition optimizer can follow the gradient and escape untrustred regions.
   return jnp.where(
       (distance < tr.trust_radius) | (tr.trust_radius > 0.5),
       acq_values,
-      -1e12 - distance,
+      -1e4 - distance,
   )
 
 
@@ -252,13 +252,13 @@ def _apply_trust_region_to_set(
     [batch_size].
   """
   distance = tr.min_linf_distance(xs)  # [batch_size, index_point_set_size]
-  # Due to output normalization, acquisition values can't be as low as -1e12.
+  # Due to output normalization, acquisition values can't be as low as -1e4.
   # We penalize the acquisition values by an amount that decreases in the
   # total distances to the trust region so that acquisition optimizer can follow
   # the gradient and escape untrustred regions.
   return acq_values + jnp.sum(
       ((distance > tr.trust_radius) & (tr.trust_radius <= 0.5))
-      * (-1e12 - distance),
+      * (-1e4 - distance),
       axis=1,
   )