From c2c35a204bd93582deb0fcd3330eef8d113fd711 Mon Sep 17 00:00:00 2001 From: Daniel Toyama Date: Wed, 22 Jul 2026 04:17:58 -0700 Subject: [PATCH] Fix bad-state simulator relaunch via `task_manager.is_healthy()`. Fix GitHub Issue #390 where `TaskManagerConfig.max_bad_states` set a write-only `_should_restart` flag that was never checked by `Coordinator`, resulting in logs claiming the simulator was restarting when no actual relaunch occurred. Add an `is_healthy()` query method to `TaskManager` that returns `False` when `max_bad_states` is exceeded, and reset bad states on `TaskManager.start()`. Update `Coordinator.rl_reset()` to check `not self._task_manager.is_healthy()` alongside simulator health before relaunching. PiperOrigin-RevId: 952027934 --- android_env/components/coordinator.py | 7 ++- android_env/components/coordinator_test.py | 62 +++++++++++++++++++++ android_env/components/task_manager.py | 16 +++++- android_env/components/task_manager_test.py | 25 +++++++++ 4 files changed, 106 insertions(+), 4 deletions(-) diff --git a/android_env/components/coordinator.py b/android_env/components/coordinator.py index c409633e..af3f0f77 100644 --- a/android_env/components/coordinator.py +++ b/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. diff --git a/android_env/components/coordinator_test.py b/android_env/components/coordinator_test.py index ac28711a..7bd86231 100644 --- a/android_env/components/coordinator_test.py +++ b/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() diff --git a/android_env/components/task_manager.py b/android_env/components/task_manager.py index 34144ecf..4b873a37 100644 --- a/android_env/components/task_manager.py +++ b/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.') diff --git a/android_env/components/task_manager_test.py b/android_env/components/task_manager_test.py index e53e5dd2..3ab43a9c 100644 --- a/android_env/components/task_manager_test.py +++ b/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()