From 0153f490eca1f0211b5507d6320db90a4c5d9f8c Mon Sep 17 00:00:00 2001 From: Tony Li Date: Fri, 10 Jul 2026 13:21:57 -0700 Subject: [PATCH] Enhance BaseWrapper to Support Generic dm_env.Environment PiperOrigin-RevId: 945862662 --- android_env/wrappers/a11y_grpc_wrapper.py | 15 ++-- .../wrappers/a11y_grpc_wrapper_test.py | 54 ++++++++++- android_env/wrappers/base_wrapper.py | 54 ++++++++--- android_env/wrappers/base_wrapper_test.py | 89 ++++++++++++++++++- android_env/wrappers/last_action_wrapper.py | 2 +- .../wrappers/last_action_wrapper_test.py | 35 +++++++- android_env/wrappers/swipe_action_wrapper.py | 2 +- .../wrappers/swipe_action_wrapper_test.py | 31 +++++++ android_env/wrappers/tap_action_wrapper.py | 2 +- .../wrappers/tap_action_wrapper_test.py | 31 +++++++ 10 files changed, 285 insertions(+), 30 deletions(-) diff --git a/android_env/wrappers/a11y_grpc_wrapper.py b/android_env/wrappers/a11y_grpc_wrapper.py index b58aad70..f35024a7 100644 --- a/android_env/wrappers/a11y_grpc_wrapper.py +++ b/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 diff --git a/android_env/wrappers/a11y_grpc_wrapper_test.py b/android_env/wrappers/a11y_grpc_wrapper_test.py index 269a2cc6..e8f2a640 100644 --- a/android_env/wrappers/a11y_grpc_wrapper_test.py +++ b/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() diff --git a/android_env/wrappers/base_wrapper.py b/android_env/wrappers/base_wrapper.py index 8868e1ca..931fe4b2 100644 --- a/android_env/wrappers/base_wrapper.py +++ b/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,7 +96,10 @@ 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 @@ -79,7 +107,7 @@ 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: diff --git a/android_env/wrappers/base_wrapper_test.py b/android_env/wrappers/base_wrapper_test.py index 9a8414a2..8afe853f 100644 --- a/android_env/wrappers/base_wrapper_test.py +++ b/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() diff --git a/android_env/wrappers/last_action_wrapper.py b/android_env/wrappers/last_action_wrapper.py index 0fa8abb9..03329644 100644 --- a/android_env/wrappers/last_action_wrapper.py +++ b/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 diff --git a/android_env/wrappers/last_action_wrapper_test.py b/android_env/wrappers/last_action_wrapper_test.py index 71911c94..12da0926 100644 --- a/android_env/wrappers/last_action_wrapper_test.py +++ b/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() diff --git a/android_env/wrappers/swipe_action_wrapper.py b/android_env/wrappers/swipe_action_wrapper.py index 2308b182..799f7507 100644 --- a/android_env/wrappers/swipe_action_wrapper.py +++ b/android_env/wrappers/swipe_action_wrapper.py @@ -66,7 +66,7 @@ def _assert_base_env(self) -> None: def stats(self) -> dict[str, Any]: """Returns a dictionary of metrics logged by the environment.""" - logs = self._env.stats() + logs = super().stats() logs.update({'env_steps': self._env_steps}) return logs diff --git a/android_env/wrappers/swipe_action_wrapper_test.py b/android_env/wrappers/swipe_action_wrapper_test.py index 20caf4bf..685e5120 100644 --- a/android_env/wrappers/swipe_action_wrapper_test.py +++ b/android_env/wrappers/swipe_action_wrapper_test.py @@ -19,6 +19,7 @@ from absl.testing import parameterized 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 swipe_action_wrapper import dm_env from dm_env import specs @@ -239,6 +240,36 @@ def test_stats(self): self.assertIn('env_steps', stats) self.assertEqual(stats['env_steps'], 0) + def test_dm_env(self): + base_env = mock.create_autospec(dm_env.Environment) + base_env.action_spec.return_value = self._base_action_spec + wrapped_env = swipe_action_wrapper.SwipeActionWrapper(base_env) + + # Test stats fallback + stats = wrapped_env.stats() + self.assertIn('env_steps', stats) + self.assertEqual(stats['env_steps'], 0) + + # 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 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 + if __name__ == '__main__': absltest.main() diff --git a/android_env/wrappers/tap_action_wrapper.py b/android_env/wrappers/tap_action_wrapper.py index 541cfa89..291230c6 100644 --- a/android_env/wrappers/tap_action_wrapper.py +++ b/android_env/wrappers/tap_action_wrapper.py @@ -52,7 +52,7 @@ def __init__( def stats(self) -> dict[str, Any]: """Returns a dictionary of metrics logged by the environment.""" - logs = self._env.stats() + logs = super().stats() logs.update({'env_steps': self._env_steps}) return logs diff --git a/android_env/wrappers/tap_action_wrapper_test.py b/android_env/wrappers/tap_action_wrapper_test.py index 2e9edeb2..6b254145 100644 --- a/android_env/wrappers/tap_action_wrapper_test.py +++ b/android_env/wrappers/tap_action_wrapper_test.py @@ -20,6 +20,7 @@ from absl.testing import parameterized 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 tap_action_wrapper import dm_env from dm_env import specs @@ -234,6 +235,36 @@ def test_step_terminal(self): self.assertIs(ts.step_type, dm_env.StepType.LAST) self.assertEqual(ts.reward, 1.0) + def test_dm_env(self): + base_env = mock.create_autospec(dm_env.Environment) + base_env.action_spec.return_value = self._base_action_spec + wrapped_env = tap_action_wrapper.TapActionWrapper(base_env) + + # Test stats fallback + stats = wrapped_env.stats() + self.assertIn('env_steps', stats) + self.assertEqual(stats['env_steps'], 0) + + # 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 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 + if __name__ == '__main__': absltest.main()