Improve Midscene a11y fallback handling

simplyme0823 · simplyme0823 · commit 7d8b8011e5b4 · 2026-05-25T14:23:50.000+08:00
diff --git a/android_world/agents/midscene.py b/android_world/agents/midscene.py
@@ -3,6 +3,7 @@
 
 
 from android_world.agents import base_agent
+from android_world.env import adb_utils
 from android_world.env import interface
 from android_world.env import representation_utils
 
@@ -160,21 +161,32 @@ def _write_response(self, status: int, content_type: str, body: str):
         self.end_headers()
         self.wfile.write(body.encode('utf-8'))
 
+      def _page_xml_from_adb_dump(self) -> str:
+        try:
+          page_xml = adb_utils.uiautomator_dump(agent_ref.env.controller)
+          if page_xml:
+            return page_xml
+        except Exception as e:
+          agent_ref._formatted_console(
+              'DOM provider adb dump fallback failed: ' + str(e)
+          )
+        return ''
+
       def _page_xml_from_state(self, state) -> str:
         if state.forest is None:
           agent_ref._formatted_console(
-              'DOM provider got empty AccessibilityForwarder forest; no page XML returned'
+              'DOM provider got empty AccessibilityForwarder forest; falling back to adb dump'
           )
-          return ''
+          return self._page_xml_from_adb_dump()
 
         page_xml = representation_utils.forest_to_raw_xml(state.forest)
         if page_xml:
           return page_xml
 
         agent_ref._formatted_console(
-            'DOM provider got empty AccessibilityForwarder XML; no page XML returned'
+            'DOM provider got empty AccessibilityForwarder XML; falling back to adb dump'
         )
-        return ''
+        return self._page_xml_from_adb_dump()
 
       def _get_state_with_retry(self):
         state = agent_ref.env.get_state()
diff --git a/android_world/env/android_world_controller.py b/android_world/env/android_world_controller.py
@@ -42,6 +42,14 @@
 # Throttle check_airplane_mode: at most once per interval to reduce ADB load
 _last_airplane_check: dict[int, float] = {}
 _AIRPLANE_CHECK_INTERVAL = 30.0
+_A11Y_FORWARDER_SERVICE = (
+    'com.google.androidenv.accessibilityforwarder/'
+    'com.google.androidenv.accessibilityforwarder.AccessibilityForwarder'
+)
+_A11Y_FORWARDER_FLAGS_RECEIVER = (
+    'com.google.androidenv.accessibilityforwarder/'
+    'com.google.androidenv.accessibilityforwarder.FlagsBroadcastReceiver'
+)
 
 
 def _has_wrapper(
@@ -109,7 +117,7 @@ def get_a11y_tree(
     try:
       forest = env.accumulate_new_extras()['accessibility_tree'][-1]  # pytype:disable=attribute-error
       return forest
-    except KeyError:
+    except (KeyError, IndexError):
       logging.warning('Could not get a11y tree, retrying.')
     time.sleep(sleep_duration)
 
@@ -454,18 +462,34 @@ def _restart_a11y_forwarder(self) -> bool:
       server_port = str(self._adb_server_port)
       device_args = ['-s', self._device_name] if self._device_name else []
 
-      # Step 1: Re-enable the accessibility service
-      cmd = [adb_path, '-P', server_port] + device_args + [
-          'shell', 'settings', 'put', 'secure',
-          'enabled_accessibility_services',
-          'com.google.androidenv.accessibilityforwarder/'
-          'com.google.androidenv.accessibilityforwarder.AccessibilityForwarder',
-      ]
-      result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
-      if result.returncode != 0:
-        logging.warning('Failed to re-enable a11y service: %s', result.stderr)
+      if self._is_remote and not self.ensure_adb_connection():
         return False
 
+      def run_shell_command(args: list[str]) -> subprocess.CompletedProcess:
+        cmd = [adb_path, '-P', server_port] + device_args + ['shell'] + args
+        return subprocess.run(cmd, capture_output=True, text=True, timeout=30)
+
+      # Step 1: Re-enable Android accessibility and the forwarder service.
+      # Some failures leave the service listed but global accessibility off.
+      for shell_args in (
+          ['settings', 'put', 'secure', 'accessibility_enabled', '1'],
+          [
+              'settings',
+              'put',
+              'secure',
+              'enabled_accessibility_services',
+              _A11Y_FORWARDER_SERVICE,
+          ],
+      ):
+        result = run_shell_command(shell_args)
+        if result.returncode != 0:
+          logging.warning(
+              'Failed to update a11y setting %s: %s',
+              ' '.join(shell_args),
+              result.stderr,
+          )
+          return False
+
       logging.info('Re-enabled AccessibilityForwarder service')
       time.sleep(2.0)  # Give the service time to start
 
@@ -480,8 +504,7 @@ def _restart_a11y_forwarder(self) -> bool:
             'shell', 'am', 'broadcast',
             '-a', 'accessibility_forwarder.intent.action.SET_GRPC',
             '--ei', 'port', str(self._a11y_port),
-            '-n', 'com.google.androidenv.accessibilityforwarder/'
-                  'com.google.androidenv.accessibilityforwarder.FlagsBroadcastReceiver',
+            '-n', _A11Y_FORWARDER_FLAGS_RECEIVER,
         ]
         subprocess.run(cmd, capture_output=True, text=True, timeout=30)
 
@@ -501,6 +524,8 @@ def get_a11y_forest(
       2. Restart AccessibilityForwarder service (handles uiautomator disruption)
       3. Full environment refresh (handles ADB disconnection / deep failures)
     """
+    self.ensure_adb_connection()
+
     try:
       return self._get_a11y_forest()
     except RuntimeError:
@@ -535,25 +560,51 @@ def get_ui_elements(self) -> list[representation_utils.UIElement]:
     self.ensure_adb_connection()
 
     if self._a11y_method == A11yMethod.A11Y_FORWARDER_APP:
-      return representation_utils.forest_to_ui_elements(
-          self.get_a11y_forest(),
-          exclude_invisible_elements=True,
-      )
+      try:
+        return representation_utils.forest_to_ui_elements(
+            self.get_a11y_forest(),
+            exclude_invisible_elements=True,
+        )
+      except RuntimeError as e:
+        logging.warning(
+            'A11y tree unavailable after recovery; falling back to '
+            'uiautomator UI elements: %s',
+            e,
+        )
+        return self._get_uiautomator_ui_elements()
     elif self._a11y_method == A11yMethod.UIAUTOMATOR:
+      return self._get_uiautomator_ui_elements()
+    else:
+      return []
+
+  def _get_uiautomator_ui_elements(self) -> list[representation_utils.UIElement]:
+    """Returns UI elements from uiautomator, or an empty list if it fails."""
+    try:
       return representation_utils.xml_dump_to_ui_elements(
           adb_utils.uiautomator_dump(self._env)
       )
-    else:
+    except Exception as e:
+      logging.warning('Failed to get UI elements via uiautomator: %s', e)
       return []
 
   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
     """Adds a11y tree info to the observation."""
     if self._a11y_method == A11yMethod.A11Y_FORWARDER_APP:
-      forest = self.get_a11y_forest()
-      ui_elements = representation_utils.forest_to_ui_elements(
-          forest,
-          exclude_invisible_elements=True,
-      )
+      try:
+        forest = self.get_a11y_forest()
+      except RuntimeError as e:
+        logging.warning(
+            'A11y tree unavailable after recovery; falling back to '
+            'uiautomator UI elements: %s',
+            e,
+        )
+        forest = None
+        ui_elements = self._get_uiautomator_ui_elements()
+      else:
+        ui_elements = representation_utils.forest_to_ui_elements(
+            forest,
+            exclude_invisible_elements=True,
+        )
     else:
       forest = None
       ui_elements = self.get_ui_elements()
diff --git a/android_world/env/android_world_controller_test.py b/android_world/env/android_world_controller_test.py
@@ -130,6 +130,36 @@ def test_process_timestep(
         exclude_invisible_elements=True,
     )
 
+  @mock.patch.object(adb_utils, 'check_airplane_mode')
+  @mock.patch.object(android_world_controller, 'get_controller')
+  @mock.patch.object(android_world_controller, '_has_wrapper')
+  @mock.patch.object(representation_utils, 'forest_to_ui_elements')
+  def test_process_timestep_continues_without_a11y_tree(
+      self,
+      mock_forest_to_ui,
+      mock_has_wrapper,
+      mock_get_controller,
+      mock_check_airplane_mode,
+  ):
+    del mock_has_wrapper, mock_get_controller, mock_check_airplane_mode
+    mock_base_env = mock.Mock(spec=env_interface.AndroidEnvInterface)
+    env = android_world_controller.AndroidWorldController(mock_base_env)
+    timestep = dm_env.TimeStep(
+        observation={}, reward=None, discount=None, step_type=None
+    )
+
+    with mock.patch.object(
+        env, 'get_a11y_forest', side_effect=RuntimeError('a11y unavailable')
+    ), mock.patch.object(
+        env, '_get_uiautomator_ui_elements', return_value=['fallback']
+    ) as mock_uiautomator_fallback:
+      processed_timestep = env._process_timestep(timestep)
+
+    self.assertIsNone(processed_timestep.observation['forest'])
+    self.assertEqual(processed_timestep.observation['ui_elements'], ['fallback'])
+    mock_uiautomator_fallback.assert_called_once()
+    mock_forest_to_ui.assert_not_called()
+
   @mock.patch.object(adb_utils, 'check_airplane_mode')
   @mock.patch.object(android_world_controller, 'get_controller')
   @mock.patch.object(android_world_controller, '_has_wrapper')
