android-env 1.2.1__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 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,12 @@
 
 """Base class for AndroidEnv wrappers."""
 
-from typing import Any, Dict
+from typing import Any
 
 from absl import logging
 from android_env import env_interface
 from android_env.proto import adb_pb2
-from android_env.proto import task_pb2
+from android_env.proto import state_pb2
 import dm_env
 from dm_env import specs
 import numpy as np
@@ -42,7 +42,7 @@ class BaseWrapper(env_interface.AndroidEnvInterface):
     action = self._process_action(action)
     return self._process_timestep(self._env.step(action))
 
-  def task_extras(self, latest_only: bool = True) -> Dict[str, np.ndarray]:
+  def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:
     return self._env.task_extras(latest_only=latest_only)
 
   def _reset_state(self):
@@ -54,31 +54,52 @@ class BaseWrapper(env_interface.AndroidEnvInterface):
   def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:
     return timestep
 
-  def observation_spec(self) -> Dict[str, specs.Array]:
+  def observation_spec(self) -> dict[str, specs.Array]:
     return self._env.observation_spec()
 
-  def action_spec(self) -> Dict[str, specs.Array]:
+  def action_spec(self) -> dict[str, specs.Array]:
     return self._env.action_spec()
 
-  def task_extras_spec(self) -> Dict[str, specs.Array]:
-    return self._env.task_extras_spec()
+  def reward_spec(self) -> specs.Array:
+    return self._env.reward_spec()
 
-  def _wrapper_stats(self) -> Dict[str, Any]:
+  def discount_spec(self) -> specs.Array:
+    return self._env.discount_spec()
+
+  def _wrapper_stats(self) -> dict[str, Any]:
     """Add wrapper specific logging here."""
     return {}
 
-  def stats(self) -> Dict[str, Any]:
+  def stats(self) -> dict[str, Any]:
     info = self._env.stats()
     info.update(self._wrapper_stats())
     return info
 
+  def load_state(
+      self, request: state_pb2.LoadStateRequest
+  ) -> state_pb2.LoadStateResponse:
+    """Loads a state."""
+    return self._env.load_state(request)
+
+  def save_state(
+      self, request: state_pb2.SaveStateRequest
+  ) -> state_pb2.SaveStateResponse:
+    """Saves a state.
+
+    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._env.save_state(request)
+
   def execute_adb_call(self,
                        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 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.
@@ -19,8 +19,8 @@ from unittest import mock
 
 from absl import logging
 from absl.testing import absltest
-from android_env import environment
-from android_env.proto import task_pb2
+from android_env import env_interface
+from android_env.proto import state_pb2
 from android_env.wrappers import base_wrapper
 
 
@@ -28,7 +28,7 @@ class BaseWrapperTest(absltest.TestCase):
 
   @mock.patch.object(logging, 'info')
   def test_base_function_forwarding(self, mock_info):
-    base_env = mock.create_autospec(environment.AndroidEnv)
+    base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     wrapped_env = base_wrapper.BaseWrapper(base_env)
     mock_info.assert_called_with('Wrapping with %s', 'BaseWrapper')
 
@@ -58,11 +58,6 @@ class BaseWrapperTest(absltest.TestCase):
     self.assertEqual(fake_action_spec, wrapped_env.action_spec())
     base_env.action_spec.assert_called_once()
 
-    fake_task_extras_spec = 'fake_task_extras_spec'
-    base_env.task_extras_spec.return_value = fake_task_extras_spec
-    self.assertEqual(fake_task_extras_spec, wrapped_env.task_extras_spec())
-    base_env.task_extras_spec.assert_called_once()
-
     fake_raw_action = 'fake_raw_action'
     type(base_env).raw_action = mock.PropertyMock(return_value=fake_raw_action)
     self.assertEqual(fake_raw_action, wrapped_env.raw_action)
@@ -72,10 +67,21 @@ class BaseWrapperTest(absltest.TestCase):
         return_value=fake_raw_observation)
     self.assertEqual(fake_raw_observation, wrapped_env.raw_observation)
 
-    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)
+    load_request = state_pb2.LoadStateRequest(args={})
+    expected_response = state_pb2.LoadStateResponse(
+        status=state_pb2.LoadStateResponse.Status.OK
+    )
+    base_env.load_state.return_value = expected_response
+    self.assertEqual(wrapped_env.load_state(load_request), expected_response)
+    base_env.load_state.assert_called_once_with(load_request)
+
+    save_request = state_pb2.SaveStateRequest(args={})
+    expected_response = state_pb2.SaveStateResponse(
+        status=state_pb2.SaveStateResponse.Status.OK
+    )
+    base_env.save_state.return_value = expected_response
+    self.assertEqual(wrapped_env.save_state(save_request), expected_response)
+    base_env.save_state.assert_called_once_with(save_request)
 
     wrapped_env.close()
     base_env.close.assert_called_once()
@@ -87,7 +93,7 @@ class BaseWrapperTest(absltest.TestCase):
       base_env.some_random_function.return_value = fake_return_value
 
   def test_multiple_wrappers(self):
-    base_env = mock.create_autospec(environment.AndroidEnv)
+    base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     wrapped_env_1 = base_wrapper.BaseWrapper(base_env)
     wrapped_env_2 = base_wrapper.BaseWrapper(wrapped_env_1)
 
@@ -101,7 +107,7 @@ class BaseWrapperTest(absltest.TestCase):
     self.assertEqual(base_env, wrapped_env_2.raw_env)
 
   def test_stats(self):
-    base_env = mock.create_autospec(environment.AndroidEnv)
+    base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     wrapped_env = base_wrapper.BaseWrapper(base_env)
     base_stats = {'base': 'stats'}
     base_env.stats.return_value = base_stats
@@ -109,7 +115,7 @@ class BaseWrapperTest(absltest.TestCase):
 
   @mock.patch.object(logging, 'info')
   def test_wrapped_stats(self, mock_info):
-    base_env = mock.create_autospec(environment.AndroidEnv)
+    base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
 
     class LoggingWrapper1(base_wrapper.BaseWrapper):
 

android_env/wrappers/discrete_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 discrete actions."""
 
-from typing import Optional, Sequence, Dict
+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,16 +58,16 @@ 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:
+  def step(self, action: dict[str, int]) -> dm_env.TimeStep:
     """Take a step in the base environment."""
 
     return self._env.step(self._process_action(action))
 
-  def _process_action(self, action: Dict[str, int]) -> Dict[str, np.ndarray]:
+  def _process_action(self, action: dict[str, int]) -> dict[str, np.ndarray]:
     """Transforms action so that it agrees with AndroidEnv's action spec."""
 
     return {
@@ -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
@@ -149,7 +150,7 @@ class DiscreteActionWrapper(base_wrapper.BaseWrapper):
 
     return [x_pos, y_pos]
 
-  def action_spec(self) -> Dict[str, specs.Array]:
+  def action_spec(self) -> dict[str, specs.Array]:
     """Action spec of the wrapped environment."""
 
     return {

android_env/wrappers/discrete_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.
@@ -18,7 +18,7 @@
 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 as action_type_lib
 from android_env.wrappers import discrete_action_wrapper
 from dm_env import specs
@@ -64,7 +64,7 @@ class DiscreteActionWrapperTest(absltest.TestCase):
         'touch_position': _make_array_spec(
             shape=(2,), dtype=np.float32, name='touch_position'),
     }
-    self.base_env = mock.create_autospec(environment.AndroidEnv)
+    self.base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     self.base_env.action_spec.return_value = self._base_action_spec
 
   def test_num_actions(self):
@@ -295,7 +295,7 @@ class DiscreteActionWrapperTest(absltest.TestCase):
         'touch_position': _make_array_spec(
             shape=(2,), dtype=np.float64, name='touch_position'),
     }
-    base_env = mock.create_autospec(environment.AndroidEnv)
+    base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     base_env.action_spec.return_value = base_action_spec
 
     wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(

android_env/wrappers/flat_interface_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 make its interface flat."""
 
-from typing import Union, Dict, Any
+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]]:
+  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.Array, Dict[str, specs.Array]]:
+  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 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,6 +15,7 @@
 
 """Tests for android_env.wrappers.flat_interface_wrapper."""
 
+from typing import cast
 from unittest import mock
 
 from absl.testing import absltest
@@ -33,10 +34,6 @@ def _make_array_spec(shape, dtype=np.float32, name=None, maximum=3, minimum=0):
       minimum=np.ones(shape) * minimum)
 
 
-def _make_discrete_array_spec(name, num_values):
-  return specs.DiscreteArray(name=name, num_values=num_values)
-
-
 def _make_timestep(observation):
   return dm_env.TimeStep(
       step_type='fake_step_type',
@@ -51,9 +48,8 @@ class FlatInterfaceWrapperTest(absltest.TestCase):
   def setUp(self):
     super().setUp()
     self.action_shape = (1,)
-    self.base_action_spec = {
-        'action_id': _make_discrete_array_spec(
-            name='action_id', num_values=4)
+    self.base_action_spec: dict[str, specs.DiscreteArray] = {
+        'action_id': specs.DiscreteArray(name='action_id', num_values=4)
     }
     self.int_obs_shape = (3, 4, 2)
     self.float_obs_shape = (2,)
@@ -148,10 +144,10 @@ class FlatInterfaceWrapperTest(absltest.TestCase):
 
   def test_action_spec(self):
     wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(self.base_env)
-    action_spec = wrapped_env.action_spec()
-    base_action_spec = self.base_action_spec['action_id']
+    action_spec = cast(specs.BoundedArray, wrapped_env.action_spec())
+    parent_action_spec = self.base_action_spec['action_id']
 
-    self.assertEqual(base_action_spec.name, action_spec.name)
+    self.assertEqual(parent_action_spec.name, action_spec.name)
     self.assertEqual((), action_spec.shape)
     self.assertEqual(np.int32, action_spec.dtype)
     self.assertEqual(0, action_spec.minimum)

android_env/wrappers/float_pixels_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,9 +15,7 @@
 
 """Converts pixel observation to from int to float32 between 0.0 and 1.0."""
 
-from typing import Dict
-
-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
@@ -34,11 +32,12 @@ class FloatPixelsWrapper(base_wrapper.BaseWrapper):
                                                       np.integer)
 
   def _process_observation(
-      self, observation: Dict[str, np.ndarray]
-  ) -> Dict[str, np.ndarray]:
+      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
 
@@ -53,10 +52,10 @@ class FloatPixelsWrapper(base_wrapper.BaseWrapper):
   def reset(self) -> dm_env.TimeStep:
     return self._process_timestep(self._env.reset())
 
-  def step(self, action: Dict[str, np.ndarray]) -> dm_env.TimeStep:
+  def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:
     return self._process_timestep(self._env.step(action))
 
-  def observation_spec(self) -> Dict[str, specs.Array]:
+  def observation_spec(self) -> dict[str, specs.Array]:
     if self._should_convert_int_to_float:
       observation_spec = self._env.observation_spec()
       observation_spec['pixels'] = specs.BoundedArray(

android_env/wrappers/float_pixels_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.
@@ -96,7 +96,7 @@ class FloatPixelsWrapperTest(absltest.TestCase):
 
   def test_float_pixels_wrapper_step(self):
     wrapped_env = float_pixels_wrapper.FloatPixelsWrapper(self.base_env)
-    ts = wrapped_env.step('fake_action')
+    ts = wrapped_env.step({'fake_action': np.array([1, 2, 3])})
 
     self.assertEqual(self.base_timestep.step_type, ts.step_type)
     self.assertEqual(self.base_timestep.reward, ts.reward)
@@ -141,7 +141,7 @@ class FloatPixelsWrapperTest(absltest.TestCase):
     # The wrapper should not touch the timestep in this case.
     fake_timestep = ('step_type', 'reward', 'discount', 'obs')
     base_env.step.return_value = fake_timestep
-    ts = wrapped_env.step('fake_action')
+    ts = wrapped_env.step({'fake_action': np.array([1, 2, 3])})
     self.assertEqual(fake_timestep, ts)
 
 

android_env/wrappers/gym_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 to expose an OpenAI Gym interface."""
 
-from typing import Any, Dict, Tuple
+from typing import Any
 
 from android_env.wrappers import base_wrapper
 import dm_env
@@ -25,16 +25,16 @@ from gym import spaces
 import numpy as np
 
 
-class GymInterfaceWrapper(base_wrapper.BaseWrapper, gym.Env):
+class GymInterfaceWrapper(gym.Env):
   """AndroidEnv with OpenAI Gym interface."""
 
   def __init__(self, env: dm_env.Environment):
-
-    base_wrapper.BaseWrapper.__init__(self, env)
+    self._env = env
     self.spec = None
     self.action_space = self._spec_to_space(self._env.action_spec())
     self.observation_space = self._spec_to_space(self._env.observation_spec())
     self.metadata = {'render.modes': ['rgb_array']}
+    self._latest_observation = None
 
   def _spec_to_space(self, spec: specs.Array) -> spaces.Space:
     """Converts dm_env specs to OpenAI Gym spaces."""
@@ -44,7 +44,8 @@ class GymInterfaceWrapper(base_wrapper.BaseWrapper, gym.Env):
 
     if isinstance(spec, dict):
       return spaces.Dict(
-          {name: self._spec_to_space(s) for name, s in spec.items()})
+          {name: self._spec_to_space(s) for name, s in spec.items()}
+      )
 
     if isinstance(spec, specs.DiscreteArray):
       return spaces.Box(
@@ -61,11 +62,13 @@ class GymInterfaceWrapper(base_wrapper.BaseWrapper, gym.Env):
           high=spec.maximum)
 
     if isinstance(spec, specs.Array):
-      return spaces.Box(
-          shape=spec.shape,
-          dtype=spec.dtype,
-          low=-np.inf,
-          high=np.inf)
+      if spec.dtype == np.uint8:
+        low = 0
+        high = 255
+      else:
+        low = -np.inf
+        high = np.inf
+      return spaces.Box(shape=spec.shape, dtype=spec.dtype, low=low, high=high)
 
     raise ValueError('Unknown type for specs: {}'.format(spec))
 
@@ -74,18 +77,21 @@ class GymInterfaceWrapper(base_wrapper.BaseWrapper, gym.Env):
     if mode == 'rgb_array':
       if self._latest_observation is None:
         return
+
       return self._latest_observation['pixels']
     else:
       raise ValueError('Only supported render mode is rgb_array.')
 
   def reset(self) -> np.ndarray:
+    self._latest_observation = None
     timestep = self._env.reset()
     return timestep.observation
 
-  def step(self, action: Dict[str, int]) -> Tuple[Any, ...]:
+  def step(self, action: dict[str, int]) -> tuple[Any, ...]:
     """Take a step in the base environment."""
-    timestep = self._env.step(self._process_action(action))
+    timestep = self._env.step(action)
     observation = timestep.observation
+    self._latest_observation = observation
     reward = timestep.reward
     done = timestep.step_type == dm_env.StepType.LAST
     info = {'discount': timestep.discount}

android_env/wrappers/gym_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.
@@ -18,7 +18,7 @@
 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 gym_wrapper
 import dm_env
 from dm_env import specs
@@ -30,7 +30,7 @@ class GymInterfaceWrapperTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self._base_env = mock.create_autospec(environment.AndroidEnv)
+    self._base_env = mock.create_autospec(env_interface.AndroidEnvInterface)
     self._base_env.action_spec.return_value = {
         'action_type':
             specs.DiscreteArray(
@@ -68,7 +68,6 @@ class GymInterfaceWrapperTest(absltest.TestCase):
   def test_render(self):
     self._base_env.step.return_value = self._fake_ts
     _ = self._wrapped_env.step(action=np.zeros(shape=(1,)))
-    self._base_env._latest_observation = {'pixels': np.ones(shape=(2, 3))}
     image = self._wrapped_env.render(mode='rgb_array')
     self.assertTrue(np.array_equal(image, np.ones(shape=(2, 3))))
 
@@ -90,7 +89,6 @@ class GymInterfaceWrapperTest(absltest.TestCase):
     self._base_env.step.return_value = self._fake_ts
     obs, _, _, _ = self._wrapped_env.step(action=np.zeros(shape=(1,)))
     self._base_env.step.assert_called_once()
-    print(obs)
     self.assertTrue(np.array_equal(obs['pixels'], np.ones(shape=(2, 3))))
 
   def test_spec_to_space(self):