android-env 1.2.2__py3-none-any.whl → 1.2.3__py3-none-any.whl

android_env/components/adb_controller.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,34 +18,20 @@
 import os
 import subprocess
 import time
-from typing import Optional
 
 from absl import logging
+from android_env.components import config_classes
 from android_env.components import errors
 
 
 class AdbController:
   """Manages communication with adb."""
 
-  def __init__(self,
-               device_name: str = '',
-               adb_path: str = 'adb',
-               adb_server_port: int = 5037,
-               default_timeout: float = 120.0):
-    """Instantiates an AdbController object.
+  def __init__(self, config: config_classes.AdbControllerConfig):
+    """Instantiates an AdbController object."""
 
-    Args:
-      device_name: Name of the device to communicate with.
-      adb_path: Path to the adb binary.
-      adb_server_port: Port for adb server.
-      default_timeout: Default timeout in seconds.
-    """
-
-    self._device_name = device_name
-    self._adb_path = os.path.expandvars(adb_path)
-    self._adb_server_port = str(adb_server_port)
-    self._default_timeout = default_timeout
-    logging.info('adb_path: %r', self._adb_path)
+    self._config = config
+    logging.info('config: %r', self._config)
 
     # Unset problematic environment variables. ADB commands will fail if these
     # are set. They are normally exported by AndroidStudio.
@@ -63,12 +49,16 @@ class AdbController:
 
   def command_prefix(self, include_device_name: bool = True) -> list[str]:
     """The command for instantiating an adb client to this server."""
-    command_prefix = [self._adb_path, '-P', self._adb_server_port]
+    command_prefix = [
+        self._config.adb_path,
+        '-P',
+        str(self._config.adb_server_port),
+    ]
     if include_device_name:
-      command_prefix.extend(['-s', self._device_name])
+      command_prefix.extend(['-s', self._config.device_name])
     return command_prefix
 
-  def init_server(self, timeout: Optional[float] = None):
+  def init_server(self, timeout: float | None = None):
     """Initialize the ADB server deamon on the given port.
 
     This function should be called immediately after initializing the first
@@ -82,7 +72,7 @@ class AdbController:
     self.execute_command(['devices'], timeout, device_specific=False)
     time.sleep(0.2)
 
-  def _restart_server(self, timeout: Optional[float] = None):
+  def _restart_server(self, timeout: float | None = None):
     """Kills and restarts the adb server.
 
     Args:
@@ -104,7 +94,7 @@ class AdbController:
   def execute_command(
       self,
       args: list[str],
-      timeout: Optional[float] = None,
+      timeout: float | None = None,
       device_specific: bool = True,
   ) -> bytes:
     """Executes an adb command.
@@ -119,13 +109,14 @@ class AdbController:
     Returns:
       The output of running such command as a binary string.
     """
-    timeout = self._default_timeout if timeout is None else timeout
+    timeout = self._config.default_timeout if timeout is None else timeout
     command = self.command_prefix(include_device_name=device_specific) + args
     command_str = 'adb ' + ' '.join(command[1:])
 
+    n_retries = 2
     n_tries = 1
     latest_error = None
-    while n_tries < 3:
+    while n_tries <= n_retries:
       try:
         logging.info('Executing ADB command: [%s]', command_str)
         cmd_output = subprocess.check_output(
@@ -150,7 +141,7 @@ class AdbController:
             logging.error('    %s', line)
         n_tries += 1
         latest_error = e
-        if device_specific:
+        if device_specific and n_tries <= n_retries:
           self._restart_server(timeout=timeout)
 
     raise errors.AdbControllerError(

android_env/components/adb_controller_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@ from unittest import mock
 
 from absl.testing import absltest
 from android_env.components import adb_controller as adb_controller_lib
+from android_env.components import config_classes
 from android_env.components import errors
 
 # Timeout to be used by default in tests below. Set to a small value to avoid
@@ -38,14 +39,24 @@ class AdbControllerTest(absltest.TestCase):
     os.environ['HOME'] = '$MY_ENV_VAR'
     self._env_before = os.environ
     self._adb_controller = adb_controller_lib.AdbController(
-        adb_path='my_adb', device_name='awesome_device', adb_server_port=9999)
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
 
   @mock.patch.object(subprocess, 'check_output', autospec=True)
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_init_server(self, mock_sleep, mock_check_output):
     # Arrange.
     adb_controller = adb_controller_lib.AdbController(
-        adb_path='my_adb', device_name='awesome_device', adb_server_port=9999)
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
 
     # Act.
     adb_controller.init_server(timeout=_TIMEOUT)
@@ -69,7 +80,12 @@ class AdbControllerTest(absltest.TestCase):
         subprocess.CalledProcessError(returncode=1, cmd='blah'),
     ] + ['fake_output'.encode('utf-8')] * 4
     adb_controller = adb_controller_lib.AdbController(
-        adb_path='my_adb', device_name='awesome_device', adb_server_port=9999)
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
 
     # Act.
     adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)
@@ -112,6 +128,73 @@ class AdbControllerTest(absltest.TestCase):
     mock_sleep.assert_has_calls(
         [mock.call(0.2), mock.call(2.0), mock.call(0.2)])
 
+  @mock.patch.object(subprocess, 'check_output', autospec=True)
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_invalid_command(self, mock_sleep, mock_check_output):
+    # Arrange.
+    restart_sequence = ['fake_output'.encode('utf-8')] * 3
+    mock_check_output.side_effect = (
+        [
+            subprocess.CalledProcessError(returncode=1, cmd='blah'),
+        ]
+        + restart_sequence
+        + [subprocess.CalledProcessError(returncode=1, cmd='blah')]
+        # Don't restart if last call fails.
+    )
+    adb_controller = adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
+
+    # Act.
+    with self.assertRaises(errors.AdbControllerError):
+      adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)
+
+    # Assert.
+    expected_env = self._env_before
+    expected_env['HOME'] = '/some/path/'
+    mock_check_output.assert_has_calls(
+        [
+            mock.call(
+                ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'kill-server'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'start-server'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'devices'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+        ],
+        any_order=False,
+    )
+    mock_sleep.assert_has_calls(
+        [mock.call(0.2), mock.call(2.0), mock.call(0.2)]
+    )
+
   @mock.patch.object(subprocess, 'check_output', autospec=True)
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_avoid_infinite_recursion(self, mock_sleep, mock_check_output):
@@ -119,7 +202,12 @@ class AdbControllerTest(absltest.TestCase):
     mock_check_output.side_effect = subprocess.CalledProcessError(
         returncode=1, cmd='blah')
     adb_controller = adb_controller_lib.AdbController(
-        adb_path='my_adb', device_name='awesome_device', adb_server_port=9999)
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
     self.assertRaises(
         errors.AdbControllerError,
         adb_controller.execute_command, ['my_command'], timeout=_TIMEOUT)
@@ -131,10 +219,13 @@ class AdbControllerInitTest(absltest.TestCase):
     os.environ['ANDROID_HOME'] = '/usr/local/Android/Sdk'
     os.environ['ANDROID_ADB_SERVER_PORT'] = '1337'
     adb_controller_lib.AdbController(
-        adb_path='my_adb',
-        device_name='awesome_device',
-        adb_server_port=9999,
-        default_timeout=_TIMEOUT)
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+            default_timeout=_TIMEOUT,
+        )
+    )
     self.assertNotIn('ANDROID_HOME', os.environ)
     self.assertNotIn('ANDROID_ADB_SERVER_PORT', os.environ)
 

android_env/components/adb_log_stream.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,8 +27,8 @@ _LOGCAT_COMMAND = ['logcat', '-v', 'epoch']
 class AdbLogStream(log_stream.LogStream):
   """Manages adb logcat process for a locally running emulator."""
 
-  def __init__(self, adb_command_prefix: list[str], *args, **kwargs):
-    super().__init__(*args, **kwargs)
+  def __init__(self, adb_command_prefix: list[str], verbose: bool = False):
+    super().__init__(verbose=verbose)
     self._adb_command_prefix = adb_command_prefix
 
   def _get_stream_output(self):

android_env/components/adb_log_stream_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

android_env/components/app_screen_checker.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,13 +15,13 @@
 
 """Determines if the current app screen matches an expected app screen."""
 
+from collections.abc import Callable, Sequence
 import enum
 import re
 import time
-from typing import Callable, Optional, Sequence, Pattern
+from typing import Self
 
 from absl import logging
-
 from android_env.components import adb_call_parser as adb_call_parser_lib
 from android_env.components import errors
 from android_env.proto import adb_pb2
@@ -31,7 +31,7 @@ from android_env.proto import task_pb2
 class _DumpsysNode:
   """A node in a dumpsys tree."""
 
-  def __init__(self, data: Optional[str] = None):
+  def __init__(self, data: str):
     self._children = []
     self._data = data
 
@@ -40,12 +40,12 @@ class _DumpsysNode:
     return self._data
 
   @property
-  def children(self) -> list['_DumpsysNode']:
+  def children(self) -> list[Self]:
     return self._children
 
-  def find_child(self,
-                 predicate: Callable[['_DumpsysNode'], bool],
-                 max_levels: int = 0) -> Optional['_DumpsysNode']:
+  def find_child(
+      self, predicate: Callable[[Self], bool], max_levels: int = 0
+  ) -> Self | None:
     """Returns the first direct child that matches `predicate`, None otherwise.
 
     Args:
@@ -126,9 +126,11 @@ def build_tree_from_dumpsys_output(dumpsys_output: str) -> _DumpsysNode:
   return root
 
 
-def matches_path(dumpsys_activity_output: str,
-                 expected_view_hierarchy_path: Sequence[Pattern[str]],
-                 max_levels: int = 0) -> bool:
+def matches_path(
+    dumpsys_activity_output: str,
+    expected_view_hierarchy_path: Sequence[re.Pattern[str]],
+    max_levels: int = 0,
+) -> bool:
   """Returns True if the current dumpsys output matches the expected path.
 
   Args:
@@ -150,7 +152,7 @@ def matches_path(dumpsys_activity_output: str,
         'view_hierarchy is None. Dumpsys activity output: %s. tree: %r',
         str(dumpsys_activity_output), root.print_tree())
     logging.error('Tree root: %s', str(root))
-    return None
+    return False
 
   current_node = view_hierarchy
   for i, regex in enumerate(expected_view_hierarchy_path):
@@ -165,7 +167,7 @@ def matches_path(dumpsys_activity_output: str,
                     regex.pattern, current_node)
       logging.error('Dumpsys activity output: %s', str(dumpsys_activity_output))
       logging.error('Tree root: %s', str(root))
-      return None
+      return False
     else:
       current_node = child
   return True

android_env/components/app_screen_checker_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

android_env/components/config_classes.py

@@ -0,0 +1,203 @@
+# coding=utf-8
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dataclass definitions used for instantiating AndroidEnv components."""
+
+import dataclasses
+
+
+@dataclasses.dataclass
+class AdbControllerConfig:
+  """Settings for instatiating an `AdbController` instance."""
+
+  # Filesystem path to the `adb` binary.
+  # NOTE: This must be a full path and must not contain environment variables
+  # or user folder shorthands (e.g. `~/some/path/to/adb`) since they will not be
+  # expanded internally by AndroidEnv.
+  adb_path: str = '~/Android/Sdk/platform-tools/adb'
+  # Port for adb server.
+  adb_server_port: int = 5037
+  # Default timeout in seconds for internal commands.
+  default_timeout: float = 120.0
+  # Name of the device to communicate with.
+  device_name: str = ''
+
+
+@dataclasses.dataclass
+class DeviceSettingsConfig:
+  """Config class for DeviceSettings."""
+
+  # Whether to show circles on the screen indicating touch position.
+  show_touches: bool = True
+  # Whether to show blue lines on the screen indicating touch position.
+  show_pointer_location: bool = True
+  # Whether or not to show the status (top) bar.
+  show_status_bar: bool = False
+  # Whether or not to show the navigation (bottom) bar.
+  show_navigation_bar: bool = False
+
+
+@dataclasses.dataclass
+class CoordinatorConfig:
+  """Config class for Coordinator."""
+
+  # Number of virtual "fingers" of the agent.
+  num_fingers: int = 1
+  # Whether to enable keyboard key events.
+  enable_key_events: bool = False
+  # Time between periodic restarts in minutes. If > 0, will trigger
+  # a simulator restart at the beginning of the next episode once the time has
+  # been reached.
+  periodic_restart_time_min: float = 0.0
+  # General Android settings.
+  device_settings: DeviceSettingsConfig = dataclasses.field(
+      default_factory=DeviceSettingsConfig
+  )
+
+
+@dataclasses.dataclass
+class SimulatorConfig:
+  """Base class for all simulator configs."""
+
+  # If true, the log stream of the simulator will be verbose.
+  verbose_logs: bool = False
+  # How often to (asynchronously) grab the screenshot from the simulator.
+  # If <= 0, stepping the environment blocks on fetching the screenshot (the
+  # environment is synchronous).
+  interaction_rate_sec: float = 0.0
+
+
+@dataclasses.dataclass
+class EmulatorLauncherConfig:
+  """Config class for EmulatorLauncher."""
+
+  # NOTE: If `adb_port`, `emulator_console_port` and `grpc_port` are defined
+  # (i.e. not all equal to 0), it is assumed that the emulator they point to
+  # exists already and EmulatorLauncher will be skipped.
+
+  # Filesystem path to the `emulator` binary.
+  emulator_path: str = '~/Android/Sdk/emulator/emulator'
+  # Filesystem path to the Android SDK root.
+  android_sdk_root: str = '~/Android/Sdk'
+  # Name of the AVD.
+  avd_name: str = ''
+  # Local directory for AVDs.
+  android_avd_home: str = '~/.android/avd'
+  # Name of the snapshot to load.
+  snapshot_name: str = ''
+  # Path to the KVM device.
+  kvm_device: str = '/dev/kvm'
+  # Path to directory which will hold temporary files.
+  tmp_dir: str = '/tmp/android_env/simulator/'
+  # GPU mode override.
+  # Please see
+  # https://developer.android.com/studio/run/emulator-acceleration#accel-graphics.
+  gpu_mode: str = 'swangle_indirect'  # Alternative: swiftshader_indirect, host
+  # Whether to run in headless mode (i.e. without a graphical window).
+  run_headless: bool = True
+  # Whether to restrict network access.
+  # If True, will disable networking on the device. This option is only
+  # available for emulator version > 31.3.9 (June 2022).
+  restrict_network: bool = False
+  # Whether to set `SHOW_PERF_STATS=1` when launching the emulator to display
+  # performance and memory statistics.
+  show_perf_stats: bool = False
+
+  # ADB port for the Android device.
+  adb_port: int = 0
+  # Port for telnet communication with the emulator.
+  emulator_console_port: int = 0
+  # Port for gRPC communication with the emulator.
+  grpc_port: int = 0
+
+
+@dataclasses.dataclass
+class EmulatorConfig(SimulatorConfig):
+  """Config class for EmulatorSimulator."""
+
+  # Configuration for launching the Android Emulator.
+  emulator_launcher: EmulatorLauncherConfig = dataclasses.field(
+      default_factory=EmulatorLauncherConfig
+  )
+  # Configuration for talking to adb.
+  adb_controller: AdbControllerConfig = dataclasses.field(
+      default_factory=AdbControllerConfig
+  )
+  # Path to file which holds emulator logs. If not provided, it will be
+  # determined by the EmulatorLauncher.
+  logfile_path: str = ''
+  # The number of times to try launching the emulator before rebooting (reboot
+  # on the n+1-st try).
+  launch_n_times_without_reboot: int = 1
+  # The number of times to try launching the emulator before reinstalling
+  # (reinstall on the n+1-st try).
+  launch_n_times_without_reinstall: int = 2
+
+
+@dataclasses.dataclass
+class FakeSimulatorConfig(SimulatorConfig):
+  """Config class for FakeSimulator."""
+
+  # The dimensions in pixels of the device screen (HxW).
+  screen_dimensions: tuple[int, int] = (0, 0)
+
+
+@dataclasses.dataclass
+class TaskManagerConfig:
+  """Config class for TaskManager."""
+
+  # If max_bad_states episodes finish in a bad state in a row, restart
+  # the simulation.
+  max_bad_states: int = 3
+  # The frequency to check for the current activity and view hierarchy.
+  # The unit is raw observation (i.e. each call to AndroidEnv.step()).
+  dumpsys_check_frequency: int = 150
+  # The maximum number of tries for extracting the current activity before
+  # forcing the episode to restart.
+  max_failed_current_activity: int = 10
+  # The maximum number of extras elements to store. If this number is exceeded,
+  # elements are dropped in the order they were received.
+  extras_max_buffer_size: int = 100
+
+
+@dataclasses.dataclass
+class TaskConfig:
+  """Base config class for loading tasks."""
+
+  # The directory for temporary task-related resources.
+  tmp_dir: str = ''
+
+
+@dataclasses.dataclass
+class FilesystemTaskConfig(TaskConfig):
+  """Config for protobuf files stored in the local filesystem."""
+
+  # Filesystem path to `.binarypb` or `.textproto` protobuf Task.
+  path: str = ''
+
+
+@dataclasses.dataclass
+class AndroidEnvConfig:
+  """Config class for AndroidEnv."""
+
+  # Configs for main components.
+  task: TaskConfig = dataclasses.field(default_factory=TaskConfig)
+  task_manager: TaskManagerConfig = dataclasses.field(
+      default_factory=TaskManagerConfig
+  )
+  coordinator: CoordinatorConfig = dataclasses.field(
+      default_factory=CoordinatorConfig
+  )
+  simulator: SimulatorConfig = dataclasses.field(default_factory=EmulatorConfig)