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

android_env/wrappers/image_rescale_wrapper.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,7 +15,7 @@
 
 """Wraps the AndroidEnv environment to rescale the observations."""
 
-from typing import Optional, Sequence, Dict
+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], (
@@ -50,16 +51,8 @@ class ImageRescaleWrapper(base_wrapper.BaseWrapper):
     # want to zoom the number of channels so we just multiply it by 1.0.
     self._zoom_factors = tuple(zoom_factors) + (1.0,)
 
-    # Save the raw image for making videos, for example.
-    self._raw_pixels = None
-
-  @property
-  def raw_pixels(self):
-    return self._raw_pixels
-
   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
     observation = timestep.observation
-    self._raw_pixels = observation['pixels'].copy()
     processed_observation = observation.copy()
     processed_observation['pixels'] = self._process_pixels(
         observation['pixels'])
@@ -78,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
@@ -98,15 +93,17 @@ class ImageRescaleWrapper(base_wrapper.BaseWrapper):
     timestep = self._env.step(action)
     return self._process_timestep(timestep)
 
-  def observation_spec(self) -> Dict[str, specs.Array]:
+  def observation_spec(self) -> dict[str, specs.Array]:
     parent_spec = self._env.observation_spec().copy()
     out_shape = np.multiply(parent_spec['pixels'].shape,
                             self._zoom_factors).astype(np.int32)
     if self._grayscale:
       # In grayscale mode we want the output shape to be [W, H, 1]
       out_shape[-1] = 1
-    parent_spec['pixels'] = specs.Array(
+    parent_spec['pixels'] = specs.BoundedArray(
         shape=out_shape,
         dtype=parent_spec['pixels'].dtype,
-        name=parent_spec['pixels'].name)
+        name=parent_spec['pixels'].name,
+        minimum=parent_spec['pixels'].minimum,
+        maximum=parent_spec['pixels'].maximum)
     return parent_spec

android_env/wrappers/image_rescale_wrapper_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.
@@ -15,42 +15,24 @@
 
 """Tests for android_env.wrappers.image_rescale_wrapper."""
 
-from typing import Any, Dict
+from typing import Any
+from unittest import mock
 
 from absl.testing import absltest
-from android_env import environment
+from android_env import env_interface
 from android_env.wrappers import image_rescale_wrapper
 import dm_env
 from dm_env import specs
 import numpy as np
 
 
-class FakeEnv(environment.AndroidEnv):
-  """A class that we can use to inject custom observations and specs."""
-
-  def __init__(self, obs_spec):
-    self._obs_spec = obs_spec
-    self._next_obs = None
-
-  def reset(self) -> dm_env.TimeStep:
-    return self._next_timestep
-
-  def step(self, action: Any) -> dm_env.TimeStep:
-    return self._next_timestep
-
-  def observation_spec(self) -> Dict[str, specs.Array]:
-    return self._obs_spec
-
-  def action_spec(self) -> Dict[str, specs.Array]:
-    assert False, 'This should not be called by tests.'
-
-  def set_next_timestep(self, timestep):
-    self._next_timestep = timestep
-
-
 def _simple_spec():
-  return specs.Array(
-      shape=np.array([300, 300, 3]), dtype=np.uint8, name='pixels')
+  return specs.BoundedArray(
+      shape=np.array([300, 300, 3]),
+      dtype=np.uint8,
+      name='pixels',
+      minimum=0,
+      maximum=255)
 
 
 def _simple_timestep():
@@ -65,9 +47,11 @@ def _simple_timestep():
 class ImageRescaleWrapperTest(absltest.TestCase):
 
   def test_100x50_grayscale(self):
-    obs_spec = {'pixels': _simple_spec()}
-    fake_env = FakeEnv(obs_spec)
-    fake_env.set_next_timestep(_simple_timestep())
+    fake_timestep = _simple_timestep()
+    fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)
+    fake_env.observation_spec.return_value = {'pixels': _simple_spec()}
+    fake_env.reset.return_value = fake_timestep
+    fake_env.step.return_value = fake_timestep
 
     wrapper = image_rescale_wrapper.ImageRescaleWrapper(
         fake_env, zoom_factors=(1.0 / 3, 1.0 / 6.0), grayscale=True)
@@ -81,9 +65,11 @@ class ImageRescaleWrapperTest(absltest.TestCase):
     self.assertEqual(step_image.shape, (100, 50, 1))
 
   def test_150x60_full_channels(self):
-    obs_spec = {'pixels': _simple_spec()}
-    fake_env = FakeEnv(obs_spec)
-    fake_env.set_next_timestep(_simple_timestep())
+    fake_timestep = _simple_timestep()
+    fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)
+    fake_env.observation_spec.return_value = {'pixels': _simple_spec()}
+    fake_env.reset.return_value = fake_timestep
+    fake_env.step.return_value = fake_timestep
 
     wrapper = image_rescale_wrapper.ImageRescaleWrapper(
         fake_env, zoom_factors=(1.0 / 2.0, 1.0 / 5.0))
@@ -97,9 +83,11 @@ class ImageRescaleWrapperTest(absltest.TestCase):
     self.assertEqual(step_image.shape, (150, 60, 3))
 
   def test_list_zoom_factor(self):
-    obs_spec = {'pixels': _simple_spec()}
-    fake_env = FakeEnv(obs_spec)
-    fake_env.set_next_timestep(_simple_timestep())
+    fake_timestep = _simple_timestep()
+    fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)
+    fake_env.observation_spec.return_value = {'pixels': _simple_spec()}
+    fake_env.reset.return_value = fake_timestep
+    fake_env.step.return_value = fake_timestep
 
     wrapper = image_rescale_wrapper.ImageRescaleWrapper(
         fake_env, zoom_factors=[0.5, 0.2])

android_env/wrappers/last_action_wrapper.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,10 +15,8 @@
 
 """Extends Android observation with the latest action taken."""
 
-from typing import Dict
-
 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
@@ -56,8 +54,8 @@ class LastActionWrapper(base_wrapper.BaseWrapper):
     return timestep._replace(observation=processed_observation)
 
   def _process_observation(
-      self, observation: Dict[str, np.ndarray]
-  ) -> Dict[str, np.ndarray]:
+      self, observation: dict[str, np.ndarray]
+  ) -> dict[str, np.ndarray]:
     """Extends observation with last_action data."""
     processed_observation = observation.copy()
     last_action_layer = self._get_last_action_layer(observation['pixels'])
@@ -78,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
@@ -92,20 +91,24 @@ class LastActionWrapper(base_wrapper.BaseWrapper):
     timestep = self._env.step(action)
     return self._process_timestep(timestep)
 
-  def observation_spec(self) -> Dict[str, specs.Array]:
+  def observation_spec(self) -> dict[str, specs.Array]:
     parent_spec = self._env.observation_spec().copy()
     shape = parent_spec['pixels'].shape
     if self._concat_to_pixels:
-      parent_spec['pixels'] = specs.Array(
+      parent_spec['pixels'] = specs.BoundedArray(
           shape=(shape[0], shape[1], shape[2] + 1),
           dtype=parent_spec['pixels'].dtype,
-          name=parent_spec['pixels'].name)
+          name=parent_spec['pixels'].name,
+          minimum=parent_spec['pixels'].minimum,
+          maximum=parent_spec['pixels'].maximum)
     else:
       parent_spec.update({
           'last_action':
-              specs.Array(
+              specs.BoundedArray(
                   shape=(shape[0], shape[1]),
                   dtype=parent_spec['pixels'].dtype,
-                  name='last_action')
+                  name='last_action',
+                  minimum=parent_spec['pixels'].minimum,
+                  maximum=parent_spec['pixels'].maximum)
       })
     return parent_spec

android_env/wrappers/last_action_wrapper_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.
@@ -15,10 +15,11 @@
 
 """Tests for android_env.wrappers.last_action_wrapper."""
 
-from typing import Any, Dict
+from typing import Any
+from unittest import mock
 
 from absl.testing import absltest
-from android_env import environment
+from android_env import env_interface
 from android_env.components import action_type
 from android_env.wrappers import last_action_wrapper
 import dm_env
@@ -26,37 +27,13 @@ from dm_env import specs
 import numpy as np
 
 
-class FakeEnv(environment.AndroidEnv):
-  """A class that we can use to inject custom observations and specs."""
-
-  def __init__(self, obs_spec):
-    self._obs_spec = obs_spec
-    self._next_obs = None
-    self._latest_action = {}
-
-  def reset(self) -> dm_env.TimeStep:
-    return self._next_timestep
-
-  def step(self, action: Any) -> dm_env.TimeStep:
-    self._latest_action = action
-    return self._next_timestep
-
-  def observation_spec(self) -> Dict[str, specs.Array]:
-    return self._obs_spec
-
-  def action_spec(self) -> Dict[str, specs.Array]:
-    assert False, 'This should not be called by tests.'
-
-  def set_next_timestep(self, timestep):
-    self._next_timestep = timestep
-
-  def close(self):
-    pass
-
-
 def _simple_spec():
-  return specs.Array(
-      shape=np.array([120, 80, 3]), dtype=np.uint8, name='pixels')
+  return specs.BoundedArray(
+      shape=np.array([120, 80, 3]),
+      dtype=np.uint8,
+      name='pixels',
+      minimum=0,
+      maximum=255)
 
 
 def _simple_timestep():
@@ -71,9 +48,11 @@ def _simple_timestep():
 class LastActionWrapperTest(absltest.TestCase):
 
   def test_concat_to_pixels(self):
-    obs_spec = {'pixels': _simple_spec()}
-    fake_env = FakeEnv(obs_spec)
-    fake_env.set_next_timestep(_simple_timestep())
+    fake_timestep = _simple_timestep()
+    fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)
+    fake_env.observation_spec.return_value = {'pixels': _simple_spec()}
+    fake_env.reset.return_value = fake_timestep
+    fake_env.step.return_value = fake_timestep
 
     wrapper = last_action_wrapper.LastActionWrapper(
         fake_env, concat_to_pixels=True)
@@ -86,10 +65,12 @@ class LastActionWrapperTest(absltest.TestCase):
     last_action_layer = reset_image[:, :, -1]
     self.assertEqual(np.sum(last_action_layer), 0)
 
-    step_timestep = wrapper.step(action={
+    action1 = {
         'action_type': action_type.ActionType.TOUCH,
         'touch_position': np.array([0.25, 0.75], dtype=np.float32),  # (W x H)
-    })
+    }
+    type(fake_env).raw_action = mock.PropertyMock(return_value=action1)
+    step_timestep = wrapper.step(action=action1)
     step_image = step_timestep.observation['pixels']
     self.assertEqual(step_image.shape, (120, 80, 4))  # (H x W)
     last_action_layer = step_image[:, :, -1]
@@ -97,19 +78,23 @@ class LastActionWrapperTest(absltest.TestCase):
     y, x = np.where(last_action_layer == 255)
     self.assertEqual((y.item(), x.item()), (90, 20))
 
-    step_timestep = wrapper.step(action={
+    action2 = {
         'action_type': action_type.ActionType.LIFT,
         'touch_position': np.array([0.25, 0.75], dtype=np.float32),
-    })
+    }
+    type(fake_env).raw_action = mock.PropertyMock(return_value=action2)
+    step_timestep = wrapper.step(action=action2)
     step_image = step_timestep.observation['pixels']
     self.assertEqual(step_image.shape, (120, 80, 4))
     last_action_layer = step_image[:, :, -1]
     self.assertEqual(np.sum(last_action_layer), 0)
 
-    step_timestep = wrapper.step(action={
+    action3 = {
         'action_type': action_type.ActionType.TOUCH,
         'touch_position': np.array([0.25, 1.0], dtype=np.float32),
-    })
+    }
+    type(fake_env).raw_action = mock.PropertyMock(return_value=action3)
+    step_timestep = wrapper.step(action=action3)
     step_image = step_timestep.observation['pixels']
     self.assertEqual(step_image.shape, (120, 80, 4))
     last_action_layer = step_image[:, :, -1]
@@ -118,9 +103,11 @@ class LastActionWrapperTest(absltest.TestCase):
     self.assertEqual((y.item(), x.item()), (119, 20))
 
   def test_no_concat_to_pixels(self):
-    obs_spec = {'pixels': _simple_spec()}
-    fake_env = FakeEnv(obs_spec)
-    fake_env.set_next_timestep(_simple_timestep())
+    fake_timestep = _simple_timestep()
+    fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)
+    fake_env.observation_spec.return_value = {'pixels': _simple_spec()}
+    fake_env.reset.return_value = fake_timestep
+    fake_env.step.return_value = fake_timestep
 
     wrapper = last_action_wrapper.LastActionWrapper(
         fake_env, concat_to_pixels=False)
@@ -134,10 +121,12 @@ class LastActionWrapperTest(absltest.TestCase):
     last_action_layer = reset_timestep.observation['last_action']
     self.assertEqual(np.sum(last_action_layer), 0)
 
-    step_timestep = wrapper.step(action={
+    action1 = {
         'action_type': action_type.ActionType.TOUCH,
         'touch_position': np.array([0.25, 0.75], dtype=np.float32),
-    })
+    }
+    type(fake_env).raw_action = mock.PropertyMock(return_value=action1)
+    step_timestep = wrapper.step(action=action1)
     step_image = step_timestep.observation['pixels']
     self.assertEqual(step_image.shape, (120, 80, 3))
     last_action_layer = step_timestep.observation['last_action']
@@ -145,19 +134,23 @@ class LastActionWrapperTest(absltest.TestCase):
     y, x = np.where(last_action_layer == 255)
     self.assertEqual((y.item(), x.item()), (90, 20))
 
-    step_timestep = wrapper.step(action={
+    action2 = {
         'action_type': action_type.ActionType.LIFT,
         'touch_position': np.array([0.25, 0.75], dtype=np.float32),
-    })
+    }
+    type(fake_env).raw_action = mock.PropertyMock(return_value=action2)
+    step_timestep = wrapper.step(action=action2)
     step_image = step_timestep.observation['pixels']
     self.assertEqual(step_image.shape, (120, 80, 3))
     last_action_layer = step_timestep.observation['last_action']
     self.assertEqual(np.sum(last_action_layer), 0)
 
-    step_timestep = wrapper.step(action={
+    action3 = {
         'action_type': action_type.ActionType.TOUCH,
         'touch_position': np.array([1.0, 0.75], dtype=np.float32),
-    })
+    }
+    type(fake_env).raw_action = mock.PropertyMock(return_value=action3)
+    step_timestep = wrapper.step(action=action3)
     step_image = step_timestep.observation['pixels']
     self.assertEqual(step_image.shape, (120, 80, 3))
     last_action_layer = step_timestep.observation['last_action']

android_env/wrappers/rate_limit_wrapper.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,7 +17,6 @@
 
 import enum
 import time
-from typing import Dict
 
 from android_env import env_interface
 from android_env.components import action_type
@@ -78,9 +77,13 @@ class RateLimitWrapper(base_wrapper.BaseWrapper):
     self._last_step_time = time.time()
     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 while maintaining a steady interaction rate."""
 
+    # If max_wait is non-positive, the wrapper has no effect.
+    if self._max_wait <= 0.0:
+      return self._env.step(action)
+
     if self._sleep_type == RateLimitWrapper.SleepType.BEFORE:
       self._wait()
 

android_env/wrappers/rate_limit_wrapper_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.
@@ -16,6 +16,7 @@
 """Tests for rate_limit_wrapper."""
 
 import time
+from typing import Any, Protocol
 from unittest import mock
 
 from absl.testing import absltest
@@ -46,6 +47,17 @@ def _get_base_env():
   return env
 
 
+class _FnWithTimestamps(Protocol):
+  """A function with `timestamp` and `timestamps` attributes."""
+
+  timestamp: float
+  timestamps: list[float]
+
+
+def _with_timestamp(fn: Any) -> _FnWithTimestamps:
+  return fn
+
+
 class RateLimitWrapperTest(parameterized.TestCase):
 
   @parameterized.named_parameters(
@@ -64,6 +76,8 @@ class RateLimitWrapperTest(parameterized.TestCase):
         'touch_position': np.array([0.123, 0.456])
     })
     mock_sleep.assert_not_called()
+    # When the wrapper is disabled, base step should only be called once.
+    env.step.assert_called_once()
 
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_enabled(self, mock_sleep):
@@ -105,6 +119,7 @@ class RateLimitWrapperTest(parameterized.TestCase):
     _ = wrapper.reset()
     mock_sleep.assert_not_called()  # It should never sleep during reset().
 
+    @_with_timestamp
     def _sleep_fn(sleep_time):
       _sleep_fn.timestamp = time.time()
       self.assertBetween(sleep_time, 0.0, 33.33)
@@ -143,6 +158,7 @@ class RateLimitWrapperTest(parameterized.TestCase):
     _ = wrapper.reset()
     mock_sleep.assert_not_called()  # It should never sleep during reset().
 
+    @_with_timestamp
     def _sleep_fn(sleep_time):
       _sleep_fn.timestamp = time.time()
       self.assertBetween(sleep_time, 0.0, 33.33)
@@ -183,12 +199,14 @@ class RateLimitWrapperTest(parameterized.TestCase):
     _ = wrapper.reset()
     mock_sleep.assert_not_called()  # It should never sleep during reset().
 
+    @_with_timestamp
     def _sleep_fn(sleep_time):
       _sleep_fn.timestamp = time.time()
       self.assertBetween(sleep_time, 0.0, 33.33)
 
     mock_sleep.side_effect = _sleep_fn
 
+    @_with_timestamp
     def _step_fn(action):
       # On even calls the action should be the actual agent action, but on odd
       # calls they should be REPEATs.
@@ -212,6 +230,9 @@ class RateLimitWrapperTest(parameterized.TestCase):
         'touch_position': np.array([0.123, 0.456])
     })
 
+    # When the wrapper is enabled, base step should be called twice.
+    self.assertEqual(env.step.call_count, 2)
+
     # `step()` should be called twice: before `sleep()` and after it.
     self.assertLen(_step_fn.timestamps, 2)
     self.assertGreaterEqual(_sleep_fn.timestamp, _step_fn.timestamps[0])

android_env/wrappers/tap_action_wrapper.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,7 +15,7 @@
 
 """Wraps the AndroidEnv environment to provide tap actions of a given duration."""
 
-from typing import Dict, 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."""
 
@@ -46,33 +43,35 @@ class TapActionWrapper(base_wrapper.BaseWrapper):
     return logs
 
   def _process_action(
-      self, action: Dict[str, np.ndarray]
-  ) -> Sequence[Dict[str, np.ndarray]]:
-
+      self, action: dict[str, np.ndarray]
+  ) -> Sequence[dict[str, np.ndarray]]:
     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)
 
     return actions
 
-  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."""
     self._env_steps += self._num_frames + 1
     actions = self._process_action(action)
@@ -93,7 +92,7 @@ class TapActionWrapper(base_wrapper.BaseWrapper):
         discount=discount,
         observation=observation)
 
-  def action_spec(self) -> Dict[str, dm_env.specs.Array]:
+  def action_spec(self) -> dict[str, dm_env.specs.Array]:
     if self._touch_only:
       return {
           'action_type':