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

android_env/components/simulators/__init__.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/simulators/base_simulator.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.
@@ -16,34 +16,23 @@
 """A base class for talking to different types of Android simulators."""
 
 import abc
+from collections.abc import Callable
+import threading
+import time
 
 from absl import logging
 from android_env.components import adb_controller
+from android_env.components import config_classes
 from android_env.components import errors
 from android_env.components import log_stream
 from android_env.proto import state_pb2
 import numpy as np
 
 
-def _print_logs_on_exception(func):
-  """Decorator function for printing simulator logs upon any exception."""
-  def wrapper(*args, **kwargs):
-    try:
-      return func(*args, **kwargs)
-    except Exception as error:
-      # Calls self.get_logs since self is the first arg.
-      for line in args[0].get_logs().splitlines():
-        logging.error(line)
-      raise errors.SimulatorError(
-          'Exception caught in simulator. Please see the simulator logs '
-          'above for more details.') from error
-  return wrapper
-
-
 class BaseSimulator(metaclass=abc.ABCMeta):
   """An interface for communicating with an Android simulator."""
 
-  def __init__(self, verbose_logs: bool = False):
+  def __init__(self, config: config_classes.SimulatorConfig):
     """Instantiates a BaseSimulator object.
 
     The simulator may be an emulator, virtual machine or even a physical device.
@@ -51,10 +40,11 @@ class BaseSimulator(metaclass=abc.ABCMeta):
     bookkeeping.
 
     Args:
-      verbose_logs: If true, the log stream of the simulator will be verbose.
+      config: Settings for this simulator.
     """
 
-    self._verbose_logs = verbose_logs
+    self._config = config
+    self._interaction_thread: InteractionThread | None = None
 
     # An increasing number that tracks the attempt at launching the simulator.
     self._num_launch_attempts: int = 0
@@ -75,12 +65,31 @@ class BaseSimulator(metaclass=abc.ABCMeta):
   def create_log_stream(self) -> log_stream.LogStream:
     """Creates a stream of logs from the simulator."""
 
-  @_print_logs_on_exception
   def launch(self) -> None:
     """Starts the simulator."""
 
+    # Stop screenshot thread if it's enabled.
+    if self._interaction_thread is not None:
+      self._interaction_thread.stop()
+      self._interaction_thread.join()
+
     self._num_launch_attempts += 1
-    self._launch_impl()
+    try:
+      self._launch_impl()
+    except Exception as error:
+      for line in self.get_logs().splitlines():
+        logging.error(line)
+      raise errors.SimulatorError(
+          'Exception caught in simulator. Please see the simulator logs '
+          'above for more details.'
+      ) from error
+
+    # Start interaction thread.
+    if self._config.interaction_rate_sec > 0:
+      self._interaction_thread = InteractionThread(
+          self._get_screenshot_impl, self._config.interaction_rate_sec
+      )
+      self._interaction_thread.start()
 
   @abc.abstractmethod
   def _launch_impl(self) -> None:
@@ -139,9 +148,18 @@ class BaseSimulator(metaclass=abc.ABCMeta):
     """
     raise NotImplementedError('This simulator does not support save_state()')
 
-  @abc.abstractmethod
   def get_screenshot(self) -> np.ndarray:
-    """Returns pixels representing the current screenshot of the simulator.
+    """Returns pixels representing the current screenshot of the simulator."""
+
+    if self._config.interaction_rate_sec > 0:
+      assert self._interaction_thread is not None
+      return self._interaction_thread.screenshot()  # Async mode.
+    else:
+      return self._get_screenshot_impl()  # Sync mode.
+
+  @abc.abstractmethod
+  def _get_screenshot_impl(self) -> np.ndarray:
+    """Actual implementation of `get_screenshot()`.
 
     The output numpy array should have shape [height, width, num_channels] and
     can be loaded into PIL using Image.fromarray(img, mode='RGB') and be saved
@@ -150,3 +168,41 @@ class BaseSimulator(metaclass=abc.ABCMeta):
 
   def close(self):
     """Frees up resources allocated by this object."""
+
+    if self._interaction_thread is not None:
+      self._interaction_thread.stop()
+      self._interaction_thread.join()
+
+
+class InteractionThread(threading.Thread):
+  """A thread that gets screenshot in the background."""
+
+  def __init__(
+      self,
+      get_screenshot_fn: Callable[[], np.ndarray],
+      interaction_rate_sec: float,
+  ):
+    super().__init__()
+    self._get_screenshot_fn = get_screenshot_fn
+    self._interaction_rate_sec = interaction_rate_sec
+    self._should_stop = threading.Event()
+    self._screenshot = self._get_screenshot_fn()
+
+  def run(self):
+    last_read = time.time()
+    while not self._should_stop.is_set():
+      self._screenshot = self._get_screenshot_fn()
+      now = time.time()
+      elapsed = now - last_read
+      last_read = now
+      sleep_time = self._interaction_rate_sec - elapsed
+      if sleep_time > 0.0:
+        time.sleep(sleep_time)
+    logging.info('InteractionThread.run() finished.')
+
+  def stop(self):
+    logging.info('Stopping InteractionThread.')
+    self._should_stop.set()
+
+  def screenshot(self) -> np.ndarray:
+    return self._screenshot

android_env/components/simulators/base_simulator_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.
@@ -13,14 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for base_simulator."""
-
+import itertools
+import time
 from unittest import mock
 
 from absl.testing import absltest
+from android_env.components import config_classes
 from android_env.components import errors
 # fake_simulator.FakeSimulator inherits from BaseSimulator, so there's no need
 # to import it here explicitly.
+from android_env.components.simulators import base_simulator
 from android_env.components.simulators.fake import fake_simulator
 import numpy as np
 
@@ -28,19 +30,25 @@ import numpy as np
 class BaseSimulatorTest(absltest.TestCase):
 
   def test_launch(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(640, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(640, 480))
+    )
     # The simulator should launch and not crash.
     simulator.launch()
 
   def test_launch_close(self):
-    simulator = fake_simulator.FakeSimulator()
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig()
+    )
     # The simulator should launch and not crash.
     simulator.launch()
     # Closing the simulator should also not crash.
     simulator.close()
 
   def test_get_screenshot(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(640, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(640, 480))
+    )
     # The simulator should launch and not crash.
     simulator.launch()
 
@@ -48,12 +56,126 @@ class BaseSimulatorTest(absltest.TestCase):
     np.testing.assert_equal(screenshot.shape, [640, 480, 3])
 
   def test_print_logs_on_exception(self):
-    simulator = fake_simulator.FakeSimulator()
-    with mock.patch.object(simulator, 'get_logs') as mock_get_logs, \
-         mock.patch.object(simulator, '_launch_impl', autospec=True) as mock_launch:
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig()
+    )
+    with mock.patch.object(
+        simulator, 'get_logs'
+    ) as mock_get_logs, mock.patch.object(
+        simulator, '_launch_impl', autospec=True
+    ) as mock_launch:
       mock_launch.side_effect = ValueError('Oh no!')
       self.assertRaises(errors.SimulatorError, simulator.launch)
       mock_get_logs.assert_called_once()
 
+  def test_get_screenshot_error_async(self):
+    """An exception in the underlying interaction thread should bubble up."""
+
+    # Arrange.
+    mock_interaction_thread = mock.create_autospec(
+        base_simulator.InteractionThread
+    )
+    mock_interaction_thread.screenshot.side_effect = (
+        errors.ReadObservationError()
+    )
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(interaction_rate_sec=0.5)
+    )
+    with mock.patch.object(
+        base_simulator,
+        'InteractionThread',
+        autospec=True,
+        return_value=mock_interaction_thread,
+    ):
+      simulator.launch()
+
+    # Act & Assert.
+    self.assertRaises(errors.ReadObservationError, simulator.get_screenshot)
+
+    # Cleanup.
+    simulator.close()
+
+  def test_get_screenshot_faster_than_screenshot_impl(self):
+    """Return same screenshot when step is faster than the interaction rate."""
+
+    # Arrange.
+    slow_rate = 0.5
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(interaction_rate_sec=slow_rate)
+    )
+
+    # Act.
+    with mock.patch.object(
+        simulator, '_get_screenshot_impl', autospec=True
+    ) as mock_get_screenshot_impl:
+      mock_get_screenshot_impl.side_effect = (
+          np.array(i, ndmin=3) for i in itertools.count(0, 1)
+      )
+      simulator.launch()
+      # Get two screenshots one after the other without pausing.
+      screenshot1 = simulator.get_screenshot()
+      screenshot2 = simulator.get_screenshot()
+
+    # Assert.
+    self.assertAlmostEqual(screenshot1[0][0][0], screenshot2[0][0][0])
+
+    # Cleanup.
+    simulator.close()
+
+  def test_get_screenshot_slower_than_screenshot_impl(self):
+    """Return different screenshots when step slower than the interaction rate."""
+
+    # Arrange.
+    fast_rate = 0.01
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(interaction_rate_sec=fast_rate)
+    )
+
+    # Act.
+    with mock.patch.object(
+        simulator, '_get_screenshot_impl', autospec=True
+    ) as mock_get_screenshot_impl:
+      mock_get_screenshot_impl.side_effect = (
+          np.array(i, ndmin=3) for i in itertools.count(0, 1)
+      )
+      simulator.launch()
+      # Sleep for 500ms between two screenshots.
+      screenshot1 = simulator.get_screenshot()
+      time.sleep(0.5)
+      screenshot2 = simulator.get_screenshot()
+
+    # Assert.
+    self.assertNotEqual(screenshot1[0][0][0], screenshot2[0][0][0])
+
+    # Cleanup.
+    simulator.close()
+
+  def test_interaction_thread_closes_upon_relaunch(self):
+    """Async interaction should kill the InteractionThread when relaunching."""
+
+    # Arrange.
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(interaction_rate_sec=0.01)
+    )
+    mock_interaction_thread = mock.create_autospec(
+        base_simulator.InteractionThread
+    )
+
+    # Act & Assert.
+    with mock.patch.object(
+        base_simulator,
+        'InteractionThread',
+        autospec=True,
+        return_value=mock_interaction_thread,
+    ):
+      simulator.launch()
+      mock_interaction_thread.stop.assert_not_called()
+      mock_interaction_thread.join.assert_not_called()
+      simulator.launch()
+      mock_interaction_thread.stop.assert_called_once()
+      mock_interaction_thread.join.assert_called_once()
+      simulator.close()
+
+
 if __name__ == '__main__':
   absltest.main()

android_env/components/simulators/emulator/__init__.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/simulators/emulator/emulator_launcher.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.
@@ -19,9 +19,9 @@ import glob
 import os
 import subprocess
 import tempfile
-from typing import Optional
 
 from absl import logging
+from android_env.components import config_classes
 
 
 class EmulatorLauncher:
@@ -29,58 +29,13 @@ class EmulatorLauncher:
 
   def __init__(
       self,
-      adb_path: str,
-      adb_port: Optional[int] = None,
-      adb_server_port: Optional[int] = None,
-      emulator_console_port: Optional[int] = None,
-      grpc_port: int = -1,
-      emulator_path: str = '',
-      android_sdk_root: str = '',
-      avd_name: str = '',
-      android_avd_home: str = '',
-      run_headless: bool = False,
-      kvm_device: str = '/dev/kvm',
-      gpu_mode: str = 'swiftshader_indirect',
-      tmp_dir: str = '',
-      snapshot_name: str = '',
-      restrict_network: bool = False,
+      config: config_classes.EmulatorLauncherConfig,
+      adb_controller_config: config_classes.AdbControllerConfig,
   ):
-    """Launches an emulator.
-
-    Args:
-      adb_path: Filesystem path to `adb` executable binary.
-      adb_port: ADB port for the Android device.
-      adb_server_port: Port of the ADB server deamon.
-      emulator_console_port: Port for telnet communication with the emulator.
-      grpc_port: Port for gRPC communication with the emulator.
-      emulator_path: Path to the emulator binary.
-      android_sdk_root: Root directory of the Android SDK.
-      avd_name: Name of the AVD.
-      android_avd_home: Local directory for AVDs.
-      run_headless: Whether to run in headless mode.
-      kvm_device: Path to the KVM device.
-      gpu_mode: GPU mode override. Supported values are listed at:
-        https://developer.android.com/studio/run/emulator-acceleration#accel-graphics
-      tmp_dir: Path to directory which will hold temporary files.
-      snapshot_name: Name of the snapshot to load.
-      restrict_network: if True, will disable networking on the device. This
-        option is only available for emulator version > 31.3.9 (June 2022).
-    """
-
-    self._adb_path = os.path.expandvars(adb_path)
-    self._adb_port = adb_port
-    self._adb_server_port = adb_server_port
-    self._emulator_console_port = emulator_console_port
-    self._grpc_port = grpc_port
-    self._emulator_path = os.path.expandvars(emulator_path)
-    self._android_sdk_root = os.path.expandvars(android_sdk_root)
-    self._avd_name = avd_name
-    self._android_avd_home = android_avd_home
-    self._run_headless = run_headless
-    self._kvm_device = kvm_device
-    self._gpu_mode = gpu_mode
-    self._snapshot_name = snapshot_name
-    self._restrict_network = restrict_network
+    """Launches an emulator."""
+
+    self._config = config
+    self._adb_controller_config = adb_controller_config
 
     self._emulator = None
     self._emulator_output = None
@@ -88,9 +43,10 @@ class EmulatorLauncher:
 
     # Create directory for tmp files.
     # Note: this will be deleted once EmulatorLauncher instance is cleaned up.
-    os.makedirs(tmp_dir, exist_ok=True)
+    os.makedirs(config.tmp_dir, exist_ok=True)
     self._local_tmp_dir_handle = tempfile.TemporaryDirectory(
-        dir=tmp_dir, prefix='simulator_instance_')
+        dir=config.tmp_dir, prefix='simulator_instance_'
+    )
     self._local_tmp_dir = self._local_tmp_dir_handle.name
     self._logfile_path = os.path.join(self._local_tmp_dir, 'emulator_output')
     logging.info('Simulator local_tmp_dir: %s', self._local_tmp_dir)
@@ -101,23 +57,28 @@ class EmulatorLauncher:
   def launch_emulator_process(self) -> None:
     """Launches the emulator."""
 
-    logging.info('Booting new emulator [%s]', self._emulator_path)
+    logging.info('Booting new emulator: %s', self._config.emulator_path)
 
     # Set necessary environment variables.
-    base_lib_dir = self._emulator_path[:-8] + 'lib64/'
+    base_lib_dir = self._config.emulator_path[:-8] + 'lib64/'
     ld_library_path = ':'.join([
         base_lib_dir + 'x11/', base_lib_dir + 'qt/lib/',
         base_lib_dir + 'gles_swiftshader/', base_lib_dir
     ])
     extra_env_vars = {
         'ANDROID_HOME': '',
-        'ANDROID_SDK_ROOT': self._android_sdk_root,
-        'ANDROID_AVD_HOME': self._android_avd_home,
-        'ANDROID_EMULATOR_KVM_DEVICE': self._kvm_device,
-        'ANDROID_ADB_SERVER_PORT': str(self._adb_server_port),
+        'ANDROID_SDK_ROOT': self._config.android_sdk_root,
+        'ANDROID_AVD_HOME': self._config.android_avd_home,
+        'ANDROID_EMULATOR_KVM_DEVICE': self._config.kvm_device,
+        'ANDROID_ADB_SERVER_PORT': str(
+            self._adb_controller_config.adb_server_port
+        ),
         'LD_LIBRARY_PATH': ld_library_path,
-        'QT_XKB_CONFIG_ROOT': str(self._emulator_path[:-8] + 'qt_config/'),
+        'QT_XKB_CONFIG_ROOT': str(
+            self._config.emulator_path[:-8] + 'qt_config/'
+        ),
         'ANDROID_EMU_ENABLE_CRASH_REPORTING': '1',
+        'SHOW_PERF_STATS': str(1 if self._config.show_perf_stats else 0),
     }
     logging.info('extra_env_vars: %s',
                  ' '.join(f'{k}={v}' for k, v in extra_env_vars.items()))
@@ -125,31 +86,51 @@ class EmulatorLauncher:
     env_vars.update(extra_env_vars)
 
     # Compile command.
-    grpc_port = ['-grpc', str(self._grpc_port)] if self._grpc_port >= 0 else []
-    run_headless = ['-no-skin', '-no-window'] if self._run_headless else []
-    ports = ['-ports', '%s,%s' % (self._emulator_console_port, self._adb_port)]
+    grpc_port = (
+        ['-grpc', str(self._config.grpc_port)]
+        if self._config.grpc_port >= 0
+        else []
+    )
+    run_headless = (
+        ['-no-skin', '-no-window'] if self._config.run_headless else []
+    )
+    ports = [
+        '-ports',
+        '%s,%s' % (self._config.emulator_console_port, self._config.adb_port),
+    ]
     snapshot = [
-        '-snapshot', self._snapshot_name, '-feature',
-        'AllowSnapshotMigration,MigratableSnapshotSave'
+        '-snapshot',
+        self._config.snapshot_name,
+        '-feature',
+        'AllowSnapshotMigration,MigratableSnapshotSave',
     ]
-    snapshot = snapshot if self._snapshot_name else ['-no-snapshot']
+    snapshot = snapshot if self._config.snapshot_name else ['-no-snapshot']
     restrict_network_args = [
         '-network-user-mode-options', 'restrict=y', '-wifi-user-mode-options',
         'restrict=y'
     ]
-    network_args = restrict_network_args if self._restrict_network else []
-    command = [
-        self._emulator_path,
-        '-adb-path',
-        self._adb_path,
-        '-gpu',
-        self._gpu_mode,
-        '-no-audio',
-        '-show-kernel',
-        '-verbose',
-        '-avd',
-        self._avd_name,
-    ] + grpc_port + run_headless + ports + snapshot + network_args
+    network_args = (
+        restrict_network_args if self._config.restrict_network else []
+    )
+    command = (
+        [
+            self._config.emulator_path,
+            '-adb-path',
+            self._adb_controller_config.adb_path,
+            '-gpu',
+            self._config.gpu_mode,
+            '-no-audio',
+            '-show-kernel',
+            '-verbose',
+            '-avd',
+            self._config.avd_name,
+        ]
+        + grpc_port
+        + run_headless
+        + ports
+        + snapshot
+        + network_args
+    )
     logging.info('Emulator launch command: %s', ' '.join(command))
     # Prepare logfile.
     self._emulator_output = open(self._logfile_path, 'wb')