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

android_env/components/task_manager_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 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,33 +55,26 @@ 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.
+  def test_start(self):
+    task_mgr = task_manager.TaskManager(task=task_pb2.Task())
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
-
-    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'])
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    self.assertIsNotNone(task_mgr._logcat_thread)
+    self.assertIsNotNone(task_mgr._dumpsys_thread)
+    self.assertIsNotNone(task_mgr._setup_step_interpreter)
 
   def test_setup_task(self):
     task_mgr = task_manager.TaskManager(task=task_pb2.Task())
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
-    self.assertIsNotNone(task_mgr._logcat_thread)
-    self.assertIsNotNone(task_mgr._setup_step_interpreter)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
+    self._setup_step_interpreter.interpret.assert_called_once()
 
   def test_step_count(self):
     task_mgr = task_manager.TaskManager(task=task_pb2.Task())
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     task_mgr.rl_reset(observation={})
     self.assertEqual(task_mgr.stats()['episode_steps'], 0)
     task_mgr.rl_step(observation={})
@@ -109,7 +102,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     timestep = task_mgr.rl_step(
         observation={
             'pixels': np.array([1, 2, 3]),
@@ -144,7 +138,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     timestep = task_mgr.rl_step(
         observation={
             'pixels': np.array([1, 2, 3]),
@@ -174,7 +169,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     timestep = task_mgr.rl_step(
         observation={
             'pixels': np.array([1, 2, 3]),
@@ -207,7 +203,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
 
     timestep = task_mgr.rl_step(
         observation={
@@ -256,7 +253,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
 
     timestep = task_mgr.rl_step(
         observation={
@@ -286,6 +284,49 @@ class TaskManagerTest(absltest.TestCase):
     np.testing.assert_equal(
         expected_extra.get('extra_dict'), extras.get('extra_dict'))
 
+  def test_get_current_extras_failed_to_parse(self):
+    # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`
+    # right away.
+    def my_add_ev_listener(event_listener: logcat_thread.EventListener):
+      # Check that the event matches what's expected.
+      event = event_listener.regexp
+      match = event.match('extra: some_extra [1, 2]')
+      if match is None:  # Ignore events that are not extras.
+        return
+
+      # Emit events.
+      fn = event_listener.handler_fn
+      fn(event, event.match('extra: extra_with_malformed_1 [1]'))
+      fn(event, event.match('extra: extra_with_malformed_1 [\'this is \\ bad]'))
+      fn(event, event.match('extra: extra_with_malformed_1 [2]'))
+      fn(event, event.match('extra: extra_with_malformed_2 [\'this is bad]'))
+      fn(event, event.match('extra: extra_with_malformed_2 [1]'))
+      fn(event, event.match('extra: extra_malformed_only [_very_bad_news]'))
+
+    # Setup the task and trigger the listener.
+    task = task_pb2.Task()
+    task.log_parsing_config.log_regexps.extra.extend([
+        '^extra: (?P<name>[^ ]*)[ ]?(?P<extra>.*)$'
+    ])
+    task_mgr = task_manager.TaskManager(task=task)
+    self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
+    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()
+
+    timestep = task_mgr.rl_step(
+        observation={
+            'pixels': np.array([1, 2, 3]),
+        })
+
+    # Check expectations.
+    self.assertIn('extras', timestep.observation)
+    extras = timestep.observation['extras']
+    np.testing.assert_almost_equal(extras.get('extra_with_malformed_1'),
+                                   [[1], [2]])
+    np.testing.assert_almost_equal(extras.get('extra_with_malformed_2'), [[1]])
+    self.assertNotIn('extra_malformed_only', extras)
+
   def test_multi_log_regexp(self):
     # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`
     # right away.
@@ -305,7 +346,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     timestep = task_mgr.rl_step(
         observation={
             'pixels': np.array([1, 2, 3]),
@@ -335,7 +377,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     timestep = task_mgr.rl_step(
         observation={
             'pixels': np.array([1, 2, 3]),
@@ -359,7 +402,8 @@ class TaskManagerTest(absltest.TestCase):
     task_mgr = task_manager.TaskManager(task=task)
     self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener
     adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-    task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)
+    task_mgr.setup_task()
     timestep = task_mgr.rl_step(
         observation={
             'pixels': np.array([1, 2, 3]),
@@ -367,38 +411,5 @@ class TaskManagerTest(absltest.TestCase):
     self.assertTrue(timestep.last())
 
 
-def test_setup_steps_skipped_if_no_change(self):
-  setup_steps = [task_pb2.SetupStep(sleep=task_pb2.Sleep(time_sec=10))]
-  task_mgr = task_manager.TaskManager(
-      task=task_pb2.Task(id='first_task', setup_steps=setup_steps))
-  adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-  task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
-  self.assertIsNotNone(task_mgr._logcat_thread)
-  self.assertIsNotNone(task_mgr._setup_step_interpreter)
-  self._setup_step_interpreter.interpret.assert_called_once_with(setup_steps)
-  task_mgr.update_task(
-      task=task_pb2.Task(id='second_task', setup_steps=setup_steps))
-  # The interpreter should not be called again.
-  task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
-  self._setup_step_interpreter.interpret.assert_called_once()
-
-
-def test_setup_steps_interpreted_if_changed(self):
-  setup_steps = [task_pb2.SetupStep(sleep=task_pb2.Sleep(time_sec=10))]
-  task_mgr = task_manager.TaskManager(
-      task=task_pb2.Task(id='first_task', setup_steps=setup_steps))
-  adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)
-  task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
-  self.assertIsNotNone(task_mgr._logcat_thread)
-  self.assertIsNotNone(task_mgr._setup_step_interpreter)
-  self._setup_step_interpreter.interpret.assert_called_once_with(setup_steps)
-  new_setup_steps = [task_pb2.SetupStep(sleep=task_pb2.Sleep(time_sec=5))]
-  task_mgr.update_task(
-      task=task_pb2.Task(id='second_task', setup_steps=new_setup_steps))
-  task_mgr.setup_task(lambda: adb_call_parser, log_stream=self._log_stream)
-  self.assertEqual(2, self._setup_step_interpreter.interpret.num_calls())
-  self._setup_step_interpreter.interpret.assert_has_calls(
-      mock.call(setup_steps), mock.call(new_setup_steps))
-
 if __name__ == '__main__':
   absltest.main()

android_env/env_interface.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 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.
@@ -20,10 +20,10 @@ extra methods that clients may use for extended functionality.
 """
 
 import abc
-from typing import Any, Dict
+from typing import Any
 
 from android_env.proto import adb_pb2
-from android_env.proto import task_pb2
+from android_env.proto import state_pb2
 import dm_env
 import numpy as np
 
@@ -34,11 +34,11 @@ class AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):
   # Methods required by dm_env.Environment.
 
   @abc.abstractmethod
-  def action_spec(self) -> Dict[str, dm_env.specs.Array]:
+  def action_spec(self) -> dict[str, dm_env.specs.Array]:
     """Returns the action specification."""
 
   @abc.abstractmethod
-  def observation_spec(self) -> Dict[str, dm_env.specs.Array]:
+  def observation_spec(self) -> dict[str, dm_env.specs.Array]:
     """Returns the observation specification."""
 
   @abc.abstractmethod
@@ -46,7 +46,7 @@ class AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):
     """Resets the current episode."""
 
   @abc.abstractmethod
-  def step(self, action: Dict[str, np.ndarray]) -> dm_env.TimeStep:
+  def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:
     """Executes `action` and returns a `TimeStep`."""
 
   @abc.abstractmethod
@@ -55,41 +55,55 @@ class AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):
 
   # Extensions provided by AndroidEnv.
 
-  @abc.abstractmethod
-  def task_extras_spec(self) -> Dict[str, dm_env.specs.Array]:
-    """Returns the specification for extra info provided by tasks."""
-
-  @abc.abstractmethod
-  def task_extras(self, latest_only: bool = True) -> Dict[str, np.ndarray]:
+  def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:
     """Returns extra info provided by tasks."""
 
+    return {}
+
   @property
-  @abc.abstractmethod
   def raw_action(self):
     """Returns the latest action."""
 
   @property
-  @abc.abstractmethod
   def raw_observation(self):
     """Returns the latest observation."""
 
-  @abc.abstractmethod
-  def stats(self) -> Dict[str, Any]:
+  def stats(self) -> dict[str, Any]:
     """Returns information generated inside the implementation."""
 
-  @abc.abstractmethod
+    return {}
+
   def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:
     """Executes `call` and returns its response."""
 
-  @abc.abstractmethod
-  def update_task(self, task: task_pb2.Task) -> bool:
-    """Replaces the current task with a new task.
+    return adb_pb2.AdbResponse()
+
+  def load_state(
+      self, request: state_pb2.LoadStateRequest
+  ) -> state_pb2.LoadStateResponse:
+    """Loads a state.
+
+    Args:
+      request: A `LoadStateRequest` containing any parameters necessary to
+        specify how/what state to load.
+
+    Returns:
+      A `LoadStateResponse` containing the status, error message (if
+      applicable), and any other relevant information.
+    """
+    raise NotImplementedError('This environment does not support loading state')
 
-    It is the caller's responsibility to call `reset()` after the task update.
+  def save_state(
+      self, request: state_pb2.SaveStateRequest
+  ) -> state_pb2.SaveStateResponse:
+    """Saves a state.
 
     Args:
-      task: A new task to replace the current one.
+      request: A `SaveStateRequest` containing any parameters necessary to
+        specify how/what state to save.
 
     Returns:
-      A bool indicating the success of the task setup.
+      A `SaveStateResponse` containing the status, error message (if
+      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 2022 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,12 +15,16 @@
 
 """Android environment implementation."""
 
-from typing import Any, Dict
+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 task_pb2
+from android_env.proto import state_pb2
 import dm_env
 import numpy as np
 
@@ -28,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 = {}
@@ -40,27 +51,17 @@ class AndroidEnv(env_interface.AndroidEnvInterface):
 
     logging.info('Action spec: %s', self.action_spec())
     logging.info('Observation spec: %s', self.observation_spec())
-    logging.info('Task extras spec: %s', self.task_extras_spec())
-
-  def action_spec(self) -> Dict[str, dm_env.specs.Array]:
-    return self._coordinator.action_spec()
 
-  def observation_spec(self) -> Dict[str, dm_env.specs.Array]:
-    return self._coordinator.observation_spec()
-
-  def task_extras_spec(self) -> Dict[str, dm_env.specs.Array]:
-    return self._coordinator.task_extras_spec()
+  def __del__(self) -> None:
+    self.close()
 
-  @property
-  def raw_action(self):
-    return self._latest_action.copy()
+  # Methods required by dm_env.Environment.
 
-  @property
-  def raw_observation(self):
-    return self._latest_observation.copy()
+  def action_spec(self) -> dict[str, dm_env.specs.Array]:
+    return self._coordinator.action_spec()
 
-  def stats(self) -> Dict[str, Any]:
-    return self._coordinator.stats()
+  def observation_spec(self) -> dict[str, dm_env.specs.Array]:
+    return self._coordinator.observation_spec()
 
   def reset(self) -> dm_env.TimeStep:
     """Resets the environment for a new RL episode."""
@@ -86,7 +87,7 @@ class AndroidEnv(env_interface.AndroidEnvInterface):
 
     return timestep
 
-  def step(self, action: Dict[str, np.ndarray]) -> dm_env.TimeStep:
+  def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:
     """Takes a step in the environment."""
 
     # Check if it's time to reset the episode.
@@ -112,47 +113,78 @@ class AndroidEnv(env_interface.AndroidEnvInterface):
 
     return timestep
 
-  def task_extras(self, latest_only: bool = True) -> Dict[str, np.ndarray]:
+  def close(self) -> None:
+    """Cleans up running processes, threads and local files."""
+    if not self._is_closed:
+      logging.info('Cleaning up AndroidEnv...')
+      if hasattr(self, '_coordinator'):
+        self._coordinator.close()
+      logging.info('Done cleaning up AndroidEnv.')
+      self._is_closed = True
+
+  # Extensions provided by AndroidEnv.
+
+  def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:
     """Returns latest task extras."""
 
-    task_extras = {}
-    for key, spec in self.task_extras_spec().items():
-      if key in self._latest_extras:
-        extra_values = self._latest_extras[key].astype(spec.dtype)
-        for extra in extra_values:
-          self._validate_type(spec, extra)
-        task_extras[key] = extra_values[-1] if latest_only else extra_values
+    task_extras = {}  # Build a copy to avoid reusing objects.
+    for k, spec in self._latest_extras.items():
+      extra_values = spec.astype(spec.dtype)
+      task_extras[k] = extra_values[-1] if latest_only else extra_values
     return task_extras
 
-  def _validate_type(self, spec: dm_env.specs.Array, data: np.ndarray):
-    # Handling of spec data type of string of arbitary length.
-    if spec.dtype.str == '<U0' and data.dtype.str.startswith(
-        '<U') and spec.shape == data.shape:
-      return
-    spec.validate(data)
+  @property
+  def raw_action(self):
+    return self._latest_action.copy()
+
+  @property
+  def raw_observation(self):
+    return self._latest_observation.copy()
+
+  def stats(self) -> dict[str, Any]:
+    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.
+  def load_state(
+      self, request: state_pb2.LoadStateRequest
+  ) -> state_pb2.LoadStateResponse:
+    """Loads a state.
 
     Args:
-      task: A new task to replace the current one.
+      request: A `LoadStateRequest` containing any parameters necessary to
+        specify how/what state to load.
 
     Returns:
-      A bool indicating the success of the task setup.
+      A `LoadStateResponse` containing the status, error message (if
+      applicable), and any other relevant information.
     """
-    return self._coordinator.update_task(task)
 
-  def close(self) -> None:
-    """Cleans up running processes, threads and local files."""
-    if not self._is_closed:
-      logging.info('Cleaning up AndroidEnv...')
-      if hasattr(self, '_coordinator'):
-        self._coordinator.close()
-      logging.info('Done cleaning up AndroidEnv.')
-      self._is_closed = True
+    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
+  ) -> state_pb2.SaveStateResponse:
+    """Saves a state.
 
-  def __del__(self) -> None:
-    self.close()
+    Args:
+      request: A `SaveStateRequest` containing any parameters necessary to
+        specify how/what state to save.
+
+    Returns:
+      A `SaveStateResponse` containing the status, error message (if
+      applicable), and any other relevant information.
+    """
+
+    return self._simulator.save_state(request)