Automated Code Change

iris/policies/gym_space_utils.py

@@ -38,8 +38,8 @@ def filter_space(space: gym.Space,
 def extend_space(space: gym.Space, key: str, value: gym.Space):
   """Adds new keys or dimensions to the space."""
   if isinstance(space, gym.spaces.Box):
-    low = np.concatenate((space.low, value.low))
-    high = np.concatenate((space.high, value.high))
+    low = np.concatenate((space.low, value.low))  # pyrefly: ignore[missing-attribute]
+    high = np.concatenate((space.high, value.high))  # pyrefly: ignore[missing-attribute]
     return gym.spaces.Box(low=low, high=high)
   elif isinstance(space, gym.spaces.Dict):
     extended_space = dict(space.spaces)
@@ -62,9 +62,9 @@ def filter_sample(
   if isinstance(x, dict):
     filtered_x = {}
     for sensor in selected:
-      filtered_x[sensor] = x[sensor]
+      filtered_x[sensor] = x[sensor]  # pyrefly: ignore[bad-index]
   else:
-    filtered_x = np.array(x).take(selected)
+    filtered_x = np.array(x).take(selected)  # pyrefly: ignore[no-matching-overload]
   return filtered_x
 
 

iris/policies/hierarchical_policy.py

@@ -78,7 +78,7 @@ def __init__(
     self._ac_space = gym.spaces.Box(-1, 1, (self._out_command_dim,))
     self._timescale = fixed_timescale
     if self._timescale is None:
-      self._timescale_low, self._timescale_high = timescale_range
+      self._timescale_low, self._timescale_high = timescale_range  # pyrefly: ignore[not-iterable]
     self._act_after_steps = 0
     self._output = np.zeros(self._out_command_dim)
     self.policy = policy(ob_space=self._ob_space, ac_space=self._ac_space)
@@ -94,7 +94,7 @@ def __call__(
     if not self._act_after_steps:
       ob = gym_space_utils.filter_sample(ob, self._selected_observations)
       ob = gym_space_utils.extend_sample(ob, "in_command", in_command)
-      self._output = self.policy.act(ob)
+      self._output = self.policy.act(ob)  # pyrefly: ignore[bad-assignment]
       if self._timescale is not None:
         self._act_after_steps = self._timescale
       else:

iris/policies/hierarchical_policy_test.py

@@ -142,7 +142,7 @@ def test_hierarchical_policy_act(self, policy_params, is_ob_dict):
     act = policy.act(ob)
     # Latent command is [1, 1, 1, 1] and low level output is
     # 2 * sum([1, 1, 1, 1, -4]) = 0
-    self.assertAlmostEqual(act, 0, places=2)
+    self.assertAlmostEqual(act, 0, places=2)  # pyrefly: ignore[no-matching-overload]
     # Check that latent command remains constant until high level activates
     interval = policy.levels[0]._act_after_steps
     ob = np.array([5, -4.5])
@@ -152,14 +152,14 @@ def test_hierarchical_policy_act(self, policy_params, is_ob_dict):
       act = policy.act(ob)
       # Latent command is still [1, 1, 1, 1] and low level output is
       # 2 * sum([1, 1, 1, 1, -4.5]) = -1
-      self.assertAlmostEqual(act, -1, places=2)
+      self.assertAlmostEqual(act, -1, places=2)  # pyrefly: ignore[no-matching-overload]
     ob = np.array([-0.5, 0.5])
     if is_ob_dict:
       ob = {"sensor_1": np.array([-0.5]), "sensor_2": np.array([0.5])}
     act = policy.act(ob)
     # Latent command has now changed to [0, 0, 0, 0] and low level output is
     # 2 * sum([0, 0, 0, 0, 0.5]) = 1
-    self.assertAlmostEqual(act, 1, places=2)
+    self.assertAlmostEqual(act, 1, places=2)  # pyrefly: ignore[no-matching-overload]
 
   def test_vision_hierarchical_policy_act(self):
     """Tests the act function for hierarchical policy with vision input."""
@@ -170,7 +170,7 @@ def test_vision_hierarchical_policy_act(self):
         "sensor_1": np.array([5]),
         "sensor_2": np.array([-4])
     }
-    act = policy.act(ob)[0]
+    act = policy.act(ob)[0]  # pyrefly: ignore[bad-index]
     self.assertAlmostEqual(act, 1, places=2)
 
 if __name__ == "__main__":

iris/policies/implicit_policy.py

@@ -439,7 +439,7 @@ def act(self, state: np.ndarray) -> np.ndarray:
       phi_state = self._energy.linearized_energy_state(state)
       if self._bootstrapped_samples == self._num_samples:
         return self._actions[np.argmax(
-            np.dot(self._lat_reps_for_actions, phi_state))]
+            np.dot(self._lat_reps_for_actions, phi_state))]  # pyrefly: ignore[bad-argument-type]
       else:
         random_indices = np.random.choice(np.arange(len(self._actions)))
         return self._actions[random_indices[np.argmax(
@@ -522,7 +522,7 @@ def act(self, state: np.ndarray) -> np.ndarray:
             base_prefix_sum = self._prefix_sum_table[seg_start_index - 1]
           prob = np.dot(
               self._prefix_sum_table[seg_end_index - 1] - base_prefix_sum,
-              phi_state)
+              phi_state)  # pyrefly: ignore[bad-argument-type]
           probs.append(prob)
           start_end_indices.append([seg_start_index, seg_end_index])
           seg_start_index = seg_end_index

iris/policies/jax_policy_test.py

@@ -43,7 +43,7 @@ def test_policy_act(self):
         init_x=init_x)
     policy.update_weights(new_weights=np.ones(6))
     act = policy.act({'a': np.array([[2, -1]])})
-    np.testing.assert_array_almost_equal(act, [0.9], 1)
+    np.testing.assert_array_almost_equal(act, [0.9], 1)  # pyrefly: ignore[bad-argument-type]
 
     # Comparing keras action output with Numpy NN policy output
     numpy_policy = nn_policy.FullyConnectedNeuralNetworkPolicy(
@@ -52,7 +52,7 @@ def test_policy_act(self):
         hidden_layer_sizes=[2])
     numpy_policy.update_weights(new_weights=np.ones(6))
     numpy_act = numpy_policy.act(np.array([2, -1]))
-    np.testing.assert_array_almost_equal(act, numpy_act)
+    np.testing.assert_array_almost_equal(act, numpy_act)  # pyrefly: ignore[bad-argument-type]
 
 
 if __name__ == '__main__':

iris/policies/keras_cnn_policy.py

@@ -40,7 +40,7 @@ def _create_vision_input_layers(self):
     for image_label in self._image_input_labels:
       vision_input_layers.append(
           tf.keras.layers.Input(
-              shape=self._ob_space[image_label].shape,
+              shape=self._ob_space[image_label].shape,  # pyrefly: ignore[bad-index]
               batch_size=1,
               dtype="float32",
               name="vision_input" + image_label,
@@ -49,7 +49,7 @@ def _create_vision_input_layers(self):
     return vision_input_layers
 
   def _create_other_input_layer(self):
-    self._other_ob_space = self._ob_space.spaces.copy()
+    self._other_ob_space = self._ob_space.spaces.copy()  # pyrefly: ignore[missing-attribute]
     for input_label in self._image_input_labels:
       del self._other_ob_space[input_label]
     self._other_ob_space = spaces.Dict(self._other_ob_space)
@@ -95,12 +95,12 @@ def _create_vision_processing_layers(
     """
     # Convolution and pooling layers.
     if pool_sizes is None:
-      pool_sizes = [None] * len(conv_filter_sizes)
+      pool_sizes = [None] * len(conv_filter_sizes)  # pyrefly: ignore[bad-assignment]
     if pool_strides is None:
-      pool_strides = [None] * len(conv_filter_sizes)
+      pool_strides = [None] * len(conv_filter_sizes)  # pyrefly: ignore[bad-assignment]
 
     for filter_size, kernel_size, pool_size, pool_stride in zip(
-        conv_filter_sizes, conv_kernel_sizes, pool_sizes, pool_strides
+        conv_filter_sizes, conv_kernel_sizes, pool_sizes, pool_strides  # pyrefly: ignore[bad-argument-type]
     ):
       x = tf.keras.layers.Conv2D(
           filter_size,
@@ -117,10 +117,10 @@ def _create_vision_processing_layers(
     if use_spatial_softmax:
       x = spatial_softmax.SpatialSoftmax(data_format="channels_last")(x)
     else:
-      x = tf.keras.layers.Flatten()(x)
+      x = tf.keras.layers.Flatten()(x)  # pyrefly: ignore[not-callable]
 
     # Encoding image into a feature vector.
-    return tf.keras.layers.Dense(
+    return tf.keras.layers.Dense(  # pyrefly: ignore[not-callable]
         image_feature_length, activation=final_vision_activation
     )(x)
 
@@ -140,8 +140,8 @@ def _create_rnn_layers(self, x, inputs):
     )
     inputs.append(lstm_h_state_input)
     inputs.append(lstm_c_state_input)
-    x = tf.keras.layers.Reshape((1, -1))(x)
-    x, h_state, c_state = tf.keras.layers.LSTM(
+    x = tf.keras.layers.Reshape((1, -1))(x)  # pyrefly: ignore[not-callable]
+    x, h_state, c_state = tf.keras.layers.LSTM(  # pyrefly: ignore[not-callable]
         units=self._rnn_units, return_state=True, stateful=True
     )(x, initial_state=[lstm_h_state_input, lstm_c_state_input])
     return x, [h_state, c_state]
@@ -235,23 +235,23 @@ def act(
       inputs.append(vision_input)
 
     if self._use_rnn:
-      inputs.extend(self._rnn_state)
+      inputs.extend(self._rnn_state)  # pyrefly: ignore[bad-argument-type]
 
     if self._other_ob_dim > 0:
       other_ob = ob.copy()
       for image_label in self._image_input_labels:
-        del other_ob[image_label]
+        del other_ob[image_label]  # pyrefly: ignore[unsupported-operation]
 
       # Flatten other observations.
       other_input = utils.flatten(self._other_ob_space, other_ob)
       inputs.append(np.array([other_input]))
 
     # Run model.
-    output = self.model(inputs)
+    output = self.model(inputs)  # pyrefly: ignore[not-callable]
 
     # Parse model output.
     if self._use_rnn:
-      num_state_objects = len(self._rnn_state)
+      num_state_objects = len(self._rnn_state)  # pyrefly: ignore[bad-argument-type]
       self._rnn_state = [output[i].numpy() for i in range(num_state_objects)]
       output = output[num_state_objects:]
     actions = output[0].numpy()

iris/policies/keras_cnn_policy_test.py

@@ -37,19 +37,19 @@ def test_policy_act(self):
     policy.reset()
     policy.update_weights(new_weights=np.ones(38))
     image = np.ones((2, 2, 1))
-    act = policy.act({
+    act = policy.act({  # pyrefly: ignore[bad-argument-type]
         'vision': image,
         'sensor1': [-3, -3],
         'sensor2': [-3, -3],
     })
-    np.testing.assert_array_almost_equal(act, np.ones((5)), 1)
+    np.testing.assert_array_almost_equal(act, np.ones((5)), 1)  # pyrefly: ignore[bad-argument-type]
     policy.update_weights(new_weights=np.zeros(38))
-    act = policy.act({
+    act = policy.act({  # pyrefly: ignore[bad-argument-type]
         'vision': image,
         'sensor1': [-3, -3],
         'sensor2': [-3, -3],
     })
-    np.testing.assert_array_almost_equal(act, np.zeros((5)), 1)
+    np.testing.assert_array_almost_equal(act, np.zeros((5)), 1)  # pyrefly: ignore[bad-argument-type]
 
   def test_lstm_state(self):
     policy = keras_cnn_policy.KerasCNNPolicy(
@@ -80,7 +80,7 @@ def test_lstm_state(self):
 
     # Checks that the LSTM state changes although the observations are the same.
     for _ in range(5):
-      policy.act(observation)
+      policy.act(observation)  # pyrefly: ignore[bad-argument-type]
       rnn_state = policy._rnn_state
       np.testing.assert_raises(AssertionError,
                                np.testing.assert_array_almost_equal,

iris/policies/keras_nn_policy_test.py

@@ -29,7 +29,7 @@ def test_policy_act(self):
         hidden_layer_sizes=[2])
     policy.update_weights(new_weights=np.ones(6))
     keras_act = policy.act(np.array([2, -1]))
-    np.testing.assert_array_almost_equal(keras_act, [0.9], 1)
+    np.testing.assert_array_almost_equal(keras_act, [0.9], 1)  # pyrefly: ignore[bad-argument-type]
 
     # Comparing keras action output with Numpy NN policy output
     numpy_policy = nn_policy.FullyConnectedNeuralNetworkPolicy(
@@ -38,7 +38,7 @@ def test_policy_act(self):
         hidden_layer_sizes=[2])
     numpy_policy.update_weights(new_weights=np.ones(6))
     numpy_act = numpy_policy.act(np.array([2, -1]))
-    np.testing.assert_array_almost_equal(keras_act, numpy_act)
+    np.testing.assert_array_almost_equal(keras_act, numpy_act)  # pyrefly: ignore[bad-argument-type]
 
   def test_policy_act_dict(self):
     """Tests act for keras NN policy with dict observation."""
@@ -54,7 +54,7 @@ def test_policy_act_dict(self):
         'sensor1': np.array([2, 2]),
         'sensor2': np.array([-1, -1])
     })
-    np.testing.assert_array_almost_equal(act, [0.9], 1)
+    np.testing.assert_array_almost_equal(act, [0.9], 1)  # pyrefly: ignore[bad-argument-type]
 
 
 if __name__ == '__main__':