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

android_env/wrappers/base_wrapper.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.
@@ -21,7 +21,6 @@ from absl import logging
 from android_env import env_interface
 from android_env.proto import adb_pb2
 from android_env.proto import state_pb2
-from android_env.proto import task_pb2
 import dm_env
 from dm_env import specs
 import numpy as np
@@ -101,9 +100,6 @@ class BaseWrapper(env_interface.AndroidEnvInterface):
                        adb_call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:
     return self._env.execute_adb_call(adb_call)
 
-  def update_task(self, task: task_pb2.Task) -> bool:
-    return self._env.update_task(task)
-
   @property
   def raw_action(self):
     return self._env.raw_action

android_env/wrappers/base_wrapper_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.
@@ -21,7 +21,6 @@ from absl import logging
 from absl.testing import absltest
 from android_env import env_interface
 from android_env.proto import state_pb2
-from android_env.proto import task_pb2
 from android_env.wrappers import base_wrapper
 
 
@@ -84,11 +83,6 @@ class BaseWrapperTest(absltest.TestCase):
     self.assertEqual(wrapped_env.save_state(save_request), expected_response)
     base_env.save_state.assert_called_once_with(save_request)
 
-    task = task_pb2.Task(id='my_task')
-    base_env.update_task.return_value = False
-    self.assertFalse(wrapped_env.update_task(task))
-    base_env.update_task.assert_called_once_with(task)
-
     wrapped_env.close()
     base_env.close.assert_called_once()
 

android_env/wrappers/discrete_action_wrapper.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.
@@ -15,7 +15,7 @@
 
 """Wraps the AndroidEnv environment to provide discrete actions."""
 
-from typing import Optional, Sequence
+from collections.abc import Sequence
 
 from android_env.components import action_type
 from android_env.wrappers import base_wrapper
@@ -24,23 +24,24 @@ from dm_env import specs
 import numpy as np
 
 
-NOISE_CLIP_VALUE = 0.4999
+_NOISE_CLIP_VALUE = 0.4999
 
 
 class DiscreteActionWrapper(base_wrapper.BaseWrapper):
   """AndroidEnv with discrete actions."""
 
-  def __init__(self,
-               env: dm_env.Environment,
-               action_grid: Optional[Sequence[int]] = (10, 10),
-               redundant_actions: bool = True,
-               noise: float = 0.1):
-
+  def __init__(
+      self,
+      env: dm_env.Environment,
+      action_grid: Sequence[int] = (10, 10),
+      redundant_actions: bool = True,
+      noise: float = 0.1,
+  ):
     super().__init__(env)
     self._parent_action_spec = self._env.action_spec()
     self._assert_base_env()
     self._action_grid = action_grid  # [height, width]
-    self._grid_size = np.product(self._action_grid)
+    self._grid_size = np.prod(self._action_grid)
     self._num_action_types = self._parent_action_spec['action_type'].num_values
     self._redundant_actions = redundant_actions
     self._noise = noise
@@ -57,9 +58,9 @@ class DiscreteActionWrapper(base_wrapper.BaseWrapper):
     """Number of discrete actions."""
 
     if self._redundant_actions:
-      return np.product(self._action_grid) * self._num_action_types
+      return self._grid_size * self._num_action_types
     else:
-      return np.product(self._action_grid) + self._num_action_types - 1
+      return self._grid_size + self._num_action_types - 1
 
   def step(self, action: dict[str, int]) -> dm_env.TimeStep:
     """Take a step in the base environment."""
@@ -133,8 +134,8 @@ class DiscreteActionWrapper(base_wrapper.BaseWrapper):
     noise_y = np.random.normal(loc=0.0, scale=self._noise)
 
     # Noise is clipped so that the action will strictly stay in the cell.
-    noise_x = max(min(noise_x, NOISE_CLIP_VALUE), -NOISE_CLIP_VALUE)
-    noise_y = max(min(noise_y, NOISE_CLIP_VALUE), -NOISE_CLIP_VALUE)
+    noise_x = max(min(noise_x, _NOISE_CLIP_VALUE), -_NOISE_CLIP_VALUE)
+    noise_y = max(min(noise_y, _NOISE_CLIP_VALUE), -_NOISE_CLIP_VALUE)
 
     x_pos = (x_pos_grid + 0.5 + noise_x) / self._action_grid[1]  # WIDTH
     y_pos = (y_pos_grid + 0.5 + noise_y) / self._action_grid[0]  # HEIGHT

android_env/wrappers/discrete_action_wrapper_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.

android_env/wrappers/flat_interface_wrapper.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.
@@ -15,7 +15,7 @@
 
 """Wraps the AndroidEnv environment to make its interface flat."""
 
-from typing import Any, Union
+from typing import Any
 
 from android_env.wrappers import base_wrapper
 import dm_env
@@ -74,7 +74,7 @@ class FlatInterfaceWrapper(base_wrapper.BaseWrapper):
     assert isinstance(base_action_spec, dict)
     assert isinstance(base_action_spec[self._action_name], specs.BoundedArray)
 
-  def _process_action(self, action: Union[int, np.ndarray, dict[str, Any]]):
+  def _process_action(self, action: int | np.ndarray | dict[str, Any]):
     if self._flat_actions:
       return {self._action_name: action}
     else:
@@ -103,7 +103,7 @@ class FlatInterfaceWrapper(base_wrapper.BaseWrapper):
     timestep = self._env.step(self._process_action(action))
     return self._process_timestep(timestep)
 
-  def observation_spec(self) -> Union[specs.Array, dict[str, specs.Array]]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
+  def observation_spec(self) -> specs.Array | dict[str, specs.Array]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
     if self._flat_observations:
       pixels_spec = self._env.observation_spec()['pixels']
       if not self._keep_action_layer:
@@ -112,7 +112,7 @@ class FlatInterfaceWrapper(base_wrapper.BaseWrapper):
     else:
       return self._env.observation_spec()
 
-  def action_spec(self) -> Union[specs.BoundedArray, dict[str, specs.Array]]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
+  def action_spec(self) -> specs.BoundedArray | dict[str, specs.Array]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
     if self._flat_actions:
       return self._env.action_spec()[self._action_name]
     else:

android_env/wrappers/flat_interface_wrapper_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.

android_env/wrappers/float_pixels_wrapper.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.
@@ -15,7 +15,7 @@
 
 """Converts pixel observation to from int to float32 between 0.0 and 1.0."""
 
-from android_env.components import utils
+from android_env.components import pixel_fns
 from android_env.wrappers import base_wrapper
 import dm_env
 from dm_env import specs
@@ -35,8 +35,9 @@ class FloatPixelsWrapper(base_wrapper.BaseWrapper):
       self, observation: dict[str, np.ndarray]
   ) -> dict[str, np.ndarray]:
     if self._should_convert_int_to_float:
-      float_pixels = utils.convert_int_to_float(observation['pixels'],
-                                                self._input_spec, np.float32)
+      float_pixels = pixel_fns.convert_int_to_float(
+          observation['pixels'], self._input_spec
+      )
       observation['pixels'] = float_pixels
     return observation
 

android_env/wrappers/float_pixels_wrapper_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.

android_env/wrappers/gym_wrapper.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/wrappers/gym_wrapper_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.

android_env/wrappers/image_rescale_wrapper.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.
@@ -15,7 +15,7 @@
 
 """Wraps the AndroidEnv environment to rescale the observations."""
 
-from typing import Optional, Sequence
+from collections.abc import Sequence
 
 from android_env.wrappers import base_wrapper
 import dm_env
@@ -37,8 +37,9 @@ class ImageRescaleWrapper(base_wrapper.BaseWrapper):
   def __init__(
       self,
       env: dm_env.Environment,
-      zoom_factors: Optional[Sequence[float]] = (0.5, 0.5),
-      grayscale: bool = False):
+      zoom_factors: Sequence[float] | None = (0.5, 0.5),
+      grayscale: bool = False,
+  ):
     super().__init__(env)
     assert 'pixels' in self._env.observation_spec()
     assert self._env.observation_spec()['pixels'].shape[-1] in [1, 3], (
@@ -70,14 +71,16 @@ class ImageRescaleWrapper(base_wrapper.BaseWrapper):
     return self._resize_image_array(image, new_shape)
 
   def _resize_image_array(
-      self,
-      grayscale_or_rbg_array: np.ndarray,
-      new_shape: Sequence[int]) -> np.ndarray:
+      self, grayscale_or_rbg_array: np.ndarray, new_shape: np.ndarray
+  ) -> np.ndarray:
     """Resize color or grayscale/action_layer array to new_shape."""
-    assert np.array(new_shape).ndim == 1
+    assert new_shape.ndim == 1
     assert len(new_shape) == 2
-    resized_array = np.array(Image.fromarray(
-        grayscale_or_rbg_array.astype('uint8')).resize(new_shape))
+    resized_array = np.array(
+        Image.fromarray(grayscale_or_rbg_array.astype('uint8')).resize(
+            tuple(new_shape)
+        )
+    )
     if resized_array.ndim == 2:
       return np.expand_dims(resized_array, axis=-1)
     return resized_array

android_env/wrappers/image_rescale_wrapper_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.

android_env/wrappers/last_action_wrapper.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,7 +16,7 @@
 """Extends Android observation with the latest action taken."""
 
 from android_env.components import action_type
-from android_env.components import utils
+from android_env.components import pixel_fns
 from android_env.wrappers import base_wrapper
 import dm_env
 from dm_env import specs
@@ -76,8 +76,9 @@ class LastActionWrapper(base_wrapper.BaseWrapper):
     if ('action_type' in last_action and
         last_action['action_type'] == action_type.ActionType.TOUCH):
       touch_position = last_action['touch_position']
-      x, y = utils.touch_position_to_pixel_position(
-          touch_position, width_height=self._screen_dimensions[::-1])
+      x, y = pixel_fns.touch_position_to_pixel_position(
+          touch_position, width_height=self._screen_dimensions[::-1]
+      )
       last_action_layer[y, x] = 255
 
     return last_action_layer

android_env/wrappers/last_action_wrapper_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.

android_env/wrappers/rate_limit_wrapper.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/wrappers/rate_limit_wrapper_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.

android_env/wrappers/tap_action_wrapper.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.
@@ -15,7 +15,7 @@
 
 """Wraps the AndroidEnv environment to provide tap actions of a given duration."""
 
-from typing import Sequence
+from collections.abc import Sequence
 
 from android_env.components import action_type
 from android_env.wrappers import base_wrapper
@@ -23,9 +23,6 @@ import dm_env
 import numpy as np
 
 
-ActionType = action_type.ActionType
-
-
 class TapActionWrapper(base_wrapper.BaseWrapper):
   """AndroidEnv with tap actions."""
 
@@ -51,20 +48,23 @@ class TapActionWrapper(base_wrapper.BaseWrapper):
     if self._touch_only:
       assert action['action_type'] == 0
       touch_action = action.copy()
-      touch_action['action_type'] = np.array(ActionType.TOUCH).astype(
-          self.action_spec()['action_type'].dtype)
+      touch_action['action_type'] = np.array(
+          action_type.ActionType.TOUCH
+      ).astype(self.action_spec()['action_type'].dtype)
       actions = [touch_action] * self._num_frames
       lift_action = action.copy()
-      lift_action['action_type'] = np.array(ActionType.LIFT).astype(
-          self.action_spec()['action_type'].dtype)
+      lift_action['action_type'] = np.array(action_type.ActionType.LIFT).astype(
+          self.action_spec()['action_type'].dtype
+      )
       actions.append(lift_action)
 
     else:
-      if action['action_type'] == ActionType.TOUCH:
+      if action['action_type'] == action_type.ActionType.TOUCH:
         actions = [action] * self._num_frames
         lift_action = action.copy()
-        lift_action['action_type'] = np.array(ActionType.LIFT).astype(
-            self.action_spec()['action_type'].dtype)
+        lift_action['action_type'] = np.array(
+            action_type.ActionType.LIFT
+        ).astype(self.action_spec()['action_type'].dtype)
         actions.append(lift_action)
       else:
         actions = [action] * (self._num_frames + 1)

android_env/wrappers/tap_action_wrapper_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.
@@ -26,9 +26,6 @@ from dm_env import specs
 import numpy as np
 
 
-ActionType = action_type.ActionType
-
-
 def _make_array_spec(shape, dtype, name):
   return specs.BoundedArray(
       name=name,
@@ -56,8 +53,8 @@ class TapActionWrapperTest(absltest.TestCase):
     wrapped_env = tap_action_wrapper.TapActionWrapper(
         self.base_env, num_frames=3)
     action = {
-        'action_type': np.array(ActionType.REPEAT, dtype=np.int32),
-        'touch_position': np.array([0.5, 0.5], dtype=np.float32)
+        'action_type': np.array(action_type.ActionType.REPEAT, dtype=np.int32),
+        'touch_position': np.array([0.5, 0.5], dtype=np.float32),
     }
     actions = wrapped_env._process_action(action)
     self.assertLen(actions, wrapped_env._num_frames + 1)
@@ -67,8 +64,8 @@ class TapActionWrapperTest(absltest.TestCase):
     wrapped_env = tap_action_wrapper.TapActionWrapper(
         self.base_env, num_frames=3)
     action = {
-        'action_type': np.array(ActionType.LIFT, dtype=np.int32),
-        'touch_position': np.array([0.5, 0.5], dtype=np.float32)
+        'action_type': np.array(action_type.ActionType.LIFT, dtype=np.int32),
+        'touch_position': np.array([0.5, 0.5], dtype=np.float32),
     }
     actions = wrapped_env._process_action(action)
     self.assertLen(actions, wrapped_env._num_frames + 1)
@@ -78,12 +75,14 @@ class TapActionWrapperTest(absltest.TestCase):
     wrapped_env = tap_action_wrapper.TapActionWrapper(
         self.base_env, num_frames=3)
     action = {
-        'action_type': np.array(ActionType.TOUCH, dtype=np.int32),
-        'touch_position': np.array([0.5, 0.5], dtype=np.float32)
+        'action_type': np.array(action_type.ActionType.TOUCH, dtype=np.int32),
+        'touch_position': np.array([0.5, 0.5], dtype=np.float32),
     }
     actions = wrapped_env._process_action(action)
     self.assertLen(actions, wrapped_env._num_frames + 1)
-    self.assertEqual(actions[-1]['action_type'], np.array(ActionType.LIFT))
+    self.assertEqual(
+        actions[-1]['action_type'], np.array(action_type.ActionType.LIFT)
+    )
 
   def test_reset(self):
     wrapped_env = tap_action_wrapper.TapActionWrapper(
@@ -95,6 +94,7 @@ class TapActionWrapperTest(absltest.TestCase):
     self.assertEqual(fake_timestep, ts)
 
   def test_step(self):
+    # Arrange.
     wrapped_env = tap_action_wrapper.TapActionWrapper(
         self.base_env, num_frames=5)
     fake_timestep = dm_env.TimeStep(
@@ -103,13 +103,21 @@ class TapActionWrapperTest(absltest.TestCase):
         discount=1.0,
         observation='fake_obs')
     self.base_env.step.return_value = fake_timestep
+    self.base_env.stats.return_value = {}
+
+    # Act.
     ts = wrapped_env.step({
-        'action_type': np.array(ActionType.REPEAT, dtype=np.int32),
-        'touch_position': np.array([0.5, 0.5], dtype=np.float32)
+        'action_type': np.array(action_type.ActionType.REPEAT, dtype=np.int32),
+        'touch_position': np.array([0.5, 0.5], dtype=np.float32),
     })
+    stats = wrapped_env.stats()
+
+    # Assert.
     self.assertEqual(wrapped_env._num_frames+1, self.base_env.step.call_count)
     self.assertIsInstance(ts, dm_env.TimeStep)
-    self.assertEqual(wrapped_env._env_steps, 6)
+    self.assertIsInstance(stats, dict)
+    self.assertIn('env_steps', stats)
+    self.assertEqual(stats['env_steps'], 6)
 
   def test_observation_spec(self):
     wrapped_env = tap_action_wrapper.TapActionWrapper(
@@ -129,5 +137,32 @@ class TapActionWrapperTest(absltest.TestCase):
     self.assertEqual(self.base_env.action_spec(),
                      action_spec)
 
+  def test_stats(self):
+    """Checks that returned stats have expected properties."""
+
+    # Arrange.
+    self.base_env.stats.return_value = {
+        'some_key': 12345,
+        'another_key': 5.4321,
+    }
+    wrapped_env = tap_action_wrapper.TapActionWrapper(
+        self.base_env, num_frames=5
+    )
+
+    # Act.
+    stats = wrapped_env.stats()
+
+    # Assert.
+    self.assertIsInstance(stats, dict)
+    # Original entries should still be present.
+    self.assertIn('some_key', stats)
+    self.assertEqual(stats['some_key'], 12345)
+    self.assertIn('another_key', stats)
+    self.assertEqual(stats['another_key'], 5.4321)
+    # TapActionWrapper inserts its own `env_steps`.
+    self.assertIn('env_steps', stats)
+    self.assertEqual(stats['env_steps'], 0)
+
+
 if __name__ == '__main__':
   absltest.main()

{android_env-1.2.2.dist-info → android_env-1.2.3.dist-info}/METADATA

@@ -1,8 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: android-env
-Version: 1.2.2
+Version: 1.2.3
 Summary: AndroidEnv environment and library for training agents.
-Home-page: https://github.com/deepmind/android_env
 Author: DeepMind
 License: 
                                          Apache License
@@ -213,22 +212,21 @@ Keywords: Android,OS,reinforcement-learning
 Requires-Python: >=3.10
 Description-Content-Type: text/plain
 License-File: LICENSE
-Requires-Dist: absl-py (>=0.1.0)
-Requires-Dist: dm-env
+Requires-Dist: absl-py>=0.1.0
+Requires-Dist: dm_env
 Requires-Dist: grpcio
-Requires-Dist: numpy (>=1.21)
-Requires-Dist: portpicker (>=1.2.0)
-Requires-Dist: protobuf (>=2.6)
+Requires-Dist: numpy>=1.21
+Requires-Dist: portpicker>=1.2.0
+Requires-Dist: protobuf>=2.6
 Requires-Dist: pygame
+Provides-Extra: testing
+Requires-Dist: gym; extra == "testing"
+Requires-Dist: pillow; extra == "testing"
+Requires-Dist: pytype; extra == "testing"
 Provides-Extra: acme
-Requires-Dist: dm-acme ; extra == 'acme'
+Requires-Dist: dm-acme; extra == "acme"
 Provides-Extra: gym
-Requires-Dist: gym ; extra == 'gym'
-Provides-Extra: testing
-Requires-Dist: attrs (==20.3.0) ; extra == 'testing'
-Requires-Dist: gym ; extra == 'testing'
-Requires-Dist: pillow ; extra == 'testing'
-Requires-Dist: pytype ; extra == 'testing'
-Requires-Dist: pytest-xdist ; extra == 'testing'
+Requires-Dist: gym; extra == "gym"
+Dynamic: license-file
 
 Read the README at https://github.com/deepmind/android_env for more information.