From 7f14aa396659d30639eb93ebf14889cf9020c68f Mon Sep 17 00:00:00 2001 From: Daniel Toyama Date: Mon, 27 Jul 2026 13:20:55 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 954795564 --- android_env/android_device_env.py | 245 ----------------- android_env/android_device_env_test.py | 354 ------------------------- android_env/loader.py | 5 +- 3 files changed, 4 insertions(+), 600 deletions(-) delete mode 100644 android_env/android_device_env.py delete mode 100644 android_env/android_device_env_test.py diff --git a/android_env/android_device_env.py b/android_env/android_device_env.py deleted file mode 100644 index 12d19716..00000000 --- a/android_env/android_device_env.py +++ /dev/null @@ -1,245 +0,0 @@ -# coding=utf-8 -# Copyright 2026 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. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Task-Free AndroidDeviceEnv implementation connecting to on-device server.""" - -import time -from typing import Any, Self - -from absl import logging -from android_env import env_interface -from android_env.components import action_type -from android_env.components import android_device -from android_env.components.specs import base_action_spec -from android_env.components.specs import base_observation_spec -from android_env.proto import adb_pb2 -from android_env.proto import device_env_service_pb2 -import dm_env -from dm_env import specs -import numpy as np - - -class AndroidDeviceEnv(env_interface.AndroidEnvInterface): - """Task-Free AndroidDeviceEnv connected to a real device via WebSockets.""" - - def __init__(self, device: android_device.AndroidDevice): - """Initializes the AndroidDeviceEnv. - - Args: - device: AndroidDevice handle instance. - """ - self._device = device - self._latest_observation_time = 0.0 - state = self._device.get_state() - h, w = 1080, 1920 - for signal in state.signals: - if ( - signal.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT - and signal.HasField('screenshot') - and signal.screenshot.HasField('decoded') - ): - h = signal.screenshot.decoded.height or 1080 - w = signal.screenshot.decoded.width or 1920 - break - self._pixel_shape = (h, w, 3) - logging.info('Detected video resolution: %s', self._pixel_shape) - - self._active_signals = { - device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT, - device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE, - device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS, - device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE, - device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS, - } - self._device.update_subscriptions(self._active_signals) - self._latest_state = None - self._latest_pixels = None - self._is_touching = False - self._last_touch_position = None - self._last_non_touch_action = None - - def action_spec(self) -> dict[str, specs.Array]: - return base_action_spec(num_fingers=1, enable_key_events=True) - - def observation_spec(self) -> dict[str, specs.Array]: - return base_observation_spec( - height=self._pixel_shape[0], width=self._pixel_shape[1] - ) - - def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse: - """Executes an ADB call using the underlying AndroidDevice.""" - return self._device.execute_adb_call(call) - - def task_extras(self, latest_only: bool = True) -> dict[str, Any]: - """Returns extras from the latest state signals.""" - extras: dict[str, Any] = {} - if self._latest_state is not None: - for sig in self._latest_state.signals: - if sig.type == device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS: - extras['logs'] = list(sig.system_logs.values) - elif ( - sig.type - == device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS - ): - extras['accessibility_events'] = list(sig.accessibility_events.events) - elif ( - sig.type == device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE - ): - extras['accessibility_tree'] = sig.accessibility_forest - return extras - - def _build_observation( - self, state: device_env_service_pb2.DeviceState, timedelta_us: int - ) -> dict[str, Any]: - pixels, active_package, audio, _ = self._parse_signals(state) - if pixels is not None: - self._latest_pixels = pixels - else: - pixels = ( - self._latest_pixels - if self._latest_pixels is not None - else np.zeros(self._pixel_shape, dtype=np.uint8) - ) - - orient_one_hot = np.zeros(4, dtype=np.uint8) - orient_val = state.orientation - if 0 <= orient_val < 4: - orient_one_hot[orient_val] = 1 - else: - orient_one_hot[0] = 1 - - obs: dict[str, Any] = { - 'pixels': pixels, - 'timedelta': np.int64(timedelta_us), - 'active_package': active_package or '', - 'orientation': orient_one_hot, - } - if audio is not None: - obs['audio'] = audio - return obs - - def reset(self) -> dm_env.TimeStep: - """Resets the environment.""" - logging.info('Resetting AndroidDeviceEnv...') - self._latest_pixels = None - self._latest_state = self._device.get_state() - obs = self._build_observation(self._latest_state, timedelta_us=0) - self._latest_observation_time = time.time() - return dm_env.restart(obs) - - def step(self, action: dict[str, Any]) -> dm_env.TimeStep: - """Takes a step in the environment by executing an action.""" - env_action_type = action.get('action_type') - - proto_action = device_env_service_pb2.Action() - if env_action_type is not None: - if env_action_type == action_type.ActionType.TOUCH: - proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN - pos = action.get('touch_position', [0.0, 0.0]) - proto_action.touch_position.x = float(pos[0]) - proto_action.touch_position.y = float(pos[1]) - self._is_touching = True - self._last_touch_position = (float(pos[0]), float(pos[1])) - elif env_action_type == action_type.ActionType.REPEAT: - if self._is_touching: - assert self._last_touch_position is not None - proto_action.action_type = ( - device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE - ) - proto_action.touch_position.x = self._last_touch_position[0] - proto_action.touch_position.y = self._last_touch_position[1] - elif env_action_type == action_type.ActionType.LIFT: - proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_UP - if 'touch_position' in action: - pos = action['touch_position'] - proto_action.touch_position.x = float(pos[0]) - proto_action.touch_position.y = float(pos[1]) - elif self._last_touch_position is not None: - proto_action.touch_position.x = self._last_touch_position[0] - proto_action.touch_position.y = self._last_touch_position[1] - self._is_touching = False - self._last_touch_position = None - elif env_action_type == action_type.ActionType.KEYDOWN: - proto_action.action_type = device_env_service_pb2.ACTION_TYPE_KEY_EVENT - if 'keycode' in action: - proto_action.keycode = int(action['keycode']) - self._last_non_touch_action = device_env_service_pb2.Action() - self._last_non_touch_action.CopyFrom(proto_action) - elif env_action_type == action_type.ActionType.KEYUP: - proto_action.action_type = ( - device_env_service_pb2.ACTION_TYPE_UNSPECIFIED - ) - self._last_non_touch_action = None - else: - proto_action.action_type = ( - device_env_service_pb2.ACTION_TYPE_UNSPECIFIED - ) - - if ( - proto_action.action_type - != device_env_service_pb2.ACTION_TYPE_UNSPECIFIED - ): - self._device.send_action(proto_action) - self._latest_state = self._device.get_state() - - now = time.time() - timedelta_us = int((now - self._latest_observation_time) * 1e6) - self._latest_observation_time = now - - obs = self._build_observation(self._latest_state, timedelta_us=timedelta_us) - return dm_env.transition(reward=0.0, observation=obs) - - def _parse_signals( - self, state: device_env_service_pb2.DeviceState | None - ) -> tuple[ - np.ndarray | None, - str | None, - np.ndarray | None, - list[str] | None, - ]: - if state is None: - return None, None, None, None - pixels = None - active_package = None - audio = None - logs = None - - for sig in state.signals: - if sig.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT: - if sig.screenshot.HasField('decoded'): - decoded = sig.screenshot.decoded - pixels = np.frombuffer(decoded.raw_pixels, dtype=np.uint8).reshape( - (decoded.height, decoded.width, 3) - ) - elif sig.type == device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE: - active_package = sig.active_package - elif sig.type == device_env_service_pb2.DEVICE_SIGNAL_AUDIO_OUTPUT: - raw_bytes = sig.audio_output.raw_bytes - if raw_bytes: - audio = np.frombuffer(raw_bytes, dtype=np.int16).reshape(-1, 2) - elif sig.type == device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS: - logs = list(sig.system_logs.values) - - return pixels, active_package, audio, logs - - def close(self) -> None: - logging.info('Closing AndroidDeviceEnv...') - self._device.close() - - def __enter__(self) -> Self: - return self - - def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: - self._device.__exit__(exc_type, exc_value, traceback) diff --git a/android_env/android_device_env_test.py b/android_env/android_device_env_test.py deleted file mode 100644 index 997b4d99..00000000 --- a/android_env/android_device_env_test.py +++ /dev/null @@ -1,354 +0,0 @@ -# coding=utf-8 -# Copyright 2026 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. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for android_device_env adapter.""" - -from unittest import mock -from absl.testing import absltest -from android_env import android_device_env -from android_env.components import action_type -from android_env.components import android_device -from android_env.components import config_classes -from android_env.components import device_connection -from android_env.proto import adb_pb2 -from android_env.proto import device_env_service_pb2 -import numpy as np - - -class FakeDeviceConnection(device_connection.DeviceConnection): - - def __init__( - self, config: config_classes.DeviceConnectionConfig | None = None - ): - config = config or config_classes.DeviceConnectionConfig() - self._config = config - self.closed: bool = False - self.subscriptions: set[int] = set() - self.injected_actions: list[device_env_service_pb2.Action] = [] - self.device_state: device_env_service_pb2.DeviceState = ( - device_env_service_pb2.DeviceState() - ) - - def connect(self): - pass - - def get_video_metadata(self) -> tuple[str, int, int]: - return ('h264', 640, 480) - - def get_device_state(self) -> device_env_service_pb2.DeviceState: - return self.device_state - - def inject_action(self, action: device_env_service_pb2.Action): - self.injected_actions.append(action) - - def update_subscriptions(self, signals: set[int]): - self.subscriptions = set(signals) - - def send_message(self, msg: device_env_service_pb2.ClientMessage): - payload_type = msg.WhichOneof('payload') - if payload_type == 'inject_action': - self.injected_actions.append(msg.inject_action.action) - elif payload_type == 'update_subscriptions': - self.subscriptions = set(msg.update_subscriptions.active_signals) - - def close(self): - self.closed = True - - -class AndroidDeviceEnvTest(absltest.TestCase): - - def test_update_subscriptions_called_on_init(self): - fake_conn = FakeDeviceConnection() - dev = android_device.AndroidDevice(connection=fake_conn) - _ = android_device_env.AndroidDeviceEnv(device=dev) - expected_signals = { - device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT, - device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE, - device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS, - device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE, - device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS, - } - self.assertEqual(fake_conn.subscriptions, expected_signals) - - def test_adapter_lifecycle(self): - fake_conn = FakeDeviceConnection() - on_close = mock.MagicMock() - dev = android_device.AndroidDevice( - connection=fake_conn, on_close_callbacks=[on_close] - ) - env = android_device_env.AndroidDeviceEnv(device=dev) - action_spec = env.action_spec() - self.assertIn('action_type', action_spec) - self.assertIn('touch_position', action_spec) - self.assertIn('keycode', action_spec) - - obs_spec = env.observation_spec() - self.assertIn('pixels', obs_spec) - self.assertIn('orientation', obs_spec) - self.assertIn('timedelta', obs_spec) - ts = env.reset() - self.assertTrue(ts.first()) - ts = env.step({}) - self.assertTrue(ts.mid()) - env.close() - self.assertTrue(fake_conn.closed) - on_close.assert_called_once() - - def test_video_resolution_detection_from_screenshot_signal(self): - fake_conn = FakeDeviceConnection() - state = device_env_service_pb2.DeviceState() - sig = state.signals.add() - sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT - sig.screenshot.decoded.height = 480 - sig.screenshot.decoded.width = 640 - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - obs_spec = env.observation_spec() - self.assertEqual(obs_spec['pixels'].shape, (480, 640, 3)) - - def test_video_resolution_fallback_default(self): - fake_conn = FakeDeviceConnection() - fake_conn.device_state = device_env_service_pb2.DeviceState() - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - obs_spec = env.observation_spec() - self.assertEqual(obs_spec['pixels'].shape, (1080, 1920, 3)) - - def test_build_observation_pixels_and_cached(self): - fake_conn = FakeDeviceConnection() - raw_pixels = np.ones((2, 2, 3), dtype=np.uint8) - state = device_env_service_pb2.DeviceState() - sig = state.signals.add() - sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT - sig.screenshot.decoded.width = 2 - sig.screenshot.decoded.height = 2 - sig.screenshot.decoded.raw_pixels = raw_pixels.tobytes() - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - ts = env.reset() - np.testing.assert_array_equal(ts.observation['pixels'], raw_pixels) - self.assertEqual(ts.observation['active_package'], '') - - # Next step with no new screenshot reuses cached pixels. - fake_conn.device_state = device_env_service_pb2.DeviceState() - ts = env.step({}) - np.testing.assert_array_equal(ts.observation['pixels'], raw_pixels) - self.assertEqual(ts.observation['active_package'], '') - - def test_build_observation_orientation_bounds(self): - fake_conn = FakeDeviceConnection() - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - - # Valid orientation 2 -> one-hot [0, 0, 1, 0] - fake_conn.device_state.orientation = 2 - ts = env.step({}) - np.testing.assert_array_equal( - ts.observation['orientation'], np.array([0, 0, 1, 0], dtype=np.uint8) - ) - - # Out-of-bounds orientation 99 -> fallback one-hot [1, 0, 0, 0] - fake_conn.device_state.orientation = 99 - ts = env.step({}) - np.testing.assert_array_equal( - ts.observation['orientation'], np.array([1, 0, 0, 0], dtype=np.uint8) - ) - - def test_task_extras(self): - fake_conn = FakeDeviceConnection() - state = device_env_service_pb2.DeviceState() - sig1 = state.signals.add() - sig1.type = device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS - sig1.system_logs.values.append('log1') - sig2 = state.signals.add() - sig2.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS - ev = sig2.accessibility_events.events.add() - sig3 = state.signals.add() - sig3.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - env.reset() - extras = env.task_extras() - self.assertEqual(extras['logs'], ['log1']) - self.assertEqual(extras['accessibility_events'], [ev]) - self.assertEqual(extras['accessibility_tree'], sig3.accessibility_forest) - - def test_step_action_types(self): - mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) - mock_dev.get_state.return_value = device_env_service_pb2.DeviceState() - env = android_device_env.AndroidDeviceEnv(device=mock_dev) - - # 1. TOUCH - env.step({ - 'action_type': action_type.ActionType.TOUCH, - 'touch_position': [0.2, 0.8], - }) - self.assertTrue(mock_dev.send_action.called) - last_act = mock_dev.send_action.call_args[0][0] - self.assertEqual( - last_act.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN - ) - self.assertAlmostEqual(last_act.touch_position.x, 0.2) - self.assertAlmostEqual(last_act.touch_position.y, 0.8) - - # 2. REPEAT when touching - mock_dev.send_action.reset_mock() - env.step({'action_type': action_type.ActionType.REPEAT}) - last_act = mock_dev.send_action.call_args[0][0] - self.assertEqual( - last_act.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE - ) - - # 3. LIFT with position - mock_dev.send_action.reset_mock() - env.step({ - 'action_type': action_type.ActionType.LIFT, - 'touch_position': [0.5, 0.5], - }) - last_act = mock_dev.send_action.call_args[0][0] - self.assertEqual( - last_act.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_UP - ) - - # 4. REPEAT when not touching (no action sent) - mock_dev.send_action.reset_mock() - env.step({'action_type': action_type.ActionType.REPEAT}) - mock_dev.send_action.assert_not_called() - - # 5. KEYDOWN - mock_dev.send_action.reset_mock() - env.step({'action_type': action_type.ActionType.KEYDOWN, 'keycode': 66}) - last_act = mock_dev.send_action.call_args[0][0] - self.assertEqual( - last_act.action_type, device_env_service_pb2.ACTION_TYPE_KEY_EVENT - ) - self.assertEqual(last_act.keycode, 66) - - # 6. KEYUP - mock_dev.send_action.reset_mock() - env.step({'action_type': action_type.ActionType.KEYUP}) - - # 7. Unrecognized action type - mock_dev.send_action.reset_mock() - env.step({'action_type': 999}) - mock_dev.send_action.assert_not_called() - - def test_parse_signals_extra_signals(self): - fake_conn = FakeDeviceConnection() - state = device_env_service_pb2.DeviceState() - sig1 = state.signals.add() - sig1.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE - sig1.active_package = 'com.example.app' - sig2 = state.signals.add() - sig2.type = device_env_service_pb2.DEVICE_SIGNAL_AUDIO_OUTPUT - pcm = np.array([[10, 20]], dtype=np.int16) - sig2.audio_output.raw_bytes = pcm.tobytes() - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - ts = env.reset() - self.assertEqual(ts.observation['active_package'], 'com.example.app') - np.testing.assert_array_equal(ts.observation['audio'], pcm) - - def test_parse_signals_none_state(self): - fake_conn = FakeDeviceConnection() - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - res = env._parse_signals(None) - self.assertEqual(res, (None, None, None, None)) - - def test_parse_signals_screenshot_undecoded(self): - fake_conn = FakeDeviceConnection() - fake_conn.subscriptions = set() - state = device_env_service_pb2.DeviceState() - sig = state.signals.add() - sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT - sig.screenshot.encoded_bytes = b'fake_compressed_jpeg_bytes' - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - pixels, _, _, _ = env._parse_signals(state) - self.assertIsNone(pixels) - self.assertTrue(sig.screenshot.HasField('encoded_bytes')) - - def test_parse_signals_system_logs(self): - fake_conn = FakeDeviceConnection() - state = device_env_service_pb2.DeviceState() - sig = state.signals.add() - sig.type = device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS - sig.system_logs.values.append('log line 1') - sig.system_logs.values.append('log line 2') - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - _, _, _, logs = env._parse_signals(state) - self.assertEqual(logs, ['log line 1', 'log line 2']) - - def test_parse_signals_screenshot_empty_signal(self): - fake_conn = FakeDeviceConnection() - fake_conn.subscriptions = set() - state = device_env_service_pb2.DeviceState() - sig = state.signals.add() - sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - pixels, _, _, _ = env._parse_signals(state) - self.assertIsNone(pixels) - - def test_parse_signals_logs_none_for_non_logs_signals(self): - fake_conn = FakeDeviceConnection() - state = device_env_service_pb2.DeviceState() - sig1 = state.signals.add() - sig1.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE - sig1.active_package = 'com.example.app' - sig2 = state.signals.add() - sig2.type = device_env_service_pb2.DEVICE_SIGNAL_AUDIO_OUTPUT - fake_conn.device_state = state - - dev = android_device.AndroidDevice(connection=fake_conn) - env = android_device_env.AndroidDeviceEnv(device=dev) - _, _, _, logs = env._parse_signals(state) - self.assertIsNone(logs) - - def test_context_manager(self): - mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) - with android_device_env.AndroidDeviceEnv(device=mock_dev) as env: - self.assertIsInstance(env, android_device_env.AndroidDeviceEnv) - mock_dev.__exit__.assert_called_once() - - def test_execute_adb_call_delegates_to_device(self): - mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) - env = android_device_env.AndroidDeviceEnv(device=mock_dev) - req = adb_pb2.AdbRequest() - mock_dev.execute_adb_call.return_value = adb_pb2.AdbResponse() - res = env.execute_adb_call(req) - mock_dev.execute_adb_call.assert_called_once_with(req) - self.assertIsInstance(res, adb_pb2.AdbResponse) - - -if __name__ == '__main__': - absltest.main() diff --git a/android_env/loader.py b/android_env/loader.py index aac854da..af5cd9eb 100644 --- a/android_env/loader.py +++ b/android_env/loader.py @@ -18,6 +18,7 @@ import os from absl import logging +from android_env import env_interface from android_env import environment from android_env.components import config_classes from android_env.components import coordinator as coordinator_lib @@ -44,7 +45,9 @@ def _load_task(task_config: config_classes.TaskConfig) -> task_pb2.Task: return task -def load(config: config_classes.AndroidEnvConfig) -> environment.AndroidEnv: +def load( + config: config_classes.AndroidEnvConfig, +) -> env_interface.AndroidEnvInterface: """Loads an AndroidEnv instance.""" task = _load_task(config.task)