Make close() functions more resilient to errors.

Failing one resource cleanup shouldn't prevent cleaning up other resources.

PiperOrigin-RevId: 778861249

Author: DeepMind

android_env/components/coordinator.py

@@ -21,12 +21,10 @@
 
 from absl import logging
 from android_env.components import action_fns
-from android_env.components import action_type as action_type_lib
 from android_env.components import adb_call_parser
 from android_env.components import config_classes
 from android_env.components import device_settings as device_settings_lib
 from android_env.components import errors
-from android_env.components import pixel_fns
 from android_env.components import specs
 from android_env.components import task_manager as task_manager_lib
 from android_env.components.simulators import base_simulator
@@ -277,6 +275,12 @@ def close(self):
     """Cleans up the state of this Coordinator."""
 
     if hasattr(self, '_task_manager'):
-      self._task_manager.stop()
+      try:
+        self._task_manager.stop()
+      except:  # pylint: disable=bare-except
+        logging.exception('Failed to stop task manager. Continuing.')
     if hasattr(self, '_simulator'):
-      self._simulator.close()
+      try:
+        self._simulator.close()
+      except:  # pylint: disable=bare-except
+        logging.exception('Failed to close simulator. Continuing.')

android_env/components/task_manager.py

@@ -23,6 +23,7 @@
 import json
 import re
 import threading
+import time
 from typing import Any
 
 from absl import logging
@@ -101,12 +102,22 @@ def setup_task(self) -> None:
 
   def stop(self) -> None:
     """Suspends task processing."""
-    self._stop_logcat_thread()
+    n_tries = 3
+    for i in range(n_tries):
+      try:
+        self._stop_logcat_thread()
+        break
+      except:  # pylint: disable=bare-except
+        logging.exception(
+            'Failed to stop logcat thread [%d/%d]. Continuing.', i + 1, n_tries
+        )
+        time.sleep(1)
 
   def start(
       self,
       adb_call_parser_factory: Callable[[], adb_call_parser_lib.AdbCallParser],
-      log_stream: log_stream_lib.LogStream) -> None:
+      log_stream: log_stream_lib.LogStream,
+  ) -> None:
     """Starts task processing."""
 
     self._start_logcat_thread(log_stream=log_stream)
@@ -150,7 +161,8 @@ def rl_reset(self, observation: dict[str, Any]) -> dm_env.TimeStep:
         step_type=dm_env.StepType.FIRST,
         reward=0.0,
         discount=0.0,
-        observation=observation)
+        observation=observation,
+    )
 
   def rl_step(self, observation: dict[str, Any]) -> dm_env.TimeStep:
     """Performs one RL step."""
@@ -202,25 +214,31 @@ def _determine_transition_fn(self) -> Callable[..., dm_env.TimeStep]:
     if self._task.max_episode_steps > 0:
       if self._stats['episode_steps'] > self._task.max_episode_steps:
         self._stats['reset_count_max_duration_reached'] += 1
-        logging.info('Maximum task duration (%r steps) reached. '
-                     'Truncating the episode.', self._task.max_episode_steps)
+        logging.info(
+            'Maximum task duration (%r steps) reached. Truncating the episode.',
+            self._task.max_episode_steps,
+        )
         return dm_env.truncation
 
     if self._task.max_episode_sec > 0.0:
       task_duration = datetime.datetime.now() - self._task_start_time
       max_episode_sec = self._task.max_episode_sec
       if task_duration > datetime.timedelta(seconds=int(max_episode_sec)):
         self._stats['reset_count_max_duration_reached'] += 1
-        logging.info('Maximum task duration (%r sec) reached. '
-                     'Truncating the episode.', max_episode_sec)
+        logging.info(
+            'Maximum task duration (%r sec) reached. Truncating the episode.',
+            max_episode_sec,
+        )
         return dm_env.truncation
 
     return dm_env.transition
 
   def _start_setup_step_interpreter(
-      self, adb_call_parser: adb_call_parser_lib.AdbCallParser):
+      self, adb_call_parser: adb_call_parser_lib.AdbCallParser
+  ):
     self._setup_step_interpreter = setup_step_interpreter.SetupStepInterpreter(
-        adb_call_parser=adb_call_parser)
+        adb_call_parser=adb_call_parser
+    )
 
   def _start_logcat_thread(self, log_stream: log_stream_lib.LogStream):
     log_stream.set_log_filters(list(self._task.log_parsing_config.filters))
@@ -229,8 +247,9 @@ def _start_logcat_thread(self, log_stream: log_stream_lib.LogStream):
     for event_listener in self._logcat_listeners():
       self._logcat_thread.add_event_listener(event_listener)
 
-  def _start_dumpsys_thread(self,
-                            adb_call_parser: adb_call_parser_lib.AdbCallParser):
+  def _start_dumpsys_thread(
+      self, adb_call_parser: adb_call_parser_lib.AdbCallParser
+  ):
     self._dumpsys_thread = dumpsys_thread.DumpsysThread(
         app_screen_checker=app_screen_checker.AppScreenChecker(
             adb_call_parser=adb_call_parser,

android_env/env_interface.py

@@ -25,6 +25,7 @@
 from android_env.proto import adb_pb2
 from android_env.proto import state_pb2
 import dm_env
+from dm_env import specs
 import numpy as np
 
 
@@ -34,11 +35,11 @@ class AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):
   # Methods required by dm_env.Environment.
 
   @abc.abstractmethod
-  def action_spec(self) -> dict[str, dm_env.specs.Array]:
+  def action_spec(self) -> dict[str, specs.Array]:
     """Returns the action specification."""
 
   @abc.abstractmethod
-  def observation_spec(self) -> dict[str, dm_env.specs.Array]:
+  def observation_spec(self) -> dict[str, specs.Array]:
     """Returns the observation specification."""
 
   @abc.abstractmethod
@@ -61,11 +62,11 @@ def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:
     return {}
 
   @property
-  def raw_action(self):
+  def raw_action(self) -> Any:
     """Returns the latest action."""
 
   @property
-  def raw_observation(self):
+  def raw_observation(self) -> Any:
     """Returns the latest observation."""
 
   def stats(self) -> dict[str, Any]:

android_env/wrappers/base_wrapper.py

@@ -29,7 +29,7 @@
 class BaseWrapper(env_interface.AndroidEnvInterface):
   """AndroidEnv wrapper."""
 
-  def __init__(self, env):
+  def __init__(self, env: env_interface.AndroidEnvInterface) -> None:
     self._env = env
     logging.info('Wrapping with %s', self.__class__.__name__)
 
@@ -96,29 +96,30 @@ def save_state(
     """
     return self._env.save_state(request)
 
-  def execute_adb_call(self,
-                       adb_call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:
+  def execute_adb_call(
+      self, adb_call: adb_pb2.AdbRequest
+  ) -> adb_pb2.AdbResponse:
     return self._env.execute_adb_call(adb_call)
 
   @property
-  def raw_action(self):
+  def raw_action(self) -> Any:
     return self._env.raw_action
 
   @property
-  def raw_observation(self):
+  def raw_observation(self) -> Any:
     return self._env.raw_observation
 
   @property
-  def raw_env(self):
+  def raw_env(self) -> env_interface.AndroidEnvInterface:
     """Recursively unwrap until we reach the true 'raw' env."""
     wrapped = self._env
     if hasattr(wrapped, 'raw_env'):
       return wrapped.raw_env
     return wrapped
 
-  def __getattr__(self, attr):
+  def __getattr__(self, attr) -> Any:
     """Delegate attribute access to underlying environment."""
     return getattr(self._env, attr)
 
-  def close(self):
+  def close(self) -> None:
     self._env.close()

android_env/wrappers/base_wrapper_test.py

@@ -101,7 +101,7 @@ def test_multiple_wrappers(self):
     base_env.close.assert_called_once()
 
   def test_raw_env(self):
-    base_env = 'fake_env'
+    base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     wrapped_env_1 = base_wrapper.BaseWrapper(base_env)
     wrapped_env_2 = base_wrapper.BaseWrapper(wrapped_env_1)
     self.assertEqual(base_env, wrapped_env_2.raw_env)

android_env/wrappers/discrete_action_wrapper.py

@@ -16,7 +16,9 @@
 """Wraps the AndroidEnv environment to provide discrete actions."""
 
 from collections.abc import Sequence
+from typing import cast
 
+from android_env import env_interface
 from android_env.components import action_type
 from android_env.wrappers import base_wrapper
 import dm_env
@@ -32,21 +34,24 @@ class DiscreteActionWrapper(base_wrapper.BaseWrapper):
 
   def __init__(
       self,
-      env: dm_env.Environment,
+      env: env_interface.AndroidEnvInterface,
       action_grid: Sequence[int] = (10, 10),
       redundant_actions: bool = True,
       noise: float = 0.1,
-  ):
+  ) -> None:
     super().__init__(env)
     self._parent_action_spec = self._env.action_spec()
     self._assert_base_env()
     self._action_grid = action_grid  # [height, width]
     self._grid_size = np.prod(self._action_grid)
-    self._num_action_types = self._parent_action_spec['action_type'].num_values
+    action_types = cast(
+        specs.DiscreteArray, self._parent_action_spec['action_type']
+    )
+    self._num_action_types = action_types.num_values
     self._redundant_actions = redundant_actions
     self._noise = noise
 
-  def _assert_base_env(self):
+  def _assert_base_env(self) -> None:
     """Checks that the wrapped env has the right action spec format."""
 
     assert len(self._parent_action_spec) == 2
@@ -142,8 +147,12 @@ def _get_touch_position(self, action_id: int) -> Sequence[float]:
 
     # Project action space to action_spec ranges. For the default case of
     # minimum = [0, 0] and maximum = [1, 1], this will not do anything.
-    x_min, y_min = self._parent_action_spec['touch_position'].minimum
-    x_max, y_max = self._parent_action_spec['touch_position'].maximum
+    x_min, y_min = cast(
+        specs.BoundedArray, self._parent_action_spec['touch_position']
+    ).minimum
+    x_max, y_max = cast(
+        specs.BoundedArray, self._parent_action_spec['touch_position']
+    ).maximum
 
     x_pos = x_min + x_pos * (x_max - x_min)
     y_pos = y_min + y_pos * (y_max - y_min)

android_env/wrappers/flat_interface_wrapper.py

@@ -15,8 +15,9 @@
 
 """Wraps the AndroidEnv environment to make its interface flat."""
 
-from typing import Any
+from typing import Any, cast
 
+from android_env import env_interface
 from android_env.wrappers import base_wrapper
 import dm_env
 from dm_env import specs
@@ -25,15 +26,17 @@
 RGB_CHANNELS = (0, 1, 2)
 
 
-def _extract_screen_pixels(obs: np.ndarray):
+def _extract_screen_pixels(obs: np.ndarray) -> np.ndarray:
   """Get only screen pixels by removing previous action layer."""
   is_grayscale_image = obs.shape[-1] == 2
   if is_grayscale_image:
     return np.expand_dims(obs[..., 0], -1)
   return obs[..., RGB_CHANNELS]
 
 
-def _get_no_action_observation_spec(obs_spec: specs.BoundedArray):
+def _get_no_action_observation_spec(
+    obs_spec: specs.BoundedArray,
+) -> specs.BoundedArray:
   """Create an observation spec without the action layer."""
   shape = np.array(obs_spec.shape)
   shape[2] -= 1
@@ -56,25 +59,29 @@ class FlatInterfaceWrapper(base_wrapper.BaseWrapper):
   space.
   """
 
-  def __init__(self,
-               env: dm_env.Environment,
-               flat_actions: bool = True,
-               flat_observations: bool = True,
-               keep_action_layer: bool = True):
+  def __init__(
+      self,
+      env: env_interface.AndroidEnvInterface,
+      flat_actions: bool = True,
+      flat_observations: bool = True,
+      keep_action_layer: bool = True,
+  ) -> None:
     super().__init__(env)
     self._flat_actions = flat_actions
     self._flat_observations = flat_observations
     self._keep_action_layer = keep_action_layer
     self._action_name = list(self._env.action_spec())[0]
     self._assert_base_env()
 
-  def _assert_base_env(self):
+  def _assert_base_env(self) -> None:
     base_action_spec = self._env.action_spec()
     assert len(base_action_spec) == 1, self._env.action_spec()
     assert isinstance(base_action_spec, dict)
     assert isinstance(base_action_spec[self._action_name], specs.BoundedArray)
 
-  def _process_action(self, action: int | np.ndarray | dict[str, Any]):
+  def _process_action(
+      self, action: int | np.ndarray | dict[str, Any]
+  ) -> int | np.ndarray | dict[str, Any]:
     if self._flat_actions:
       return {self._action_name: action}
     else:
@@ -85,13 +92,15 @@ def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
       step_type, reward, discount, observation = timestep
       # Keep only the pixels.
       pixels = observation['pixels']
-      pixels = pixels if self._keep_action_layer else _extract_screen_pixels(
-          pixels)
+      pixels = (
+          pixels if self._keep_action_layer else _extract_screen_pixels(pixels)
+      )
       return dm_env.TimeStep(
           step_type=step_type,
           reward=reward,
           discount=discount,
-          observation=pixels)
+          observation=pixels,
+      )
     else:
       return timestep
 
@@ -105,7 +114,9 @@ def step(self, action: int) -> dm_env.TimeStep:
 
   def observation_spec(self) -> specs.Array | dict[str, specs.Array]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
     if self._flat_observations:
-      pixels_spec = self._env.observation_spec()['pixels']
+      pixels_spec = cast(
+          specs.BoundedArray, self._env.observation_spec()['pixels']
+      )
       if not self._keep_action_layer:
         return _get_no_action_observation_spec(pixels_spec)
       return pixels_spec
@@ -114,6 +125,6 @@ def observation_spec(self) -> specs.Array | dict[str, specs.Array]:  # pytype: d
 
   def action_spec(self) -> specs.BoundedArray | dict[str, specs.Array]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
     if self._flat_actions:
-      return self._env.action_spec()[self._action_name]
+      return self._env.action_spec()[self._action_name]  # pytype: disable=bad-return-type
     else:
       return self._env.action_spec()

android_env/wrappers/float_pixels_wrapper.py

@@ -15,6 +15,7 @@
 
 """Converts pixel observation to from int to float32 between 0.0 and 1.0."""
 
+from android_env import env_interface
 from android_env.components import pixel_fns
 from android_env.wrappers import base_wrapper
 import dm_env
@@ -25,7 +26,7 @@
 class FloatPixelsWrapper(base_wrapper.BaseWrapper):
   """Wraps AndroidEnv for Panultimate agent."""
 
-  def __init__(self, env: dm_env.Environment):
+  def __init__(self, env: env_interface.AndroidEnvInterface) -> None:
     super().__init__(env)
     self._input_spec = self._env.observation_spec()['pixels']
     self._should_convert_int_to_float = np.issubdtype(self._input_spec.dtype,