diff --git a/iris/algorithms/piars_algorithm_test.py b/iris/algorithms/piars_algorithm_test.py
index cbbdaf3..5979bd3 100644
--- a/iris/algorithms/piars_algorithm_test.py
+++ b/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,
diff --git a/iris/coordinator_rl_test.py b/iris/coordinator_rl_test.py
index 4ab6d3e..3498b5f 100644
--- a/iris/coordinator_rl_test.py
+++ b/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))
diff --git a/iris/normalizer.py b/iris/normalizer.py
index 9067c34..e855b98 100644
--- a/iris/normalizer.py
+++ b/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]
diff --git a/iris/policies/gym_space_utils.py b/iris/policies/gym_space_utils.py
index 8a03baf..22e6aea 100644
--- a/iris/policies/gym_space_utils.py
+++ b/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)
diff --git a/iris/policies/implicit_policy.py b/iris/policies/implicit_policy.py
index da2e8ea..37ab0f9 100644
--- a/iris/policies/implicit_policy.py
+++ b/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
diff --git a/iris/policies/implicit_policy_test.py b/iris/policies/implicit_policy_test.py
index 65f84c8..451fdda 100644
--- a/iris/policies/implicit_policy_test.py
+++ b/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]
diff --git a/iris/policies/jax_policy.py b/iris/policies/jax_policy.py
index e46a743..630bb40 100644
--- a/iris/policies/jax_policy.py
+++ b/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],
diff --git a/iris/policies/keras_cnn_policy.py b/iris/policies/keras_cnn_policy.py
index 630a640..87deca4 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)
diff --git a/iris/policies/keras_pi_policy.py b/iris/policies/keras_pi_policy.py
index 3bd54e1..a2d2aa3 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)
diff --git a/iris/policies/keras_policy.py b/iris/policies/keras_policy.py
index a797e12..337bd56 100644
--- a/iris/policies/keras_policy.py
+++ b/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
diff --git a/iris/policies/linear_policy.py b/iris/policies/linear_policy.py
index becfb4c..6df2230 100644
--- a/iris/policies/linear_policy.py
+++ b/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)
diff --git a/iris/policies/nas_policy.py b/iris/policies/nas_policy.py
index 7054a06..cb3a5fb 100644
--- a/iris/policies/nas_policy.py
+++ b/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])
diff --git a/iris/policies/nn_policy.py b/iris/policies/nn_policy.py
index 9ba0623..79ef3b4 100644
--- a/iris/policies/nn_policy.py
+++ b/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):
diff --git a/iris/workers/maml_worker_test.py b/iris/workers/maml_worker_test.py
index 58dc3dc..02a56df 100644
--- a/iris/workers/maml_worker_test.py
+++ b/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
     )
diff --git a/iris/workers/multi_agent_rl_worker.py b/iris/workers/multi_agent_rl_worker.py
index 67f42e6..12e42d6 100644
--- a/iris/workers/multi_agent_rl_worker.py
+++ b/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:
diff --git a/iris/workers/pyglove_rl_worker_test.py b/iris/workers/pyglove_rl_worker_test.py
index 177558b..f1ac2fd 100644
--- a/iris/workers/pyglove_rl_worker_test.py
+++ b/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,
diff --git a/iris/workers/rl_representation_worker.py b/iris/workers/rl_representation_worker.py
index 703abd3..ca3a500 100644
--- a/iris/workers/rl_representation_worker.py
+++ b/iris/workers/rl_representation_worker.py
@@ -113,7 +113,7 @@ def work(  # pyrefly: ignore[bad-override]
     if obs_norm_state is not None:
       self._observation_normalizer.state = obs_norm_state  # pyrefly: ignore[bad-argument-type]
 
-    if env_seed is not None:
+    if env_seed is not None and hasattr(self._env, "seed"):
       self._env.seed(env_seed)
 
     video = None
@@ -122,7 +122,19 @@ def work(  # pyrefly: ignore[bad-override]
 
     reward = 0.0
     metrics = collections.defaultdict(float)
-    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)
+    ):
+      obs = reset_res[0]
+    else:
+      obs = reset_res
     time_step = ts.TimeStep(
         step_type=ts.StepType.FIRST,
         reward=np.array(0.0, dtype=np.float32),
@@ -136,7 +148,12 @@ def work(  # pyrefly: ignore[bad-override]
       action = self._policy.act(normalized_obs)
       action_step = policy_step.PolicyStep(action)  # pyrefly: ignore[missing-argument]
       action = self._action_denormalizer(action)
-      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
       reward += r
       next_time_step = ts.TimeStep(
           step_type=ts.StepType.LAST if done else ts.StepType.MID,
@@ -166,13 +183,18 @@ def work(  # pyrefly: ignore[bad-override]
       time_step = next_time_step
 
       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: %f, Done: %d", st, r, done)
         if mdict:
           logging.info("Metrics: %s", mdict)
       if done:
-        next_obs, r, _, _ = self._env.step(action)
+        step_res = self._env.step(action)
+        next_obs = step_res[0]
         next_time_step = ts.TimeStep(
             step_type=ts.StepType.FIRST,
             reward=np.array(0.0, dtype=np.float32),
diff --git a/iris/workers/rl_representation_worker_test.py b/iris/workers/rl_representation_worker_test.py
index c96e349..c5a5f15 100644
--- a/iris/workers/rl_representation_worker_test.py
+++ b/iris/workers/rl_representation_worker_test.py
@@ -24,7 +24,7 @@
 class RlRepresentationWorkerTest(absltest.TestCase):
 
   def test_rl_representation_worker(self):
-    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
     )
diff --git a/iris/workers/rl_worker.py b/iris/workers/rl_worker.py
index f2ff307..45cebee 100644
--- a/iris/workers/rl_worker.py
+++ b/iris/workers/rl_worker.py
@@ -59,7 +59,7 @@ def __init__(
       rollout_length: int = 500,
       send_training_state_to_env: bool = False,
       metrics_fn: Optional[
-          Callable[[gym.Env, Dict[str, Any]], Dict[str, worker.FloatLike]]
+          Callable[[gym.Env, Dict[str, Any]], Mapping[str, worker.FloatLike]]
       ] = lambda a, b: {},
       stats_fn: Optional[
           Callable[
@@ -104,6 +104,7 @@ def __init__(
     if policy_args is None:
       policy_args = {}
 
+    self._env: Any
     if not isinstance(env, gym.Env):
       self._env = env(**env_args)
     else:
@@ -208,6 +209,7 @@ def work(  # pyrefly: ignore[bad-override]
       num_rollouts += 1
 
     if self._retry_rollout and not valid_rollout:
+      workerlock.release()
       raise RolloutRetryError(
           f"Rollout error: max_rollout_retries: {self._max_rollout_retries}"
       )
@@ -262,9 +264,21 @@ def _run_rollout(
 
     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
@@ -274,7 +288,12 @@ def _run_rollout(
           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
       rewards.append(r)
 
       step_output = {
@@ -293,7 +312,11 @@ def _run_rollout(
       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: %f, Done: %d", self._step, r, done)
         if mdict:
diff --git a/iris/workers/rl_worker_test.py b/iris/workers/rl_worker_test.py
index 19766db..f678076 100644
--- a/iris/workers/rl_worker_test.py
+++ b/iris/workers/rl_worker_test.py
@@ -26,7 +26,7 @@
 class RlWorkerTest(parameterized.TestCase):
 
   def test_rl_worker(self):
-    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
     )
@@ -68,8 +68,8 @@ def stats_fn(name, values):
     self.assertIn('Num invalid rollouts: 0', logs.output[-1])
     self.assertLen(logs.output, 401)
     self.assertLessEqual(result1.value, 0)
-    self.assertEqual(result1.metrics['extra_metric'], 1.0)
-    self.assertLessEqual(result1.metrics['summed_extra_metric'], 200.0)
+    self.assertEqual(result1.metrics['extra_metric'], 1.0)  # pyrefly: ignore[unsupported-operation]
+    self.assertLessEqual(result1.metrics['summed_extra_metric'], 200.0)  # pyrefly: ignore[unsupported-operation]
 
     result2 = worker_obj.work(
         params_to_eval=np.ones(3),
@@ -103,8 +103,13 @@ def __init__(self, env: gym.Env):
         self._max_rollouts = max_rollouts
         self._num_rollouts = 0
 
-      def step(self, action):
-        obs, reward, done, info = super().step(action)
+      def step(self, action):  # pyrefly: ignore[bad-override]
+        step_res = super().step(action)
+        if len(step_res) == 5:
+          obs, reward, terminated, truncated, info = step_res
+          done = terminated or truncated
+        else:
+          obs, reward, done, info = step_res
         if done:
           self._num_rollouts += 1
         if done and self._retry_rollout_info_flag:
@@ -113,7 +118,7 @@ def step(self, action):
           )
         return obs, reward, done, info
 
-    env = gym.make(id='Pendulum-v0')
+    env = gym.make(id='Pendulum-v1')
     env = _RolloutEnv(env)
     policy = linear_policy.LinearPolicy(
         ob_space=env.observation_space, ac_space=env.action_space
@@ -156,7 +161,7 @@ def step(self, action):
       )
       self.assertLen(logs.output, 400 * expected_rollouts + 1)
       self.assertLessEqual(result1.value, 0)
-      self.assertLessEqual(result1.metrics['extra_metric'], 500.0)
+      self.assertLessEqual(result1.metrics['extra_metric'], 500.0)  # pyrefly: ignore[unsupported-operation]
       self.assertEqual(env._num_rollouts, expected_rollouts)
 
       result2 = worker_obj.work(
@@ -168,7 +173,7 @@ def step(self, action):
 
   @parameterized.parameters((None, 0), (0, 0), (1, 1), (2, 2), (3, 3))
   def test_rl_worker_with_iteration(self, iteration, expected_iteration):
-    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
     )
@@ -219,7 +224,7 @@ def act(
       ) -> np.ndarray | dict[str, np.ndarray]:
         return self._ac_space.sample()
 
-    env = gym.make(id='Pendulum-v0')
+    env = gym.make(id='Pendulum-v1')
     policy = _DummyPolicy(
         ob_space=env.observation_space, ac_space=env.action_space
     )