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

android_env/components/coordinator_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.
@@ -23,7 +23,9 @@ from absl.testing import absltest
 from absl.testing import parameterized
 from android_env.components import action_type
 from android_env.components import adb_call_parser
+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 errors
 from android_env.components import task_manager
 from android_env.components.simulators import base_simulator
@@ -34,15 +36,6 @@ import dm_env
 import numpy as np
 
 
-class MockScreenshotGetter:
-  def __init__(self):
-    self._screenshot_index = 0
-
-  def get_screenshot(self):
-    self._screenshot_index += 1
-    return np.array(self._screenshot_index, ndmin=3)
-
-
 class CoordinatorTest(parameterized.TestCase):
 
   def setUp(self):
@@ -64,8 +57,8 @@ class CoordinatorTest(parameterized.TestCase):
     self._coordinator = coordinator_lib.Coordinator(
         simulator=self._simulator,
         task_manager=self._task_manager,
-        num_fingers=1,
-        periodic_restart_time_min=0)
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+    )
 
   def tearDown(self):
     super().tearDown()
@@ -80,15 +73,31 @@ class CoordinatorTest(parameterized.TestCase):
 
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_reset(self, unused_mock_sleep):
+    """'relaunch_count_execute_action' should be zero if there's no error."""
+
+    self._coordinator.rl_reset()
+    stats = self._coordinator.stats()
+    self.assertIn('relaunch_count_execute_action', stats)
+    self.assertEqual(stats['relaunch_count_execute_action'], 0)
+
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_reset_error_sending_action(self, unused_mock_sleep):
+    """'relaunch_count_execute_action' should be positive if there's an error."""
+
+    self._simulator.send_touch.side_effect = errors.SendActionError()
     self._coordinator.rl_reset()
+    stats = self._coordinator.stats()
+    self.assertIn('relaunch_count_execute_action', stats)
+    self.assertEqual(stats['relaunch_count_execute_action'], 1)
 
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_lift_all_fingers(self, unused_mock_sleep):
     self._coordinator = coordinator_lib.Coordinator(
         simulator=self._simulator,
         task_manager=self._task_manager,
-        num_fingers=3,
-        periodic_restart_time_min=0)
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+        config=config_classes.CoordinatorConfig(num_fingers=3),
+    )
     self._coordinator.rl_reset()
     expected_actions = [
         # (x, y, is_down, identifier).
@@ -147,136 +156,6 @@ class CoordinatorTest(parameterized.TestCase):
     self.assertEqual(timestep.reward, 0.0)
     self.assertTrue(timestep.last())
 
-  @mock.patch.object(time, 'sleep', autospec=True)
-  def test_process_action_error_async(self, unused_mock_sleep):
-    mock_interaction_thread = mock.create_autospec(
-        coordinator_lib.InteractionThread)
-    with mock.patch.object(
-        coordinator_lib,
-        'InteractionThread',
-        autospec=True,
-        return_value=mock_interaction_thread):
-      coordinator = coordinator_lib.Coordinator(
-          simulator=self._simulator,
-          task_manager=self._task_manager,
-          num_fingers=1,
-          interaction_rate_sec=0.016)
-
-      def fake_rl_step(agent_action, simulator_signals):
-        del agent_action
-        self.assertFalse(simulator_signals['simulator_healthy'])
-        return dm_env.truncation(reward=0.0, observation=None)
-
-      self._task_manager.rl_step.side_effect = fake_rl_step
-      mock_interaction_thread.screenshot.side_effect = errors.ReadObservationError(
-      )
-      timestep = coordinator.rl_step(
-          agent_action={
-              'action_type': np.array(action_type.ActionType.LIFT),
-              'touch_position': np.array([0.5, 0.5]),
-          })
-      self.assertIsNone(timestep.observation)
-      self.assertEqual(timestep.reward, 0.0)
-      self.assertTrue(timestep.last())
-      coordinator.close()
-
-  def test_async_step_faster_than_screenshot(self):
-    """Return same screenshot when step is faster than the interaction rate."""
-    screenshot_getter = MockScreenshotGetter()
-    self._simulator.get_screenshot.side_effect = screenshot_getter.get_screenshot
-    def fake_rl_step(simulator_signals):
-      return dm_env.transition(
-          reward=10.0,
-          observation={
-              'pixels': simulator_signals['pixels'],
-              'orientation': simulator_signals['orientation'],
-              'timedelta': simulator_signals['timedelta'],
-              'extras': {
-                  'extra': [0.0]
-              }
-          })
-    self._task_manager.rl_step.side_effect = fake_rl_step
-    coordinator = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        num_fingers=1,
-        interaction_rate_sec=0.5)
-    ts1 = coordinator.rl_step(
-        agent_action={
-            'action_type': np.array(action_type.ActionType.LIFT),
-            'touch_position': np.array([0.5, 0.5]),
-        })
-    ts2 = coordinator.rl_step(
-        agent_action={
-            'action_type': np.array(action_type.ActionType.LIFT),
-            'touch_position': np.array([0.5, 0.5]),
-        })
-    np.testing.assert_almost_equal(ts2.observation['pixels'],
-                                   ts1.observation['pixels'])
-    coordinator.close()
-
-  def test_async_step_slower_than_screenshot(self):
-    """Return different screenshots when step slower than the interaction rate."""
-    screenshot_getter = MockScreenshotGetter()
-    self._simulator.get_screenshot.side_effect = screenshot_getter.get_screenshot
-
-    def fake_rl_step(simulator_signals):
-      return dm_env.transition(
-          reward=10.0,
-          observation={
-              'pixels': simulator_signals['pixels'],
-              'orientation': simulator_signals['orientation'],
-              'timedelta': simulator_signals['timedelta'],
-              'extras': {
-                  'extra': [0.0]
-              }
-          })
-
-    self._task_manager.rl_step.side_effect = fake_rl_step
-    coordinator = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        num_fingers=1,
-        interaction_rate_sec=0.01)
-    ts1 = coordinator.rl_step(
-        agent_action={
-            'action_type': np.array(action_type.ActionType.LIFT),
-            'touch_position': np.array([0.5, 0.5]),
-        })
-    time.sleep(0.5)
-    ts2 = coordinator.rl_step(
-        agent_action={
-            'action_type': np.array(action_type.ActionType.LIFT),
-            'touch_position': np.array([0.5, 0.5]),
-        })
-    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal,
-                             ts2.observation['pixels'],
-                             ts1.observation['pixels'])
-    coordinator.close()
-
-  def test_interaction_thread_closes_upon_relaunch(self):
-    """Async coordinator should kill the InteractionThread when relaunching."""
-    mock_interaction_thread = mock.create_autospec(
-        coordinator_lib.InteractionThread)
-    with mock.patch.object(
-        coordinator_lib,
-        'InteractionThread',
-        autospec=True,
-        return_value=mock_interaction_thread):
-      coordinator = coordinator_lib.Coordinator(
-          simulator=self._simulator,
-          task_manager=self._task_manager,
-          num_fingers=1,
-          periodic_restart_time_min=1E-6,
-          interaction_rate_sec=0.5)
-      mock_interaction_thread.stop.assert_not_called()
-      mock_interaction_thread.join.assert_not_called()
-      time.sleep(0.1)
-      coordinator.rl_reset()
-      mock_interaction_thread.stop.assert_called_once()
-      mock_interaction_thread.join.assert_called_once()
-      coordinator.close()
-
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_execute_action_touch(self, unused_mock_sleep):
 
@@ -308,8 +187,9 @@ class CoordinatorTest(parameterized.TestCase):
     self._coordinator = coordinator_lib.Coordinator(
         simulator=self._simulator,
         task_manager=self._task_manager,
-        num_fingers=3,
-        periodic_restart_time_min=0)
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+        config=config_classes.CoordinatorConfig(num_fingers=3),
+    )
 
     def fake_rl_step(simulator_signals):
       return dm_env.transition(
@@ -398,143 +278,6 @@ class CoordinatorTest(parameterized.TestCase):
     self.assertEqual(response, expected_response)
     self._adb_call_parser.parse.assert_called_with(call)
 
-  @mock.patch.object(time, 'sleep', autospec=True)
-  @mock.patch.object(tempfile, 'gettempdir', autospec=True)
-  def test_with_tmp_dir_no_tempfile_call(self, mock_gettempdir,
-                                         unused_mock_sleep):
-    """If passing a `tmp_dir`, `tempfile.gettempdir()` should not be called."""
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        periodic_restart_time_min=0,
-        tmp_dir=absltest.get_default_test_tmpdir())
-    mock_gettempdir.assert_not_called()
-
-  @mock.patch.object(time, 'sleep', autospec=True)
-  @mock.patch.object(tempfile, 'gettempdir', autospec=True)
-  def test_no_tmp_dir_calls_tempfile(self, mock_gettempdir, unused_mock_sleep):
-    """If not passing a `tmp_dir`, `tempfile.gettempdir()` should be called."""
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        periodic_restart_time_min=0)
-    mock_gettempdir.assert_called_once()
-
-  @parameterized.parameters(
-      (True, '1'),
-      (False, '0'),
-  )
-  @mock.patch.object(time, 'sleep', autospec=True)
-  def test_touch_indicator(self, show, expected_value, unused_mock_sleep):
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        show_touches=show)
-    self._adb_call_parser.parse.assert_any_call(
-        adb_pb2.AdbRequest(
-            settings=adb_pb2.AdbRequest.SettingsRequest(
-                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
-                put=adb_pb2.AdbRequest.SettingsRequest.Put(
-                    key='show_touches', value=expected_value))))
-
-  @parameterized.parameters(
-      (True, '1'),
-      (False, '0'),
-  )
-  @mock.patch.object(time, 'sleep', autospec=True)
-  def test_pointer_location(self, show, expected_value, unused_mock_sleep):
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        show_pointer_location=show)
-    self._adb_call_parser.parse.assert_any_call(
-        adb_pb2.AdbRequest(
-            settings=adb_pb2.AdbRequest.SettingsRequest(
-                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
-                put=adb_pb2.AdbRequest.SettingsRequest.Put(
-                    key='pointer_location', value=expected_value))))
-
-  @parameterized.parameters(
-      (True, True, 'null*'),
-      (True, False, 'immersive.status=*'),
-      (False, True, 'immersive.navigation=*'),
-      (False, False, 'immersive.full=*'),
-      (None, None, 'immersive.full=*'),  # Defaults to hiding both.
-  )
-  @mock.patch.object(time, 'sleep', autospec=True)
-  def test_bar_visibility(self, show_navigation_bar, show_status_bar,
-                          expected_value, unused_mock_sleep):
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        show_navigation_bar=show_navigation_bar,
-        show_status_bar=show_status_bar)
-    self._adb_call_parser.parse.assert_any_call(
-        adb_pb2.AdbRequest(
-            settings=adb_pb2.AdbRequest.SettingsRequest(
-                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,
-                put=adb_pb2.AdbRequest.SettingsRequest.Put(
-                    key='policy_control', value=expected_value))))
-
-  def test_load_state(self):
-    expected_response = state_pb2.LoadStateResponse(
-        status=state_pb2.LoadStateResponse.Status.OK
-    )
-    request = state_pb2.LoadStateRequest(args={'foo': 'bar'})
-    self._simulator.load_state.return_value = expected_response
-    stop_call_count = self._task_manager.stop.call_count
-    start_call_count = self._task_manager.start.call_count
-    response = self._coordinator.load_state(request)
-    self.assertEqual(response, expected_response)
-    self._simulator.load_state.assert_called_once_with(request)
-    self.assertEqual(self._task_manager.stop.call_count, stop_call_count + 1)
-    self.assertEqual(self._task_manager.start.call_count, start_call_count + 1)
-
-  def test_save_state(self):
-    expected_response = state_pb2.SaveStateResponse(
-        status=state_pb2.SaveStateResponse.Status.OK
-    )
-    request = state_pb2.SaveStateRequest(args={'foo': 'bar'})
-    self._simulator.save_state.return_value = expected_response
-    response = self._coordinator.save_state(request)
-    self.assertEqual(response, expected_response)
-    self._simulator.save_state.assert_called_once_with(request)
-
-  @mock.patch.object(time, 'sleep', autospec=True)
-  def test_update_task_succeeds(self, unused_mock_sleep):
-    task = task_pb2.Task(id='my_task')
-    stop_call_count = self._task_manager.stop.call_count
-    start_call_count = self._task_manager.start.call_count
-    setup_call_count = self._task_manager.setup_task.call_count
-    success = self._coordinator.update_task(task)
-    self.assertEqual(1,
-                     self._task_manager.stop.call_count - stop_call_count)
-    self.assertEqual(
-        1, self._task_manager.setup_task.call_count - setup_call_count)
-    self.assertEqual(1, self._task_manager.start.call_count - start_call_count)
-    self._task_manager.update_task.assert_called_once_with(task)
-    self.assertTrue(success)
-    self._task_manager.stats.return_value = {}
-    self.assertEqual(0, self._coordinator.stats()['failed_task_updates'])
-
-  @mock.patch.object(time, 'sleep', autospec=True)
-  def test_update_task_fails(self, unused_mock_sleep):
-    task = task_pb2.Task(id='my_task')
-    self._task_manager.setup_task.side_effect = errors.StepCommandError
-    stop_call_count = self._task_manager.stop.call_count
-    start_call_count = self._task_manager.start.call_count
-    setup_call_count = self._task_manager.setup_task.call_count
-    success = self._coordinator.update_task(task)
-    self.assertEqual(1,
-                     self._task_manager.stop.call_count - stop_call_count)
-    self.assertEqual(
-        1, self._task_manager.setup_task.call_count - setup_call_count)
-    self.assertEqual( 1, self._task_manager.start.call_count - start_call_count)
-    self._task_manager.update_task.assert_called_once_with(task)
-    self.assertFalse(success)
-    self._task_manager.stats.return_value = {}
-    self.assertEqual(1, self._coordinator.stats()['failed_task_updates'])
-
 
 if __name__ == '__main__':
   absltest.main()

android_env/components/device_settings.py

@@ -0,0 +1,174 @@
+# 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.
+
+"""Sets and gets some global settings on an Android device."""
+
+from typing import Final
+from unittest import mock
+
+from absl import logging
+from android_env.components import adb_call_parser
+from android_env.components import config_classes
+from android_env.components.simulators import base_simulator
+from android_env.proto import adb_pb2
+import numpy as np
+
+
+# The internal `AdbCallParser` instance is lazily instantiated within
+# `DeviceSettings`. If we make it optional (i.e. `| None`), pytype will think
+# that it could be `None`, requiring either explicit runtime checks or escape
+# hatches in every actual call, even if it's never actually `None` if reached
+# via the public API.
+# The trick here is to create this dummy instance of the right type that's used
+# as a sentinel to indicate that it hasn't been initialized yet.
+_PLACEHOLDER_ADB_CALL_PARSER: Final[adb_call_parser.AdbCallParser] = (
+    mock.create_autospec(adb_call_parser.AdbCallParser)
+)
+
+
+class DeviceSettings:
+  """An abstraction for general properties and settings of an Android device."""
+
+  def __init__(self, simulator: base_simulator.BaseSimulator):
+    self._simulator = simulator
+    self._adb_call_parser = _PLACEHOLDER_ADB_CALL_PARSER
+
+    # The size of the device screen in pixels.
+    self._screen_width: int = 0
+    self._screen_height: int = 0
+    # The device orientation.
+    self._orientation = np.zeros(4, dtype=np.uint8)
+
+  def update(self, config: config_classes.DeviceSettingsConfig) -> None:
+    """Sets the configuration of the device according to `config`."""
+
+    if self._adb_call_parser is _PLACEHOLDER_ADB_CALL_PARSER:
+      self._adb_call_parser = adb_call_parser.AdbCallParser(
+          adb_controller=self._simulator.create_adb_controller()
+      )
+
+    self._update_screen_size()
+    self._set_show_touches(config.show_touches)
+    self._set_show_pointer_location(config.show_pointer_location)
+    self._set_status_navigation_bars(
+        config.show_navigation_bar, config.show_status_bar
+    )
+
+  def screen_width(self) -> int:
+    """The screen width in pixels. Only valid after `update()` is called."""
+
+    return self._screen_width
+
+  def screen_height(self) -> int:
+    """The screen height in pixels. Only valid after `update()` is called."""
+
+    return self._screen_height
+
+  def get_orientation(self) -> np.ndarray:
+    """Returns the device orientation. Please see specs.py for details."""
+
+    if self._adb_call_parser is _PLACEHOLDER_ADB_CALL_PARSER:
+      self._adb_call_parser = adb_call_parser.AdbCallParser(
+          adb_controller=self._simulator.create_adb_controller()
+      )
+
+    self._update_orientation()
+    return self._orientation
+
+  def _update_screen_size(self) -> None:
+    """Sets the screen size from a screenshot ignoring the color channel."""
+
+    screenshot = self._simulator.get_screenshot()
+    self._screen_height = screenshot.shape[0]
+    self._screen_width = screenshot.shape[1]
+
+  def _set_show_touches(self, show: bool) -> None:
+    """Whether to display circles indicating the touch position."""
+
+    self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            settings=adb_pb2.AdbRequest.SettingsRequest(
+                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
+                put=adb_pb2.AdbRequest.SettingsRequest.Put(
+                    key='show_touches', value='1' if show else '0'
+                ),
+            )
+        )
+    )
+
+  def _set_show_pointer_location(self, show: bool) -> None:
+    """Whether to display blue lines on the screen indicating touch position."""
+
+    self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            settings=adb_pb2.AdbRequest.SettingsRequest(
+                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
+                put=adb_pb2.AdbRequest.SettingsRequest.Put(
+                    key='pointer_location', value='1' if show else '0'
+                ),
+            )
+        )
+    )
+
+  def _set_status_navigation_bars(
+      self, show_navigation: bool, show_status: bool
+  ) -> None:
+    """Whether to display the status (top) and navigation (bottom) bars."""
+
+    if show_navigation and show_status:
+      policy_control_value = 'null*'
+    elif show_navigation and not show_status:
+      policy_control_value = 'immersive.status=*'
+    elif not show_navigation and show_status:
+      policy_control_value = 'immersive.navigation=*'
+    else:
+      policy_control_value = 'immersive.full=*'
+
+    self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            settings=adb_pb2.AdbRequest.SettingsRequest(
+                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,
+                put=adb_pb2.AdbRequest.SettingsRequest.Put(
+                    key='policy_control', value=policy_control_value
+                ),
+            )
+        )
+    )
+
+  def _update_orientation(self) -> None:
+    """Updates the current device orientation."""
+
+    # Skip fetching the orientation if we already have it.
+    if not np.all(self._orientation == np.zeros(4)):
+      return
+
+    orientation_response = self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            get_orientation=adb_pb2.AdbRequest.GetOrientationRequest()
+        )
+    )
+    if orientation_response.status != adb_pb2.AdbResponse.Status.OK:
+      logging.error('Got bad orientation: %r', orientation_response)
+      return
+
+    orientation = orientation_response.get_orientation.orientation
+    if orientation not in {0, 1, 2, 3}:
+      logging.error('Got bad orientation: %r', orientation)
+      return
+
+    # Transform into one-hot format.
+    orientation_onehot = np.zeros([4], dtype=np.uint8)
+    orientation_onehot[orientation] = 1
+    self._orientation = orientation_onehot