From e30d8099f0c582a9786b67095b0d1e6de2421207 Mon Sep 17 00:00:00 2001 From: The android_world Authors Date: Tue, 22 Oct 2024 11:11:43 -0700 Subject: [PATCH] Fix that UIAutomator may not return output dump on a real device if the device screen is off. PiperOrigin-RevId: 688618828 --- android_world/agents/random_agent.py | 2 + android_world/env/adb_utils.py | 76 +++++++++++++++++-- android_world/env/android_world_controller.py | 43 +++++++---- android_world/env/env_launcher.py | 11 ++- run.py | 7 ++ 5 files changed, 118 insertions(+), 21 deletions(-) diff --git a/android_world/agents/random_agent.py b/android_world/agents/random_agent.py index 785ea0c..d045052 100644 --- a/android_world/agents/random_agent.py +++ b/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) diff --git a/android_world/env/adb_utils.py b/android_world/env/adb_utils.py index 29afc82..3608874 100644 --- a/android_world/env/adb_utils.py +++ b/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') diff --git a/android_world/env/android_world_controller.py b/android_world/env/android_world_controller.py index bd11ac5..01fa187 100644 --- a/android_world/env/android_world_controller.py +++ b/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) diff --git a/android_world/env/env_launcher.py b/android_world/env/env_launcher.py index 168ad00..64e056c 100644 --- a/android_world/env/env_launcher.py +++ b/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 diff --git a/run.py b/run.py index dcaa4e0..33139e5 100644 --- a/run.py +++ b/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)