From 0588153dd3757ef32f756ba1fdb3e1a99fa3182d Mon Sep 17 00:00:00 2001
From: Hana Joo <hanajoo@google.com>
Date: Thu, 2 Jul 2026 05:57:54 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 941655541
---
 iris/policies/gym_space_utils.py              |  8 +--
 iris/policies/hierarchical_policy.py          |  4 +-
 iris/policies/hierarchical_policy_test.py     |  8 +--
 iris/policies/implicit_policy.py              |  4 +-
 iris/policies/jax_policy_test.py              |  4 +-
 iris/policies/keras_cnn_policy.py             | 26 +++++-----
 iris/policies/keras_cnn_policy_test.py        | 10 ++--
 iris/policies/keras_nn_policy_test.py         |  6 +--
 iris/policies/keras_pi_policy.py              | 52 +++++++++----------
 iris/policies/keras_pi_policy_test.py         |  8 +--
 iris/policies/keras_policy.py                 |  2 +-
 iris/policies/keras_toeplitz_policy_test.py   |  4 +-
 iris/policies/linear_policy.py                |  2 +-
 iris/policies/linear_policy_test.py           |  4 +-
 iris/policies/nn_policy.py                    |  2 +-
 iris/policies/nn_policy_test.py               |  4 +-
 iris/policies/option_hierarchical_policy.py   |  2 +-
 .../option_hierarchical_policy_test.py        | 10 ++--
 iris/policies/pct_policy_test.py              |  8 +--
 19 files changed, 84 insertions(+), 84 deletions(-)

diff --git a/iris/policies/gym_space_utils.py b/iris/policies/gym_space_utils.py
index a47ee15..69322e7 100644
--- a/iris/policies/gym_space_utils.py
+++ b/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
 
 
diff --git a/iris/policies/hierarchical_policy.py b/iris/policies/hierarchical_policy.py
index 662d925..49cd315 100644
--- a/iris/policies/hierarchical_policy.py
+++ b/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:
diff --git a/iris/policies/hierarchical_policy_test.py b/iris/policies/hierarchical_policy_test.py
index 65b8ae1..115cb2d 100644
--- a/iris/policies/hierarchical_policy_test.py
+++ b/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__":
diff --git a/iris/policies/implicit_policy.py b/iris/policies/implicit_policy.py
index 5643589..41f709d 100644
--- a/iris/policies/implicit_policy.py
+++ b/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
diff --git a/iris/policies/jax_policy_test.py b/iris/policies/jax_policy_test.py
index f0077ed..da2d0fb 100644
--- a/iris/policies/jax_policy_test.py
+++ b/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__':
diff --git a/iris/policies/keras_cnn_policy.py b/iris/policies/keras_cnn_policy.py
index 8812ab3..d21cbcd 100644
--- a/iris/policies/keras_cnn_policy.py
+++ b/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()
diff --git a/iris/policies/keras_cnn_policy_test.py b/iris/policies/keras_cnn_policy_test.py
index 821828c..e6a4fc9 100644
--- a/iris/policies/keras_cnn_policy_test.py
+++ b/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,
diff --git a/iris/policies/keras_nn_policy_test.py b/iris/policies/keras_nn_policy_test.py
index ca954c1..039e392 100644
--- a/iris/policies/keras_nn_policy_test.py
+++ b/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__':
diff --git a/iris/policies/keras_pi_policy.py b/iris/policies/keras_pi_policy.py
index ae2fe66..f828dc7 100644
--- a/iris/policies/keras_pi_policy.py
+++ b/iris/policies/keras_pi_policy.py
@@ -56,7 +56,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,
@@ -68,7 +68,7 @@ def _create_other_input_layer(self):
     if isinstance(self._ob_space, gym.spaces.Box):
       self._other_ob_space = self._ob_space
     else:
-      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)
@@ -116,12 +116,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,
@@ -138,10 +138,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)
 
@@ -204,7 +204,7 @@ def build_h(
 
     # state: fully connected layers.
     for h_fc_layer_size in h_fc_layer_sizes:
-      x = tf.keras.layers.Dense(h_fc_layer_size, activation="tanh")(x)
+      x = tf.keras.layers.Dense(h_fc_layer_size, activation="tanh")(x)  # pyrefly: ignore[not-callable]
     outputs = [x, vision_output] if vision_output is not None else x
     self.h_model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
 
@@ -220,9 +220,9 @@ def build_f(
     )
     x = state_input
     for f_fc_layer_size in f_fc_layer_sizes:
-      x = tf.keras.layers.Dense(f_fc_layer_size, activation="tanh")(x)
-    p = tf.keras.layers.Dense(self._ac_dim, activation="tanh")(x)
-    v = tf.keras.layers.Dense(num_supports)(x)
+      x = tf.keras.layers.Dense(f_fc_layer_size, activation="tanh")(x)  # pyrefly: ignore[not-callable]
+    p = tf.keras.layers.Dense(self._ac_dim, activation="tanh")(x)  # pyrefly: ignore[not-callable]
+    v = tf.keras.layers.Dense(num_supports)(x)  # pyrefly: ignore[not-callable]
     self.f_model = tf.keras.models.Model(inputs=state_input, outputs=[p, v])
 
   def build_g(self, state_dim: int, g_fc_layer_sizes: Sequence[int], **kwargs):
@@ -238,9 +238,9 @@ def build_g(self, state_dim: int, g_fc_layer_sizes: Sequence[int], **kwargs):
 
     x = tf.keras.layers.concatenate([state_input, action_input])
     for g_fc_layer_size in g_fc_layer_sizes:
-      x = tf.keras.layers.Dense(g_fc_layer_size, activation="tanh")(x)
-    u_next = tf.keras.layers.Dense(1)(x)
-    s_next = tf.keras.layers.Dense(state_dim, activation="tanh")(x)
+      x = tf.keras.layers.Dense(g_fc_layer_size, activation="tanh")(x)  # pyrefly: ignore[not-callable]
+    u_next = tf.keras.layers.Dense(1)(x)  # pyrefly: ignore[not-callable]
+    s_next = tf.keras.layers.Dense(state_dim, activation="tanh")(x)  # pyrefly: ignore[not-callable]
     self.g_model = tf.keras.models.Model(
         inputs=[state_input, action_input], outputs=[u_next, s_next]
     )
@@ -251,8 +251,8 @@ def build_px(self, state_dim: int, **kwargs):
     )
 
     x = state_input
-    x = tf.keras.layers.Dense(64, activation="tanh")(x)
-    z = tf.keras.layers.Dense(state_dim)(x)
+    x = tf.keras.layers.Dense(64, activation="tanh")(x)  # pyrefly: ignore[not-callable]
+    z = tf.keras.layers.Dense(state_dim)(x)  # pyrefly: ignore[not-callable]
     self.px_model = tf.keras.models.Model(inputs=state_input, outputs=z)
 
   def build_py(self, state_dim: int, image_feature_length: int, **kwargs):
@@ -264,8 +264,8 @@ def build_py(self, state_dim: int, image_feature_length: int, **kwargs):
     )
 
     x = state_input
-    x = tf.keras.layers.Dense(64, activation="tanh")(x)
-    z = tf.keras.layers.Dense(state_dim)(x)
+    x = tf.keras.layers.Dense(64, activation="tanh")(x)  # pyrefly: ignore[not-callable]
+    z = tf.keras.layers.Dense(state_dim)(x)  # pyrefly: ignore[not-callable]
     self.py_model = tf.keras.models.Model(inputs=state_input, outputs=z)
 
   def act(
@@ -288,15 +288,15 @@ def act(
     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.
-    s, _ = self.h_model(inputs)
-    output = self.model(s)
+    s, _ = self.h_model(inputs)  # pyrefly: ignore[not-callable]
+    output = self.model(s)  # pyrefly: ignore[not-callable]
 
     # Parse model output.
     actions = output.numpy()
@@ -317,7 +317,7 @@ def rollout(
 
     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)
@@ -326,13 +326,13 @@ def rollout(
     inputs.append(np.array([other_input]))
 
     # Run model.
-    s, _ = self.h_model(inputs)
+    s, _ = self.h_model(inputs)  # pyrefly: ignore[not-callable]
     reward = 0.0
     for _ in range(rollout_length):
-      action = self.model(s)
-      u_next, s = self.g_model([s, action])
+      action = self.model(s)  # pyrefly: ignore[not-callable]
+      u_next, s = self.g_model([s, action])  # pyrefly: ignore[not-callable]
       reward += u_next
-    _, z = self.f_model(s)
+    _, z = self.f_model(s)  # pyrefly: ignore[not-callable]
     vd = tf.nn.softmax(z)
     supports = tf.linspace(-10.0, 10.0, 51)
     v = tf.reduce_sum(vd * supports[None, ...], axis=-1)
diff --git a/iris/policies/keras_pi_policy_test.py b/iris/policies/keras_pi_policy_test.py
index 193fed7..f5b16d3 100644
--- a/iris/policies/keras_pi_policy_test.py
+++ b/iris/policies/keras_pi_policy_test.py
@@ -41,20 +41,20 @@ def test_policy_act(self):
     policy.update_weights(np.ones(21))
     policy.update_representation_weights(np.ones(1001))
     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(np.zeros(21))
     policy.update_representation_weights(np.zeros(1001))
-    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]
 
 
 if __name__ == '__main__':
diff --git a/iris/policies/keras_policy.py b/iris/policies/keras_policy.py
index a797e12..f3873a3 100644
--- a/iris/policies/keras_policy.py
+++ b/iris/policies/keras_policy.py
@@ -119,6 +119,6 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
       The actions in reinforcement learning.
     """
     ob = utils.flatten(self._ob_space, ob)
-    actions = self.model(np.array([ob])).numpy()[0]
+    actions = self.model(np.array([ob])).numpy()[0]  # pyrefly: ignore[not-callable]
     actions = utils.unflatten(self._ac_space, actions)
     return actions
diff --git a/iris/policies/keras_toeplitz_policy_test.py b/iris/policies/keras_toeplitz_policy_test.py
index e4633dc..4efe3be 100644
--- a/iris/policies/keras_toeplitz_policy_test.py
+++ b/iris/policies/keras_toeplitz_policy_test.py
@@ -28,7 +28,7 @@ def test_policy_act(self):
         hidden_layer_sizes=[3])
     policy.update_weights(new_weights=np.ones(8))
     act = policy.act(np.array([2, -1]))
-    np.testing.assert_array_almost_equal(act, np.array([0.9, 0.9]), 1)
+    np.testing.assert_array_almost_equal(act, np.array([0.9, 0.9]), 1)  # pyrefly: ignore[bad-argument-type]
 
   def test_policy_act_dict(self):
     """Tests the act function for Toeplitz policy with dict observation."""
@@ -44,7 +44,7 @@ def test_policy_act_dict(self):
         'sensor1': np.array([2, 2]),
         'sensor2': np.array([-1, -1])
     })
-    np.testing.assert_array_almost_equal(act, np.array([0.9, 0.9, 0.9, 0.9]), 1)
+    np.testing.assert_array_almost_equal(act, np.array([0.9, 0.9, 0.9, 0.9]), 1)  # pyrefly: ignore[bad-argument-type]
 
 
 if __name__ == '__main__':
diff --git a/iris/policies/linear_policy.py b/iris/policies/linear_policy.py
index 6276295..4719d65 100644
--- a/iris/policies/linear_policy.py
+++ b/iris/policies/linear_policy.py
@@ -51,6 +51,6 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
     """
     ob = utils.flatten(self._ob_space, ob)
     matrix_weights = np.reshape(self._weights, (self._ac_dim, self._ob_dim))
-    actions = self._activation(np.dot(matrix_weights, ob))
+    actions = self._activation(np.dot(matrix_weights, ob))  # pyrefly: ignore[not-callable]
     actions = utils.unflatten(self._ac_space, actions)
     return actions
diff --git a/iris/policies/linear_policy_test.py b/iris/policies/linear_policy_test.py
index 50c2b45..ccc89a3 100644
--- a/iris/policies/linear_policy_test.py
+++ b/iris/policies/linear_policy_test.py
@@ -26,7 +26,7 @@ def test_policy_act(self):
         ob_space=spaces.Box(low=-10, high=10, shape=(2,)),
         ac_space=spaces.Box(low=-10, high=10, shape=(1,)))
     policy.update_weights(new_weights=np.ones(2))
-    act = policy.act(np.array([2, -1]))[0]
+    act = policy.act(np.array([2, -1]))[0]  # pyrefly: ignore[bad-index]
     self.assertEqual(act, 1)
 
   def test_policy_act_dict(self):
@@ -41,7 +41,7 @@ def test_policy_act_dict(self):
     act = policy.act({
         'sensor1': np.array([2, 2]),
         'sensor2': np.array([-1, -1])
-    })[0]
+    })[0]  # pyrefly: ignore[bad-index]
     self.assertEqual(act, 1)
 
 
diff --git a/iris/policies/nn_policy.py b/iris/policies/nn_policy.py
index 79e4764..9e8d1ca 100644
--- a/iris/policies/nn_policy.py
+++ b/iris/policies/nn_policy.py
@@ -71,7 +71,7 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
           self._weights[start:end],
           (self._layer_sizes[ith_layer + 1], self._layer_sizes[ith_layer]))
       ith_layer_result = np.dot(mat_weight, ith_layer_result)
-      ith_layer_result = self._activation(ith_layer_result)
+      ith_layer_result = self._activation(ith_layer_result)  # pyrefly: ignore[not-callable]
     actions = ith_layer_result
     actions = utils.unflatten(self._ac_space, actions)
     return actions
diff --git a/iris/policies/nn_policy_test.py b/iris/policies/nn_policy_test.py
index 68eb539..763b017 100644
--- a/iris/policies/nn_policy_test.py
+++ b/iris/policies/nn_policy_test.py
@@ -28,7 +28,7 @@ def test_policy_act(self):
         hidden_layer_sizes=[2],
         activation='clip')
     policy.update_weights(new_weights=np.ones(6))
-    act = policy.act(np.array([2, -1]))[0]
+    act = policy.act(np.array([2, -1]))[0]  # pyrefly: ignore[bad-index]
     self.assertEqual(act, 1)
 
   def test_policy_act_dict(self):
@@ -45,7 +45,7 @@ def test_policy_act_dict(self):
     act = policy.act({
         'sensor1': np.array([2, 2]),
         'sensor2': np.array([-1, -1])
-    })[0]
+    })[0]  # pyrefly: ignore[bad-index]
     self.assertEqual(act, 1)
 
 
diff --git a/iris/policies/option_hierarchical_policy.py b/iris/policies/option_hierarchical_policy.py
index 403db6b..1f3a63a 100644
--- a/iris/policies/option_hierarchical_policy.py
+++ b/iris/policies/option_hierarchical_policy.py
@@ -85,7 +85,7 @@ def __call__(self,
     if terminated:
       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.idx > 0:
         self._terminate = bool(np.random.binomial(1, (self._output[0] + 1.)/2.))
         self._output = self._output[1:]
diff --git a/iris/policies/option_hierarchical_policy_test.py b/iris/policies/option_hierarchical_policy_test.py
index d8525ee..080a515 100644
--- a/iris/policies/option_hierarchical_policy_test.py
+++ b/iris/policies/option_hierarchical_policy_test.py
@@ -143,7 +143,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 low level doesn't terminate and latent command remains constant
     # until low level terminates
     ob = np.array([5, -4.5])
@@ -153,21 +153,21 @@ 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([5, -3.5])
     if is_ob_dict:
       ob = {"sensor_1": np.array([5]), "sensor_2": np.array([-3.5])}
     act = policy.act(ob)
     # Latent command is still [1, 1, 1, 1] and low level output is
     # 2 * sum([1, 1, 1, 1, -3.5]) = 1 and low level terminates
-    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."""
@@ -179,7 +179,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)
 
 
diff --git a/iris/policies/pct_policy_test.py b/iris/policies/pct_policy_test.py
index d26cea4..722deb7 100644
--- a/iris/policies/pct_policy_test.py
+++ b/iris/policies/pct_policy_test.py
@@ -30,7 +30,7 @@ def test_pct_encoder(self):
     pct_enc = pct_policy.PCTEncoder(emb_dim=out_dim, attention_type='perf-relu')
     params = pct_enc.init(jax.random.PRNGKey(0), pc)
     emb = pct_enc.apply(params, pc)
-    self.assertEqual(emb.shape, (batch_size, out_dim))
+    self.assertEqual(emb.shape, (batch_size, out_dim))  # pyrefly: ignore[missing-attribute]
 
   def test_pct_encoder_with_mask(self):
     batch_size = 10
@@ -46,10 +46,10 @@ def test_pct_encoder_with_mask(self):
     pct_enc = pct_policy.PCTEncoder(emb_dim=out_dim)
     params = pct_enc.init(jax.random.PRNGKey(0), pc)
     emb = pct_enc.apply(params, pc, mask)
-    self.assertEqual(emb.shape, (batch_size, out_dim))
+    self.assertEqual(emb.shape, (batch_size, out_dim))  # pyrefly: ignore[missing-attribute]
 
     emb_no_mask = pct_enc.apply(params, pc[:, :-masked_points], None)
-    np.testing.assert_allclose(emb_no_mask, emb, atol=1e-1, rtol=1e-1)
+    np.testing.assert_allclose(emb_no_mask, emb, atol=1e-1, rtol=1e-1)  # pyrefly: ignore[no-matching-overload]
 
   def test_policy_act(self):
     """Tests the act function for PCT policy."""
@@ -76,7 +76,7 @@ def test_policy_act(self):
     x = ob_space.sample()
 
     jax_act = policy.act(x)
-    np.testing.assert_array_almost_equal(jax_act, np.zeros(7), 2)
+    np.testing.assert_array_almost_equal(jax_act, np.zeros(7), 2)  # pyrefly: ignore[bad-argument-type]
 
 
 if __name__ == '__main__':