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

android_env/components/task_manager.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,16 +16,18 @@
 """TaskManager handles all events and information related to the task."""
 
 import ast
+from collections.abc import Callable
 import copy
 import datetime
 import json
 import re
 import threading
-from typing import Any, Callable
+from typing import Any
 
 from absl import logging
 from android_env.components import adb_call_parser as adb_call_parser_lib
 from android_env.components import app_screen_checker
+from android_env.components import config_classes
 from android_env.components import dumpsys_thread
 from android_env.components import log_stream as log_stream_lib
 from android_env.components import logcat_thread
@@ -41,28 +43,18 @@ class TaskManager:
   def __init__(
       self,
       task: task_pb2.Task,
-      max_bad_states: int = 3,
-      dumpsys_check_frequency: int = 150,
-      max_failed_current_activity: int = 10,
+      config: config_classes.TaskManagerConfig | None = None,
   ):
     """Controls task-relevant events and information.
 
     Args:
       task: A task proto defining the RL task.
-      max_bad_states: How many bad states in a row are allowed before a restart
-        of the simulator is triggered.
-      dumpsys_check_frequency: Frequency, in steps, at which to check
-        current_activity and view hierarchy
-      max_failed_current_activity: The maximum number of tries for extracting
-        the current activity before forcing the episode to restart.
+      config: Configuration for this instance.
     """
-    self._task = task
-    self._max_bad_states = max_bad_states
-    self._dumpsys_check_frequency = dumpsys_check_frequency
-    self._max_failed_current_activity = max_failed_current_activity
 
+    self._task = task
+    self._config = config or config_classes.TaskManagerConfig()
     self._lock = threading.Lock()
-    self._extras_max_buffer_size = 100
     self._logcat_thread = None
     self._dumpsys_thread = None
     self._setup_step_interpreter = None
@@ -92,13 +84,6 @@ class TaskManager:
 
     logging.info('Task config: %s', self._task)
 
-  def task(self) -> task_pb2.Task:
-    return self._task
-
-  def update_task(self, task: task_pb2.Task) -> None:
-    self._stats['task_updates'] += 1
-    self._task = task
-
   def stats(self) -> dict[str, Any]:
     """Returns a dictionary of stats.
 
@@ -248,9 +233,11 @@ class TaskManager:
     self._dumpsys_thread = dumpsys_thread.DumpsysThread(
         app_screen_checker=app_screen_checker.AppScreenChecker(
             adb_call_parser=adb_call_parser,
-            expected_app_screen=self._task.expected_app_screen),
-        check_frequency=self._dumpsys_check_frequency,
-        max_failed_current_activity=self._max_failed_current_activity)
+            expected_app_screen=self._task.expected_app_screen,
+        ),
+        check_frequency=self._config.dumpsys_check_frequency,
+        max_failed_current_activity=self._config.max_failed_current_activity,
+    )
 
   def _stop_logcat_thread(self):
     if self._logcat_thread is not None:
@@ -266,11 +253,11 @@ class TaskManager:
     to a good state.
     """
     logging.warning('Bad state detected.')
-    if self._max_bad_states:
+    if self._config.max_bad_states:
       self._is_bad_episode = True
       self._bad_state_counter += 1
       logging.warning('Bad state counter: %d.', self._bad_state_counter)
-      if self._bad_state_counter >= self._max_bad_states:
+      if self._bad_state_counter >= self._config.max_bad_states:
         logging.error('Too many consecutive bad states. Restarting simulator.')
         self._stats['restart_count_max_bad_states'] += 1
         self._should_restart = True
@@ -342,8 +329,13 @@ class TaskManager:
       if extra:
         try:
           extra = ast.literal_eval(extra)
-        # Except all to avoid unnecessary crashes, only log error.
-        except Exception:  # pylint: disable=broad-except
+        except (
+            ValueError,
+            TypeError,
+            SyntaxError,
+            MemoryError,
+            RecursionError,
+        ):
           logging.exception('Could not parse extra: %s', extra)
           # Don't try to process the extra as text; that would probably crash.
           return
@@ -380,7 +372,10 @@ class TaskManager:
         latest_extras = self._latest_values['extra']
         if extra_name in latest_extras:
           # If latest extra is not flushed, append.
-          if len(latest_extras[extra_name]) >= self._extras_max_buffer_size:
+          if (
+              len(latest_extras[extra_name])
+              >= self._config.extras_max_buffer_size
+          ):
             latest_extras[extra_name].pop(0)
           latest_extras[extra_name].append(extra)
         else:

android_env/components/task_manager_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.
@@ -55,23 +55,6 @@ class TaskManagerTest(absltest.TestCase):
         log_stream, 'LogStream',
         return_value=self._log_stream).start()
 
-  def test_update_task(self):
-    init_task = task_pb2.Task(id='initial_task')
-    new_task = task_pb2.Task(id='updated_task')
-    task_mgr = task_manager.TaskManager(task=init_task)
-
-    # Setting up the initial task so that the setup_step_interpreter
-    # is properly initialized.
-    adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
-    task_mgr.setup_task()
-
-    self.assertEqual(init_task, task_mgr.task())
-    self.assertEqual(0, task_mgr.stats()['task_updates'])
-    task_mgr.update_task(new_task)
-    self.assertEqual(new_task, task_mgr.task())
-    self.assertEqual(1, task_mgr.stats()['task_updates'])
-
   def test_start(self):
     task_mgr = task_manager.TaskManager(task=task_pb2.Task())
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)

android_env/env_interface.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.
@@ -24,7 +24,6 @@ from typing import Any
 
 from android_env.proto import adb_pb2
 from android_env.proto import state_pb2
-from android_env.proto import task_pb2
 import dm_env
 import numpy as np
 
@@ -79,20 +78,6 @@ class AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):
 
     return adb_pb2.AdbResponse()
 
-  def update_task(self, task: task_pb2.Task) -> bool:
-    """Replaces the current task with a new task.
-
-    It is the caller's responsibility to call `reset()` after the task update.
-
-    Args:
-      task: A new task to replace the current one.
-
-    Returns:
-      A bool indicating the success of the task setup.
-    """
-
-    return True
-
   def load_state(
       self, request: state_pb2.LoadStateRequest
   ) -> state_pb2.LoadStateResponse:
@@ -122,4 +107,3 @@ class AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):
       applicable), and any other relevant information.
     """
     raise NotImplementedError('This environment does not support saving state')
-

android_env/environment.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,10 +19,12 @@ from typing import Any
 
 from absl import logging
 from android_env import env_interface
+from android_env.components import adb_call_parser
 from android_env.components import coordinator as coordinator_lib
+from android_env.components import task_manager as task_manager_lib
+from android_env.components.simulators import base_simulator
 from android_env.proto import adb_pb2
 from android_env.proto import state_pb2
-from android_env.proto import task_pb2
 import dm_env
 import numpy as np
 
@@ -30,10 +32,17 @@ import numpy as np
 class AndroidEnv(env_interface.AndroidEnvInterface):
   """An RL environment that interacts with Android apps."""
 
-  def __init__(self, coordinator: coordinator_lib.Coordinator):
+  def __init__(
+      self,
+      simulator: base_simulator.BaseSimulator,
+      coordinator: coordinator_lib.Coordinator,
+      task_manager: task_manager_lib.TaskManager,
+  ):
     """Initializes the state of this AndroidEnv object."""
 
+    self._simulator = simulator
     self._coordinator = coordinator
+    self._task_manager = task_manager
     self._latest_action = {}
     self._latest_observation = {}
     self._latest_extras = {}
@@ -133,22 +142,13 @@ class AndroidEnv(env_interface.AndroidEnvInterface):
     return self._latest_observation.copy()
 
   def stats(self) -> dict[str, Any]:
-    return self._coordinator.stats()
+    coordinator_stats = self._coordinator.stats()
+    task_manager_stats = self._task_manager.stats()
+    return coordinator_stats | task_manager_stats
 
   def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:
     return self._coordinator.execute_adb_call(call)
 
-  def update_task(self, task: task_pb2.Task) -> bool:
-    """Replaces the current task with a new task.
-
-    Args:
-      task: A new task to replace the current one.
-
-    Returns:
-      A bool indicating the success of the task setup.
-    """
-    return self._coordinator.update_task(task)
-
   def load_state(
       self, request: state_pb2.LoadStateRequest
   ) -> state_pb2.LoadStateResponse:
@@ -162,7 +162,16 @@ class AndroidEnv(env_interface.AndroidEnvInterface):
       A `LoadStateResponse` containing the status, error message (if
       applicable), and any other relevant information.
     """
-    return self._coordinator.load_state(request)
+
+    self._task_manager.stop()
+    response = self._simulator.load_state(request)
+    self._task_manager.start(
+        adb_call_parser_factory=lambda: adb_call_parser.AdbCallParser(
+            self._simulator.create_adb_controller()
+        ),
+        log_stream=self._simulator.create_log_stream(),
+    )
+    return response
 
   def save_state(
       self, request: state_pb2.SaveStateRequest
@@ -177,4 +186,5 @@ class AndroidEnv(env_interface.AndroidEnvInterface):
       A `SaveStateResponse` containing the status, error message (if
       applicable), and any other relevant information.
     """
-    return self._coordinator.save_state(request)
+
+    return self._simulator.save_state(request)

android_env/environment_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.
@@ -19,10 +19,13 @@ from unittest import mock
 
 from absl.testing import absltest
 from android_env import environment
+from android_env.components import config_classes
 from android_env.components import coordinator as coordinator_lib
+from android_env.components import task_manager as task_manager_lib
+from android_env.components.simulators import base_simulator
+from android_env.components.simulators.fake import fake_simulator
 from android_env.proto import adb_pb2
 from android_env.proto import state_pb2
-from android_env.proto import task_pb2
 import dm_env
 import numpy as np
 
@@ -44,10 +47,21 @@ def _create_mock_coordinator() -> coordinator_lib.Coordinator:
   return coordinator
 
 
+def _create_fake_simulator() -> fake_simulator.FakeSimulator:
+  return fake_simulator.FakeSimulator(
+      config=config_classes.FakeSimulatorConfig(screen_dimensions=(123, 456))
+  )
+
+
 class AndroidEnvTest(absltest.TestCase):
 
   def test_specs(self):
-    env = environment.AndroidEnv(_create_mock_coordinator())
+    simulator = _create_fake_simulator()
+    coordinator = _create_mock_coordinator()
+    task_manager = mock.create_autospec(task_manager_lib.TaskManager)
+    env = environment.AndroidEnv(
+        simulator=simulator, coordinator=coordinator, task_manager=task_manager
+    )
 
     # Check action spec.
     self.assertNotEmpty(env.action_spec())
@@ -77,7 +91,9 @@ class AndroidEnvTest(absltest.TestCase):
     self.assertEqual(env.observation_spec()['orientation'].shape, (4,))
 
   def test_reset_and_step(self):
-    coordinator = mock.create_autospec(coordinator_lib.Coordinator)
+    simulator = _create_fake_simulator()
+    coordinator = _create_mock_coordinator()
+    task_manager = mock.create_autospec(task_manager_lib.TaskManager)
     coordinator.action_spec.return_value = {
         'action_type':
             dm_env.specs.DiscreteArray(num_values=3),
@@ -90,7 +106,9 @@ class AndroidEnvTest(absltest.TestCase):
         'timedelta': dm_env.specs.Array(shape=(), dtype=np.int64),
         'orientation': dm_env.specs.Array(shape=(4,), dtype=np.uint8),
     }
-    env = environment.AndroidEnv(coordinator)
+    env = environment.AndroidEnv(
+        simulator=simulator, coordinator=coordinator, task_manager=task_manager
+    )
     coordinator.rl_reset.return_value = dm_env.TimeStep(
         step_type=dm_env.StepType.FIRST,
         reward=0.0,
@@ -125,9 +143,8 @@ class AndroidEnvTest(absltest.TestCase):
     self.assertIn('click', extras)
     self.assertEqual(extras['click'], np.array([246], dtype=np.int64))
 
-    coordinator.stats.return_value = {
-        'my_measurement': 135,
-    }
+    coordinator.stats.return_value = {'my_measurement': 135}
+    task_manager.stats.return_value = {'another_measurement': 79}
 
     # Step again in the environment and check expectations again.
     pixels = np.random.rand(987, 654, 3)
@@ -189,8 +206,12 @@ class AndroidEnvTest(absltest.TestCase):
     np.testing.assert_equal(obs['orientation'], (1, 0, 0, 0))
 
   def test_adb_call(self):
+    simulator = _create_fake_simulator()
     coordinator = _create_mock_coordinator()
-    env = environment.AndroidEnv(coordinator)
+    task_manager = mock.create_autospec(task_manager_lib.TaskManager)
+    env = environment.AndroidEnv(
+        simulator=simulator, coordinator=coordinator, task_manager=task_manager
+    )
     call = adb_pb2.AdbRequest(
         force_stop=adb_pb2.AdbRequest.ForceStop(package_name='blah'))
     expected_response = adb_pb2.AdbResponse(
@@ -203,41 +224,46 @@ class AndroidEnvTest(absltest.TestCase):
     coordinator.execute_adb_call.assert_called_once_with(call)
 
   def test_load_state(self):
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
     coordinator = _create_mock_coordinator()
-    env = environment.AndroidEnv(coordinator)
+    task_manager = mock.create_autospec(task_manager_lib.TaskManager)
+    env = environment.AndroidEnv(
+        simulator=simulator, coordinator=coordinator, task_manager=task_manager
+    )
     expected_response = state_pb2.LoadStateResponse(
         status=state_pb2.LoadStateResponse.Status.OK
     )
     request = state_pb2.LoadStateRequest(args={'foo': 'bar'})
-    coordinator.load_state.return_value = expected_response
+    simulator.load_state.return_value = expected_response
     response = env.load_state(request)
     self.assertEqual(response, expected_response)
-    coordinator.load_state.assert_called_once_with(request)
+    simulator.load_state.assert_called_once_with(request)
+    task_manager.stop.assert_called_once()
+    task_manager.start.assert_called_once()
 
   def test_save_state(self):
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
     coordinator = _create_mock_coordinator()
-    env = environment.AndroidEnv(coordinator)
+    task_manager = mock.create_autospec(task_manager_lib.TaskManager)
+    env = environment.AndroidEnv(
+        simulator=simulator, coordinator=coordinator, task_manager=task_manager
+    )
     expected_response = state_pb2.SaveStateResponse(
         status=state_pb2.SaveStateResponse.Status.OK
     )
     request = state_pb2.SaveStateRequest(args={'foo': 'bar'})
-    coordinator.save_state.return_value = expected_response
+    simulator.save_state.return_value = expected_response
     response = env.save_state(request)
     self.assertEqual(response, expected_response)
-    coordinator.save_state.assert_called_once_with(request)
-
-  def test_update_task(self):
-    coordinator = _create_mock_coordinator()
-    env = environment.AndroidEnv(coordinator)
-    task = task_pb2.Task()
-    coordinator.update_task.return_value = True
-    response = env.update_task(task)
-    self.assertEqual(response, True)
-    coordinator.update_task.assert_called_once_with(task)
+    simulator.save_state.assert_called_once_with(request)
 
   def test_double_close(self):
+    simulator = _create_fake_simulator()
     coordinator = _create_mock_coordinator()
-    env = environment.AndroidEnv(coordinator)
+    task_manager = mock.create_autospec(task_manager_lib.TaskManager)
+    env = environment.AndroidEnv(
+        simulator=simulator, coordinator=coordinator, task_manager=task_manager
+    )
     env.close()
     env.close()
     coordinator.close.assert_called_once()

android_env/loader.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.
@@ -17,59 +17,73 @@
 
 import os
 
+from absl import logging
 from android_env import environment
+from android_env.components import config_classes
 from android_env.components import coordinator as coordinator_lib
+from android_env.components import device_settings as device_settings_lib
 from android_env.components import task_manager as task_manager_lib
 from android_env.components.simulators.emulator import emulator_simulator
+from android_env.components.simulators.fake import fake_simulator
 from android_env.proto import task_pb2
 
 from google.protobuf import text_format
 
 
-def load(task_path: str,
-         avd_name: str,
-         android_avd_home: str = '~/.android/avd',
-         android_sdk_root: str = '~/Android/Sdk',
-         emulator_path: str = '~/Android/Sdk/emulator/emulator',
-         adb_path: str = '~/Android/Sdk/platform-tools/adb',
-         run_headless: bool = False) -> environment.AndroidEnv:
-  """Loads an AndroidEnv instance.
-
-  Args:
-    task_path: Path to the task textproto file.
-    avd_name: Name of the AVD (Android Virtual Device).
-    android_avd_home: Path to the AVD (Android Virtual Device).
-    android_sdk_root: Root directory of the SDK.
-    emulator_path: Path to the emulator binary.
-    adb_path: Path to the ADB (Android Debug Bridge).
-    run_headless: If True, the emulator display is turned off.
-  Returns:
-    env: An AndroidEnv instance.
-  """
-
-  # Create simulator.
-  simulator = emulator_simulator.EmulatorSimulator(
-      adb_controller_args=dict(
-          adb_path=os.path.expanduser(adb_path),
-          adb_server_port=5037,
-      ),
-      emulator_launcher_args=dict(
-          avd_name=avd_name,
-          android_avd_home=os.path.expanduser(android_avd_home),
-          android_sdk_root=os.path.expanduser(android_sdk_root),
-          emulator_path=os.path.expanduser(emulator_path),
-          run_headless=run_headless,
-          gpu_mode='swiftshader_indirect',
-      ),
-  )
+def _load_task(task_config: config_classes.TaskConfig) -> task_pb2.Task:
+  """Returns the task according to `task_config`."""
 
-  # Prepare task.
   task = task_pb2.Task()
-  with open(task_path, 'r') as proto_file:
-    text_format.Parse(proto_file.read(), task)
+  match task_config:
+    case config_classes.FilesystemTaskConfig():
+      with open(task_config.path, 'r') as proto_file:
+        text_format.Parse(proto_file.read(), task)
+    case _:
+      logging.error('Unsupported TaskConfig: %r', task_config)
+
+  return task
+
+
+def load(config: config_classes.AndroidEnvConfig) -> environment.AndroidEnv:
+  """Loads an AndroidEnv instance."""
 
+  task = _load_task(config.task)
   task_manager = task_manager_lib.TaskManager(task)
-  coordinator = coordinator_lib.Coordinator(simulator, task_manager)
 
-  # Load environment.
-  return environment.AndroidEnv(coordinator=coordinator)
+  match config.simulator:
+    case config_classes.EmulatorConfig():
+      _process_emulator_launcher_config(config.simulator)
+      simulator = emulator_simulator.EmulatorSimulator(config=config.simulator)
+    case config_classes.FakeSimulatorConfig():
+      simulator = fake_simulator.FakeSimulator(config=config.simulator)
+    case _:
+      raise ValueError('Unsupported simulator config: {config.simulator}')
+
+  device_settings = device_settings_lib.DeviceSettings(simulator)
+  coordinator = coordinator_lib.Coordinator(
+      simulator, task_manager, device_settings
+  )
+  return environment.AndroidEnv(
+      simulator=simulator, coordinator=coordinator, task_manager=task_manager
+  )
+
+
+def _process_emulator_launcher_config(
+    emulator_config: config_classes.EmulatorConfig,
+) -> None:
+  """Adjusts the configuration of the emulator depending on some conditions."""
+
+  # Expand the user directory if specified.
+  launcher_config = emulator_config.emulator_launcher
+  launcher_config.android_avd_home = os.path.expanduser(
+      launcher_config.android_avd_home
+  )
+  launcher_config.android_sdk_root = os.path.expanduser(
+      launcher_config.android_sdk_root
+  )
+  launcher_config.emulator_path = os.path.expanduser(
+      launcher_config.emulator_path
+  )
+  emulator_config.adb_controller.adb_path = os.path.expanduser(
+      emulator_config.adb_controller.adb_path
+  )