Automated Code Change

iris/algorithms/ars_algorithm.py

@@ -189,7 +189,7 @@ def get_param_suggestions(
         param_suggestions = self._np_random_state.normal(
             0, 1, (self._num_suggestions, dimensions)
         )
-      self._last_std_used = self._std
+      self._last_std_used = self._std  # pyrefly: ignore[bad-assignment]
       if callable(self._std):
         self._last_std_used = self._std(self._iteration)
       param_suggestions = np.vstack([
@@ -213,7 +213,7 @@ def state(self) -> Dict[str, Any]:
   def _get_state(self) -> Dict[str, Any]:
     state = {"params_to_eval": self._opt_params}
     if self._obs_norm_data_buffer is not None:
-      state["obs_norm_state"] = self._obs_norm_data_buffer.state
+      state["obs_norm_state"] = self._obs_norm_data_buffer.state  # pyrefly: ignore[bad-assignment]
     return state
 
   @state.setter

iris/algorithms/ars_algorithm_test.py

@@ -92,7 +92,7 @@ def test_restore_state_from_checkpoint(self, expected_obs_norm_state):
     )
     init_state = {'init_params': np.array([10.0, 10.0])}
     if expected_obs_norm_state:
-      init_state['obs_norm_buffer_data'] = {
+      init_state['obs_norm_buffer_data'] = {  # pyrefly: ignore[bad-assignment]
           'mean': np.asarray([0.0, 0.0]),
           'std': np.asarray([1.0, 1.0]),
           'n': 0,

iris/algorithms/cma_algorithm.py

@@ -103,7 +103,7 @@ def process_evaluations(self,
     # Update the observation buffer
     if self._obs_norm_data_buffer is not None:
       for r in filtered_eval_results:
-        self._obs_norm_data_buffer.merge(r.obs_norm_buffer_data)
+        self._obs_norm_data_buffer.merge(r.obs_norm_buffer_data)  # pyrefly: ignore[bad-argument-type]
 
   def get_param_suggestions(self,
                             evaluate: bool = False
@@ -141,7 +141,7 @@ def state(self) -> Dict[str, Any]:
   def _get_state(self) -> Dict[str, Any]:
     state = {"params_to_eval": self._opt_params}
     if self._obs_norm_data_buffer is not None:
-      state["obs_norm_state"] = self._obs_norm_data_buffer.state
+      state["obs_norm_state"] = self._obs_norm_data_buffer.state  # pyrefly: ignore[bad-assignment]
     return state
 
   @state.setter

iris/algorithms/learnable_ars_algorithm.py

@@ -169,7 +169,7 @@ def get_param_suggestions(
         param_suggestions = self._np_random_state.normal(
             0, 1, (self._num_suggestions, dimensions)
         )
-      self._last_std_used = self._std
+      self._last_std_used = self._std  # pyrefly: ignore[bad-assignment]
       param_suggestions = np.vstack([
           self._opt_params,
           self._opt_params + self._last_std_used * param_suggestions,
@@ -194,7 +194,7 @@ def process_evaluations(
     model_input = np.concatenate([[self._iteration], rewards])
 
     if self._tree_weights is None:
-      self._model_state = self._restore_state_from_checkpoint(self._model_path)
+      self._model_state = self._restore_state_from_checkpoint(self._model_path)  # pyrefly: ignore[bad-argument-type]
       self._tree_weights = self._model.init(
           jax.random.PRNGKey(seed=self._seed), model_input, self._model_state
       )

iris/algorithms/multi_agent_ars_algorithm.py

@@ -110,7 +110,7 @@ def restore_state_from_checkpoint(self, new_state: Dict[str, Any]) -> None:
           )
       }
       if self._obs_norm_data_buffer is not None:
-        duplicated_state["obs_norm_state"] = {}
+        duplicated_state["obs_norm_state"] = {}  # pyrefly: ignore[bad-assignment]
         duplicated_state["obs_norm_state"]["mean"] = np.tile(
             new_state["obs_norm_state"]["mean"], self._num_agents
         )
@@ -193,10 +193,10 @@ def _get_top_evaluation_results(
       neg_eval_results: Sequence[worker_util.EvaluationResult],
   ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     pos_evals = np.array(
-        [r.metrics[f"reward_{agent_key}"] for r in pos_eval_results]
+        [r.metrics[f"reward_{agent_key}"] for r in pos_eval_results]  # pyrefly: ignore[unsupported-operation]
     )
     neg_evals = np.array(
-        [r.metrics[f"reward_{agent_key}"] for r in neg_eval_results]
+        [r.metrics[f"reward_{agent_key}"] for r in neg_eval_results]  # pyrefly: ignore[unsupported-operation]
     )
     if self._top_sort_type == "max":
       max_evals = np.max(np.vstack([pos_evals, neg_evals]), axis=0)

iris/algorithms/multi_agent_ars_algorithm_test.py

@@ -224,7 +224,7 @@ def test_restore_state_from_checkpoint(
     self.assertEqual(algo._num_agents, num_agents)
     init_state = {'init_params': np.array([10.0, 10.0])}
     if state['obs_norm_state'] is not None:
-      init_state['obs_norm_buffer_data'] = {
+      init_state['obs_norm_buffer_data'] = {  # pyrefly: ignore[bad-assignment]
           'mean': np.asarray([0.0, 0.0]),
           'std': np.asarray([1.0, 1.0]),
           'n': 0,

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])
+    res_list.append(result[i])  # pyrefly: ignore[unsupported-operation]
   return np.array(res_list)
 
 
@@ -129,7 +129,7 @@ def general_jacobian_decoder(atranspose, yprime, optimization_parameters,
     list_res = []
     for j in range(n):
       list_res.append(res[j])
-    final_solutions.append(np.float32(list_res))
+    final_solutions.append(np.float32(list_res))  # pyrefly: ignore[bad-argument-type]
   return np.array(final_solutions)
 
 

iris/algorithms/pes_algorithm.py

@@ -111,7 +111,7 @@ def process_evaluations(
         pos_directions.append((params - self._opt_params) / self._std)
         pos_directions[-1] = self._positive_cumulative_perturbations[
             i] + pos_directions[-1]
-        if pos_eval_results[i].metrics["current_step"] == 0:
+        if pos_eval_results[i].metrics["current_step"] == 0:  # pyrefly: ignore[unsupported-operation]
           self._positive_cumulative_perturbations[i] = 0
         else:
           self._positive_cumulative_perturbations[i] = pos_directions[-1]
@@ -120,7 +120,7 @@ def process_evaluations(
         neg_directions.append((params - self._opt_params) / self._std)
         neg_directions[-1] = self._negative_cumulative_perturbations[
             i] + neg_directions[-1]
-        if neg_eval_results[i].metrics["current_step"] == 0:
+        if neg_eval_results[i].metrics["current_step"] == 0:  # pyrefly: ignore[unsupported-operation]
           self._negative_cumulative_perturbations[i] = 0
         else:
           self._negative_cumulative_perturbations[i] = neg_directions[-1]
@@ -136,7 +136,7 @@ def process_evaluations(
       max_evals = np.max(np.vstack([pos_evals, neg_evals]), axis=0)
     elif self._top_sort_type == "diff":
       max_evals = np.abs(pos_evals - neg_evals)
-    idx = (-max_evals).argsort()[:self._num_top]
+    idx = (-max_evals).argsort()[:self._num_top]  # pyrefly: ignore[unbound-name]
     pos_evals = pos_evals[idx]
     neg_evals = neg_evals[idx]
     all_top_evals = np.hstack([pos_evals, neg_evals])
@@ -212,7 +212,7 @@ def state(self) -> Dict[str, Any]:
   def _get_state(self) -> Dict[str, Any]:
     state = {"params_to_eval": self._opt_params}
     if self._obs_norm_data_buffer is not None:
-      state["obs_norm_state"] = self._obs_norm_data_buffer.state
+      state["obs_norm_state"] = self._obs_norm_data_buffer.state  # pyrefly: ignore[bad-assignment]
     return state
 
   @state.setter

iris/algorithms/piars_algorithm.py

@@ -207,10 +207,10 @@ def __init__(
     else:
       self.policy = policy
 
-    obs_spec = gym_wrapper.spec_from_gym_space(self._env.observation_space)
-    action_spec = gym_wrapper.spec_from_gym_space(self._env.action_space)
+    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]
     time_step_spec = ts.time_step_spec(observation_spec=obs_spec)
-    policy_step_spec = policy_step.PolicyStep(action=action_spec)
+    policy_step_spec = policy_step.PolicyStep(action=action_spec)  # pyrefly: ignore[missing-argument]
     collect_data_spec = trajectory.from_transition(
         time_step_spec, policy_step_spec, time_step_spec
     )
@@ -293,7 +293,7 @@ def train(self, obs_norm_state=None):
     if self.global_step % self.reverb_checkpoint_period == 0:
       logging.info("Start checkpointing reverb data.")
       self.reverb_rb.py_client.checkpoint()
-    print("train/loss: {}".format(np.mean(loss.numpy())))
+    print("train/loss: {}".format(np.mean(loss.numpy())))  # pyrefly: ignore[unbound-name]
 
   @tf.function
   def train_single_step(self, obs, reward, action, discount):
@@ -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)
+    s, _ = self.policy.h_model(obs)  # pyrefly: ignore[not-callable]
     outputs = []
     for i in range(self._rollout_length):
-      p, v = self.policy.f_model(s)
-      u_next, s_next = self.policy.g_model([s, actions[:, i, ...]])
+      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]
       outputs.append((p, v, u_next, s))
       s = s_next
-    p, v = self.policy.f_model(s)
+    p, v = self.policy.f_model(s)  # pyrefly: ignore[not-callable]
     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)
+      _, hy_vision = self.policy.h_model(obs_k)  # pyrefly: ignore[not-callable]
       # A trick from https://arxiv.org/abs/2011.10566
       hy_vision = tf.stop_gradient(hy_vision)
-      zx = self.policy.px_model(hx)
-      zy = self.policy.py_model(hy_vision)
+      zx = self.policy.px_model(hx)  # pyrefly: ignore[not-callable]
+      zy = self.policy.py_model(hy_vision)  # pyrefly: ignore[not-callable]
       iyz, _, _ = infonce(zx, zy, temperature=0.1)
       loss_pi = -iyz
 
@@ -406,11 +406,11 @@ def infonce(hidden_x, hidden_y, temperature=0.1):
         loss_v += self.distributional_value_loss(
             value_logits=z,
             value_supports=self.supports,
-            target_value_logits=last_value_distribution,
-            target_value_supports=target_value_supports[i],
+            target_value_logits=last_value_distribution,  # pyrefly: ignore[unbound-name]
+            target_value_supports=target_value_supports[i],  # pyrefly: ignore[unbound-name]
         )
         vd = tf.nn.softmax(z)
-        pred_value_sum += tf.reduce_sum(vd * self.supports[None, ...], axis=-1)
+        pred_value_sum += tf.reduce_sum(vd * self.supports[None, ...], axis=-1)  # pyrefly: ignore[unbound-name]
       # reward loss
       loss_r += tf.reduce_sum(
           tf.math.square(u_next - tf.stop_gradient(rewards[:, i : i + 1])), -1
@@ -424,9 +424,9 @@ def infonce(hidden_x, hidden_y, temperature=0.1):
         "loss_pi": tf.reduce_mean(loss_pi),
     }
     if self.use_value_loss:
-      metrics["value"] = tf.reduce_mean(pred_value_sum) / self._rollout_length
+      metrics["value"] = tf.reduce_mean(pred_value_sum) / self._rollout_length  # pyrefly: ignore[unbound-name]
     if self.use_pi_loss:
-      metrics["iyz"] = tf.reduce_mean(iyz)
+      metrics["iyz"] = tf.reduce_mean(iyz)  # pyrefly: ignore[unbound-name]
     return loss, metrics
 
   def distributional_value_loss(
@@ -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)
+    inner_dims = list(space.shape)  # pyrefly: ignore[bad-argument-type]
     outer_dims = list(x.shape)[: -len(inner_dims)]
     x = np.reshape(x, outer_dims + [np.prod(inner_dims)])
     return x

iris/algorithms/pyglove_algorithm.py

@@ -97,7 +97,7 @@ def get_param_suggestions(self,
 
     for metadata in metadata_list:
       suggestion = {"params_to_eval": np.empty((), dtype=np.float64)}
-      suggestion["metadata"] = metadata
+      suggestion["metadata"] = metadata  # pyrefly: ignore[bad-assignment]
       vanilla_suggestions.append(suggestion)
 
     return vanilla_suggestions

iris/algorithms/pyribs_algorithm.py

@@ -181,7 +181,7 @@ def get_param_suggestions(
           buffer_lib.STD: elite[_OBS_NORM_STD],
       }
     else:
-      param_suggestions = self._scheduler.ask()
+      param_suggestions = self._scheduler.ask()  # pyrefly: ignore[missing-attribute]
       buffer = self._obs_norm_data_buffer.state
 
     return [
@@ -202,12 +202,12 @@ def process_evaluations(
     obs_norm_std = []
     obs_norm_mean = []
     for result in eval_results:
-      self._obs_norm_data_buffer.merge(result.obs_norm_buffer_data)
+      self._obs_norm_data_buffer.merge(result.obs_norm_buffer_data)  # pyrefly: ignore[bad-argument-type]
       objective.append(result.value)
-      measures.append([result.metrics[name] for name in self._measure_names])
-      obs_norm_n.append(result.obs_norm_buffer_data[buffer_lib.N])
-      obs_norm_std.append(result.obs_norm_buffer_data[buffer_lib.STD])
-      obs_norm_mean.append(result.obs_norm_buffer_data[buffer_lib.MEAN])
+      measures.append([result.metrics[name] for name in self._measure_names])  # pyrefly: ignore[unsupported-operation]
+      obs_norm_n.append(result.obs_norm_buffer_data[buffer_lib.N])  # pyrefly: ignore[unsupported-operation]
+      obs_norm_std.append(result.obs_norm_buffer_data[buffer_lib.STD])  # pyrefly: ignore[unsupported-operation]
+      obs_norm_mean.append(result.obs_norm_buffer_data[buffer_lib.MEAN])  # pyrefly: ignore[unsupported-operation]
 
     # Store the state of the obs_norm_buffer for each solution so that it can be
     # reproduced later when evaluating the policy, similar to other algorithms
@@ -218,7 +218,7 @@ def process_evaluations(
         _OBS_NORM_N: obs_norm_n,
     }
 
-    self._scheduler.tell(
+    self._scheduler.tell(  # pyrefly: ignore[missing-attribute]
         objective=objective,
         measures=measures,
         **extra_fields,

iris/algorithms/pyribs_algorithm_test.py

@@ -102,7 +102,7 @@ def test_get_param_suggestions_for_eval(self):
     # Give the first evaluation a high score so it is the elite.
     evaluations[0].value = 1000
     if evaluations[0].obs_norm_buffer_data is not None:
-      evaluations[0].obs_norm_buffer_data[buffer.N] = 1000
+      evaluations[0].obs_norm_buffer_data[buffer.N] = 1000  # pyrefly: ignore[unsupported-operation]
     self.test_algorithm.process_evaluations(evaluations)
 
     eval_suggestions = self.test_algorithm.get_param_suggestions(evaluate=True)
@@ -115,7 +115,7 @@ def test_get_param_suggestions_for_eval(self):
       )
       np.testing.assert_equal(
           eval_suggestion[algorithm.OBS_NORM_BUFFER_STATE][buffer.N],
-          evaluations[0].obs_norm_buffer_data[buffer.N],
+          evaluations[0].obs_norm_buffer_data[buffer.N],  # pyrefly: ignore[unsupported-operation]
       )
       self.assertFalse(eval_suggestion[algorithm.UPDATE_OBS_NORM_BUFFER])
 
@@ -208,7 +208,7 @@ def test_process_evaluations(self):
         worker_util.EvaluationResult(
             params_evaluated=np.ones((13,)),
             value=1,
-            obs_norm_buffer_data={
+            obs_norm_buffer_data={  # pyrefly: ignore[bad-argument-type]
                 buffer.N: 1,
                 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={
+            obs_norm_buffer_data={  # pyrefly: ignore[bad-argument-type]
                 buffer.N: 2,
                 buffer.STD: np.ones((8,)) * 2,
                 buffer.MEAN: np.ones((8,)) * 2,

iris/buffer.py

@@ -197,7 +197,7 @@ def state(self, new_state: Dict[str, Any]) -> None:
   @property
   def _var(self) -> np.ndarray:
     return (
-        self._data[UNNORM_VAR] / (self._data[N] - 1)
+        self._data[UNNORM_VAR] / (self._data[N] - 1)  # pyrefly: ignore[bad-return]
         if self._data[N] > 1
         else np.ones_like(self._data[MEAN])
     )

iris/checkpoint_evaluator.py

@@ -49,7 +49,7 @@ def main(argv):
     worker_config.worker_args.write_to_replay = False
   worker = worker_config["worker_class"](
       worker_id=0, **worker_config["worker_args"])
-  state = checkpoint_util.load_checkpoint_state(_CHECKPOINT_FILE.value)
+  state = checkpoint_util.load_checkpoint_state(_CHECKPOINT_FILE.value)  # pyrefly: ignore[bad-argument-type]
   returns = []
   times = []
   metric_dict = collections.defaultdict(list)
@@ -61,7 +61,7 @@ def main(argv):
         gfile.MakeDirs(
             _VIDEO_PATH.value, mode=gfile.LEGACY_GROUP_WRITABLE_WORLD_READABLE
         )
-      video_path = os.path.join(_VIDEO_PATH.value, "video_" + str(i) + ".mp4")
+      video_path = os.path.join(_VIDEO_PATH.value, "video_" + str(i) + ".mp4")  # pyrefly: ignore[no-matching-overload]
       result = worker.work(
           **state,
           enable_logging=True,

iris/coordinator.py

@@ -476,7 +476,7 @@ def _restore_checkpoint(
           int(checkpoint_paths_sorted[0].split("_")[-1]) + 1
       )
       if latest_checkpoint_num_iterations > max_allowed_iteration_for_restart:
-        raise checkpoint_load_error
+        raise checkpoint_load_error  # pyrefly: ignore[bad-raise]
     return None, 0
 
   def evaluate(

iris/coordinator_rl_test.py

@@ -94,7 +94,7 @@ def make_bb_program(
       workers.append(worker_handle)
 
   if warmstartdir:
-    warmstartdir = pathlib.Path(warmstartdir)
+    warmstartdir = pathlib.Path(warmstartdir)  # pyrefly: ignore[bad-assignment]
   algo = algo_config["algorithm_class"](**algo_config["algorithm_args"])
 
   # Launches eval worker instances if there is at least one num_eval_workers.

iris/coordinator_test.py

@@ -69,7 +69,7 @@ def make_bb_program(
       workers.append(worker_handle)
 
   if warmstartdir:
-    warmstartdir = pathlib.Path(warmstartdir)
+    warmstartdir = pathlib.Path(warmstartdir)  # pyrefly: ignore[bad-assignment]
   algo = algo_config["algorithm_class"](**algo_config["algorithm_args"])
 
   # Launches eval worker instances if there is at least one num_eval_workers.