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

android_env/components/simulators/fake/fake_simulator_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.
@@ -17,6 +17,7 @@
 
 import re
 from absl.testing import absltest
+from android_env.components import config_classes
 from android_env.components.simulators.fake import fake_simulator
 import numpy as np
 
@@ -24,18 +25,24 @@ import numpy as np
 class FakeSimulatorTest(absltest.TestCase):
 
   def test_device_name(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     self.assertEqual(simulator.adb_device_name(), 'fake_simulator')
 
   def test_launch_close(self):
     # The simulator should launch and not crash.
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     simulator.launch()
     # Closing the simulator should also not crash.
     simulator.close()
 
   def test_get_screenshot(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     simulator.launch()
 
     screenshot = simulator.get_screenshot()
@@ -43,7 +50,9 @@ class FakeSimulatorTest(absltest.TestCase):
     np.testing.assert_equal(screenshot.dtype, np.uint8)
 
   def test_log_stream(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     simulator.launch()
     log_stream = simulator.create_log_stream()
     # Start yielding lines from LogStream.
@@ -61,12 +70,16 @@ class FakeSimulatorTest(absltest.TestCase):
         break
 
   def test_adb_output(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     simulator.launch()
     adb_controller = simulator.create_adb_controller()
     line = adb_controller.execute_command(['shell', 'dumpsys', 'input'])
     line = line.decode('utf-8')
-    orientation = re.match(r'\s+SurfaceOrientation:\s+(\d)', line).group(1)
+    matches = re.match(r'\s+SurfaceOrientation:\s+(\d)', line)
+    self.assertIsNotNone(matches)
+    orientation = matches.group(1)
     self.assertEqual(orientation, '0')
     line = adb_controller.execute_command(['shell', 'service', 'check', 'foo'])
     line = line.decode('utf-8')
@@ -77,7 +90,9 @@ class FakeSimulatorTest(absltest.TestCase):
                      'topActivity=ComponentInfo{fake_activity}')
 
   def test_send_touch(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     simulator.launch()
     simulator.send_touch([(0, 1, True, 0)])
     simulator.send_touch([(0, 1, False, 0)])
@@ -85,11 +100,13 @@ class FakeSimulatorTest(absltest.TestCase):
     # without crashing anything.
 
   def test_send_key(self):
-    simulator = fake_simulator.FakeSimulator(screen_dimensions=(320, 480))
+    simulator = fake_simulator.FakeSimulator(
+        config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))
+    )
     simulator.launch()
-    simulator.send_key(123, 'keydown')
-    simulator.send_key(123, 'keyup')
-    simulator.send_key(123, 'keypress')
+    simulator.send_key(np.int32(123), 'keydown')
+    simulator.send_key(np.int32(123), 'keyup')
+    simulator.send_key(np.int32(123), 'keypress')
     # No assertions, we just want to ensure that `send_key()` can be called
     # without crashing anything.
 

android_env/components/specs.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,11 +15,8 @@
 
 """Base specs for AndroidEnv."""
 
-from typing import Dict
-
 from android_env.components import action_type
 from android_env.proto import task_pb2
-import dm_env
 from dm_env import specs
 import numpy as np
 
@@ -41,11 +38,13 @@ _PROTO_DTYPE_TO_NUMPY_DTYPE = {
     task_pb2.ArraySpec.DataType.STRING_U25: np.dtype(('<U25')),
     task_pb2.ArraySpec.DataType.STRING_U250: np.dtype(('<U250')),
     task_pb2.ArraySpec.DataType.STRING: np.dtype(('<U0')),
+    task_pb2.ArraySpec.DataType.OBJECT: np.dtype('O'),
 }
 
 
-def base_action_spec(num_fingers: int = 1,
-                     enable_key_events: bool = False) -> Dict[str, specs.Array]:
+def base_action_spec(
+    num_fingers: int = 1, enable_key_events: bool = False
+) -> dict[str, specs.Array]:
   """Default action spec for AndroidEnv.
 
   Args:
@@ -58,6 +57,8 @@ def base_action_spec(num_fingers: int = 1,
     touch_position: Position [x, y] of the touch action, where x, y are float
       values between 0.0 and 1.0 corresponding to the relative position on the
       screen. IGNORED when (action_type != ActionType.TOUCH).
+    keycode: code representing the desired key press in XKB format. See the
+      emulator_controller_pb2 for details.
     action_type_i: Action type for additional fingers (i>1).
     touch_position_i: Touch position for additional fingers (i>1).
   """
@@ -98,7 +99,7 @@ def base_action_spec(num_fingers: int = 1,
   return action_spec
 
 
-def base_observation_spec(height: int, width: int) -> Dict[str, specs.Array]:
+def base_observation_spec(height: int, width: int) -> dict[str, specs.Array]:
   """Default observation spec for AndroidEnv.
 
   Args:
@@ -119,30 +120,19 @@ def base_observation_spec(height: int, width: int) -> Dict[str, specs.Array]:
 
   return {
       'pixels':
-          specs.Array(
+          specs.BoundedArray(
               shape=(height, width, 3),
               dtype=np.uint8,
-              name='pixels'),
+              name='pixels',
+              minimum=0,
+              maximum=255),
       'timedelta':
           specs.Array(shape=(), dtype=np.int64, name='timedelta'),
       'orientation':
-          specs.Array(shape=np.array([4]), dtype=np.uint8, name='orientation'),
-  }
-
-
-def base_task_extras_spec(task: task_pb2.Task) -> Dict[str, dm_env.specs.Array]:
-  """Task extras spec for AndroidEnv, as read from a task_pb2.Task."""
-
-  return {
-      spec.name: _convert_spec(spec)
-      for spec in task.extras_spec
+          specs.BoundedArray(
+              shape=np.array([4]),
+              dtype=np.uint8,
+              name='orientation',
+              minimum=0,
+              maximum=1),
   }
-
-
-def _convert_spec(array_spec: task_pb2.ArraySpec) -> specs.Array:
-  """Converts ArraySpec proto to dm_env specs.Array."""
-
-  return specs.Array(
-      shape=array_spec.shape,
-      dtype=_PROTO_DTYPE_TO_NUMPY_DTYPE[array_spec.dtype],
-      name=array_spec.name)

android_env/components/specs_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.
@@ -79,49 +79,6 @@ class SpecsTest(parameterized.TestCase):
     self.assertEqual(observation_spec['orientation'].shape, (4,))
     self.assertEqual(observation_spec['orientation'].dtype, np.uint8)
 
-  def test_base_task_extras_spec(self):
-    array_spec_1 = task_pb2.ArraySpec()
-    array_spec_1.name = 'my_extra_1'
-    array_spec_1.shape.extend([10, 10])
-    array_spec_1.dtype = task_pb2.ArraySpec.FLOAT
-
-    array_spec_2 = task_pb2.ArraySpec()
-    array_spec_2.name = 'my_extra_2'
-    array_spec_2.shape.extend([1])
-    array_spec_2.dtype = task_pb2.ArraySpec.UINT8
-
-    fake_task = task_pb2.Task()
-    fake_task.extras_spec.extend([array_spec_1, array_spec_2])
-    task_extras_spec = specs.base_task_extras_spec(fake_task)
-    for spec in task_extras_spec.values():
-      self.assertIsInstance(spec, dm_env_specs.Array)
-
-    self.assertEqual(task_extras_spec['my_extra_1'].shape, (10, 10))
-    self.assertEqual(task_extras_spec['my_extra_2'].shape, (1,))
-    self.assertEqual(task_extras_spec['my_extra_1'].dtype, np.float32)
-    self.assertEqual(task_extras_spec['my_extra_2'].dtype, np.uint8)
-
-  @parameterized.parameters(
-      ('name_1', [480, 320, 3], task_pb2.ArraySpec.FLOAT, np.float32),
-      ('name_2', [100, 100, 3], task_pb2.ArraySpec.INT32, np.int32),
-      ('name_3', [123, 456, 3], task_pb2.ArraySpec.UINT8, np.uint8),
-      ('name_4', [480, 320, 1], task_pb2.ArraySpec.BOOL, np.bool_),
-      ('', [480, 320], task_pb2.ArraySpec.STRING_U25, np.dtype(('<U25'))),
-  )
-  def test_convert_spec(self, name, shape, dtype, expected_dtype):
-    fake_array_spec = task_pb2.ArraySpec()
-    fake_array_spec.name = name
-    fake_array_spec.dtype = dtype
-    fake_array_spec.shape.extend(shape)
-    fake_task = task_pb2.Task()
-    fake_task.extras_spec.extend([fake_array_spec])
-    task_extras_spec = specs.base_task_extras_spec(fake_task)
-    for spec in task_extras_spec.values():
-      self.assertIsInstance(spec, dm_env_specs.Array)
-
-    self.assertEqual(task_extras_spec[name].shape, tuple(shape))
-    self.assertEqual(task_extras_spec[name].dtype, expected_dtype)
-
 
 if __name__ == '__main__':
   absltest.main()

android_env/components/task_manager.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.
@@ -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, Dict
+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
@@ -35,34 +37,24 @@ import dm_env
 import numpy as np
 
 
-class TaskManager():
+class TaskManager:
   """Handles all events and information related to the task."""
 
   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,14 +84,7 @@ 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]:
+  def stats(self) -> dict[str, Any]:
     """Returns a dictionary of stats.
 
     This method is expected to be called after setup_task() has been called.
@@ -109,21 +94,24 @@ class TaskManager():
       output.update(self._setup_step_interpreter.stats())
     return output
 
-  def setup_task(
+  def setup_task(self) -> None:
+    """Performs one-off task setup.."""
+    self._setup_step_interpreter.interpret(self._task.setup_steps)
+
+  def stop(self) -> None:
+    """Suspends task processing."""
+    self._stop_logcat_thread()
+
+  def start(
       self,
       adb_call_parser_factory: Callable[[], adb_call_parser_lib.AdbCallParser],
       log_stream: log_stream_lib.LogStream) -> None:
-    """Performs one-off task setup.."""
+    """Starts task processing."""
 
     self._start_logcat_thread(log_stream=log_stream)
+    self._logcat_thread.resume()
     self._start_dumpsys_thread(adb_call_parser_factory())
     self._start_setup_step_interpreter(adb_call_parser_factory())
-    self._setup_step_interpreter.interpret(self._task.setup_steps)
-
-  def stop_task(self) -> None:
-    """Suspends task processing."""
-
-    self._stop_logcat_thread()
 
   def reset_task(self) -> None:
     """Resets a task for a new run."""
@@ -146,7 +134,7 @@ class TaskManager():
           'episode_end': False,
       }
 
-  def rl_reset(self, observation: Dict[str, Any]) -> dm_env.TimeStep:
+  def rl_reset(self, observation: dict[str, Any]) -> dm_env.TimeStep:
     """Performs one RL step."""
 
     self._stats['episode_steps'] = 0
@@ -163,7 +151,7 @@ class TaskManager():
         discount=0.0,
         observation=observation)
 
-  def rl_step(self, observation: Dict[str, Any]) -> dm_env.TimeStep:
+  def rl_step(self, observation: dict[str, Any]) -> dm_env.TimeStep:
     """Performs one RL step."""
 
     self._stats['episode_steps'] += 1
@@ -184,7 +172,7 @@ class TaskManager():
     self._latest_values['reward'] = 0.0
     return reward
 
-  def _get_current_extras(self) -> Dict[str, Any]:
+  def _get_current_extras(self) -> dict[str, Any]:
     """Returns task extras accumulated since the last step."""
     extras = {}
     for name, values in self._latest_values['extra'].items():
@@ -213,8 +201,8 @@ class TaskManager():
     if self._task.max_episode_steps > 0:
       if self._stats['episode_steps'] > self._task.max_episode_steps:
         self._stats['reset_count_max_duration_reached'] += 1
-        logging.info('Maximum task duration (steps) reached. '
-                     'Truncating the episode.')
+        logging.info('Maximum task duration (%r steps) reached. '
+                     'Truncating the episode.', self._task.max_episode_steps)
         return dm_env.truncation
 
     if self._task.max_episode_sec > 0.0:
@@ -222,8 +210,8 @@ class TaskManager():
       max_episode_sec = self._task.max_episode_sec
       if task_duration > datetime.timedelta(seconds=int(max_episode_sec)):
         self._stats['reset_count_max_duration_reached'] += 1
-        logging.info('Maximum task duration (sec) reached. '
-                     'Truncating the episode.')
+        logging.info('Maximum task duration (%r sec) reached. '
+                     'Truncating the episode.', max_episode_sec)
         return dm_env.truncation
 
     return dm_env.transition
@@ -245,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:
@@ -263,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
@@ -339,9 +329,16 @@ 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
       else:
         # No extra value provided for boolean extra. Setting value to True.
         extra = 1
@@ -375,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: