Update callers for Gym 0.26.2 compatibility.

iris/algorithms/piars_algorithm_test.py

@@ -28,7 +28,7 @@ class AlgorithmTest(parameterized.TestCase):
       (False, False),
   )
   def test_ars_gradient(self, orthogonal_suggestions, quasirandom_suggestions):
-    env = gym.make(id='Pendulum-v0')
+    env = gym.make(id='Pendulum-v1')
     policy = keras_pi_policy.KerasPIPolicy(
         ob_space=env.observation_space,
         ac_space=env.action_space,

iris/coordinator_rl_test.py

@@ -50,10 +50,11 @@ def __init__(self):
 
   def step(self, action):
     del action
-    return np.zeros(self._ob_dim), 1.0, False, {}
+    return np.zeros(self._ob_dim), 1.0, False, False, {}
 
-  def reset(self):
-    return np.zeros(self._ob_dim)
+  def reset(self, *, seed=None, options=None):
+    del seed, options
+    return np.zeros(self._ob_dim), {}
 
   def render(self, mode: str = "rgb_array"):
     return np.zeros((16, 16))

iris/normalizer.py

@@ -119,8 +119,8 @@ def __init__(
       self, space: gym.Space, ignored_keys: Optional[Sequence[str]] = None
   ) -> None:
     super().__init__(space, ignored_keys)
-    low = self._flat_space.low
-    high = self._flat_space.high
+    low = self._flat_space.low  # pyrefly: ignore[missing-attribute]
+    high = self._flat_space.high  # pyrefly: ignore[missing-attribute]
     self._state["mid"] = (low + high) / 2.0
     self._state["half_range"] = (high - low) / 2.0
 
@@ -145,9 +145,9 @@ def __call__(
     del update_buffer  # No buffer to update
     action = action.copy()
     ignored_action = self._filter_ignored_input(action)  # pyrefly: ignore[bad-argument-type]
-    action = utils.flatten(self._space, action)
+    action = utils.flatten(self._space, action)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     action = (action * self._state["half_range"]) + self._state["mid"]
-    action = utils.unflatten(self._space, action)
+    action = utils.unflatten(self._space, action)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     action = self._add_ignored_input(action, ignored_action)  # pyrefly: ignore[bad-argument-type]
     return action
 
@@ -159,8 +159,8 @@ def __init__(
       self, space: gym.Space, ignored_keys: Optional[Sequence[str]] = None
   ) -> None:
     super().__init__(space, ignored_keys)
-    low = self._flat_space.low
-    high = self._flat_space.high
+    low = self._flat_space.low  # pyrefly: ignore[missing-attribute]
+    high = self._flat_space.high  # pyrefly: ignore[missing-attribute]
     self._state["mid"] = (low + high) / 2.0
     self._state["half_range"] = (high - low) / 2.0
 
@@ -186,9 +186,9 @@ def __call__(
     observation = observation.copy()
     ignored_observation = self._filter_ignored_input(observation)  # pyrefly: ignore[bad-argument-type]
 
-    observation = utils.flatten(self._space, observation)
+    observation = utils.flatten(self._space, observation)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     observation = (observation - self._state["mid"]) / self._state["half_range"]
-    observation = utils.unflatten(self._space, observation)
+    observation = utils.unflatten(self._space, observation)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     observation = self._add_ignored_input(observation, ignored_observation)  # pyrefly: ignore[bad-argument-type]
     return observation
 
@@ -201,9 +201,9 @@ def __init__(
   ) -> None:
     super().__init__(space, ignored_keys)
     shape = self._flat_space.shape
-    self._state[buffer.MEAN] = np.zeros(shape, dtype=np.float64)
-    self._state[buffer.STD] = np.ones(shape, dtype=np.float64)
-    self._buffer = buffer.MeanStdBuffer(shape)
+    self._state[buffer.MEAN] = np.zeros(shape, dtype=np.float64)  # pyrefly: ignore[no-matching-overload]
+    self._state[buffer.STD] = np.ones(shape, dtype=np.float64)  # pyrefly: ignore[no-matching-overload]
+    self._buffer = buffer.MeanStdBuffer(shape)  # pyrefly: ignore[bad-argument-type]
 
   @property
   def buffer(self) -> buffer.MeanStdBuffer:
@@ -216,12 +216,12 @@ def __call__(
   ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
     observation = observation.copy()
     ignored_observation = self._filter_ignored_input(observation)  # pyrefly: ignore[bad-argument-type]
-    observation = utils.flatten(self._space, observation)
+    observation = utils.flatten(self._space, observation)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     if update_buffer:
       self._buffer.push(observation)  # pyrefly: ignore[bad-argument-type]
     observation -= self._state[buffer.MEAN]
     observation /= self._state[buffer.STD] + _EPSILON
-    observation = utils.unflatten(self._space, observation)
+    observation = utils.unflatten(self._space, observation)  # pyrefly: ignore[bad-argument-type, bad-assignment]
     observation = self._add_ignored_input(observation, ignored_observation)  # pyrefly: ignore[bad-argument-type]
     return observation
 
@@ -232,7 +232,7 @@ class RunningMeanStdAgentVsAgentNormalizer(RunningMeanStdNormalizer):
   def __init__(self, space: gym.Space) -> None:
     # We use the "ignored_keys" to split the agent obs to process individually.
     super().__init__(space, ignored_keys=["opp"])
-    self._buffer = buffer.MeanStdBuffer(shape=self._flat_space.shape)
+    self._buffer = buffer.MeanStdBuffer(shape=self._flat_space.shape)  # pyrefly: ignore[bad-argument-type]
 
   @property
   def buffer(self) -> buffer.MeanStdBuffer:
@@ -246,14 +246,14 @@ def __call__(
     observation = observation.copy()
 
     opp_observation = self._filter_ignored_input(observation)  # pyrefly: ignore[bad-argument-type]
-    opp_observation = utils.flatten(self._space_ignored, opp_observation)
+    opp_observation = utils.flatten(self._space_ignored, opp_observation)  # pyrefly: ignore[bad-argument-type]
 
-    arm_observation = utils.flatten(self._space, observation)
+    arm_observation = utils.flatten(self._space, observation)  # pyrefly: ignore[bad-argument-type]
     if update_buffer:
-      self._buffer.push(arm_observation)
+      self._buffer.push(arm_observation)  # pyrefly: ignore[bad-argument-type]
       # Adding the opponent observations to the buffer greatly speeds learning
       # in self-play scenarios.
-      self._buffer.push(opp_observation)
+      self._buffer.push(opp_observation)  # pyrefly: ignore[bad-argument-type]
 
     def _normalized(obs, unflatten_space):
       obs -= self._state[buffer.MEAN]

iris/policies/gym_space_utils.py

@@ -23,13 +23,13 @@ def filter_space(space: gym.Space,
                  selected: Sequence[Union[str, int]]) -> gym.Space:
   """Filters gym space based on list of selected keys or dimensions."""
   if isinstance(space, gym.spaces.Box):
-    low = space.low.take(selected)
-    high = space.high.take(selected)
+    low = space.low.take(selected)  # pyrefly: ignore[no-matching-overload]
+    high = space.high.take(selected)  # pyrefly: ignore[no-matching-overload]
     return gym.spaces.Box(low=low, high=high)
   elif isinstance(space, gym.spaces.Dict):
     filtered_space = {}
     for sensor in selected:
-      filtered_space[sensor] = space[sensor]
+      filtered_space[sensor] = space[sensor]  # pyrefly: ignore[bad-index]
     return gym.spaces.Dict(filtered_space)
   else:
     raise NotImplementedError
@@ -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/implicit_policy.py

@@ -563,7 +563,7 @@ def act(
     Returns:
       The actions in reinforcement learning.
     """
-    ob = utils.flatten(self._ob_space, ob)
+    ob = utils.flatten(self._ob_space, ob)  # pyrefly: ignore[bad-assignment]
     action = self._action_calculator.act(ob)  # pyrefly: ignore[bad-argument-type]
     action = utils.unflatten(self._ac_space, action)
     return action

iris/policies/implicit_policy_test.py

@@ -29,9 +29,9 @@ def test_policy_act(self):
         'sensor2': spaces.Box(low=-10, high=10, shape=(2,)),
     })
     ac_space = spaces.Box(low=-10, high=10, shape=(3,))
-    action_low = utils.flatten_space(ac_space).low
-    action_high = utils.flatten_space(ac_space).high
-    state_dim = len(utils.flatten_space(ob_space).low)
+    action_low = utils.flatten_space(ac_space).low  # pyrefly: ignore[missing-attribute]
+    action_high = utils.flatten_space(ac_space).high  # pyrefly: ignore[missing-attribute]
+    state_dim = len(utils.flatten_space(ob_space).low)  # pyrefly: ignore[missing-attribute]
     action_dim = len(action_low)
     latent_dim = 30
     st_hidden_layer_sizes = [100, 100]
@@ -67,9 +67,9 @@ def test_policy_update_get_weights(self):
     })
     ac_space = spaces.Box(low=15, high=16, shape=(10,))
 
-    action_low = utils.flatten_space(ac_space).low
-    action_high = utils.flatten_space(ac_space).high
-    state_dim = len(utils.flatten_space(ob_space).low)
+    action_low = utils.flatten_space(ac_space).low  # pyrefly: ignore[missing-attribute]
+    action_high = utils.flatten_space(ac_space).high  # pyrefly: ignore[missing-attribute]
+    state_dim = len(utils.flatten_space(ob_space).low)  # pyrefly: ignore[missing-attribute]
     action_dim = len(action_low)
     latent_dim = 54
     st_hidden_layer_sizes = [45, 55]

iris/policies/jax_policy.py

@@ -87,7 +87,7 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
     Returns:
       The action in reinforcement learning.
     """
-    ob = utils.flatten(self._ob_space, ob)
+    ob = utils.flatten(self._ob_space, ob)  # pyrefly: ignore[bad-assignment]
     action = self.model.apply(
         self._tree_weights,
         [ob],

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/policies/keras_policy.py

@@ -118,7 +118,7 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
     Returns:
       The actions in reinforcement learning.
     """
-    ob = utils.flatten(self._ob_space, ob)
-    actions = self.model(np.array([ob])).numpy()[0]
+    ob = utils.flatten(self._ob_space, ob)  # pyrefly: ignore[bad-assignment]
+    actions = self.model(np.array([ob])).numpy()[0]  # pyrefly: ignore[not-callable]
     actions = utils.unflatten(self._ac_space, actions)
     return actions

iris/policies/linear_policy.py

@@ -49,7 +49,7 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
     Returns:
       The actions in reinforcement learning.
     """
-    ob = utils.flatten(self._ob_space, ob)
+    ob = utils.flatten(self._ob_space, ob)  # pyrefly: ignore[bad-assignment]
     matrix_weights = np.reshape(self._weights, (self._ac_dim, self._ob_dim))
     actions = self._activation(np.dot(matrix_weights, ob))  # pyrefly: ignore[no-matching-overload, not-callable]
     actions = utils.unflatten(self._ac_space, actions)

iris/policies/nas_policy.py

@@ -79,10 +79,10 @@ def __init__(
   def act(
       self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
   ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
-    ob = utils.flatten(self._ob_space, ob)
+    ob = utils.flatten(self._ob_space, ob)  # pyrefly: ignore[bad-assignment]
     values = [0.0] * self._total_nb_nodes
     for i in range(self._ob_dim):
-      values[i] = ob[i]  # pyrefly: ignore[bad-index]
+      values[i] = ob[i]  # pyrefly: ignore[bad-index, unsupported-operation]
     for i in range(self._total_nb_nodes):
       if (i > self._ob_dim) and (i < self._total_nb_nodes - self._ac_dim):
         values[i] = np.tanh(values[i] + self._biases[i])

iris/policies/nn_policy.py

@@ -61,7 +61,7 @@ def act(self, ob: Union[np.ndarray, Dict[str, np.ndarray]]
     Returns:
       The actions in reinforcement learning.
     """
-    ob = utils.flatten(self._ob_space, ob)
+    ob = utils.flatten(self._ob_space, ob)  # pyrefly: ignore[bad-assignment]
     ith_layer_result = ob
     num_layers = len(self._layer_sizes)
     for ith_layer in range(num_layers - 1):

iris/workers/maml_worker_test.py

@@ -158,7 +158,7 @@ class MamlWorkerTest(parameterized.TestCase):
       (maml_worker.HillClimbAdaptation,), (maml_worker.GradientAdaptation,)
   )
   def test_maml_rl_worker(self, adaptation_cls):
-    env = gym.make(id='Pendulum-v0')
+    env = gym.make(id='Pendulum-v1')
     policy = linear_policy.LinearPolicy(
         ob_space=env.observation_space, ac_space=env.action_space
     )

iris/workers/multi_agent_rl_worker.py

@@ -85,9 +85,21 @@ def work(  # pyrefly: ignore[bad-override]
 
     if self._step == 0:
       self._policy.reset()
-      if env_seed is not None:
+      if env_seed is not None and hasattr(self._env, "seed"):
         self._env.seed(env_seed)
-      self._obs = self._env.reset()
+      reset_res = (
+          self._env.reset(seed=env_seed)
+          if env_seed is not None and not hasattr(self._env, "seed")
+          else self._env.reset()
+      )
+      if (
+          isinstance(reset_res, tuple)
+          and len(reset_res) == 2
+          and isinstance(reset_res[1], dict)
+      ):
+        self._obs = reset_res[0]
+      else:
+        self._obs = reset_res
 
     if partial_rollout_length is None:
       partial_rollout_length = self._rollout_length
@@ -96,7 +108,12 @@ def work(  # pyrefly: ignore[bad-override]
           self._obs, update_obs_norm_buffer
       )
       action = self._action_denormalizer(self._policy.act(normalized_obs))
-      next_obs, r, done, info = self._env.step(action)
+      step_res = self._env.step(action)
+      if len(step_res) == 5:
+        next_obs, r, terminated, truncated, info = step_res
+        done = terminated or truncated
+      else:
+        next_obs, r, done, info = step_res
 
       for rkey, rval in r.items():
         if rkey not in reward_dict:
@@ -124,7 +141,11 @@ def work(  # pyrefly: ignore[bad-override]
       self._obs = next_obs
 
       if record_video and video is not None:
-        video.add_frame(self._env.render(mode="rgb_array"))
+        try:
+          frame = self._env.render(mode="rgb_array")
+        except TypeError:
+          frame = self._env.render()
+        video.add_frame(frame)
       if enable_logging:
         logging.info("Step: %d, Reward: %s, Done: %d", self._step, r, done)
         if mdict:

iris/workers/pyglove_rl_worker_test.py

@@ -24,7 +24,7 @@
 class PygloveRlWorkerTest(absltest.TestCase):
 
   def test_pyglove_rl_worker(self):
-    env = gym.make(id='Pendulum-v0')
+    env = gym.make(id='Pendulum-v1')
     policy = nas_policy.NumpyEdgeSparsityPolicy(
         ob_space=env.observation_space,
         ac_space=env.action_space,