Enhance BaseWrapper to Support Generic dm_env.Environment

android_env/wrappers/a11y_grpc_wrapper.py

@@ -187,7 +187,7 @@ def _start_a11y_services(self) -> None:
             ),
         )
     )
-    start_service_response = self._env.execute_adb_call(start_service_request)
+    start_service_response = self.execute_adb_call(start_service_request)
     if start_service_response.status != adb_pb2.AdbResponse.Status.OK:
       raise RuntimeError(
           'Could not start accessibility forwarder '
@@ -204,7 +204,7 @@ def _install_a11y_forwarding_apk(self) -> None:
             blob=adb_pb2.AdbRequest.InstallApk.Blob(contents=a11y_fwd_apk),
         )
     )
-    install_response = self._env.execute_adb_call(install_request)
+    install_response = self.execute_adb_call(install_request)
     if install_response.status != adb_pb2.AdbResponse.Status.OK:
       raise ValueError(
           f'Could not install accessibility_forwarder.apk: {install_response}.'
@@ -222,9 +222,7 @@ def _enable_a11y_tree_logs(self) -> None:
             ),
         )
     )
-    enable_tree_logs_response = self._env.execute_adb_call(
-        enable_tree_logs_request
-    )
+    enable_tree_logs_response = self.execute_adb_call(enable_tree_logs_request)
     if enable_tree_logs_response.status != adb_pb2.AdbResponse.Status.OK:
       raise ValueError(
           'Could not enable accessibility tree logging: '
@@ -405,9 +403,10 @@ def reset(self) -> dm_env.TimeStep:
     self._servicer.pause_and_clear()
     timestep = self._env.reset()
     self._servicer.resume()
-    if self._env.stats()['relaunch_count'] > self._relaunch_count:
+    relaunch_count = self.stats().get('relaunch_count', 0)
+    if relaunch_count > self._relaunch_count:
       self._configure_grpc()
-      self._relaunch_count = self._env.stats()['relaunch_count']
+      self._relaunch_count = relaunch_count
     self._accumulated_extras = {}
     timeout = self._a11y_info_timeout or 0.0
     new_observation = self._fetch_task_extras_and_update_observation(
@@ -456,7 +455,7 @@ def _fetch_task_extras(self) -> dict[str, Any]:
       EnableNetworkingError: after a fixed number of attempts to revive the a11y
         services by re-enabling the network connection.
     """
-    base_extras = self._env.task_extras(latest_only=False).copy()
+    base_extras = super().task_extras(latest_only=False).copy()
     # If the previous future is done, reset it to the initial state.
     if (
         self._enabling_networking_future is not None

android_env/wrappers/a11y_grpc_wrapper_test.py

@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for a11y_grpc_wrapper."""
-
 import time
 from unittest import mock
 
@@ -962,6 +960,58 @@ def test_broadcast_sent_port(self):
         base_env.execute_adb_call.call_args[0][0].send_broadcast.action,
     )
 
+  @mock.patch.object(
+      a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True
+  )
+  @mock.patch.object(grpc, 'server', autospec=True)
+  def test_dm_env(self, mock_server, mock_add_servicer):
+    del mock_server, mock_add_servicer
+    base_env = mock.create_autospec(dm_env.Environment)
+    base_env.action_spec.return_value = {
+        'action_type': dm_env.specs.DiscreteArray(
+            num_values=3, name='action_type'
+        )
+    }
+    base_env.reset.return_value = dm_env.restart(
+        observation={'pixels': np.zeros((10, 10, 3))}
+    )
+    base_env.step.return_value = dm_env.transition(
+        observation={'pixels': np.zeros((10, 10, 3))}, reward=0.0
+    )
+
+    wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(
+        base_env,
+        install_a11y_forwarding=False,
+        start_a11y_service=False,
+        enable_a11y_tree_info=False,
+    )
+
+    # Test stats fallback
+    self.assertEqual({}, wrapped_env.stats())
+
+    # Test task_extras fallback
+    wrapped_env.reset()
+    self.assertEqual({}, wrapped_env.task_extras())
+
+    # Test execute_adb_call fallback
+    adb_request = adb_pb2.AdbRequest()
+    self.assertEqual(
+        adb_pb2.AdbResponse(), wrapped_env.execute_adb_call(adb_request)
+    )
+
+    # Test missing attributes
+    with self.assertRaisesRegex(
+        AttributeError, 'does not have attribute raw_action'
+    ):
+      _ = wrapped_env.raw_action
+
+    with self.assertRaisesRegex(
+        AttributeError, 'does not have attribute raw_observation'
+    ):
+      _ = wrapped_env.raw_observation
+
+    wrapped_env.close()
+
 
 if __name__ == '__main__':
   absltest.main()

android_env/wrappers/base_wrapper.py

@@ -27,10 +27,18 @@
 
 
 class BaseWrapper(env_interface.AndroidEnvInterface):
-  """AndroidEnv wrapper."""
-
-  def __init__(self, env: env_interface.AndroidEnvInterface) -> None:
-    self._env = env
+  """AndroidEnv wrapper.
+
+  Supports wrapping standard `dm_env.Environment` instances in addition to
+  `AndroidEnvInterface` instances.
+  """
+
+  def __init__(self, env: dm_env.Environment) -> None:
+    # `self._env` can be a standard `dm_env.Environment` or an
+    # `AndroidEnvInterface`. We use the `|` operator to support both, and use
+    # `hasattr` checks before calling AndroidEnv-specific methods that are not
+    # present in standard `dm_env.Environment`.
+    self._env: env_interface.AndroidEnvInterface | dm_env.Environment = env
     logging.info('Wrapping with %s', self.__class__.__name__)
 
   def reset(self) -> dm_env.TimeStep:
@@ -42,8 +50,25 @@ def step(self, action: Any) -> dm_env.TimeStep:
     action = self._process_action(action)
     return self._process_timestep(self._env.step(action))
 
+  def _get_env_attr(self, name: str) -> Any:
+    """Safely gets an attribute from the wrapped environment."""
+    if hasattr(self._env, name):
+      return getattr(self._env, name)
+    raise AttributeError(
+        f'Underlying environment {type(self._env).__name__} does not have'
+        f' attribute {name}.'
+    )
+
+  def _delegate_to_env(self, method_name: str, *args, **kwargs) -> Any:
+    """Delegates to the underlying env if the method exists, else uses super()."""
+    method = getattr(self._env, method_name, None)
+    if method is not None and callable(method):
+      return method(*args, **kwargs)
+    # Fallback to super() to call the default AndroidEnvInterface implementation
+    return getattr(super(), method_name)(*args, **kwargs)
+
   def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:
-    return self._env.task_extras(latest_only=latest_only)
+    return self._delegate_to_env('task_extras', latest_only=latest_only)
 
   def _reset_state(self):
     pass
@@ -71,15 +96,18 @@ def _wrapper_stats(self) -> dict[str, Any]:
     return {}
 
   def stats(self) -> dict[str, Any]:
-    info = self._env.stats()
+    if hasattr(self._env, 'stats'):
+      info = self._env.stats()
+    else:
+      info = super().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)
+    return self._delegate_to_env('load_state', request)
 
   def save_state(
       self, request: state_pb2.SaveStateRequest
@@ -94,23 +122,23 @@ def save_state(
       A `SaveStateResponse` containing the status, error message (if
       applicable), and any other relevant information.
     """
-    return self._env.save_state(request)
+    return self._delegate_to_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)
+    return self._delegate_to_env('execute_adb_call', adb_call)
 
   @property
   def raw_action(self) -> Any:
-    return self._env.raw_action
+    return self._get_env_attr('raw_action')
 
   @property
   def raw_observation(self) -> Any:
-    return self._env.raw_observation
+    return self._get_env_attr('raw_observation')
 
   @property
-  def raw_env(self) -> env_interface.AndroidEnvInterface:
+  def raw_env(self) -> dm_env.Environment:
     """Recursively unwrap until we reach the true 'raw' env."""
     wrapped = self._env
     if hasattr(wrapped, 'raw_env'):
@@ -119,6 +147,8 @@ def raw_env(self) -> env_interface.AndroidEnvInterface:
 
   def __getattr__(self, attr) -> Any:
     """Delegate attribute access to underlying environment."""
+    if attr in ('raw_action', 'raw_observation'):
+      return self._get_env_attr(attr)
     return getattr(self._env, attr)
 
   def close(self) -> None:

android_env/wrappers/base_wrapper_test.py

@@ -13,15 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for android_env.wrappers.base_wrapper."""
-
 from unittest import mock
 
 from absl import logging
 from absl.testing import absltest
 from android_env import env_interface
+from android_env.proto import adb_pb2
 from android_env.proto import state_pb2
 from android_env.wrappers import base_wrapper
+import dm_env
+from dm_env import specs
 
 
 class BaseWrapperTest(absltest.TestCase):
@@ -150,5 +151,89 @@ def _wrapper_stats(self):
     self.assertEqual(expected_stats, wrapped_env.stats())
 
 
+class BaseWrapperDmEnvTest(absltest.TestCase):
+
+  def test_dm_env_forwarding(self):
+    base_env = mock.create_autospec(dm_env.Environment, instance=True)
+    wrapped_env = base_wrapper.BaseWrapper(base_env)
+
+    fake_ts = dm_env.TimeStep(
+        step_type=dm_env.StepType.MID, reward=0.0, discount=1.0, observation={}
+    )
+    base_env.reset.return_value = fake_ts
+    self.assertEqual(fake_ts, wrapped_env.reset())
+    base_env.reset.assert_called_once()
+
+    fake_action = 'fake_action'
+    base_env.step.return_value = fake_ts
+    self.assertEqual(fake_ts, wrapped_env.step(fake_action))
+    base_env.step.assert_called_once_with(fake_action)
+
+    fake_obs_spec = 'fake_obs_spec'
+    base_env.observation_spec.return_value = fake_obs_spec
+    self.assertEqual(fake_obs_spec, wrapped_env.observation_spec())
+    base_env.observation_spec.assert_called_once()
+
+    fake_action_spec = 'fake_action_spec'
+    base_env.action_spec.return_value = fake_action_spec
+    self.assertEqual(fake_action_spec, wrapped_env.action_spec())
+    base_env.action_spec.assert_called_once()
+
+    fake_reward_spec = mock.create_autospec(specs.Array)
+    fake_discount_spec = mock.create_autospec(specs.Array)
+    base_env.reward_spec.return_value = fake_reward_spec
+    base_env.discount_spec.return_value = fake_discount_spec
+
+    self.assertEqual(fake_reward_spec, wrapped_env.reward_spec())
+    self.assertEqual(fake_discount_spec, wrapped_env.discount_spec())
+    base_env.reward_spec.assert_called_once()
+    base_env.discount_spec.assert_called_once()
+
+    wrapped_env.close()
+    base_env.close.assert_called_once()
+
+  def test_safe_fallbacks(self):
+    base_env = mock.create_autospec(dm_env.Environment, instance=True)
+    wrapped_env = base_wrapper.BaseWrapper(base_env)
+
+    # task_extras should return default value {} without crashing
+    self.assertEqual({}, wrapped_env.task_extras(latest_only=True))
+
+    # stats should return default value {} without crashing
+    self.assertEqual({}, wrapped_env.stats())
+
+    # execute_adb_call should return default value AdbResponse without crashing
+    self.assertEqual(
+        adb_pb2.AdbResponse(),
+        wrapped_env.execute_adb_call(adb_pb2.AdbRequest()),
+    )
+
+  def test_raw_action_observation_raises_attribute_error(self):
+    base_env = mock.create_autospec(dm_env.Environment, instance=True)
+    wrapped_env = base_wrapper.BaseWrapper(base_env)
+
+    with self.assertRaisesRegex(
+        AttributeError,
+        'Underlying environment .*Mock does not have attribute raw_action',
+    ):
+      _ = wrapped_env.raw_action
+
+    with self.assertRaisesRegex(
+        AttributeError,
+        'Underlying environment .*Mock does not have attribute raw_observation',
+    ):
+      _ = wrapped_env.raw_observation
+
+  def test_getattr_delegation(self):
+    base_env = mock.create_autospec(dm_env.Environment, instance=True)
+    base_env.custom_attr = 'custom_value'
+    wrapped_env = base_wrapper.BaseWrapper(base_env)
+
+    self.assertEqual('custom_value', wrapped_env.custom_attr)
+
+    with self.assertRaises(AttributeError):
+      _ = wrapped_env.non_existent_attribute
+
+
 if __name__ == '__main__':
   absltest.main()

android_env/wrappers/last_action_wrapper.py

@@ -72,7 +72,7 @@ def _process_observation(
   def _get_last_action_layer(self, pixels: np.ndarray) -> np.ndarray:
     """Makes sure the rescaling doesn't distort the last_action layer."""
 
-    last_action = self._env.raw_action
+    last_action = self.raw_action
     last_action_layer = np.zeros(self._screen_dimensions, dtype=pixels.dtype)
 
     if ('action_type' in last_action and

android_env/wrappers/last_action_wrapper_test.py

@@ -13,14 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for android_env.wrappers.last_action_wrapper."""
-
-from typing import Any
 from unittest import mock
 
 from absl.testing import absltest
 from android_env import env_interface
 from android_env.components import action_type
+from android_env.proto import adb_pb2
 from android_env.wrappers import last_action_wrapper
 import dm_env
 from dm_env import specs
@@ -158,5 +156,36 @@ def test_no_concat_to_pixels(self):
     y, x = np.where(last_action_layer == 255)
     self.assertEqual((y.item(), x.item()), (90, 79))
 
+  def test_dm_env(self):
+    base_env = mock.create_autospec(dm_env.Environment)
+    base_env.observation_spec.return_value = {'pixels': _simple_spec()}
+    wrapped_env = last_action_wrapper.LastActionWrapper(base_env)
+
+    # Test stats fallback
+    self.assertEqual({}, wrapped_env.stats())
+
+    # Test task_extras fallback (should return {})
+    self.assertEqual({}, wrapped_env.task_extras())
+
+    # Test execute_adb_call fallback (should return default AdbResponse)
+    adb_request = adb_pb2.AdbRequest()
+    self.assertEqual(
+        adb_pb2.AdbResponse(), wrapped_env.execute_adb_call(adb_request)
+    )
+
+    # Test that reset and step raise AttributeError because they need raw_action
+    base_env.reset.return_value = _simple_timestep()
+    with self.assertRaisesRegex(
+        AttributeError, 'does not have attribute raw_action'
+    ):
+      wrapped_env.reset()
+
+    base_env.step.return_value = _simple_timestep()
+    with self.assertRaisesRegex(
+        AttributeError, 'does not have attribute raw_action'
+    ):
+      wrapped_env.step(action={})
+
+
 if __name__ == '__main__':
   absltest.main()