diff --git a/iris/algorithms/ars_algorithm_test.py b/iris/algorithms/ars_algorithm_test.py
index 8b55d69..c297cc8 100644
--- a/iris/algorithms/ars_algorithm_test.py
+++ b/iris/algorithms/ars_algorithm_test.py
@@ -41,12 +41,12 @@ def test_ars_gradient(self, orthogonal_suggestions, quasirandom_suggestions):
     suggestions = algo.get_param_suggestions()
     self.assertLen(suggestions, 6)
     eval_results = [
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.empty(0), 0),
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.empty(0), 0),
     ]
     algo.process_evaluations(eval_results)
     np.testing.assert_array_equal(algo._opt_params, np.array([10, 11]))
@@ -63,12 +63,12 @@ def test_ars_gradient_with_schedule(self):
     suggestions = algo.get_param_suggestions()
     self.assertLen(suggestions, 6)
     eval_results = [
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.empty(0), 0),
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.empty(0), 0),
     ]
     algo.process_evaluations(eval_results)
     np.testing.assert_array_equal(algo._opt_params, np.array([10, 11]))
diff --git a/iris/algorithms/cma_algorithm_test.py b/iris/algorithms/cma_algorithm_test.py
index 0addc70..4785c7d 100644
--- a/iris/algorithms/cma_algorithm_test.py
+++ b/iris/algorithms/cma_algorithm_test.py
@@ -54,7 +54,7 @@ def test_cma_optimization(self):
                 np.array(suggestion['params_to_eval']),
                 test_fn(np.array(suggestion['params_to_eval']))))
       if i%10 == 0:
-        eval_results[0] = worker_util.EvaluationResult(np.empty(0), 0)  # pytype: disable=wrong-arg-types  # numpy-scalars
+        eval_results[0] = worker_util.EvaluationResult(np.empty(0), 0)
       self.algo.process_evaluations(eval_results)
     np.testing.assert_almost_equal(self.algo._opt_params, _TRUE_OPTIMAL)
 
diff --git a/iris/algorithms/es_enas_algorithm.py b/iris/algorithms/es_enas_algorithm.py
index 829ffe6..bf7cac0 100644
--- a/iris/algorithms/es_enas_algorithm.py
+++ b/iris/algorithms/es_enas_algorithm.py
@@ -90,7 +90,7 @@ def proper_unserialize(metadata: str) -> pg.DNA:
         return dna
 
       if self._multithreading:
-        dna_list = self._pool.map(proper_unserialize, eval_metadatas)  # pytype:disable=attribute-error
+        dna_list = self._pool.map(proper_unserialize, eval_metadatas)
       else:
         dna_list = map(proper_unserialize, eval_metadatas)
       dna_list = list(dna_list)
@@ -108,7 +108,7 @@ def get_param_suggestions(self,
       # Note that for faster serialization, DNASpec is removed from DNA.
       dna_list = [self._controller.propose_dna() for _ in vanilla_suggestions]
       if self._multithreading:
-        metadata_list = self._pool.map(pg.to_json_str, dna_list)  # pytype:disable=attribute-error
+        metadata_list = self._pool.map(pg.to_json_str, dna_list)
       else:
         metadata_list = map(pg.to_json_str, dna_list)
       metadata_list = list(metadata_list)
@@ -130,7 +130,7 @@ def _get_state(self) -> Dict[str, Any]:
     return vanilla_state
 
   def _set_state(self, new_state: Dict[str, Any]) -> None:
-    super()._set_state(new_state)  # pytype: disable=attribute-error
+    super()._set_state(new_state)
     self._interval_counter = new_state["interval_counter"]
     self._dna_spec = pg.from_json_str(new_state["serialized_dna_spec"])
     self._controller = self._controller_fn(dna_spec=self._dna_spec)
diff --git a/iris/algorithms/es_enas_algorithm_test.py b/iris/algorithms/es_enas_algorithm_test.py
index a020503..09b7d82 100644
--- a/iris/algorithms/es_enas_algorithm_test.py
+++ b/iris/algorithms/es_enas_algorithm_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pytype: disable=attribute-error
 from gym import spaces
 from iris.algorithms import es_enas_algorithm
 from iris.policies import nas_policy
@@ -44,7 +43,7 @@ def make_evaluation_results(suggestion_list):
         value=np.random.uniform(),
         metadata=suggestion['metadata'])
     eval_results.append(evaluation_result)
-  eval_results.append(worker_util.EvaluationResult(np.empty(0), 0))  # pytype: disable=wrong-arg-types  # numpy-scalars
+  eval_results.append(worker_util.EvaluationResult(np.empty(0), 0))
   return eval_results
 
 
diff --git a/iris/algorithms/multi_agent_ars_algorithm_test.py b/iris/algorithms/multi_agent_ars_algorithm_test.py
index 25c5b28..e72658a 100644
--- a/iris/algorithms/multi_agent_ars_algorithm_test.py
+++ b/iris/algorithms/multi_agent_ars_algorithm_test.py
@@ -49,42 +49,42 @@ def test_init(self, agent_keys, expected_agent_keys, expected_num_agents):
 
   def _build_evaluation_results(self) -> list[worker_util.EvaluationResult]:
     eval_results = [
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.array([10.0, 11.0, 12.0, 13.0]),
             value=10,
             metrics={'reward_arm': 10, 'reward_opp': -5},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.array([10.0, 11.0, 14.0, 15.0]),
             value=10,
             metrics={'reward_arm': 10, 'reward_opp': -10},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.empty(0),
             value=0,
             metrics={'reward_arm': 0, 'reward_opp': 0},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.array([1.0, 2.0, 3.0, 4.0]),
             value=10,
             metrics={'reward_arm': 10, 'reward_opp': -10},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.array([10.0, 11.0, 12.0, 13.0]),
             value=-10,
             metrics={'reward_arm': -10, 'reward_opp': 5},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.array([10.0, 11.0, 14.0, 15.0]),
             value=-10,
             metrics={'reward_arm': -10, 'reward_opp': 10},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.array([5.0, 6.0, 7.0, 8.0]),
             value=-10,
             metrics={'reward_arm': -10, 'reward_opp': 10},
         ),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(
             params_evaluated=np.empty(0),
             value=0,
             metrics={'reward_arm': 0, 'reward_opp': 0},
diff --git a/iris/algorithms/optimizers.py b/iris/algorithms/optimizers.py
index 5b27c99..462acb4 100644
--- a/iris/algorithms/optimizers.py
+++ b/iris/algorithms/optimizers.py
@@ -87,7 +87,7 @@ def vector_decoding_function(A, b, optimization_parameters, loss_function):
   result = x.value
   res_list = []
   for i in range(n):
-    res_list.append(result[i])  # pyrefly: ignore[unsupported-operation]
+    res_list.append(result[i])
   return np.array(res_list)
 
 
diff --git a/iris/algorithms/pes_algorithm_test.py b/iris/algorithms/pes_algorithm_test.py
index 71f3b67..6056d24 100644
--- a/iris/algorithms/pes_algorithm_test.py
+++ b/iris/algorithms/pes_algorithm_test.py
@@ -37,18 +37,24 @@ def test_pes_gradient(self, orthogonal_suggestions, quasirandom_suggestions):
     init_state = {'init_params': np.array([10., 10.])}
     algo.initialize(init_state)
     eval_results = [
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
-            np.array([10., 11.]), 10, metrics={'current_step': 5}),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
-            np.empty(0), 0, metrics={'current_step': 5}),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
-            np.array([10., 11.]), 10, metrics={'current_step': 5}),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
-            np.array([10., 9.]), -10, metrics={'current_step': 5}),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
-            np.array([10., 9.]), -10, metrics={'current_step': 5}),
-        worker_util.EvaluationResult(  # pytype: disable=wrong-arg-types  # numpy-scalars
-            np.empty(0), 0, metrics={'current_step': 5}),
+        worker_util.EvaluationResult(
+            np.array([10.0, 11.0]), 10, metrics={'current_step': 5}
+        ),
+        worker_util.EvaluationResult(
+            np.empty(0), 0, metrics={'current_step': 5}
+        ),
+        worker_util.EvaluationResult(
+            np.array([10.0, 11.0]), 10, metrics={'current_step': 5}
+        ),
+        worker_util.EvaluationResult(
+            np.array([10.0, 9.0]), -10, metrics={'current_step': 5}
+        ),
+        worker_util.EvaluationResult(
+            np.array([10.0, 9.0]), -10, metrics={'current_step': 5}
+        ),
+        worker_util.EvaluationResult(
+            np.empty(0), 0, metrics={'current_step': 5}
+        ),
     ]
     algo.process_evaluations(eval_results)
     np.testing.assert_array_equal(algo._opt_params, np.array([10, 11]))
diff --git a/iris/algorithms/piars_algorithm.py b/iris/algorithms/piars_algorithm.py
index 7c4b9f9..8ff13b7 100644
--- a/iris/algorithms/piars_algorithm.py
+++ b/iris/algorithms/piars_algorithm.py
@@ -207,8 +207,8 @@ def __init__(
     else:
       self.policy = policy
 
-    obs_spec = gym_wrapper.spec_from_gym_space(self._env.observation_space)  # pyrefly: ignore[bad-argument-type]
-    action_spec = gym_wrapper.spec_from_gym_space(self._env.action_space)  # pyrefly: ignore[bad-argument-type]
+    obs_spec = gym_wrapper.spec_from_gym_space(self._env.observation_space)
+    action_spec = gym_wrapper.spec_from_gym_space(self._env.action_space)
     time_step_spec = ts.time_step_spec(observation_spec=obs_spec)
     policy_step_spec = policy_step.PolicyStep(action=action_spec)  # pyrefly: ignore[missing-argument]
     collect_data_spec = trajectory.from_transition(
@@ -325,14 +325,14 @@ def train_single_step(self, obs, reward, action, discount):
   @tf.function
   def rollout(self, obs, actions):
     """Latent rollout."""
-    s, _ = self.policy.h_model(obs)  # pyrefly: ignore[not-callable]
+    s, _ = self.policy.h_model(obs)
     outputs = []
     for i in range(self._rollout_length):
-      p, v = self.policy.f_model(s)  # pyrefly: ignore[not-callable]
-      u_next, s_next = self.policy.g_model([s, actions[:, i, ...]])  # pyrefly: ignore[not-callable]
+      p, v = self.policy.f_model(s)
+      u_next, s_next = self.policy.g_model([s, actions[:, i, ...]])
       outputs.append((p, v, u_next, s))
       s = s_next
-    p, v = self.policy.f_model(s)  # pyrefly: ignore[not-callable]
+    p, v = self.policy.f_model(s)
     outputs.append((p, v, None, s))
     return outputs
 
@@ -368,11 +368,11 @@ def infonce(hidden_x, hidden_y, temperature=0.1):
       # Latent state (from visual + other observations) for the first time step
       hx = latent_traj[0][-1]
       # Latent state (from visual observations) for the last time step
-      _, hy_vision = self.policy.h_model(obs_k)  # pyrefly: ignore[not-callable]
+      _, hy_vision = self.policy.h_model(obs_k)
       # A trick from https://arxiv.org/abs/2011.10566
       hy_vision = tf.stop_gradient(hy_vision)
-      zx = self.policy.px_model(hx)  # pyrefly: ignore[not-callable]
-      zy = self.policy.py_model(hy_vision)  # pyrefly: ignore[not-callable]
+      zx = self.policy.px_model(hx)
+      zy = self.policy.py_model(hy_vision)
       iyz, _, _ = infonce(zx, zy, temperature=0.1)
       loss_pi = -iyz
 
@@ -455,7 +455,7 @@ def flatten_nested(space, x):
   """Flatten nested."""
   if isinstance(space, spaces.Box):
     x = np.asarray(x, dtype=np.float32)
-    inner_dims = list(space.shape)  # pyrefly: ignore[bad-argument-type]
+    inner_dims = list(space.shape)
     outer_dims = list(x.shape)[: -len(inner_dims)]
     x = np.reshape(x, outer_dims + [np.prod(inner_dims)])
     return x
diff --git a/iris/algorithms/pyglove_algorithm.py b/iris/algorithms/pyglove_algorithm.py
index f379c7c..9b571ab 100644
--- a/iris/algorithms/pyglove_algorithm.py
+++ b/iris/algorithms/pyglove_algorithm.py
@@ -70,7 +70,7 @@ def proper_unserialize(metadata: str) -> pg.DNA:
       return dna
 
     if self._multithreading:
-      dna_list = self._pool.map(proper_unserialize, eval_metadatas)  # pytype:disable=attribute-error
+      dna_list = self._pool.map(proper_unserialize, eval_metadatas)
     else:
       dna_list = map(proper_unserialize, eval_metadatas)
     dna_list = list(dna_list)
@@ -89,7 +89,7 @@ def get_param_suggestions(self,
     ]
     # Note that for faster serialization, DNASpec is removed from DNA.
     if self._multithreading:
-      metadata_list = self._pool.map(pg.to_json_str, dna_list)  # pytype:disable=attribute-error
+      metadata_list = self._pool.map(pg.to_json_str, dna_list)
     else:
       metadata_list = map(pg.to_json_str, dna_list)
 
@@ -104,8 +104,8 @@ def get_param_suggestions(self,
 
   def _get_state(self) -> Dict[str, Any]:
     vanilla_state = {}
-    vanilla_state["serialized_dna_spec"] = pg.to_json_str(self._dna_spec)  # pytype:disable=attribute-error
-    vanilla_state["controller_alg_state"] = self._controller.get_state()  # pytype:disable=attribute-error
+    vanilla_state["serialized_dna_spec"] = pg.to_json_str(self._dna_spec)
+    vanilla_state["controller_alg_state"] = self._controller.get_state()
     return vanilla_state
 
   def _set_state(self, new_state: Dict[str, Any]) -> None:
diff --git a/iris/algorithms/pyribs_algorithm_test.py b/iris/algorithms/pyribs_algorithm_test.py
index d025e16..2684d0b 100644
--- a/iris/algorithms/pyribs_algorithm_test.py
+++ b/iris/algorithms/pyribs_algorithm_test.py
@@ -208,7 +208,7 @@ def test_process_evaluations(self):
         worker_util.EvaluationResult(
             params_evaluated=np.ones((13,)),
             value=1,
-            obs_norm_buffer_data={  # pyrefly: ignore[bad-argument-type]
+            obs_norm_buffer_data={
                 buffer.N: 1,  # pyrefly: ignore[bad-assignment]
                 buffer.STD: np.ones((8,)),
                 buffer.MEAN: np.ones((8,)),
@@ -219,7 +219,7 @@ def test_process_evaluations(self):
         worker_util.EvaluationResult(
             params_evaluated=np.ones((13,) * 2),
             value=2,
-            obs_norm_buffer_data={  # pyrefly: ignore[bad-argument-type]
+            obs_norm_buffer_data={
                 buffer.N: 2,  # pyrefly: ignore[bad-assignment]
                 buffer.STD: np.ones((8,)) * 2,
                 buffer.MEAN: np.ones((8,)) * 2,
diff --git a/iris/algorithms/rbo_algorithm_test.py b/iris/algorithms/rbo_algorithm_test.py
index 2bc0b25..01fd100 100644
--- a/iris/algorithms/rbo_algorithm_test.py
+++ b/iris/algorithms/rbo_algorithm_test.py
@@ -32,14 +32,14 @@ def test_rbo_gradient(self):
     init_state = {'init_params': np.array([10., 10.])}
     algo.initialize(init_state)
     eval_results = [
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.empty(0), 0),
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.empty(0), 0),
     ]
     algo.process_evaluations(eval_results)
     np.testing.assert_array_almost_equal(
@@ -67,14 +67,14 @@ def test_rbo_gradient_2(self, regression_method, orthogonal_suggestions,
     init_state = {'init_params': np.array([10., 10.])}
     algo.initialize(init_state)
     eval_results = [
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 11.]), 10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.array([10., 9.]), -10),  # pytype: disable=wrong-arg-types  # numpy-scalars
-        worker_util.EvaluationResult(np.empty(0), 0),  # pytype: disable=wrong-arg-types  # numpy-scalars
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.empty(0), 0),
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 11.0]), 10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.array([10.0, 9.0]), -10),
+        worker_util.EvaluationResult(np.empty(0), 0),
     ]
     algo.process_evaluations(eval_results)
     np.testing.assert_equal(len(algo._opt_params), 2)
diff --git a/iris/policies/gym_space_utils.py b/iris/policies/gym_space_utils.py
index 69322e7..8a03baf 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))  # pyrefly: ignore[missing-attribute]
-    high = np.concatenate((space.high, value.high))  # pyrefly: ignore[missing-attribute]
+    low = np.concatenate((space.low, value.low))
+    high = np.concatenate((space.high, value.high))
     return gym.spaces.Box(low=low, high=high)
   elif isinstance(space, gym.spaces.Dict):
     extended_space = dict(space.spaces)
diff --git a/iris/policies/implicit_policy.py b/iris/policies/implicit_policy.py
index a2777e8..da2e8ea 100644
--- a/iris/policies/implicit_policy.py
+++ b/iris/policies/implicit_policy.py
@@ -438,8 +438,9 @@ def act(self, state: np.ndarray) -> np.ndarray:
     else:
       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))]  # pyrefly: ignore[bad-argument-type, no-matching-overload]
+        return self._actions[
+            np.argmax(np.dot(self._lat_reps_for_actions, phi_state))  # pyrefly: ignore[no-matching-overload]
+        ]
       else:
         random_indices = np.random.choice(np.arange(len(self._actions)))
         return self._actions[random_indices[np.argmax(
@@ -522,7 +523,8 @@ def act(self, state: np.ndarray) -> np.ndarray:
             base_prefix_sum = self._prefix_sum_table[seg_start_index - 1]
           prob = np.dot(  # pyrefly: ignore[no-matching-overload]
               self._prefix_sum_table[seg_end_index - 1] - base_prefix_sum,
-              phi_state)  # pyrefly: ignore[bad-argument-type]
+              phi_state,
+          )
           probs.append(prob)
           start_end_indices.append([seg_start_index, seg_end_index])
           seg_start_index = seg_end_index
diff --git a/iris/policies/keras_cnn_policy.py b/iris/policies/keras_cnn_policy.py
index d21cbcd..630a640 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,  # pyrefly: ignore[bad-index]
+              shape=self._ob_space[image_label].shape,
               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()  # pyrefly: ignore[missing-attribute]
+    self._other_ob_space = self._ob_space.spaces.copy()
     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)
@@ -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)  # pyrefly: ignore[not-callable]
+      x = tf.keras.layers.Flatten()(x)
 
     # Encoding image into a feature vector.
-    return tf.keras.layers.Dense(  # pyrefly: ignore[not-callable]
+    return tf.keras.layers.Dense(
         image_feature_length, activation=final_vision_activation
     )(x)
 
@@ -140,14 +140,13 @@ 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)  # pyrefly: ignore[not-callable]
-    x, h_state, c_state = tf.keras.layers.LSTM(  # pyrefly: ignore[not-callable]
+    x = tf.keras.layers.Reshape((1, -1))(x)
+    x, h_state, c_state = tf.keras.layers.LSTM(
         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]
 
-  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
-  def _build_model(
+  def _build_model(  # pyrefly: ignore[bad-override]
       self,
       fc_layer_sizes: Sequence[int],
       use_rnn: bool = False,
@@ -209,8 +208,6 @@ def _build_model(
 
     self.model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
 
-  # pytype: enable=signature-mismatch  # overriding-parameter-count-checks
-
   def reset(self) -> None:
     """Resets the policy's internal state (default LSTM)."""
     lstm_h_state = np.zeros(shape=(1, self._rnn_units), dtype="float32")
@@ -247,7 +244,7 @@ def act(
       inputs.append(np.array([other_input]))
 
     # Run model.
-    output = self.model(inputs)  # pyrefly: ignore[not-callable]
+    output = self.model(inputs)
 
     # Parse model output.
     if self._use_rnn:
diff --git a/iris/policies/keras_cnn_policy_test.py b/iris/policies/keras_cnn_policy_test.py
index 62bc788..577af65 100644
--- a/iris/policies/keras_cnn_policy_test.py
+++ b/iris/policies/keras_cnn_policy_test.py
@@ -37,14 +37,14 @@ def test_policy_act(self):
     policy.reset()
     policy.update_weights(new_weights=np.ones(38))
     image = np.ones((2, 2, 1))
-    act = policy.act({  # pyrefly: ignore[bad-argument-type]
+    act = policy.act({
         'vision': image,
         'sensor1': [-3, -3],  # pyrefly: ignore[bad-assignment]
         'sensor2': [-3, -3],  # pyrefly: ignore[bad-assignment]
     })
     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({  # pyrefly: ignore[bad-argument-type]
+    act = policy.act({
         'vision': image,
         'sensor1': [-3, -3],  # pyrefly: ignore[bad-assignment]
         'sensor2': [-3, -3],  # pyrefly: ignore[bad-assignment]
@@ -82,22 +82,32 @@ def test_lstm_state(self):
     for _ in range(5):
       policy.act(observation)  # pyrefly: ignore[bad-argument-type]
       rnn_state = policy._rnn_state
-      np.testing.assert_raises(AssertionError,
-                               np.testing.assert_array_almost_equal,
-                               prev_h_state, rnn_state[0])  #  pytype: disable=unsupported-operands
-      np.testing.assert_raises(AssertionError,
-                               np.testing.assert_array_almost_equal,
-                               prev_c_state, rnn_state[1])  #  pytype: disable=unsupported-operands
-      prev_h_state = rnn_state[0]  #  pytype: disable=unsupported-operands
-      prev_c_state = rnn_state[1]  #  pytype: disable=unsupported-operands
+      np.testing.assert_raises(
+          AssertionError,
+          np.testing.assert_array_almost_equal,
+          prev_h_state,
+          rnn_state[0],  # pyrefly: ignore[unsupported-operation]
+      )
+      np.testing.assert_raises(
+          AssertionError,
+          np.testing.assert_array_almost_equal,
+          prev_c_state,
+          rnn_state[1],  # pyrefly: ignore[unsupported-operation]
+      )
+      prev_h_state = rnn_state[0]  # pyrefly: ignore[unsupported-operation]
+      prev_c_state = rnn_state[1]  # pyrefly: ignore[unsupported-operation]
 
     # Checks that the LSTM state is reset to zero.
     policy.reset()
     new_rnn_state = policy._rnn_state
-    np.testing.assert_array_almost_equal(new_rnn_state[0],  #  pytype: disable=unsupported-operands
-                                         np.zeros(shape=(1, 2)))
-    np.testing.assert_array_almost_equal(new_rnn_state[0],  #  pytype: disable=unsupported-operands
-                                         np.zeros(shape=(1, 2)))
+    np.testing.assert_array_almost_equal(
+        new_rnn_state[0],  # pyrefly: ignore[unsupported-operation]
+        np.zeros(shape=(1, 2)),
+    )
+    np.testing.assert_array_almost_equal(
+        new_rnn_state[0],  # pyrefly: ignore[unsupported-operation]
+        np.zeros(shape=(1, 2)),
+    )
 
 
 if __name__ == '__main__':
diff --git a/iris/policies/keras_nn_policy.py b/iris/policies/keras_nn_policy.py
index c821f76..9b74556 100644
--- a/iris/policies/keras_nn_policy.py
+++ b/iris/policies/keras_nn_policy.py
@@ -22,8 +22,7 @@
 class KerasNNPolicy(keras_policy.KerasPolicy):
   """Policy class that computes action by running feed fwd neural network."""
 
-  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
-  def _build_model(
+  def _build_model(  # pyrefly: ignore[bad-override]
       self,
       hidden_layer_sizes: Sequence[int],
       activation: str = "tanh",
@@ -59,5 +58,3 @@ def _build_model(
     self.model = tf.keras.models.Model(
         inputs=[input_layer], outputs=[output_layer]
     )
-
-  # pytype: enable=signature-mismatch  # overriding-parameter-count-checks
diff --git a/iris/policies/keras_pi_policy.py b/iris/policies/keras_pi_policy.py
index f828dc7..3bd54e1 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,  # pyrefly: ignore[bad-index]
+              shape=self._ob_space[image_label].shape,
               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()  # pyrefly: ignore[missing-attribute]
+      self._other_ob_space = self._ob_space.spaces.copy()
       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)
@@ -138,15 +138,14 @@ 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)  # pyrefly: ignore[not-callable]
+      x = tf.keras.layers.Flatten()(x)
 
     # Encoding image into a feature vector.
-    return tf.keras.layers.Dense(  # pyrefly: ignore[not-callable]
+    return tf.keras.layers.Dense(
         image_feature_length, activation=final_vision_activation
     )(x)
 
-  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
-  def _build_model(
+  def _build_model(  # pyrefly: ignore[bad-override]
       self, state_dim: int, fc_layer_sizes: Sequence[int], **kwargs
   ) -> None:
     # hidden state input
@@ -163,8 +162,6 @@ def _build_model(
         inputs=state_input, outputs=[action_output]
     )
 
-  # pytype: enable=signature-mismatch  # overriding-parameter-count-checks
-
   def build_h(
       self,
       h_fc_layer_sizes: Sequence[int],
@@ -204,7 +201,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)  # pyrefly: ignore[not-callable]
+      x = tf.keras.layers.Dense(h_fc_layer_size, activation="tanh")(x)
     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 +217,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)  # 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]
+      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)
     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 +235,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)  # 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]
+      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)
     self.g_model = tf.keras.models.Model(
         inputs=[state_input, action_input], outputs=[u_next, s_next]
     )
@@ -251,8 +248,8 @@ def build_px(self, state_dim: int, **kwargs):
     )
 
     x = state_input
-    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]
+    x = tf.keras.layers.Dense(64, activation="tanh")(x)
+    z = tf.keras.layers.Dense(state_dim)(x)
     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 +261,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)  # pyrefly: ignore[not-callable]
-    z = tf.keras.layers.Dense(state_dim)(x)  # pyrefly: ignore[not-callable]
+    x = tf.keras.layers.Dense(64, activation="tanh")(x)
+    z = tf.keras.layers.Dense(state_dim)(x)
     self.py_model = tf.keras.models.Model(inputs=state_input, outputs=z)
 
   def act(
@@ -295,8 +292,8 @@ def act(
       inputs.append(np.array([other_input]))
 
     # Run model.
-    s, _ = self.h_model(inputs)  # pyrefly: ignore[not-callable]
-    output = self.model(s)  # pyrefly: ignore[not-callable]
+    s, _ = self.h_model(inputs)
+    output = self.model(s)
 
     # Parse model output.
     actions = output.numpy()
@@ -326,13 +323,13 @@ def rollout(
     inputs.append(np.array([other_input]))
 
     # Run model.
-    s, _ = self.h_model(inputs)  # pyrefly: ignore[not-callable]
+    s, _ = self.h_model(inputs)
     reward = 0.0
     for _ in range(rollout_length):
-      action = self.model(s)  # pyrefly: ignore[not-callable]
-      u_next, s = self.g_model([s, action])  # pyrefly: ignore[not-callable]
+      action = self.model(s)
+      u_next, s = self.g_model([s, action])
       reward += u_next
-    _, z = self.f_model(s)  # pyrefly: ignore[not-callable]
+    _, z = self.f_model(s)
     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 d532c98..3b88f4e 100644
--- a/iris/policies/keras_pi_policy_test.py
+++ b/iris/policies/keras_pi_policy_test.py
@@ -41,7 +41,7 @@ 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({  # pyrefly: ignore[bad-argument-type]
+    act = policy.act({
         'vision': image,
         'sensor1': [-3, -3],  # pyrefly: ignore[bad-assignment]
         'sensor2': [-3, -3],  # pyrefly: ignore[bad-assignment]
@@ -49,7 +49,7 @@ def test_policy_act(self):
     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({  # pyrefly: ignore[bad-argument-type]
+    act = policy.act({
         'vision': image,
         'sensor1': [-3, -3],  # pyrefly: ignore[bad-assignment]
         'sensor2': [-3, -3],  # pyrefly: ignore[bad-assignment]
diff --git a/iris/policies/keras_policy.py b/iris/policies/keras_policy.py
index f3873a3..a797e12 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]  # pyrefly: ignore[not-callable]
+    actions = self.model(np.array([ob])).numpy()[0]
     actions = utils.unflatten(self._ac_space, actions)
     return actions
diff --git a/iris/policies/keras_toeplitz_policy.py b/iris/policies/keras_toeplitz_policy.py
index 5efa917..8f15e9d 100644
--- a/iris/policies/keras_toeplitz_policy.py
+++ b/iris/policies/keras_toeplitz_policy.py
@@ -135,8 +135,7 @@ def call(self, inputs: tf.Tensor) -> tf.Tensor:
 class KerasToeplitzPolicy(keras_policy.KerasPolicy):
   """Policy class that computes action by running toeplitz network."""
 
-  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
-  def _build_model(
+  def _build_model(  # pyrefly: ignore[bad-override]
       self,
       hidden_layer_sizes: Sequence[int],
       activation: str = "tanh",
@@ -172,5 +171,3 @@ def _build_model(
     self.model = tf.keras.models.Model(
         inputs=[input_layer], outputs=[output_layer]
     )
-
-  # pytype: enable=signature-mismatch  # overriding-parameter-count-checks
diff --git a/iris/policies/spatial_softmax_test.py b/iris/policies/spatial_softmax_test.py
index adb0dae..0f7a241 100644
--- a/iris/policies/spatial_softmax_test.py
+++ b/iris/policies/spatial_softmax_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pytype: disable=attribute-error
 from iris.policies import spatial_softmax
 import tensorflow as tf2
 import tensorflow.compat.v1 as tf
@@ -32,13 +31,13 @@
 class SpatialSoftmaxTest(tf.test.TestCase):
 
   def test_with_default(self):
-    out = test_utils.layer_test(
+    out = test_utils.layer_test(  # pyrefly: ignore[missing-attribute]
         spatial_softmax.SpatialSoftmax, input_shape=_INPUT_SHAPE
     )
     self.assertAllEqual(out.shape, (_INPUT_SHAPE[0], _INPUT_SHAPE[3] * 2))
 
   def test_with_preset_temperature(self):
-    test_utils.layer_test(
+    test_utils.layer_test(  # pyrefly: ignore[missing-attribute]
         spatial_softmax.SpatialSoftmax,
         kwargs={'temperature': _TEMPERATURE},
         input_shape=_INPUT_SHAPE,
@@ -50,7 +49,7 @@ def test_get_weights(self):
 
   def test_with_channels_first(self):
     input_shape = (16, 128, 32, 32)
-    out = test_utils.layer_test(
+    out = test_utils.layer_test(  # pyrefly: ignore[missing-attribute]
         spatial_softmax.SpatialSoftmax,
         kwargs={'data_format': 'channels_first'},
         input_shape=input_shape,