Fix bad-state simulator relaunch via task_manager.is_healthy().

android_env/components/coordinator.py

@@ -191,8 +191,11 @@ def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:
   def rl_reset(self) -> dm_env.TimeStep:
     """Resets the RL episode."""
 
-    # Relaunch the simulator if necessary.
-    if not self._simulator_healthy or self._should_periodic_relaunch():
+    if (
+        not self._simulator_healthy
+        or not self._task_manager.is_healthy()
+        or self._should_periodic_relaunch()
+    ):
       self._launch_simulator()
 
     # Reset counters.

android_env/components/coordinator_test.py

@@ -26,7 +26,9 @@
 from android_env.components import config_classes
 from android_env.components import coordinator as coordinator_lib
 from android_env.components import device_settings as device_settings_lib
+from android_env.components import dumpsys_thread
 from android_env.components import errors
+from android_env.components import logcat_thread
 from android_env.components import task_manager
 from android_env.components.simulators import base_simulator
 from android_env.proto import adb_pb2
@@ -278,6 +280,66 @@ def test_execute_adb_call(self, unused_mock_sleep):
     self.assertEqual(response, expected_response)
     self._adb_call_parser.parse.assert_called_with(call)
 
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_reset_unhealthy_task_manager(self, unused_mock_sleep):
+    """rl_reset should relaunch simulator if task_manager is unhealthy."""
+    self._task_manager.is_healthy.return_value = False
+    relaunch_count = self._coordinator.stats()['relaunch_count']
+    self._coordinator.rl_reset()
+    self.assertEqual(
+        self._coordinator.stats()['relaunch_count'], relaunch_count + 1
+    )
+
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_max_bad_states_triggers_relaunch(self, unused_mock_sleep):
+    """Verifies that consecutive bad states cause Coordinator to relaunch simulator."""
+    dumpsys = mock.create_autospec(dumpsys_thread.DumpsysThread)
+    dumpsys.check_user_exited.return_value = True
+    self.enter_context(
+        mock.patch.object(
+            dumpsys_thread, 'DumpsysThread', return_value=dumpsys
+        )
+    )
+    self.enter_context(
+        mock.patch.object(
+            logcat_thread, 'LogcatThread', autospec=True
+        )
+    )
+
+    real_task_manager = task_manager.TaskManager(
+        task=task_pb2.Task(),
+        config=config_classes.TaskManagerConfig(max_bad_states=2),
+    )
+    coordinator = coordinator_lib.Coordinator(
+        simulator=self._simulator,
+        task_manager=real_task_manager,
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+    )
+    self.addCleanup(coordinator.close)
+
+    # Initial launch.
+    coordinator._launch_simulator()
+    self._simulator.launch.reset_mock()
+
+    # Episode 1 bad state: user exits task.
+    coordinator.rl_reset()
+    coordinator.rl_step({
+        'action_type': np.array(action_type.ActionType.TOUCH),
+        'touch_position': np.array([0.5, 0.5]),
+    })
+    # After 1st bad state, simulator is NOT relaunched yet.
+    coordinator.rl_reset()
+    self._simulator.launch.assert_not_called()
+
+    # Episode 2 bad state: user exits task again (2 >= max_bad_states=2).
+    coordinator.rl_step({
+        'action_type': np.array(action_type.ActionType.TOUCH),
+        'touch_position': np.array([0.5, 0.5]),
+    })
+    # Now max_bad_states is exceeded. Reset MUST trigger _launch_simulator().
+    coordinator.rl_reset()
+    self._simulator.launch.assert_called_once()
+
 
 if __name__ == '__main__':
   absltest.main()

android_env/components/task_manager.py

@@ -104,6 +104,12 @@ def stats(self) -> dict[str, Any]:
       output.update(self._setup_step_interpreter.stats())
     return output
 
+  def is_healthy(self) -> bool:
+    """Returns True if task manager has not exceeded max bad states."""
+    if self._config.max_bad_states:
+      return self._bad_state_counter < self._config.max_bad_states
+    return True
+
   def setup_task(self) -> None:
     """Performs one-off task setup.."""
     assert self._setup_step_interpreter is not None, (
@@ -132,6 +138,8 @@ def start(
   ) -> None:
     """Starts task processing."""
 
+    self._bad_state_counter = 0
+    self._is_bad_episode = False
     self._start_logcat_thread(log_stream=log_stream)
     assert (
         self._logcat_thread is not None
@@ -316,9 +324,13 @@ def _increment_bad_state(self) -> None:
       self._bad_state_counter += 1
       logging.warning('Bad state counter: %d.', self._bad_state_counter)
       if self._bad_state_counter >= self._config.max_bad_states:
-        logging.error('Too many consecutive bad states. Restarting simulator.')
+        logging.error(
+            'Too many consecutive bad states (%d >= %d). Task manager is'
+            ' unhealthy.',
+            self._bad_state_counter,
+            self._config.max_bad_states,
+        )
         self._stats['restart_count_max_bad_states'] += 1
-        self._should_restart = True
     else:
       logging.warning('Max bad states not set, bad states will be ignored.')
 

android_env/components/task_manager_test.py

@@ -530,6 +530,31 @@ def my_add_ev_listener(event_listener: logcat_thread.EventListener):
     extras = timestep.observation['extras']
     np.testing.assert_almost_equal([2, 3], extras['overflow_extra'])
 
+  def test_is_healthy_and_max_bad_states(self):
+    config = config_classes.TaskManagerConfig(max_bad_states=2)
+    task_mgr = task_manager.TaskManager(task=task_pb2.Task(), config=config)
+    adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
+
+    # Initially healthy.
+    self.assertTrue(task_mgr.is_healthy())
+
+    # User exits once (1 bad state).
+    self._dumpsys_thread.check_user_exited.return_value = True
+    task_mgr.rl_reset(observation={})
+    task_mgr.rl_step(observation={})
+    self.assertTrue(task_mgr.is_healthy())
+
+    # User exits second time (2 bad states >= max_bad_states).
+    task_mgr.rl_reset(observation={})
+    task_mgr.rl_step(observation={})
+    self.assertFalse(task_mgr.is_healthy())
+
+    # Start clears bad state counter, restoring health.
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    self.assertTrue(task_mgr.is_healthy())
+
 
 if __name__ == '__main__':
   absltest.main()