Restore pyrefly typecheck ignore comments for Gym 0.19.0 compatibility and update actions/setup-python to v5 in Iris GitHub workflow.

.github/workflows/core_test.yml

@@ -9,7 +9,7 @@
       - main
 
 jobs:
   test-ubuntu:
     name: "pytest on ${{ matrix.python-version }} on ${{ matrix.os }}"
     runs-on: "${{ matrix.os }}"
     strategy:
@@ -17,9 +17,9 @@
         python-version: ['3.10']
         os: [ubuntu-latest]
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install essential dependencies

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)
-    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)  # pyrefly: ignore[missing-argument]
     collect_data_spec = trajectory.from_transition(
@@ -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/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)

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)

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)

iris/workers/rl_representation_worker.py

@@ -56,8 +56,8 @@ def __init__(
     if reverb_client is not None:
       self._init_state["reverb_server_addr"] = reverb_client.server_address
 
-    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)  # pyrefly: ignore[missing-argument]
     collect_data_spec = trajectory.from_transition(

iris/workers/rl_worker.py

@@ -124,13 +124,13 @@ def __init__(
 
     if not isinstance(observation_normalizer, normalizer.Normalizer):
       self._observation_normalizer = observation_normalizer(
-          self._env.observation_space
+          self._env.observation_space  # pyrefly: ignore[bad-argument-type]
       )
     else:
       self._observation_normalizer = observation_normalizer
 
     if not isinstance(action_denormalizer, normalizer.Normalizer):
-      self._action_denormalizer = action_denormalizer(self._env.action_space)
+      self._action_denormalizer = action_denormalizer(self._env.action_space)  # pyrefly: ignore[bad-argument-type]
     else:
       self._action_denormalizer = action_denormalizer