Fix that UIAutomator may not return output dump on a real device if t…

…he device screen is off.

PiperOrigin-RevId: 688618828

android_world/agents/random_agent.py

@@ -68,6 +68,8 @@ def _generate_random_action(
       action_details['text'] = ''.join(
           random.choices(text_characters, k=10)
       )  # Random text of length 10
+    elif action_type == json_action.SWIPE:
+      action_details['direction'] = random.choice(scroll_directions)
   elif action_type == json_action.SCROLL:
     action_details['direction'] = random.choice(scroll_directions)
 

android_world/env/adb_utils.py

@@ -136,6 +136,8 @@
     'gallery': 'content://media/external/images/media/',
 }
 
+_DEFAULT_DUMP_FILE = '/sdcard/window_dump.xml'
+
 
 def check_ok(response: adb_pb2.AdbResponse, message=None) -> None:
   """Check an ADB response and raise RuntimeError if not OK.
@@ -1633,12 +1635,76 @@ def set_root_if_needed(
   return issue_generic_request(['root'], env, timeout_sec)
 
 
+def _dump_ui_tree(
+    env: env_interface.AndroidEnvInterface,
+) -> adb_pb2.AdbResponse:
+  """Dumps the UI tree from the device."""
+  dump_request = adb_pb2.AdbRequest(
+      uiautomator=adb_pb2.AdbRequest.UIAutomatorRequest(file=_DEFAULT_DUMP_FILE)
+  )
+  logging.info('Dumping a11y tree from UIAutomator.')
+  dump_response = env.execute_adb_call(dump_request)
+  logging.debug('Dump response: %s', dump_response)
+  return dump_response
+
+
+def switch_on_device_screen(
+    env: env_interface.AndroidEnvInterface,
+) -> adb_pb2.AdbResponse:
+  """Switches on the device screen."""
+  screen_on_request = adb_pb2.AdbRequest(
+      generic=adb_pb2.AdbRequest.GenericRequest(
+          args=['shell', 'input', 'keyevent', 'KEYCODE_POWER']
+      )
+  )
+  logging.info('Switching on the device screen.')
+  screen_on_response = env.execute_adb_call(screen_on_request)
+  return screen_on_response
+
+
+def _read_ui_tree_from_device(
+    env: env_interface.AndroidEnvInterface,
+) -> adb_pb2.AdbResponse:
+  """Reads the UI tree from the device."""
+  tree_request = adb_pb2.AdbRequest(
+      generic=adb_pb2.AdbRequest.GenericRequest(
+          args=['shell', 'cat', _DEFAULT_DUMP_FILE]
+      )
+  )
+  logging.info(
+      'Reading a11y tree from the dump file: %s.', _DEFAULT_DUMP_FILE
+  )
+  tree_response = env.execute_adb_call(tree_request)
+  logging.debug('Tree response: %s', tree_response)
+  return tree_response
+
+
 def uiautomator_dump(env) -> str:
   """Issues a uiautomator dump request and returns the UI hierarchy."""
-  dump_args = 'shell uiautomator dump /sdcard/window_dump.xml'
-  issue_generic_request(dump_args, env)
 
-  read_args = 'shell cat /sdcard/window_dump.xml'
-  response = issue_generic_request(read_args, env)
+  # UIAutomator may not return output if the device screen is off.
+  retries = 1
+  dump_response = _dump_ui_tree(env)
+  for _ in range(retries):
+    if dump_response.status != adb_pb2.AdbResponse.Status.OK:
+      # If fails, try to switch on the device screen.
+      _ = switch_on_device_screen(env)
+      dump_response = _dump_ui_tree(env)
+    else:
+      break
+
+  if dump_response.status != adb_pb2.AdbResponse.Status.OK:
+    raise ValueError(
+        f'Could not dump a11y tree from UIAutomator: {dump_response}.'
+    )
+
+  # If the ui is successfully dumped, it will be in a file
+  # /sdcard/window_dump.xml. Run another adb shell command to get the file
+  # content.
+  tree_response = _read_ui_tree_from_device(env)
+  if tree_response.status != adb_pb2.AdbResponse.Status.OK:
+    raise ValueError(
+        f'Could not read a11y tree from dump file: {tree_response}.'
+    )
 
-  return response.generic.output.decode('utf-8')
+  return tree_response.generic.output.decode('utf-8')

android_world/env/android_world_controller.py

@@ -287,22 +287,39 @@ def get_controller(
     console_port: int = 5554,
     adb_path: str = DEFAULT_ADB_PATH,
     grpc_port: int = 8554,
+    real_device_name: str | None = None,
 ) -> AndroidWorldController:
   """Creates a controller by connecting to an existing Android environment."""
 
-  config = config_classes.AndroidEnvConfig(
-      task=config_classes.FilesystemTaskConfig(
-          path=_write_default_task_proto()
-      ),
-      simulator=config_classes.EmulatorConfig(
-          emulator_launcher=config_classes.EmulatorLauncherConfig(
-              emulator_console_port=console_port,
-              adb_port=console_port + 1,
-              grpc_port=grpc_port,
-          ),
-          adb_controller=config_classes.AdbControllerConfig(adb_path=adb_path),
-      ),
-  )
+  if real_device_name:
+    config = config_classes.AndroidEnvConfig(
+        task=config_classes.FilesystemTaskConfig(
+            path=_write_default_task_proto()
+        ),
+        simulator=config_classes.RealDeviceConfig(
+            device_name=real_device_name,
+            adb_controller_config=config_classes.AdbControllerConfig(
+                adb_path=adb_path,
+                device_name=real_device_name,
+            ),
+        ),
+    )
+  else:
+    config = config_classes.AndroidEnvConfig(
+        task=config_classes.FilesystemTaskConfig(
+            path=_write_default_task_proto()
+        ),
+        simulator=config_classes.EmulatorConfig(
+            emulator_launcher=config_classes.EmulatorLauncherConfig(
+                emulator_console_port=console_port,
+                adb_port=console_port + 1,
+                grpc_port=grpc_port,
+            ),
+            adb_controller=config_classes.AdbControllerConfig(
+                adb_path=adb_path
+            ),
+        ),
+    )
   android_env_instance = loader.load(config)
   logging.info('Setting up AndroidWorldController.')
   return AndroidWorldController(android_env_instance)

android_world/env/env_launcher.py

@@ -30,11 +30,14 @@
 
 
 def _get_env(
-    console_port: int, adb_path: str, grpc_port: int
+    console_port: int,
+    adb_path: str,
+    grpc_port: int,
+    real_device_name: str | None = None,
 ) -> interface.AsyncEnv:
   """Creates an AsyncEnv by connecting to an existing Android environment."""
   controller = android_world_controller.get_controller(
-      console_port, adb_path, grpc_port
+      console_port, adb_path, grpc_port, real_device_name
   )
   return interface.AsyncAndroidEnv(controller)
 
@@ -99,6 +102,7 @@ def load_and_setup_env(
     freeze_datetime: bool = True,
     adb_path: str = android_world_controller.DEFAULT_ADB_PATH,
     grpc_port: int = 8554,
+    real_device_name: str | None = None,
 ) -> interface.AsyncEnv:
   """Create environment with `get_env()` and perform env setup and validation.
 
@@ -118,10 +122,11 @@ def load_and_setup_env(
       2023, to ensure consistent benchmarking.
     adb_path: The location of the adb binary.
     grpc_port: The port for gRPC communication with the emulator.
+    real_device_name: The name of the real device to use.
 
   Returns:
     An interactable Android environment.
   """
-  env = _get_env(console_port, adb_path, grpc_port)
+  env = _get_env(console_port, adb_path, grpc_port, real_device_name)
   setup_env(env, emulator_setup, freeze_datetime)
   return env

run.py

@@ -149,6 +149,12 @@ def _find_adb_directory() -> str:
     ),
 ]
 
+_REAL_DEVICE_NAME = flags.DEFINE_string(
+    'real_device_name',
+    None,
+    'The real device name as shown in `adb devices` to run the agent on.',
+)
+
 
 def _get_agent(
     env: interface.AsyncEnv,
@@ -200,6 +206,7 @@ def _main() -> None:
       console_port=_DEVICE_CONSOLE_PORT.value,
       emulator_setup=_EMULATOR_SETUP.value,
       adb_path=_ADB_PATH.value,
+      real_device_name=_REAL_DEVICE_NAME.value,
   )
   env_launcher.verify_api_level(env)