Fix bug introduced by bug fix in sklearn's GP.

PiperOrigin-RevId: 495767145

balloon_learning_environment/env/balloon/acs.py

@@ -21,14 +21,14 @@
 from scipy import interpolate
 
 
-_PRESSURE_RATIO_TO_POWER: interpolate.interpolate.interp1d = (
+_PRESSURE_RATIO_TO_POWER: interpolate.interp1d = (
     interpolate.interp1d(
         np.array([1.0, 1.05, 1.2, 1.25, 1.35]),  # pressure_ratio
         np.array([100.0, 100.0, 300.0, 400.0, 400.0]),  # power
         fill_value='extrapolate'))
 
 
-_PRESSURE_RATIO_POWER_TO_EFFICIENCY: interpolate.interpolate.interp2d = (
+_PRESSURE_RATIO_POWER_TO_EFFICIENCY: interpolate.interp2d = (
     interpolate.interp2d(
         np.linspace(1.05, 1.35, 13),  # pressure_ratio
         np.linspace(100.0, 400.0, 4),  # power

balloon_learning_environment/env/features.py

@@ -138,7 +138,8 @@ def observe(self, observation: simulator_data.SimulatorObservation) -> None:
   def get_features(self) -> np.ndarray:
     """Gets the current feature vector given all observations."""
 
-  @abc.abstractproperty
+  @property
+  @abc.abstractmethod
   def observation_space(self) -> gym.Space:
     """Gets the observation space specification for the feature vector."""
 

balloon_learning_environment/env/wind_gp.py

@@ -19,9 +19,6 @@
 This lets us query any point (x, y, p, t) in the wind field for its value,
 as well as the model confidence's in this value.
 
----- Open issues
-* The forecast is not used
-
 """
 
 import datetime as dt
@@ -190,8 +187,6 @@ def query_batch(self, locations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
 
       # Output should be a N x 2 set of predictions about local measurements,
       # and a N-sized vector of standard deviations.
-      # TODO(bellemare): Determine why deviations is a single number per sample,
-      # instead of two (since we have two values being predicted).
       means, deviations = self.model.predict(locations, return_std=True)
 
       # Deviations are std.dev., convert to variance and normalize.
@@ -200,6 +195,17 @@ def query_batch(self, locations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
       # like 0.07 from the GP, but that doesn't seem to match the Loon code.
       deviations = deviations**2 / _SIGMA_EXP_SQUARED
 
+      # Previously, there was a bug in sklearn which meant that the GP
+      # only returned 1 deviation per location when it should have
+      # returned 2. This has since been fixed, but people running with
+      # older code won't need to apply the next block because the shapes
+      # will work as expected.
+      assert deviations.ndim in (1, 2)
+      if deviations.ndim == 2:
+        # We use the standard deviation as a measure of uncertainty,
+        # and take the average of the two components as a proxy for this.
+        deviations = np.mean(deviations, axis=1)
+
       # TODO(bellemare): Sal says this needs normalizing so that the lower bound
       # is really zero.