[
{
"path": ".github/workflows/tests.yml",
"content": "name: tests\n\non:\n push:\n branches: [ main ]\n pull_request:\n branches: [ main ]\n workflow_dispatch:\n inputs:\n git-ref:\n description: Git Ref (Optional)\n required: false\n\njobs:\n build:\n runs-on: ubuntu-latest\n env:\n TEST_TMPDIR: '/tmp'\n strategy:\n matrix:\n python-version: [\"3.11\", \"3.12\", \"3.13\"]\n steps:\n - uses: actions/checkout@v6\n\n - name: Set up Python ${{ matrix.python-version }}\n uses: actions/setup-python@v6\n with:\n python-version: ${{ matrix.python-version }}\n\n - name: Install dependencies\n run: |\n pip install --upgrade pip setuptools\n python setup.py install\n pip install .[testing]\n\n - name: Run tests\n run: |\n # Find all test files, print their names and execute them in parallel\n # with a maximum of 20 proccesses.\n find . -type f -name \"*_test.py\" -print0 | xargs -t -0 -n1 -P 20 python3\n"
},
{
"path": "CONTRIBUTING.md",
"content": "\n# How to Contribute\n\n# Pull Requests\n\nPlease send in fixes or feature additions through Pull Requests.\n\n## Contributor License Agreement\n\nContributions to this project must be accompanied by a Contributor License\nAgreement. You (or your employer) retain the copyright to your contribution,\nthis simply gives us permission to use and redistribute your contributions as\npart of the project. Head over to to see\nyour current agreements on file or to sign a new one.\n\nYou generally only need to submit a CLA once, so if you've already submitted one\n(even if it was for a different project), you probably don't need to do it\nagain.\n\n## Code reviews\n\nAll submissions, including submissions by project members, require review. We\nuse GitHub pull requests for this purpose. Consult\n[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more\ninformation on using pull requests.\n"
},
{
"path": "LICENSE",
"content": "\n Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License."
},
{
"path": "README.md",
"content": "\n# AndroidEnv - The Android Learning Environment\n\n![](\"docs/images/device_control.gif\")
\n\n[AndroidEnv](https://github.com/deepmind/android_env) is a Python library that\nexposes an [Android](https://www.android.com/) device as a Reinforcement\nLearning (RL) environment. The library provides a flexible platform for defining\ncustom tasks on top of the Android Operating System, including any Android\napplication. Agents interact with the device through a universal action\ninterface - the touchscreen - by sending localized touch and lift events to the\nsystem. The library processes these events and returns pixel observations and\nrewards as provided by specific [task definitions](docs/tasks_guide.md). For\nexample, rewards might be given for events such as successfully scrolling down a\npage, sending an email, or achieving some score in a game, depending on the\nresearch purpose and how the user configures the task.\n\n[](https://github.com/deepmind/android_env/actions/workflows/tests.yml)\n[](https://badge.fury.io/py/android-env)\n[](https://pepy.tech/project/android-env)\n\n## Index\n\n* [Environment details](docs/environment.md)\n* [Running AndroidEnv](docs/instructions.md)\n* [Setting up a virtual Android device](docs/emulator_guide.md)\n* [Defining a task in AndroidEnv](docs/tasks_guide.md)\n* [Example tasks available for download](docs/example_tasks.md)\n\n## Environment features\n\nThere are a number of aspects that make AndroidEnv a challenging yet suitable\nenvironment for Reinforcement Learning research:\n\n* Allowing agents to interact with a system used daily by billions of users\n around the world, AndroidEnv offers a platform for RL agents to navigate,\n learn tasks and have direct impact in **real-world contexts**. The\n environment wraps a simulated Android device, which runs independently from\n the environment, completely unaltered, and works in exactly the same way as\n the devices that humans use, exposing exactly the same features and\n services.\n\n* The platform offers a virtually infinite **range of possible tasks**, all\n sharing a common action interface. The library facilitates the design of\n Reinforcement Learning tasks for any existing or custom built Android\n application. For example, it exposes the broad world of Android games,\n ranging from card games, puzzle games, time reactive games, all requiring a\n diverse set of action combinations and interaction types.\n\n* The environment runs on top of a **real-time simulation** of an Android\n device. In other words, the environment dynamics does not wait for the agent\n to deliberate, and the speed of the simulation cannot be increased.\n\n* The observation is a collection of **RGB values** corresponding to the\n displayed pixels on the screen. The exact screen resolution depends on the\n simulated device, but in general it will be considered relatively large in\n an RL context. However, users have the option of downsampling each\n observation.\n\n* The learning environment has an interesting, **complex action space** unique\n to the touchscreen interface of Android.\n\n * The raw, **hybrid action space** consists of a continuous tuple\n signifying the action location, and a discrete signal determining\n whether the agent wants to touch the screen or lift its virtual finger.\n * Raw actions are highly **composable**: the Android UI and most\n applications were designed so that they could be intuitively navigated\n via common\n [touchscreen gestures](https://developer.android.com/training/gestures/detector)\n such as tapping, scrolling, swiping, pinching, drag & drop etc. This is\n still the case in AndroidEnv: to trigger meaningful changes in the\n environment, the agent often has to perform carefully timed and\n positioned sequences of raw actions. For example, in order to navigate\n to the next image in a photo gallery, the agent would have to perform a\n *swipe*, touching the screen multiple times, gradually shifting the\n actions' positions to the right. Thus, in most contexts raw actions do\n not trigger changes in the state of the environment unless correctly\n chained together to make up a human gesture.\n * The action interface is **closely related to the observation space**, as\n meaningful touch and lift events are often either co-localized or\n strongly correlated to the location or movement of salient objects in\n the observation. For example, the position of a button on the screen\n aligns with the location of the actions that trigger the button press.\n * The library provides tools for flexibly **altering the action\n interface** if needed for particular studies, such as discretization or\n hard-coding gesture skills. Still, we believe that the real challenge\n remains in devising agents that are capable of dealing with a large\n suite of diverse tasks, through acting and learning in the complex\n unifying action interface.\n\n# Getting started\n\n### Installation\n\nThe easiest way to get AndroidEnv is with pip:\n\n```shell\n$ python3 -m pip install android-env\n```\n\nPlease note that `/examples` are not included in this package.\n\nAlternatively, you can clone the repository from git's `main` branch:\n\n```shell\n$ git clone https://github.com/deepmind/android_env/\n$ cd android_env\n$ python3 setup.py install\n```\n\nUpdate: the environment now runs on Windows, but please keep in mind that this\noption is not well-maintained or widely supported, as Unix-based systems are the\nprimary target platforms of this project.\n\n### Create a simulator\n\nBefore running the environment, you will need access to an emulated Android\ndevice. For instructions on creating a virtual Android device, see the\n[Emulator guide](docs/emulator_guide.md).\n\n### Define a task\n\nThen, you will want to define what the agent's *task* is. At this point, the\nagent will be able to communicate with the emulated device, but it will not yet\nhave an objective, or access to signals such as rewards or RL episode ends.\nLearn [how to define an RL task](docs/tasks_guide.md) of your own, or use one of\nthe [existing task definitions](docs/example_tasks.md) for training.\n\n### Load and run\n\nTo find out how to run and train agents on AndroidEnv, see these\n[detailed instructions](docs/instructions.md). Here you can also find example\nscripts demonstrating how to run a random agent, an\n[acme](https://github.com/deepmind/acme) agent, or a human agent on AndroidEnv.\n\n## About\n\nThis library is developed and maintained by [DeepMind](http://deepmind.com). \\\nYou can find the [technical report](https://arxiv.org/abs/2105.13231) on Arxiv,\nas well as an introductory\n[blog\npost](https://www.deepmind.com/publications/androidenv-the-android-learning-environment)\non DeepMind's website.\n\nIf you use AndroidEnv in your research, you can cite the paper using the\nfollowing BibTeX:\n\n```\n@article{ToyamaEtAl2021AndroidEnv,\n title = {{AndroidEnv}: A Reinforcement Learning Platform for Android},\n author = {Daniel Toyama and Philippe Hamel and Anita Gergely and\n Gheorghe Comanici and Amelia Glaese and Zafarali Ahmed and Tyler\n Jackson and Shibl Mourad and Doina Precup},\n year = {2021},\n eprint = {2105.13231},\n archivePrefix = {arXiv},\n primaryClass = {cs.LG},\n volume = {abs/2105.13231},\n url = {http://arxiv.org/abs/2105.13231},\n}\n```\n\nDisclaimer: This is not an official Google product.\n"
},
{
"path": "android_env/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/apps/MODULE.bazel",
"content": "# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Bazel dependencies for building Catch.\nmodule(\n name = \"catch\",\n version = \"1.0\",\n)\n\nbazel_dep(name = \"rules_android\", version = \"0.6.6\")\nbazel_dep(name = \"rules_kotlin\", version = \"2.1.8\")\nbazel_dep(name = \"rules_jvm_external\", version = \"6.7\")\nbazel_dep(name = \"rules_robolectric\", version = \"4.16\", repo_name = \"robolectric\")\nbazel_dep(name = \"rules_java\", version = \"9.0.3\")\nbazel_dep(name = \"protobuf\", version = \"30.0\")\n\n# To avoid conflict with different protobuf versions.\nsingle_version_override(\n module_name = \"protobuf\",\n version = \"30.0\",\n)\n\nmaven = use_extension(\"@rules_jvm_external//:extensions.bzl\", \"maven\")\n\n# Need to set testonly = True because the package depends on testonly targets.\nmaven.artifact(\n testonly = True,\n artifact = \"runner\",\n group = \"androidx.test\",\n version = \"1.7.0\",\n)\nmaven.artifact(\n testonly = True,\n artifact = \"junit\",\n group = \"androidx.test.ext\",\n version = \"1.3.0\",\n)\nmaven.artifact(\n testonly = True,\n artifact = \"mockito-kotlin\",\n group = \"org.mockito.kotlin\",\n version = \"6.1.0\",\n)\nmaven.install(\n artifacts = [\n \"androidx.test.ext:junit:1.3.0\",\n \"androidx.test:runner:1.7.0\",\n \"com.google.guava:guava:32.0.1-jre\",\n \"com.google.truth:truth:1.4.0\",\n \"org.mockito.kotlin:mockito-kotlin:6.1.0\",\n \"org.mockito:mockito-core:5.20.0\",\n \"org.robolectric:robolectric:4.16\",\n \"org.yaml:snakeyaml:2.5\",\n ],\n repositories = [\n \"https://maven.google.com\",\n \"https://repo1.maven.org/maven2\",\n ],\n)\nuse_repo(maven, \"maven\")\n\nremote_android_extensions = use_extension(\n \"@rules_android//bzlmod_extensions:android_extensions.bzl\",\n \"remote_android_tools_extensions\",\n)\nuse_repo(remote_android_extensions, \"android_tools\")\n\nandroid_sdk_repository_extension = use_extension(\"@rules_android//rules/android_sdk_repository:rule.bzl\", \"android_sdk_repository_extension\")\nuse_repo(android_sdk_repository_extension, \"androidsdk\")\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/AccessibilityForwarder.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.accessibilityservice.AccessibilityService\nimport android.util.Log\nimport android.view.accessibility.AccessibilityEvent\nimport android.view.accessibility.AccessibilityNodeInfo\nimport android.view.accessibility.AccessibilityWindowInfo\nimport com.google.androidenv.accessibilityforwarder.A11yServiceGrpcKt.A11yServiceCoroutineStub\nimport io.grpc.ManagedChannel\nimport io.grpc.ManagedChannelBuilder\nimport io.grpc.ProxyDetector\nimport io.grpc.StatusException\nimport kotlinx.coroutines.TimeoutCancellationException\nimport kotlinx.coroutines.runBlocking\nimport kotlinx.coroutines.withTimeout\n\n/**\n * An Android service that listens to accessibility events and forwards them via gRPC.\n *\n * This service also logs the accessibility tree if [LogFlags.logAccessibilityTree] is set and if\n * [LogFlags.grpcPort] is positive.\n *\n * Please see\n * https://developer.android.com/reference/android/view/accessibility/AccessibilityEvent#getEventType()\n * for a comprehensive list of events emitted by Android.\n */\nclass AccessibilityForwarder(\n private val channelFactory: (host: String, port: Int) -> ManagedChannel = { host, port ->\n ManagedChannelBuilder.forAddress(host, port)\n .proxyDetector(ProxyDetector { _ -> null })\n .usePlaintext()\n .build()\n }\n) : AccessibilityService() {\n\n init {\n // Spawn long-running thread for periodically logging the tree.\n Thread(\n Runnable {\n while (LogFlags.a11yTreePeriodMs > 0) {\n try {\n logAccessibilityTree()\n } catch (e: ConcurrentModificationException) {\n continue\n }\n\n Thread.sleep(/* millis= */ LogFlags.a11yTreePeriodMs)\n }\n }\n )\n .start()\n }\n\n // grpcStub has a backing property that can be reset to null.\n private var _grpcStub: A11yServiceCoroutineStub? = null\n val grpcStub: A11yServiceCoroutineStub\n get() {\n if (_grpcStub == null) {\n Log.i(TAG, \"Building channel on ${LogFlags.grpcHost}:${LogFlags.grpcPort}.\")\n _grpcStub = A11yServiceCoroutineStub(channelFactory(LogFlags.grpcHost, LogFlags.grpcPort))\n }\n return _grpcStub!!\n }\n\n private fun resetGrpcStub() {\n _grpcStub = null\n }\n\n override fun onInterrupt() {\n LogFlags.a11yTreePeriodMs = 0 // Turn off periodic tree forwarding.\n }\n\n override fun onAccessibilityEvent(event: AccessibilityEvent?) {\n if (event == null) {\n Log.i(TAG, \"`event` is null.\")\n return\n }\n\n logExtrasForEvent(event)\n val eventType = event.eventType\n val eventTypeStr: String = AccessibilityEvent.eventTypeToString(eventType)\n if (eventTypeStr.isNotEmpty()) {\n Log.i(TAG, eventTypeStr)\n }\n }\n\n private fun logAccessibilityTree() {\n if (!LogFlags.logAccessibilityTree) {\n Log.i(TAG, \"Not logging accessibility tree\")\n return\n }\n\n val windows = getWindowsOrNull()\n\n if (windows == null) {\n Log.i(TAG, \"windows is null.\")\n return\n }\n\n // Check gRPC port before actually building the forest.\n if (LogFlags.grpcPort <= 0) {\n Log.w(TAG, \"Can't log accessibility tree because gRPC port has not been set.\")\n return\n }\n\n val forest = creator.buildForest(windows)\n try {\n val grpcTimeoutMillis = 1000L\n val response: ForestResponse =\n with(grpcStub) {\n Log.i(TAG, \"sending (blocking) gRPC request for tree.\")\n runBlocking { withTimeout(grpcTimeoutMillis) { sendForest(forest) } }\n }\n if (response.error.isNotEmpty()) {\n Log.w(TAG, \"gRPC response.error: ${response.error}\")\n } else {\n Log.i(TAG, \"gRPC request for tree succeeded.\")\n }\n } catch (e: StatusException) {\n Log.w(TAG, \"gRPC StatusException; are you sure networking is turned on?\")\n Log.i(TAG, \"extra: exception ['$e']\")\n resetGrpcStub()\n } catch (e: TimeoutCancellationException) {\n Log.w(TAG, \"gRPC TimeoutCancellationException; are you sure networking is turned on?\")\n Log.i(TAG, \"extra: exception ['$e']\")\n resetGrpcStub()\n }\n }\n\n private fun getWindowsOrNull(): List? =\n try {\n windows\n } catch (e: NullPointerException) {\n null\n }\n\n /** Logs extras for all event types. */\n private fun logExtrasForEvent(event: AccessibilityEvent) {\n\n val events: MutableMap = mutableMapOf()\n\n val sourceDescription = event.source?.contentDescription()\n if (!sourceDescription.isNullOrEmpty()) {\n events.put(\"source_content_description\", sourceDescription)\n }\n\n // Output the event text.\n val eventText = event.text.joinToString(\", \")\n if (eventText.isNotEmpty()) {\n events.put(\"event_text\", eventText)\n }\n\n // Output the source text.\n val sourceText = event.source?.text?.toString()\n if (!sourceText.isNullOrEmpty()) {\n events.put(\"source_text\", sourceText)\n }\n\n val eventTypeStr: String = AccessibilityEvent.eventTypeToString(event.eventType)\n if (eventTypeStr.isNotEmpty()) {\n events.put(\"event_type\", eventTypeStr)\n }\n\n val className = event.source?.className?.toString()\n if (!className.isNullOrEmpty()) {\n events.put(\"source_class_name\", className)\n }\n\n val packageName = event.packageName?.toString()\n if (!packageName.isNullOrEmpty()) {\n events.put(\"event_package_name\", packageName)\n }\n\n // Text editing properties.\n val beforeText = event.beforeText?.toString()\n if (!beforeText.isNullOrEmpty()) {\n events.put(\"before_text\", beforeText)\n }\n\n val fromIndex = event.fromIndex\n if (fromIndex != -1) {\n events.put(\"from_index\", fromIndex.toString())\n }\n\n val toIndex = event.toIndex\n if (toIndex != -1) {\n events.put(\"to_index\", toIndex.toString())\n }\n\n val addedCount = event.addedCount\n if (addedCount != -1) {\n events.put(\"added_count\", addedCount.toString())\n }\n\n val removedCount = event.removedCount\n if (removedCount != -1) {\n events.put(\"removed_count\", removedCount.toString())\n }\n\n // Text traversal properties\n val movementGranularity = event.movementGranularity\n if (movementGranularity != 0) {\n events.put(\"movement_granularity\", movementGranularity.toString())\n }\n\n val action = event.action\n if (action != 0) {\n events.put(\"action\", action.toString())\n }\n\n // Scrolling properties.\n if (eventTypeStr == \"TYPE_VIEW_SCROLLED\") {\n events.put(\"scroll_delta_x\", event.scrollDeltaX.toString())\n events.put(\"scroll_delta_y\", event.scrollDeltaY.toString())\n }\n\n // Report viewID so we know exactly where the event came from.\n val viewId = event.source?.viewIdResourceName?.toString()\n if (!viewId.isNullOrEmpty()) {\n events.put(\"view_id\", viewId)\n }\n\n // Format [events] as a Python dict.\n if (events.isNotEmpty()) {\n events.put(\"event_timestamp_ms\", event.eventTime.toString(10))\n // Check if we want to use gRPC.\n if (LogFlags.grpcPort > 0) {\n try {\n val grpcTimeoutMillis = 1000L\n val request = eventRequest { this.event.putAll(events) }\n val response: EventResponse =\n with(grpcStub) {\n Log.i(TAG, \"sending (blocking) gRPC request for event.\")\n runBlocking { withTimeout(grpcTimeoutMillis) { sendEvent(request) } }\n }\n if (response.error.isNotEmpty()) {\n Log.w(TAG, \"gRPC response.error: ${response.error}\")\n } else {\n Log.i(TAG, \"gRPC request for event succeeded.\")\n }\n } catch (e: StatusException) {\n Log.w(TAG, \"gRPC StatusException; are you sure networking is turned on?\")\n Log.i(TAG, \"extra: exception ['$e']\")\n resetGrpcStub()\n } catch (e: TimeoutCancellationException) {\n Log.w(TAG, \"gRPC TimeoutCancellationException; are you sure networking is turned on?\")\n Log.i(TAG, \"extra: exception ['$e']\")\n resetGrpcStub()\n }\n } else {\n Log.w(TAG, \"Can't log accessibility event because gRPC port has not been set.\")\n }\n }\n }\n\n /** Recursively climbs the accessibility tree until the root, collecting descriptions. */\n private fun AccessibilityNodeInfo?.contentDescription(): String {\n if (this == null) {\n return \"\"\n }\n\n val descriptions = mutableListOf()\n var current: AccessibilityNodeInfo? = this\n while (current != null) {\n val description = current.contentDescription\n if (description != null) {\n descriptions.add(description.toString())\n }\n\n current = current.parent\n }\n return descriptions.joinToString(\", \")\n }\n\n companion object {\n private const val TAG = \"AndroidRLTask\"\n private val creator = AccessibilityTreeCreator()\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/AccessibilityForwarderTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.view.accessibility.AccessibilityEvent\nimport android.view.accessibility.AccessibilityNodeInfo\nimport android.view.accessibility.AccessibilityWindowInfo\nimport com.google.common.truth.Truth.assertThat\nimport io.grpc.Status\nimport io.grpc.StatusException\nimport io.grpc.inprocess.InProcessChannelBuilder\nimport io.grpc.inprocess.InProcessServerBuilder\nimport io.grpc.testing.GrpcCleanupRule\nimport org.junit.Assert.assertFalse\nimport org.junit.Rule\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.robolectric.RobolectricTestParameterInjector\nimport org.robolectric.Shadows.shadowOf\n\n@RunWith(RobolectricTestParameterInjector::class)\nclass AccessibilityForwarderTest {\n\n @get:Rule(order = 1) val cleanupRule = GrpcCleanupRule()\n\n class FakeAccessibilityService : A11yServiceGrpcKt.A11yServiceCoroutineImplBase() {\n var sendForestChecker: (AndroidAccessibilityForest) -> String = { _ -> \"\" }\n var sendEventChecker: (EventRequest) -> String = { _ -> \"\" }\n\n override suspend fun sendForest(request: AndroidAccessibilityForest) = forestResponse {\n error = sendForestChecker(request)\n }\n\n override suspend fun sendEvent(request: EventRequest) = eventResponse {\n error = sendEventChecker(request)\n }\n }\n\n protected lateinit var forwarder: AccessibilityForwarder\n protected val fakeA11yService = FakeAccessibilityService()\n protected val channel by lazy {\n val serverName: String = InProcessServerBuilder.generateName()\n cleanupRule.register(\n InProcessServerBuilder.forName(serverName)\n .directExecutor()\n .addService(fakeA11yService)\n .build()\n .start()\n )\n cleanupRule.register(InProcessChannelBuilder.forName(serverName).directExecutor().build())\n }\n\n /** Initializes [forwarder] and [LogFlags] from the given args. */\n fun createForwarder(\n logAccessibilityTree: Boolean = false,\n a11yTreePeriodMs: Long = 0,\n grpcHost: String = \"10.0.2.2\",\n grpcPort: Int = 0,\n a11yWindows: MutableList? = null,\n ) {\n LogFlags.logAccessibilityTree = logAccessibilityTree\n LogFlags.a11yTreePeriodMs = a11yTreePeriodMs\n LogFlags.grpcHost = grpcHost\n LogFlags.grpcPort = grpcPort\n forwarder = AccessibilityForwarder({ _, _ -> channel })\n if (a11yWindows == null) {\n shadowOf(forwarder).setWindows(mutableListOf(AccessibilityWindowInfo.obtain()))\n } else {\n shadowOf(forwarder).setWindows(a11yWindows)\n }\n }\n\n @Test\n fun onInterrupt_doesNotCrash() {\n // Arrange.\n createForwarder(logAccessibilityTree = false)\n fakeA11yService.sendEventChecker = { _: EventRequest ->\n assertFalse(true) // This should not be called.\n \"\" // This should be unreachable\n }\n\n // Act.\n forwarder.onInterrupt()\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_nullEventShouldBeIgnored() {\n // Arrange.\n createForwarder(logAccessibilityTree = false)\n fakeA11yService.sendEventChecker = { _: EventRequest ->\n assertFalse(true) // This should not be called.\n \"\" // This should be unreachable\n }\n\n // Act.\n forwarder.onAccessibilityEvent(null)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_knownEventWithNoInformationShouldNotBeEmitted() {\n // Arrange.\n createForwarder(logAccessibilityTree = false)\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"\")\n var event = AccessibilityEvent()\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { _: EventRequest ->\n assertFalse(true) // This should not be called.\n \"\" // This should be unreachable\n }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_typeViewClicked_sendEventViaGrpc() {\n // Arrange.\n createForwarder(logAccessibilityTree = false, grpcPort = 1234)\n forwarder = AccessibilityForwarder({ _, _ -> channel })\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"My Content Description\")\n nodeInfo.setText(\"My Source Text\")\n nodeInfo.setClassName(\"AwesomeClass\")\n var event = AccessibilityEvent()\n event.setEventTime(1357924680)\n event.setEventType(AccessibilityEvent.TYPE_VIEW_CLICKED)\n event.getText().add(\"Some text!\")\n event.setPackageName(\"some.loooong.package.name\")\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { request: EventRequest ->\n // Check that all fields are consistent with how they were set above.\n assertThat(request.eventMap.get(\"event_type\")).isEqualTo(\"TYPE_VIEW_CLICKED\")\n assertThat(request.eventMap.get(\"event_package_name\")).isEqualTo(\"some.loooong.package.name\")\n assertThat(request.eventMap.get(\"source_content_description\"))\n .isEqualTo(\"My Content Description\")\n assertThat(request.eventMap.get(\"source_text\")).isEqualTo(\"My Source Text\")\n assertThat(request.eventMap.get(\"source_class_name\")).isEqualTo(\"AwesomeClass\")\n assertThat(request.eventMap.get(\"event_text\")).isEqualTo(\"Some text!\")\n assertThat(request.eventMap.get(\"event_timestamp_ms\")).isEqualTo(\"1357924680\")\n // No error message\n \"\"\n }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_typeViewTextChanged_ensureAllFieldsForwarded() {\n // Arrange.\n createForwarder(logAccessibilityTree = false, grpcPort = 1234)\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"My Content Description\")\n nodeInfo.setText(\"My Source Text\")\n nodeInfo.setClassName(\"AwesomeClass\")\n var event = AccessibilityEvent()\n event.setEventTime(1357924680)\n event.setEventType(AccessibilityEvent.TYPE_VIEW_TEXT_CHANGED)\n event.getText().add(\"Some text!\")\n event.fromIndex = 7\n event.beforeText = \"Old words\"\n event.addedCount = 12\n event.removedCount = 9\n event.setPackageName(\"some.loooong.package.name\")\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { request: EventRequest ->\n // Check that all fields are consistent with how they were set above.\n assertThat(request.eventMap.get(\"event_type\")).isEqualTo(\"TYPE_VIEW_TEXT_CHANGED\")\n assertThat(request.eventMap.get(\"event_package_name\")).isEqualTo(\"some.loooong.package.name\")\n assertThat(request.eventMap.get(\"source_content_description\"))\n .isEqualTo(\"My Content Description\")\n assertThat(request.eventMap.get(\"source_text\")).isEqualTo(\"My Source Text\")\n assertThat(request.eventMap.get(\"source_class_name\")).isEqualTo(\"AwesomeClass\")\n assertThat(request.eventMap.get(\"event_text\")).isEqualTo(\"Some text!\")\n assertThat(request.eventMap.get(\"event_timestamp_ms\")).isEqualTo(\"1357924680\")\n assertThat(request.eventMap.get(\"from_index\")).isEqualTo(\"7\")\n assertThat(request.eventMap.get(\"before_text\")).isEqualTo(\"Old words\")\n assertThat(request.eventMap.get(\"added_count\")).isEqualTo(\"12\")\n assertThat(request.eventMap.get(\"removed_count\")).isEqualTo(\"9\")\n assertFalse(request.eventMap.containsKey(\"to_index\"))\n assertFalse(request.eventMap.containsKey(\"view_id\"))\n assertFalse(request.eventMap.containsKey(\"action\"))\n assertFalse(request.eventMap.containsKey(\"movement_granularity\"))\n assertFalse(request.eventMap.containsKey(\"scroll_delta_x\"))\n assertFalse(request.eventMap.containsKey(\"scroll_delta_y\"))\n // No error message\n \"\"\n }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_typeViewScrolled_ensureAllFieldsForwarded() {\n // Arrange.\n createForwarder(logAccessibilityTree = false, grpcPort = 1234)\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"My Content Description\")\n nodeInfo.setText(\"My Source Text\")\n nodeInfo.setClassName(\"AwesomeClass\")\n var event = AccessibilityEvent()\n event.setEventTime(1357924680)\n event.setEventType(AccessibilityEvent.TYPE_VIEW_SCROLLED)\n event.getText().add(\"Some text!\")\n event.scrollDeltaX = 13\n event.scrollDeltaY = 27\n event.setPackageName(\"some.loooong.package.name\")\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { request: EventRequest ->\n // Check that all fields are consistent with how they were set above.\n assertThat(request.eventMap.get(\"event_type\")).isEqualTo(\"TYPE_VIEW_SCROLLED\")\n assertThat(request.eventMap.get(\"event_package_name\")).isEqualTo(\"some.loooong.package.name\")\n assertThat(request.eventMap.get(\"source_content_description\"))\n .isEqualTo(\"My Content Description\")\n assertThat(request.eventMap.get(\"source_text\")).isEqualTo(\"My Source Text\")\n assertThat(request.eventMap.get(\"source_class_name\")).isEqualTo(\"AwesomeClass\")\n assertThat(request.eventMap.get(\"event_text\")).isEqualTo(\"Some text!\")\n assertThat(request.eventMap.get(\"event_timestamp_ms\")).isEqualTo(\"1357924680\")\n assertThat(request.eventMap.get(\"scroll_delta_x\")).isEqualTo(\"13\")\n assertThat(request.eventMap.get(\"scroll_delta_y\")).isEqualTo(\"27\")\n assertFalse(request.eventMap.containsKey(\"from_index\"))\n assertFalse(request.eventMap.containsKey(\"to_index\"))\n assertFalse(request.eventMap.containsKey(\"before_text\"))\n assertFalse(request.eventMap.containsKey(\"added_count\"))\n assertFalse(request.eventMap.containsKey(\"removed_count\"))\n // No error message\n \"\"\n }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_typeViewTextTraversedAtMovementGranularity_ensureAllFieldsForwarded() {\n // Arrange.\n createForwarder(logAccessibilityTree = false, grpcPort = 1234)\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"My Content Description\")\n nodeInfo.setText(\"My Source Text\")\n nodeInfo.setClassName(\"AwesomeClass\")\n nodeInfo.viewIdResourceName = \"this.big.old.view.id\"\n var event = AccessibilityEvent()\n event.setEventTime(1357924680)\n event.setEventType(AccessibilityEvent.TYPE_VIEW_TEXT_TRAVERSED_AT_MOVEMENT_GRANULARITY)\n event.getText().add(\"Some text!\")\n event.setPackageName(\"some.loooong.package.name\")\n event.movementGranularity = 5\n event.fromIndex = 6\n event.toIndex = 8\n event.action = 23\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { request: EventRequest ->\n // Check that all fields are consistent with how they were set above.\n assertThat(request.eventMap.get(\"event_type\"))\n .isEqualTo(\"TYPE_VIEW_TEXT_TRAVERSED_AT_MOVEMENT_GRANULARITY\")\n assertThat(request.eventMap.get(\"event_package_name\")).isEqualTo(\"some.loooong.package.name\")\n assertThat(request.eventMap.get(\"source_content_description\"))\n .isEqualTo(\"My Content Description\")\n assertThat(request.eventMap.get(\"source_text\")).isEqualTo(\"My Source Text\")\n assertThat(request.eventMap.get(\"source_class_name\")).isEqualTo(\"AwesomeClass\")\n assertThat(request.eventMap.get(\"event_text\")).isEqualTo(\"Some text!\")\n assertThat(request.eventMap.get(\"event_timestamp_ms\")).isEqualTo(\"1357924680\")\n assertThat(request.eventMap.get(\"movement_granularity\")).isEqualTo(\"5\")\n assertThat(request.eventMap.get(\"from_index\")).isEqualTo(\"6\")\n assertThat(request.eventMap.get(\"to_index\")).isEqualTo(\"8\")\n assertThat(request.eventMap.get(\"view_id\")).isEqualTo(\"this.big.old.view.id\")\n assertThat(request.eventMap.get(\"action\")).isEqualTo(\"23\")\n // No error message\n \"\"\n }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_sendingevent_grpcTimeout() {\n // Arrange.\n createForwarder(\n logAccessibilityTree = false,\n a11yTreePeriodMs = 0,\n grpcHost = \"amazing.host\",\n grpcPort = 4321,\n )\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"My Content Description\")\n nodeInfo.setText(\"My Source Text\")\n nodeInfo.setClassName(\"AwesomeClass\")\n var event = AccessibilityEvent()\n event.setEventTime(1357924680)\n event.setEventType(AccessibilityEvent.TYPE_VIEW_CLICKED)\n event.getText().add(\"Some text!\")\n event.setPackageName(\"some.loooong.package.name\")\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { _ ->\n // Delay the request to prompt a timeout\n Thread.sleep(1500L)\n \"\" // Return no error.\n }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Run a second request to ensure that the channel gets rebuilt.\n fakeA11yService.sendEventChecker = { _ -> \"\" }\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun onAccessibilityEvent_sendingevent_grpcStatusException() {\n // Arrange.\n createForwarder(logAccessibilityTree = false, grpcHost = \"amazing.host\", grpcPort = 4321)\n var nodeInfo = AccessibilityNodeInfo()\n nodeInfo.setContentDescription(\"My Content Description\")\n nodeInfo.setText(\"My Source Text\")\n nodeInfo.setClassName(\"AwesomeClass\")\n var event = AccessibilityEvent()\n event.setEventTime(1357924680)\n event.setEventType(AccessibilityEvent.TYPE_VIEW_CLICKED)\n event.getText().add(\"Some text!\")\n event.setPackageName(\"some.loooong.package.name\")\n shadowOf(event).setSourceNode(nodeInfo)\n fakeA11yService.sendEventChecker = { _ -> throw StatusException(Status.UNAVAILABLE) }\n\n // Act.\n forwarder.onAccessibilityEvent(event)\n\n // Run a second request to ensure that the channel gets rebuilt.\n fakeA11yService.sendEventChecker = { _ -> \"\" }\n forwarder.onAccessibilityEvent(event)\n\n // Assert.\n // See `sendEventChecker` above.\n }\n\n @Test\n fun logAccessibilityTreeFalse_doesNotLogAccessibilityTree() {\n // Arrange.\n createForwarder(logAccessibilityTree = false, a11yTreePeriodMs = 10, grpcPort = 13579)\n fakeA11yService.sendForestChecker = { _: AndroidAccessibilityForest ->\n assertFalse(true) // This should not be called.\n \"\" // This should be unreachable\n }\n\n // Act.\n Thread.sleep(1000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Assert.\n // See `sendForestChecker` above.\n }\n\n @Test\n fun grpcPortZero_doesNotSendTree() {\n // Arrange.\n createForwarder(logAccessibilityTree = true, a11yTreePeriodMs = 10, grpcPort = 0)\n fakeA11yService.sendForestChecker = { _: AndroidAccessibilityForest ->\n assertFalse(true) // This should not be called.\n \"\" // This should be unreachable\n }\n\n // Act.\n Thread.sleep(1000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Assert.\n // See `sendForestChecker` above.\n }\n\n @Test\n fun grpcPortPositive_shouldSendTreeViaGrpc() {\n // Arrange.\n val window = AccessibilityWindowInfo()\n shadowOf(window).setType(AccessibilityWindowInfo.TYPE_SYSTEM)\n createForwarder(\n logAccessibilityTree = true,\n a11yTreePeriodMs = 10,\n grpcPort = 1234,\n a11yWindows = mutableListOf(window),\n )\n fakeA11yService.sendForestChecker = { request: AndroidAccessibilityForest ->\n // Check that we get only a single window.\n assertThat(request.windowsList.size).isEqualTo(1)\n // And that its type is what we set above.\n assertThat(request.windowsList[0].windowType)\n .isEqualTo(AndroidAccessibilityWindowInfo.WindowType.TYPE_SYSTEM)\n // The error message\n \"Something went wrong!\"\n }\n\n // Act.\n Thread.sleep(1000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Assert.\n // See `sendForestChecker` above.\n }\n\n @Test\n fun grpcPortPositiveAndHost_shouldSendTreeViaGrpc() {\n // Arrange.\n fakeA11yService.sendForestChecker = { request: AndroidAccessibilityForest ->\n // Check that we get only a single window.\n assertThat(request.windowsList.size).isEqualTo(1)\n // And that its type is what we set above.\n assertThat(request.windowsList[0].windowType)\n .isEqualTo(AndroidAccessibilityWindowInfo.WindowType.TYPE_ACCESSIBILITY_OVERLAY)\n \"\" // Return no error.\n }\n val window = AccessibilityWindowInfo()\n shadowOf(window).setType(AccessibilityWindowInfo.TYPE_ACCESSIBILITY_OVERLAY)\n createForwarder(\n logAccessibilityTree = true,\n a11yTreePeriodMs = 500,\n grpcHost = \"amazing.host\",\n grpcPort = 4321,\n a11yWindows = mutableListOf(window),\n )\n\n // Act.\n Thread.sleep(1000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Assert.\n // See `sendForestChecker` above.\n }\n\n @Test\n fun sendingForest_grpcTimeout() {\n // Arrange.\n fakeA11yService.sendForestChecker = { _ ->\n // Delay the request to prompt a timeout\n Thread.sleep(1500L)\n \"\" // Return no error.\n }\n val window = AccessibilityWindowInfo()\n shadowOf(window).setType(AccessibilityWindowInfo.TYPE_ACCESSIBILITY_OVERLAY)\n createForwarder(\n logAccessibilityTree = true,\n a11yTreePeriodMs = 10,\n grpcHost = \"amazing.host\",\n grpcPort = 4321,\n a11yWindows = mutableListOf(window),\n )\n\n // Act.\n Thread.sleep(2000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Run a second request to ensure that the channel gets rebuilt.\n fakeA11yService.sendForestChecker = { _ -> \"\" }\n Thread.sleep(2000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Assert.\n // See `sendForestChecker` above.\n }\n\n @Test\n fun sendingForest_grpcStatusException() {\n // Arrange.\n val window = AccessibilityWindowInfo()\n shadowOf(window).setType(AccessibilityWindowInfo.TYPE_ACCESSIBILITY_OVERLAY)\n createForwarder(\n logAccessibilityTree = true,\n a11yTreePeriodMs = 10,\n grpcHost = \"amazing.host\",\n grpcPort = 4321,\n a11yWindows = mutableListOf(window),\n )\n fakeA11yService.sendForestChecker = { _ -> throw StatusException(Status.UNAVAILABLE) }\n\n // Act.\n Thread.sleep(1000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Run a second request to ensure that the channel gets rebuilt.\n fakeA11yService.sendForestChecker = { _ -> \"\" }\n Thread.sleep(1000) // Sleep a bit to give time to trigger the tree logging function.\n\n // Assert.\n // See `sendForestChecker` above.\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/AccessibilityTreeCreator.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.graphics.Rect\nimport android.util.Log\nimport android.view.accessibility.AccessibilityNodeInfo\nimport android.view.accessibility.AccessibilityWindowInfo\nimport com.google.androidenv.accessibilityforwarder.AndroidAccessibilityWindowInfo.WindowType\nimport java.util.concurrent.ConcurrentHashMap\nimport java.util.stream.Collectors\nimport kotlin.collections.mutableListOf\nimport kotlinx.coroutines.Deferred\nimport kotlinx.coroutines.async\nimport kotlinx.coroutines.awaitAll\nimport kotlinx.coroutines.runBlocking\n\n/** Helper methods for creating the android accessibility info extra. */\nclass AccessibilityTreeCreator() {\n\n /** Creates an accessibility forest proto. */\n fun buildForest(windowInfos: List): AndroidAccessibilityForest {\n val sourcesMap: ConcurrentHashMap =\n ConcurrentHashMap()\n val windows: List =\n processWindowsAndBlock(windowInfos, sourcesMap)\n return androidAccessibilityForest { this.windows += windows }\n }\n\n private fun processWindowsAndBlock(\n windowInfos: List,\n sourcesMap: ConcurrentHashMap,\n ): List {\n val windows: List\n runBlocking { windows = processWindows(windowInfos, sourcesMap) }\n return windows\n }\n\n private suspend fun processWindows(\n windowInfos: List,\n sourcesMap: ConcurrentHashMap,\n ): List {\n var windowInfoProtos = mutableListOf()\n for (i in windowInfos.size - 1 downTo 0) {\n val windowInfoProto = processWindow(windowInfos.get(i), sourcesMap)\n windowInfoProto?.let { windowInfoProtos.add(windowInfoProto) }\n }\n return windowInfoProtos.toList()\n }\n\n private suspend fun processWindow(\n windowInfo: AccessibilityWindowInfo,\n sources: ConcurrentHashMap,\n ): AndroidAccessibilityWindowInfo? {\n val bounds = Rect()\n windowInfo.getBoundsInScreen(bounds)\n val root: AccessibilityNodeInfo? = windowInfo.root\n if (root == null) {\n Log.i(TAG, \"window root is null\")\n return androidAccessibilityWindowInfo {\n this.tree = androidAccessibilityTree {}\n this.isActive = windowInfo.isActive\n this.id = windowInfo.id\n this.layer = windowInfo.layer\n this.isAccessibilityFocused = windowInfo.isAccessibilityFocused\n this.isFocused = windowInfo.isFocused\n this.boundsInScreen = convertToRectProto(bounds)\n this.windowType = toWindowType(windowInfo.type)\n }\n }\n val treeDeferred: Deferred\n runBlocking { treeDeferred = async { processNodesInWindow(root, sources) } }\n return androidAccessibilityWindowInfo {\n this.tree = treeDeferred.await()\n this.isActive = windowInfo.isActive\n this.id = windowInfo.id\n this.layer = windowInfo.layer\n this.isAccessibilityFocused = windowInfo.isAccessibilityFocused\n this.isFocused = windowInfo.isFocused\n this.boundsInScreen = convertToRectProto(bounds)\n this.windowType = toWindowType(windowInfo.type)\n }\n }\n\n private suspend fun processNodesInWindow(\n root: AccessibilityNodeInfo,\n sources: ConcurrentHashMap,\n ): AndroidAccessibilityTree {\n Log.d(TAG, \"processNodesInWindow()\")\n val traversalQueue = ArrayDeque()\n traversalQueue.add(ParentChildNodePair.builder().child(root).build())\n val uniqueIdsCache: UniqueIdsGenerator = UniqueIdsGenerator()\n var currentDepth = 0\n val nodesDeferred = mutableListOf>()\n val seenNodes: HashSet = HashSet()\n seenNodes.add(root)\n runBlocking {\n while (!traversalQueue.isEmpty()) {\n // Traverse the tree layer-by-layer.\n // The first layer has only the root and depth 0.\n // The second layer has all the root's children and depth 1.\n for (nodesAtCurrentDepth in traversalQueue.size downTo 1) {\n val nodePair: ParentChildNodePair = traversalQueue.removeFirst()\n for (i in 0 until nodePair.child().childCount) {\n val childNode: AccessibilityNodeInfo? = nodePair.child().getChild(i)\n if (childNode != null && !seenNodes.contains(childNode)) {\n traversalQueue.add(\n ParentChildNodePair.builder().child(childNode).parent(nodePair.child()).build()\n )\n seenNodes.add(childNode)\n }\n }\n val thisDepth = currentDepth\n var deferred = async { processNode(nodePair, sources, uniqueIdsCache, thisDepth) }\n nodesDeferred.add(deferred)\n }\n currentDepth++\n }\n }\n return androidAccessibilityTree { this.nodes += nodesDeferred.awaitAll() }\n }\n\n companion object {\n private const val TAG = \"AndroidRLTask\"\n }\n}\n\nprivate fun processNode(\n nodePair: ParentChildNodePair,\n sourceBuilder: ConcurrentHashMap,\n uniqueIdsCache: UniqueIdsGenerator,\n nodeDepth: Int,\n): AndroidAccessibilityNodeInfo {\n val node: AccessibilityNodeInfo = nodePair.child()\n val immutableNode: AndroidAccessibilityNodeInfo =\n createAndroidAccessibilityNode(\n node,\n uniqueIdsCache.getUniqueId(node),\n nodeDepth,\n getChildUniqueIds(node, uniqueIdsCache),\n )\n sourceBuilder.put(immutableNode, node)\n return immutableNode\n}\n\nprivate fun createAndroidAccessibilityNode(\n node: AccessibilityNodeInfo,\n nodeId: Int,\n depth: Int,\n childIds: List,\n): AndroidAccessibilityNodeInfo {\n val bounds = Rect()\n node.getBoundsInScreen(bounds)\n val actions = node.getActionList().stream().map(::createAction).collect(Collectors.toList())\n return androidAccessibilityNodeInfo {\n this.actions += actions\n this.boundsInScreen = convertToRectProto(bounds)\n this.isCheckable = node.isCheckable\n this.isChecked = node.isChecked\n this.className = stringFromNullableCharSequence(node.getClassName())\n this.isClickable = node.isClickable\n this.contentDescription = stringFromNullableCharSequence(node.getContentDescription())\n this.isEditable = node.isEditable\n this.isEnabled = node.isEnabled\n this.isFocusable = node.isFocusable\n this.hintText = stringFromNullableCharSequence(node.getHintText())\n this.isLongClickable = node.isLongClickable\n this.packageName = stringFromNullableCharSequence(node.getPackageName())\n this.isPassword = node.isPassword\n this.isScrollable = node.isScrollable\n this.isSelected = node.isSelected\n this.text = stringFromNullableCharSequence(node.getText())\n this.textSelectionEnd = node.getTextSelectionEnd().toLong()\n this.textSelectionStart = node.getTextSelectionStart().toLong()\n this.viewIdResourceName = node.getViewIdResourceName() ?: \"\"\n this.isVisibleToUser = node.isVisibleToUser\n this.windowId = node.windowId\n this.uniqueId = nodeId\n this.childIds += childIds\n this.drawingOrder = node.drawingOrder\n this.tooltipText = stringFromNullableCharSequence(node.getTooltipText())\n this.depth = depth\n }\n}\n\nprivate fun createAction(\n action: AccessibilityNodeInfo.AccessibilityAction\n): AndroidAccessibilityAction =\n AndroidAccessibilityAction.newBuilder()\n .setId(action.id)\n .setLabel(stringFromNullableCharSequence(action.label))\n .build()\n\nprivate fun getChildUniqueIds(\n node: AccessibilityNodeInfo,\n uniqueIdsCache: UniqueIdsGenerator,\n): List {\n val ids = mutableListOf()\n for (childId in 0 until node.getChildCount()) {\n val child: AccessibilityNodeInfo = node.getChild(childId) ?: continue\n ids.add(uniqueIdsCache.getUniqueId(child))\n }\n return ids.toList()\n}\n\nfun stringFromNullableCharSequence(cs: CharSequence?): String = cs?.toString() ?: \"\"\n\nfun convertToRectProto(rect: Rect) = protoRect {\n left = rect.left\n top = rect.top\n right = rect.right\n bottom = rect.bottom\n}\n\nprivate fun toWindowType(type: Int): WindowType =\n when (type) {\n AccessibilityWindowInfo.TYPE_ACCESSIBILITY_OVERLAY -> WindowType.TYPE_ACCESSIBILITY_OVERLAY\n AccessibilityWindowInfo.TYPE_APPLICATION -> WindowType.TYPE_APPLICATION\n AccessibilityWindowInfo.TYPE_INPUT_METHOD -> WindowType.TYPE_INPUT_METHOD\n AccessibilityWindowInfo.TYPE_SYSTEM -> WindowType.TYPE_SYSTEM\n AccessibilityWindowInfo.TYPE_SPLIT_SCREEN_DIVIDER -> WindowType.TYPE_SPLIT_SCREEN_DIVIDER\n else -> WindowType.UNKNOWN_TYPE\n }\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/AccessibilityTreeCreatorTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.view.accessibility.AccessibilityNodeInfo\nimport android.view.accessibility.AccessibilityWindowInfo\nimport kotlin.test.assertEquals\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.robolectric.RobolectricTestRunner\nimport org.robolectric.Shadows.shadowOf\n\n@RunWith(RobolectricTestRunner::class)\nclass AccessibilityTreeCreatorTest {\n\n @Test\n fun buildForest_buildsAccessibilityForestCorrectly() {\n val creator = AccessibilityTreeCreator()\n\n val forest = creator.buildForest(mutableListOf(createWindowInfo()))\n\n assertEquals(forest.windowsCount, 1)\n assertEquals(forest.getWindows(0).tree.nodesCount, 3)\n var rootNode: AndroidAccessibilityNodeInfo? = null\n var checkableNode: AndroidAccessibilityNodeInfo? = null\n val nodes = forest.getWindows(0).tree.nodesList\n for (i in nodes.size - 1 downTo 0) {\n if (nodes[i].text == \"root node\") {\n rootNode = nodes[i]\n }\n if (nodes[i].isCheckable == true) {\n checkableNode = nodes[i]\n }\n }\n assertEquals(rootNode?.childIdsCount, 2)\n assertEquals(checkableNode?.text, \"Check box\")\n }\n\n @Test\n fun buildForest_noRootInWindow_returnsEmptyTree() {\n val creator = AccessibilityTreeCreator()\n val windowInfo = AccessibilityWindowInfo.obtain()\n shadowOf(windowInfo).setType(AccessibilityWindowInfo.TYPE_ACCESSIBILITY_OVERLAY)\n\n val forest = creator.buildForest(mutableListOf(windowInfo))\n\n assertEquals(0, forest.getWindows(0).tree.nodesList.size)\n }\n\n private fun createAccessibilityNodeInfo(): AccessibilityNodeInfo {\n val root = AccessibilityNodeInfo.obtain()\n root.text = \"root node\"\n root.isClickable = true\n val accessibilityNodeInfo = AccessibilityNodeInfo.obtain()\n accessibilityNodeInfo.viewIdResourceName = \"test\"\n accessibilityNodeInfo.isClickable = true\n accessibilityNodeInfo.isEditable = true\n accessibilityNodeInfo.hintText = \"Please enter your address\"\n shadowOf(root).addChild(accessibilityNodeInfo)\n val anotherChildNode = AccessibilityNodeInfo.obtain()\n anotherChildNode.isCheckable = true\n anotherChildNode.text = \"Check box\"\n shadowOf(root).addChild(anotherChildNode)\n return root\n }\n\n private fun createWindowInfo(): AccessibilityWindowInfo {\n val windowInfo = AccessibilityWindowInfo.obtain()\n shadowOf(windowInfo).setType(AccessibilityWindowInfo.TYPE_ACCESSIBILITY_OVERLAY)\n shadowOf(windowInfo).setRoot(createAccessibilityNodeInfo())\n return windowInfo\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/AndroidManifest.xml",
"content": "\n\n\n\n\n\n\n \n\n \n \n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/AndroidManifest_lite.xml",
"content": "\n\n\n\n\n\n\n \n\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/FlagsBroadcastReceiver.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.content.BroadcastReceiver\nimport android.content.Context\nimport android.content.Intent\nimport android.util.Log\n\n/** Broadcast receiver responsible for enabling or disabling flags. */\nclass FlagsBroadcastReceiver() : BroadcastReceiver() {\n\n override fun onReceive(context: Context?, intent: Intent?) {\n val action = intent?.action\n Log.i(TAG, \"Received broadcast intent with action: \" + action)\n when (action) {\n ACTION_ENABLE_ACCESSIBILITY_TREE_LOGS -> {\n Log.i(TAG, \"Enabling Accessibility Tree logging.\")\n LogFlags.logAccessibilityTree = true\n }\n ACTION_DISABLE_ACCESSIBILITY_TREE_LOGS -> {\n Log.i(TAG, \"Disabling Accessibility Tree logging.\")\n LogFlags.logAccessibilityTree = false\n }\n ACTION_SET_GRPC -> {\n // The Android Emulator uses 10.0.2.2 as a redirect to the workstation's IP. Most often the\n // gRPC server will be running locally so it makes sense to use this as the default value.\n // See https://developer.android.com/studio/run/emulator-networking#networkaddresses.\n val host = intent.getStringExtra(\"host\") ?: \"10.0.2.2\"\n // The TCP port to connect. If <=0 gRPC is disabled.\n val port = intent.getIntExtra(\"port\", 0)\n Log.i(TAG, \"Setting gRPC endpoint to ${host}:${port}.\")\n LogFlags.grpcHost = host\n LogFlags.grpcPort = port\n }\n else -> Log.w(TAG, \"Unknown action: ${action}\")\n }\n }\n\n companion object {\n private const val TAG = \"FlagsBroadcastReceiver\"\n private const val ACTION_ENABLE_ACCESSIBILITY_TREE_LOGS =\n \"accessibility_forwarder.intent.action.ENABLE_ACCESSIBILITY_TREE_LOGS\"\n private const val ACTION_DISABLE_ACCESSIBILITY_TREE_LOGS =\n \"accessibility_forwarder.intent.action.DISABLE_ACCESSIBILITY_TREE_LOGS\"\n private const val ACTION_SET_GRPC = \"accessibility_forwarder.intent.action.SET_GRPC\"\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/FlagsBroadcastReceiverTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.content.Intent\nimport com.google.common.truth.Truth.assertThat\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.robolectric.RobolectricTestRunner\n\n@RunWith(RobolectricTestRunner::class)\nclass FlagsBroadcastReceiverTest {\n\n @Test\n fun onReceive_nullIntent_shouldNotLogAnything() {\n // Arrange.\n LogFlags.logAccessibilityTree = false\n val receiver = FlagsBroadcastReceiver()\n\n // Act.\n receiver.onReceive(context = null, intent = null)\n\n // Assert.\n assertThat(LogFlags.logAccessibilityTree).isFalse()\n }\n\n @Test\n fun onReceive_nullIntent_actionShouldNotLogAnything() {\n // Arrange.\n LogFlags.logAccessibilityTree = false\n val receiver = FlagsBroadcastReceiver()\n val intent = Intent()\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.logAccessibilityTree).isFalse()\n }\n\n @Test\n fun onReceive_unknownIntent_actionShouldIssueWarning() {\n // Arrange.\n LogFlags.logAccessibilityTree = false\n val receiver = FlagsBroadcastReceiver()\n val intent = Intent(\"SOME_WEIRD_ACTION\")\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.logAccessibilityTree).isFalse()\n }\n\n @Test\n fun onReceive_intentWithDisableAction_shouldDisableTreeLogging() {\n // Arrange.\n LogFlags.logAccessibilityTree = true\n val receiver = FlagsBroadcastReceiver()\n val intent = Intent(\"accessibility_forwarder.intent.action.DISABLE_ACCESSIBILITY_TREE_LOGS\")\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.logAccessibilityTree).isFalse()\n }\n\n @Test\n fun onReceive_intentWithEnableAction_shouldEnableTreeLogging() {\n // Arrange.\n LogFlags.logAccessibilityTree = false\n val receiver = FlagsBroadcastReceiver()\n val intent = Intent(\"accessibility_forwarder.intent.action.ENABLE_ACCESSIBILITY_TREE_LOGS\")\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.logAccessibilityTree).isTrue()\n }\n\n @Test\n fun onReceive_intentWithSetGrpcActionNoArgs_shouldDefaultToEmuIpAndPortZero() {\n // Arrange.\n LogFlags.grpcHost = \"some_host\"\n LogFlags.grpcPort = 9999\n val receiver = FlagsBroadcastReceiver()\n val intent = Intent(\"accessibility_forwarder.intent.action.SET_GRPC\")\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.grpcHost).isEqualTo(\"10.0.2.2\")\n assertThat(LogFlags.grpcPort).isEqualTo(0)\n }\n\n @Test\n fun onReceive_intentWithSetGrpcActionWithHostNoPort_shouldDefaultPortToZero() {\n // Arrange.\n LogFlags.grpcHost = \"some_host\"\n LogFlags.grpcPort = 9999\n val receiver = FlagsBroadcastReceiver()\n val intent =\n Intent(\"accessibility_forwarder.intent.action.SET_GRPC\").apply {\n putExtra(\"host\", \"awesome.server.ca\")\n }\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.grpcHost).isEqualTo(\"awesome.server.ca\")\n assertThat(LogFlags.grpcPort).isEqualTo(0)\n }\n\n @Test\n fun onReceive_intentWithSetGrpcActionWithPortNoHost_shouldDefaultHostToEmuIp() {\n // Arrange.\n LogFlags.grpcHost = \"some_host\"\n LogFlags.grpcPort = 9999\n val receiver = FlagsBroadcastReceiver()\n val intent =\n Intent(\"accessibility_forwarder.intent.action.SET_GRPC\").apply { putExtra(\"port\", 54321) }\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.grpcHost).isEqualTo(\"10.0.2.2\")\n assertThat(LogFlags.grpcPort).isEqualTo(54321)\n }\n\n @Test\n fun onReceive_intentWithSetGrpcActionWithHostAndPort_shouldSetBoth() {\n // Arrange.\n LogFlags.grpcHost = \"some_host\"\n LogFlags.grpcPort = 9999\n val receiver = FlagsBroadcastReceiver()\n val intent =\n Intent(\"accessibility_forwarder.intent.action.SET_GRPC\").apply {\n putExtra(\"host\", \"grpc.ca\")\n putExtra(\"port\", 54321)\n }\n\n // Act.\n receiver.onReceive(context = null, intent = intent)\n\n // Assert.\n assertThat(LogFlags.grpcHost).isEqualTo(\"grpc.ca\")\n assertThat(LogFlags.grpcPort).isEqualTo(54321)\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/LogFlags.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\n/**\n * Controls global settings in AccessibilityForwarder.\n *\n * Please note that this class is not thread safe.\n */\nobject LogFlags {\n // Whether to log the accessibility tree.\n var logAccessibilityTree: Boolean = false\n // How frequent to emit a11y trees (in milliseconds).\n var a11yTreePeriodMs: Long = 100\n\n // The gRPC server to connect to. (Only available if grpcPort>0).\n var grpcHost: String = \"\"\n // If >0 this represents the gRPC port number to connect to.\n var grpcPort: Int = 0\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/ParentChildNodePair.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport android.view.accessibility.AccessibilityNodeInfo\nimport com.google.auto.value.AutoValue\n\n/** Parent and child [AccessibilityNodeInfo] relationship. */\n@AutoValue\ninternal abstract class ParentChildNodePair {\n abstract fun parent(): AccessibilityNodeInfo?\n\n abstract fun child(): AccessibilityNodeInfo\n\n /** [ParentChildNodePair] builder. */\n @AutoValue.Builder\n abstract class Builder {\n abstract fun parent(parent: AccessibilityNodeInfo?): Builder\n\n abstract fun child(child: AccessibilityNodeInfo): Builder\n\n abstract fun build(): ParentChildNodePair\n }\n\n companion object {\n @JvmStatic fun builder(): Builder = AutoValue_ParentChildNodePair.Builder()\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/UniqueIdsGenerator.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.accessibilityforwarder\n\nimport java.util.concurrent.ConcurrentHashMap\nimport java.util.concurrent.atomic.AtomicInteger\nimport java.util.function.Function\n\n/** Thread-safe helper class for assigning a unique ID to an object. */\ninternal class UniqueIdsGenerator {\n private val nextId = AtomicInteger(0)\n private val uniqueIdsByNode = ConcurrentHashMap()\n\n fun getUniqueId(a: A): Int {\n return uniqueIdsByNode.computeIfAbsent(a, Function { _: A -> nextId.getAndIncrement() })!!\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/accessibilityforwarder/res/xml/accessibility_forwarder_service.xml",
"content": "\n\n\n\n\n\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/AndroidManifest.xml",
"content": "\n\n\n\n\n \n \n \n \n \n \n \n \n \n\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/BUILD.bazel",
"content": "# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Classic RL task implemented as an Android app.\nload(\"@rules_android//rules:rules.bzl\", \"android_binary\")\nload(\"@rules_kotlin//kotlin:android.bzl\", \"kt_android_library\")\n\npackage(\n default_visibility = [\":catch_packages\"],\n)\n\npackage_group(\n name = \"catch_packages\",\n packages = [\n \"//java/com/google/androidenv/catch/...\",\n \"//javatests/com/google/androidenv/catch/...\",\n ],\n)\n\nlicenses([\"notice\"])\n\nandroid_binary(\n name = \"app\",\n manifest = \"AndroidManifest.xml\",\n multidex = \"native\",\n deps = [\":MainActivity\"],\n)\n\nkt_android_library(\n name = \"GameLogic\",\n srcs = [\"GameLogic.kt\"],\n deps = [\n \"//java/com/google/androidenv/catch/sprite:Background\",\n \"//java/com/google/androidenv/catch/sprite:Ball\",\n \"//java/com/google/androidenv/catch/sprite:LineSegment\",\n \"//java/com/google/androidenv/catch/sprite:Paddle\",\n ],\n)\n\nkt_android_library(\n name = \"GameLogicThread\",\n srcs = [\"GameLogicThread.kt\"],\n deps = [\n \":GameLogic\",\n ],\n)\n\nkt_android_library(\n name = \"MainActivity\",\n srcs = [\"MainActivity.kt\"],\n manifest = \"AndroidManifest.xml\",\n resource_files = glob([\"res/**\"]),\n deps = [\n \":GameLogic\",\n \":GameLogicThread\",\n \":RenderThread\",\n \"//java/com/google/androidenv/catch/sprite:Background\",\n \"//java/com/google/androidenv/catch/sprite:Ball\",\n \"//java/com/google/androidenv/catch/sprite:Paddle\",\n ],\n)\n\nkt_android_library(\n name = \"RenderThread\",\n srcs = [\"RenderThread.kt\"],\n deps = [\n \":GameLogic\",\n ],\n)\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/GameLogic.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.graphics.Canvas\nimport android.view.MotionEvent\nimport com.google.androidenv.catch.sprite.Background\nimport com.google.androidenv.catch.sprite.Ball\nimport com.google.androidenv.catch.sprite.LineSegment\nimport com.google.androidenv.catch.sprite.Paddle\nimport java.time.Duration\nimport java.time.Instant\nimport kotlin.random.Random\n\n/** The class that contains the game logic. */\nopen class GameLogic(\n // Expected number of frames per second.\n fps: Int = 60,\n // Pseudo random number generator.\n private val rand: Random = Random.Default,\n // Width and height of the game in pixels.\n private val width: Int,\n private val height: Int,\n // UI objects in the game.\n private var background: Background = Background(),\n private var ball: Ball = Ball(maxX = width, maxY = height, rand = rand),\n private var paddle: Paddle = Paddle(maxX = width, y = height),\n) {\n\n private val sleepTime: Duration = Duration.ofMillis((1000.0 / fps).toLong())\n\n /** Reinitializes the state of the game. */\n // Need to make this open to allow for testing.\n open fun reset() {\n this.ball.reset()\n }\n\n /** Runs one \"throw\" of a [ball] that needs to be caught by the [paddle]. */\n // Need to make this open to allow for testing.\n open fun run(): Boolean {\n var lastTimestamp = Instant.now()\n do {\n Thread.sleep(sleepTime.toMillis())\n val now = Instant.now()\n val interval = Duration.between(lastTimestamp, now)\n lastTimestamp = now\n ball.update(interval)\n } while (!ball.isOutOfBounds())\n\n return ball.intersects(LineSegment(paddle.topLeft(), paddle.topRight()))\n }\n\n /** Processes a user event (e.g. a touchscreen event) and updates the [paddle] accordingly. */\n fun handleTouch(event: MotionEvent) {\n paddle.x = event.x.toInt()\n }\n\n /** Renders the game on [c]. */\n open fun render(c: Canvas) {\n background.draw(c)\n ball.draw(c)\n paddle.draw(c)\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/GameLogicThread.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.util.Log\n\n/** A thread that continuously runs the game logic, resetting after each internal [run()]. */\nclass GameLogicThread(private val game: GameLogic, private val loggingTag: String) : Thread() {\n\n /** Whether this thread should continuously run. */\n private var shouldRun: Boolean = true\n /** A counter of game runs. */\n private var counter: Int = 0\n\n /**\n * Lets the current [run()] iteration complete then break exit this [Thread].\n *\n * Notice that [shouldRun] cannot have a private getter with a public setter (please see\n * https://youtrack.jetbrains.com/issue/KT-3110 for details), hence this public function. Also\n * notice that we cannot call this function [stop()] since it would shadow [Thread.stop()].\n */\n public fun finish() {\n shouldRun = false\n }\n\n /** Continuously runs the [game] until [finish()] is called. */\n public override fun run() {\n while (shouldRun) {\n game.reset()\n Log.i(loggingTag, \"${counter++} - ${game.run()}\")\n }\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/MainActivity.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.app.Activity\nimport android.content.Intent\nimport android.graphics.Color\nimport android.os.Bundle\nimport android.util.Log\nimport android.view.SurfaceHolder\nimport android.view.SurfaceView\nimport android.view.View\nimport android.view.Window\nimport com.google.androidenv.catch.sprite.Background\nimport com.google.androidenv.catch.sprite.Ball\nimport com.google.androidenv.catch.sprite.Paddle\n\n/** The activity that allows users to play the RL game of Catch. */\nclass MainActivity() : Activity(), SurfaceHolder.Callback {\n\n private var surfaceView: SurfaceView? = null\n private var renderThread: RenderThread? = null\n private var gameLogicThread: GameLogicThread? = null\n\n private val fps: Int = 60\n private var gameCounter: Int = 0\n private var width: Int = -1\n private var height: Int = -1\n\n private var extras: Bundle? = null\n\n // [Activity] overrides.\n\n /** Initializes the Android [View] and sets up callbacks. */\n override fun onCreate(savedInstanceState: Bundle?) {\n super.onCreate(savedInstanceState)\n Log.i(TAG, \"MainActivity::onCreate()\")\n requestWindowFeature(Window.FEATURE_NO_TITLE)\n setContentView(R.layout.main)\n val surface: SurfaceView? = findViewById(R.id.surfaceView)\n if (surface == null) throw Exception(\"Could not create SurfaceView. Aborting...\")\n\n surface.visibility = View.VISIBLE\n surface.holder.addCallback(this)\n surfaceView = surface\n extras = intent?.extras\n }\n\n override fun onNewIntent(intent: Intent?) {\n super.onNewIntent(intent)\n Log.i(TAG, \"MainActivity::onNewIntent()\")\n extras = intent?.extras\n startGame()\n }\n\n // [SurfaceHolder.Callback] overrides.\n\n override fun surfaceCreated(holder: SurfaceHolder) {\n Log.i(TAG, \"MainActivity::surfaceCreated()\")\n renderThread = RenderThread(surfaceHolder = holder, fps = fps).also { it.start() }\n }\n\n override fun surfaceChanged(holder: SurfaceHolder, format: Int, width: Int, height: Int) {\n Log.i(TAG, \"MainActivity::surfaceChanged()\")\n this.width = width\n this.height = height\n startGame()\n }\n\n override fun surfaceDestroyed(holder: SurfaceHolder) {\n Log.i(TAG, \"MainActivity::surfaceDestroyed()\")\n renderThread?.finish()\n renderThread?.join()\n gameLogicThread?.finish()\n gameLogicThread?.join()\n }\n\n private fun startGame() {\n Log.i(TAG, \"MainActivity::startGame()\")\n if (width <= 0 || height <= 0) {\n Log.e(TAG, \"MainActivity::startGame() - Width or height not initialized yet.\")\n return\n }\n val backgroundColor = Color.parseColor(extras?.getString(\"backgroundColor\") ?: \"BLACK\")\n val ballColor = Color.parseColor(extras?.getString(\"ballColor\") ?: \"WHITE\")\n val ballRadius = extras?.getFloat(\"ballRadius\", 10.0f) ?: 10.0f\n val ballSpeed = extras?.getFloat(\"ballSpeed\", 0.2f) ?: 0.2f\n val paddleColor = Color.parseColor(extras?.getString(\"paddleColor\") ?: \"WHITE\")\n val paddleWidth = extras?.getInt(\"paddleWidth\", 80) ?: 80\n val paddleHeight = extras?.getInt(\"paddleHeight\", 10) ?: 10\n Log.i(TAG, \"MainActivity::startGame() - extras bundle: $extras\")\n val game =\n GameLogic(\n width = width,\n height = height,\n fps = fps,\n background = Background(color = backgroundColor),\n ball =\n Ball(\n maxX = width,\n maxY = height,\n color = ballColor,\n radius = ballRadius,\n speed = ballSpeed,\n ),\n paddle =\n Paddle(\n color = paddleColor,\n width = paddleWidth,\n height = paddleHeight,\n maxX = width,\n y = (height - paddleHeight / 2),\n ),\n )\n\n // Stop the previous game logic thread if it's running.\n gameLogicThread?.finish()\n gameLogicThread?.join()\n\n // Create and start the new GameLogicThread, passing the game instance.\n gameLogicThread = GameLogicThread(game, TAG).also { it.start() }\n\n // Pass the same game instance to the render thread.\n renderThread?.game = game\n\n surfaceView?.setOnTouchListener(\n // Suppress warning for ClickableViewAccessibility since click handling\n // is not within an OnTouchListener.\n @SuppressWarnings(\"ClickableViewAccessibility\")\n View.OnTouchListener { _, motionEvent ->\n game.handleTouch(motionEvent)\n true\n }\n )\n }\n\n companion object {\n private const val TAG = \"AndroidRLTask\"\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/RenderThread.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.graphics.Canvas\nimport android.view.SurfaceHolder\nimport java.time.Duration\n\n/** A thread that continuously renders the game logic onto a surface. */\nclass RenderThread(private val surfaceHolder: SurfaceHolder, private val fps: Int = 60) : Thread() {\n\n /** Whether this thread should continuously run. */\n private var shouldRun: Boolean = true\n /** How long to sleep at each [run()] iteration. */\n private val sleepTime: Duration = Duration.ofMillis((1000.0 / fps).toLong())\n /** The class responsible for issuing rendering commands to the canvas. */\n var game: GameLogic? = null\n\n /**\n * Runs the current game logic [run()] to completion.\n *\n * Notice that [shouldRun] cannot have a private getter with a public setter (please see\n * https://youtrack.jetbrains.com/issue/KT-3110 for details), hence this public function. Also\n * notice that we cannot call this function [stop()] since it would shadow [Thread.stop()].\n */\n public fun finish() {\n shouldRun = false\n }\n\n /** Continuously renders the [game] onto [surfaceHolder]. */\n public override fun run() {\n while (shouldRun) {\n if (surfaceHolder.surface?.isValid() ?: false) {\n val c: Canvas = surfaceHolder.lockCanvas()\n game?.render(c)\n surfaceHolder.unlockCanvasAndPost(c)\n }\n Thread.sleep(sleepTime.toMillis())\n }\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/res/layout/main.xml",
"content": "\n\n\n\n\n \n\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/res/values/strings.xml",
"content": "\n\n\n\n\n Catch\n\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/BUILD.bazel",
"content": "# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Sprites for the app.\n\nload(\"@rules_kotlin//kotlin:android.bzl\", \"kt_android_library\")\n\npackage(\n default_visibility = [\"//java/com/google/androidenv/catch:catch_packages\"],\n)\n\nlicenses([\"notice\"])\n\nkt_android_library(\n name = \"Background\",\n srcs = [\"Background.kt\"],\n deps = [\":Sprite\"],\n)\n\nkt_android_library(\n name = \"Ball\",\n srcs = [\"Ball.kt\"],\n deps = [\n \":LineSegment\",\n \":Point\",\n \":Sprite\",\n ],\n)\n\nkt_android_library(\n name = \"LineSegment\",\n srcs = [\"LineSegment.kt\"],\n deps = [\":Point\"],\n)\n\nkt_android_library(\n name = \"Paddle\",\n srcs = [\"Paddle.kt\"],\n deps = [\n \":Point\",\n \":Sprite\",\n ],\n)\n\nkt_android_library(\n name = \"Point\",\n srcs = [\"Point.kt\"],\n)\n\nkt_android_library(\n name = \"Sprite\",\n srcs = [\"Sprite.kt\"],\n)\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/Background.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport android.graphics.Color\n\n/** Represents the static background behind all objects. */\nopen class Background(private val color: Int = Color.BLACK) : Sprite() {\n /** Paints the canvas with the color given in the constructor. */\n override fun draw(c: Canvas) {\n c.drawColor(color)\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/Ball.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport android.graphics.Color\nimport android.graphics.Paint\nimport java.time.Duration\nimport kotlin.math.ceil\nimport kotlin.math.sqrt\nimport kotlin.random.Random\n\n/** Represents a ball that travels down in space with constant speed. */\nopen class Ball(\n private val maxX: Int,\n private val maxY: Int,\n private val color: Int = Color.WHITE,\n private val radius: Float = 10.0f,\n // `speed`'s unit is in pixels/ms.\n private val speed: Float = 1.0f,\n private val rand: Random = Random.Default,\n) : Sprite() {\n\n // `x` and `y` represent the position of the center of this ball.\n //\n // Valid range [0, maxX]. 0==left, maxX==right.\n private var x: Int = rand.nextInt(maxX)\n // Valid range [0, maxY]. 0==top, maxY==bottom.\n private var y: Int = ceil(radius).toInt()\n\n private val paint: Paint =\n Paint(Paint.ANTI_ALIAS_FLAG).apply {\n style = Paint.Style.FILL\n color = (this@Ball).color\n }\n\n /** Returns `true` if this ball intersects the given line [segment]. */\n fun intersects(segment: LineSegment): Boolean {\n\n /** A vector with two components. */\n data class Vector2D(val u: Int, val v: Int) {\n /** Returns the dot product between two 2D vectors. */\n fun dot(other: Vector2D): Int = u * other.u + v * other.v\n }\n\n /** Returns the vector representing [p] minus [q]. */\n fun pointDiff(p: Point, q: Point): Vector2D = Vector2D(p.x - q.x, p.y - q.y)\n\n val direction = pointDiff(segment.p1, segment.p0) // p0 -> p1.\n val centerToP = pointDiff(segment.p0, Point(x, y)) // Ball center -> p0.\n\n // The `(centerToP + m * direction)` function models all the points in the line segment where\n // the independent variable `m` is a real number in [0,1]. Putting this function into the\n // formula for the circle (x ^ 2 + y ^ 2 = radius ^ 2) gives a quadratic equation\n // (am^2 + bm + c = 0) where:\n // [a] = direction · direction\n // [b] = 2 centerToP · direction\n // [c] = centerToP · centerToP - radius ^ 2\n val a = direction.dot(direction)\n val b = 2 * centerToP.dot(direction)\n val c = centerToP.dot(centerToP) - radius * radius\n\n val delta = b * b - 4 * a * c\n if (delta < 0)\n return false // No real roots means the (infinite) line does not intersect the ball.\n\n val d = sqrt(delta)\n val m1 = (-b - d) / (2 * a)\n val m2 = (-b + d) / (2 * a)\n\n // If a root is in [0,1], the line segment intersects the circumference.\n // If [m1] < 0 and [m2] > 1, the line segment is \"within\" the circle meaning the circle\n // intersects the infinite line, but not the line segment. In this case, we consider that it\n // touched the ball.\n return (m1 >= 0 && m1 <= 1) || (m2 >= 0 && m2 <= 1) || (m1 < 0 && m2 > 1)\n }\n\n /** Places the ball at the top of the screen at a random x-coordinate. */\n fun reset() {\n x = rand.nextInt(maxX)\n y = ceil(radius).toInt()\n }\n\n /** Moves the ball down by [timeDeltaMs]. */\n open fun update(timeDelta: Duration) {\n y += (speed * timeDelta.toMillis()).toInt()\n }\n\n /** Returns whether the ball is over [maxY]. */\n fun isOutOfBounds(): Boolean = y + radius > maxY || y - radius < 0\n\n /** Draws this ball in `c`. */\n override fun draw(c: Canvas) {\n c.drawCircle(x.toFloat(), y.toFloat(), radius, paint)\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/LineSegment.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\n/** Represents a finite line segment in 2D connected by two points [p0] and [p1]. */\ndata class LineSegment(val p0: Point, val p1: Point)\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/Paddle.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport android.graphics.Color\nimport android.graphics.Paint\nimport android.graphics.Rect\nimport kotlin.ranges.coerceIn\n\n/** Represents a paddle to hit/catch a falling ball. */\nopen class Paddle(\n private val color: Int = Color.WHITE,\n // Width and height in pixels.\n private val width: Int = 80,\n private val height: Int = 10,\n // maxX is the maximum X value for the center of the paddle.\n private val maxX: Int = 100,\n // The vertical position of the center of this paddle in pixels.\n val y: Int = 100,\n) : Sprite() {\n\n // Memoize a few things to make [draw()] a bit faster.\n private val halfH = height / 2\n private val halfW = width / 2\n private val paint =\n Paint(Paint.ANTI_ALIAS_FLAG).apply {\n style = Paint.Style.FILL\n color = (this@Paddle).color\n }\n\n // The horizontal center of the paddle.\n var x: Int = maxX / 2 // Start in the middle.\n set(value) {\n field = value.coerceIn(0, maxX)\n }\n\n /** Returns the (x,y) coordinates of the top-left corner. */\n fun topLeft(): Point = Point(x - halfW, y - halfH)\n\n /** Returns the (x,y) coordinates of the top-right corner. */\n fun topRight(): Point = Point(x + halfW, y - halfH)\n\n fun move(deltaX: Int) {\n x += deltaX\n }\n\n override fun draw(c: Canvas) {\n val rect =\n Rect().apply {\n bottom = y + halfH\n top = y - halfH\n left = x - halfW\n right = x + halfW\n }\n c.drawRect(rect, paint)\n }\n}\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/Point.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\n/** Represents a cartesian point in 2D. */\ndata class Point(val x: Int, val y: Int)\n"
},
{
"path": "android_env/apps/java/com/google/androidenv/catch/sprite/Sprite.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\n\n/** Represents something that can be drawn on the screen. */\nopen class Sprite {\n\n /** Draws the Sprite in the given canvas. */\n open fun draw(c: Canvas) {}\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/AndroidManifest.xml",
"content": "\n\n\n\n\n \n \n \n \n \n \n \n \n \n\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/BUILD.bazel",
"content": "# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Tests for the Android version of the RL Catch game.\nload(\"@rules_kotlin//kotlin:android.bzl\", \"kt_android_local_test\")\nload(\"@rules_kotlin//kotlin:core.bzl\", \"kt_kotlinc_options\")\n\nkt_kotlinc_options(\n name = \"kt_kotlinc_options\",\n jvm_target = \"11\", # Need to override default 1.8.\n x_no_param_assertions = True,\n)\n\nkt_android_local_test(\n name = \"GameLogicTest\",\n srcs = [\"GameLogicTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n deps = [\n \"//java/com/google/androidenv/catch:GameLogic\",\n \"//java/com/google/androidenv/catch/sprite:Background\",\n \"//java/com/google/androidenv/catch/sprite:Ball\",\n \"//java/com/google/androidenv/catch/sprite:Paddle\",\n \"@maven//:androidx_test_ext_junit\",\n \"@maven//:androidx_test_runner\",\n \"@maven//:com_google_truth_truth\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_robolectric_robolectric\",\n \"@robolectric//bazel:android-all\",\n ],\n)\n\nkt_android_local_test(\n name = \"GameLogicThreadTest\",\n srcs = [\"GameLogicThreadTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n deps = [\n \"//java/com/google/androidenv/catch:GameLogic\",\n \"//java/com/google/androidenv/catch:GameLogicThread\",\n \"@maven//:androidx_test_ext_junit\",\n \"@maven//:com_google_truth_truth\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_robolectric_robolectric\",\n \"@robolectric//bazel:android-all\",\n ],\n)\n\nkt_android_local_test(\n name = \"MainActivityTest\",\n srcs = [\n \"MainActivityTest.kt\",\n ],\n kotlinc_opts = \":kt_kotlinc_options\",\n manifest = \"AndroidManifest.xml\",\n deps = [\n \"//java/com/google/androidenv/catch:MainActivity\",\n \"@maven//:androidx_test_ext_junit\",\n \"@maven//:junit_junit\",\n \"@maven//:org_robolectric_robolectric\",\n \"@robolectric//bazel:android-all\",\n ],\n)\n\nkt_android_local_test(\n name = \"RenderThreadTest\",\n srcs = [\"RenderThreadTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n deps = [\n \"//java/com/google/androidenv/catch:GameLogic\",\n \"//java/com/google/androidenv/catch:RenderThread\",\n \"@maven//:androidx_test_ext_junit\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_mockito_mockito_core\",\n \"@maven//:org_robolectric_robolectric\",\n \"@robolectric//bazel:android-all\",\n ],\n)\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/GameLogicTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.graphics.Canvas\nimport androidx.test.core.view.MotionEventBuilder\nimport androidx.test.ext.junit.runners.AndroidJUnit4\nimport com.google.androidenv.catch.sprite.Background\nimport com.google.androidenv.catch.sprite.Ball\nimport com.google.androidenv.catch.sprite.Paddle\nimport com.google.common.truth.Truth.assertThat\nimport java.time.Duration\nimport java.time.Instant\nimport kotlin.random.Random\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.mockito.kotlin.any\nimport org.mockito.kotlin.atLeast\nimport org.mockito.kotlin.atMost\nimport org.mockito.kotlin.doReturn\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.spy\nimport org.mockito.kotlin.times\nimport org.mockito.kotlin.verify\n\n@RunWith(AndroidJUnit4::class)\nclass GameLogicTest {\n\n @Test\n fun run_ballIsMissed() {\n // Arrange.\n val width = 123\n val height = 33\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 37 }\n val game =\n GameLogic(\n rand = mockRandom,\n width = width,\n height = height,\n ball = Ball(maxX = width, maxY = height, radius = 5.0f, rand = mockRandom),\n paddle = Paddle(maxX = width, y = height, width = 3, height = 2),\n )\n game.reset()\n game.handleTouch(\n MotionEventBuilder.newBuilder().setPointer(/* x= */ 12.0f, /* y= */ 31.0f).build()\n )\n\n // Act.\n val outcome = game.run() // Ball falls at x==37, ev.x==12 so ball is missed.\n\n // Assert.\n assertThat(outcome).isEqualTo(false)\n }\n\n @Test\n fun run_ballIsCaught() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 53 }\n val game = GameLogic(rand = mockRandom, width = 321, height = 47)\n game.reset()\n game.handleTouch(\n MotionEventBuilder.newBuilder().setPointer(/* x= */ 53.0f, /* y= */ 43.0f).build()\n )\n\n // Act.\n val outcome = game.run() // Ball falls at x==53, ev.x==53 so ball is caught.\n\n // Assert.\n assertThat(outcome).isEqualTo(true)\n }\n\n @Test\n fun run_resetAllowsMultipleGamesToBePlayedWithASingleObjectAndDoesNotHang() {\n // Arrange.\n val mockRandom: Random = mock()\n val game = GameLogic(width = 101, height = 59, rand = mockRandom)\n\n // Act.\n repeat(17) {\n game.reset()\n val unused = game.run() // Ignore the outcome since we only care about run() terminating.\n }\n\n // Assert.\n // [rand.nextInt()] should be called once at construction and then 17 times for [reset()].\n verify(mockRandom, times(18)).nextInt(any())\n }\n\n @Test\n fun run_inASeparateThread() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 23 }\n val game = GameLogic(rand = mockRandom, width = 321, height = 89)\n game.reset()\n game.handleTouch(\n MotionEventBuilder.newBuilder().setPointer(/* x= */ 23.0f, /* y= */ 29.0f).build()\n )\n var outcome: Boolean = false\n\n class MyThread(val g: GameLogic, var outcome: Boolean) : Thread() {\n public override fun run() {\n outcome = g.run()\n }\n }\n val someThread = MyThread(game, outcome)\n\n // Act.\n someThread.start() // Ball falls at x==23, ev.x==23 so ball is caught.\n someThread.join()\n\n // Assert.\n assertThat(outcome).isEqualTo(true)\n }\n\n @Test\n fun run_fpsLeadstoApproximatelyNumberOfElapsedTimeAndUpdateCalls() {\n // Arrange.\n val width = 123\n val height = 300\n val ball = spy(Ball(maxX = width, maxY = height, speed = 2.0f, radius = 1.0f))\n val game = GameLogic(fps = 100, width = width, height = height, ball = ball)\n game.reset()\n\n // Act.\n val start = Instant.now()\n val unused = game.run()\n val end = Instant.now()\n\n // Assert.\n val elapsed = Duration.between(start, end)\n // The ball should take around `height / speed = 150` milliseconds to reach the bottom. Due to\n // timing non-determinism, we accept values between 100 and 200.\n assertThat(elapsed.toMillis()).isAtLeast(100L)\n assertThat(elapsed.toMillis()).isAtMost(200L)\n // At fps==100, we expect [update()] to be called every `1000 / 100 = 10` milliseconds. We\n // expect [elapsed] to be around 150ms (checked above) which should be around `150 / 10 = 15`\n // calls, so to account for timing non-determinism we accept between 5 and 25 calls.\n verify(ball, atLeast(5)).update(any())\n verify(ball, atMost(25)).update(any())\n }\n\n @Test\n fun render_drawCanBeCalledMultipleTimesWithinASingleRun() {\n // Arrange.\n val width = 321\n val height = 89\n val mockCanvas: Canvas = mock()\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 23 }\n val background = spy(Background())\n val paddle = spy(Paddle())\n val ball = spy(Ball(maxX = width, maxY = height))\n val game =\n GameLogic(\n rand = mockRandom,\n width = width,\n height = height,\n background = background,\n ball = ball,\n paddle = paddle,\n )\n game.reset()\n game.handleTouch(\n MotionEventBuilder.newBuilder().setPointer(/* x= */ 23.0f, /* y= */ 29.0f).build()\n )\n\n class MyThread(val g: GameLogic) : Thread() {\n public override fun run() {\n val unused = g.run()\n }\n }\n val someThread = MyThread(game)\n\n // Act.\n someThread.start()\n repeat(11) { game.render(mockCanvas) }\n someThread.join()\n\n // Assert.\n verify(background, times(11)).draw(mockCanvas)\n verify(ball, times(11)).draw(mockCanvas)\n verify(paddle, times(11)).draw(mockCanvas)\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/GameLogicThreadTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.util.Log\nimport androidx.test.ext.junit.runners.AndroidJUnit4\nimport com.google.common.truth.Truth.assertThat\nimport org.junit.Rule\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.mockito.kotlin.atLeastOnce\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.verify\nimport org.robolectric.junit.rules.ExpectedLogMessagesRule\n\n@RunWith(AndroidJUnit4::class)\nclass GameLogicThreadTest {\n\n // Rule to assert log messages, taken as a reference from MainActivityTest.kt\n @get:Rule val expectedLogMessagesRule = ExpectedLogMessagesRule()\n\n private val mockGame: GameLogic = mock()\n private val testTag = \"TestAndroidRLTask\"\n\n @Test\n fun run_iteratesGameAndLogs() {\n // Arrange\n val gameLogicThread = GameLogicThread(mockGame, testTag)\n\n // Act\n gameLogicThread.start()\n Thread.sleep(100) // Allow time for the thread to execute at least once.\n gameLogicThread.finish()\n gameLogicThread.join() // Wait for the thread to terminate.\n\n // Assert\n // Verify that the game's core methods were called at least once.\n verify(mockGame, atLeastOnce()).reset()\n verify(mockGame, atLeastOnce()).run()\n // Expect the log message from the run() loop.\n // The mock 'game.run()' returns false by default.\n expectedLogMessagesRule.expectLogMessage(Log.INFO, testTag, \"0 - false\")\n }\n\n @Test\n fun finish_stopsTheThread() {\n // Arrange\n val gameLogicThread = GameLogicThread(mockGame, testTag)\n\n // Act\n gameLogicThread.start()\n // Let it run for a moment before stopping it.\n Thread.sleep(50)\n gameLogicThread.finish()\n gameLogicThread.join()\n\n // Assert\n assertThat(gameLogicThread.isAlive).isFalse()\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/MainActivityTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.content.Intent\nimport android.util.Log\nimport androidx.test.ext.junit.rules.ActivityScenarioRule\nimport androidx.test.ext.junit.runners.AndroidJUnit4\nimport java.lang.reflect.Method\nimport org.junit.Before\nimport org.junit.Rule\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.robolectric.junit.rules.ExpectedLogMessagesRule\n\n@RunWith(AndroidJUnit4::class)\nclass MainActivityTest {\n @get:Rule(order = 0) val activityScenarioRule = ActivityScenarioRule(MainActivity::class.java)\n @get:Rule(order = 1) val expectedLogMessagesRule = ExpectedLogMessagesRule()\n\n @Before\n fun setUp() {\n expectedLogMessagesRule.expectLogMessage(Log.INFO, TAG, \"MainActivity::onCreate()\")\n }\n\n @Test\n fun surfaceChanged_logsStartsGame() {\n activityScenarioRule.scenario.onActivity { activity ->\n // Arrange.\n val surfaceView = activity.findViewById(R.id.surfaceView)\n val surfaceHolder = surfaceView.holder\n\n // Act - Trigger the surfaceChanged callback with positive width and height.\n activity.surfaceChanged(surfaceHolder, 0, 100, 200)\n\n // Assert.\n expectedLogMessagesRule.expectLogMessage(Log.INFO, TAG, \"MainActivity::surfaceChanged()\")\n expectedLogMessagesRule.expectLogMessage(Log.INFO, TAG, \"MainActivity::startGame()\")\n }\n }\n\n @Test\n fun onNewIntent_logsStartsGame_errorsOnUninitializedWidthOrHeight() {\n // Arrange.\n val newIntent = Intent()\n // Find the onNewIntent method using reflection\n val onNewIntentMethod: Method =\n MainActivity::class.java.getDeclaredMethod(\"onNewIntent\", Intent::class.java)\n // Enable access to protected method\n onNewIntentMethod.isAccessible = true\n\n activityScenarioRule.scenario.onActivity { activity ->\n // Act - Invoke the onNewIntent method using reflection.\n onNewIntentMethod.invoke(activity, newIntent)\n\n // Assert.\n expectedLogMessagesRule.expectLogMessage(Log.INFO, TAG, \"MainActivity::onNewIntent()\")\n expectedLogMessagesRule.expectLogMessage(Log.INFO, TAG, \"MainActivity::startGame()\")\n // In this test case where we don't call surfaceChanged(), default width and height\n // are -1 and should trigger this error to prevent Ball from initializing\n // with invalid negative values, since nextInt() expects a positive number.\n expectedLogMessagesRule.expectLogMessage(\n Log.ERROR,\n TAG,\n \"MainActivity::startGame() - Width or height not initialized yet.\",\n )\n }\n }\n\n companion object {\n private const val TAG = \"AndroidRLTask\"\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/RenderThreadTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch\n\nimport android.graphics.Canvas\nimport android.view.Surface\nimport android.view.SurfaceHolder\nimport androidx.test.ext.junit.runners.AndroidJUnit4\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.mockito.Mockito.verifyNoInteractions\nimport org.mockito.kotlin.any\nimport org.mockito.kotlin.atLeast\nimport org.mockito.kotlin.atMost\nimport org.mockito.kotlin.doReturn\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.verify\n\n@RunWith(AndroidJUnit4::class)\nclass RenderThreadTest {\n\n @Test\n fun run_finishBeforeStartResultsInNoRendering() {\n // Arrange.\n val surfaceHolder: SurfaceHolder = mock()\n val renderThread = RenderThread(surfaceHolder = surfaceHolder, fps = 1000)\n val game: GameLogic = mock()\n renderThread.game = game\n\n // Act.\n renderThread.finish()\n renderThread.start()\n\n // Assert.\n verifyNoInteractions(game)\n verifyNoInteractions(surfaceHolder)\n }\n\n @Test\n fun run_startResultsInSomeRendering() {\n // Arrange.\n val canvas: Canvas = mock()\n val surface: Surface = mock() { on { isValid() } doReturn true }\n val surfaceHolder: SurfaceHolder =\n mock() {\n on { getSurface() } doReturn surface\n on { lockCanvas() } doReturn canvas\n }\n val renderThread = RenderThread(surfaceHolder = surfaceHolder, fps = 1000)\n val game: GameLogic = mock()\n renderThread.game = game\n\n // Act.\n renderThread.start()\n Thread.sleep(/* millis= */ 500) // Sleep for at least one loop iteration.\n renderThread.finish()\n\n // Assert.\n verify(surfaceHolder, atLeast(1)).surface\n verify(surfaceHolder, atLeast(1)).lockCanvas()\n verify(surfaceHolder, atLeast(1)).unlockCanvasAndPost(any())\n verify(game, atLeast(1)).render(canvas)\n }\n\n @Test\n fun run_finishStopsRendering() {\n // Arrange.\n val canvas: Canvas = mock()\n val surface: Surface = mock() { on { isValid() } doReturn true }\n val surfaceHolder: SurfaceHolder =\n mock() {\n on { getSurface() } doReturn surface\n on { lockCanvas() } doReturn canvas\n }\n val renderThread = RenderThread(surfaceHolder = surfaceHolder, fps = 20)\n val game: GameLogic = mock()\n renderThread.game = game\n\n // Act.\n renderThread.start()\n Thread.sleep(/* millis= */ 500) // Sleep for around 10 iterations\n renderThread.finish()\n Thread.sleep(/* millis= */ 500) // Sleep some more to ensure nothing runs after.\n\n // Assert.\n verify(surfaceHolder, atLeast(1)).surface\n verify(surfaceHolder, atLeast(1)).lockCanvas()\n verify(surfaceHolder, atLeast(1)).unlockCanvasAndPost(any())\n // We expect [game.render()] to be executed for around 500 / (1000 / 20 = 50) = 10 times. To\n // allow for some timing non-determinism we allow it to execute up to 15 times, but not more\n // than that since [renderThread.finish()] should stop the thread from calling it.\n verify(game, atLeast(1)).render(canvas)\n verify(game, atMost(15)).render(canvas)\n }\n\n @Test\n fun run_expectedFramesPerSecond() {\n // Arrange.\n val canvas: Canvas = mock()\n val surface: Surface = mock() { on { isValid() } doReturn true }\n val surfaceHolder: SurfaceHolder =\n mock() {\n on { getSurface() } doReturn surface\n on { lockCanvas() } doReturn canvas\n }\n val renderThread = RenderThread(surfaceHolder = surfaceHolder, fps = 5)\n val game: GameLogic = mock()\n renderThread.game = game\n\n // Act.\n renderThread.start()\n Thread.sleep(/* millis= */ 2000) // Sleep for around 10 loop iterations.\n renderThread.finish()\n\n // Assert.\n verify(surfaceHolder, atLeast(1)).surface\n verify(surfaceHolder, atLeast(1)).lockCanvas()\n verify(surfaceHolder, atLeast(1)).unlockCanvasAndPost(any())\n // We expect [game.render()] to be called around 2000ms / 5fps = 10 times but to account for\n // timing non-determinism we allow ±4 iterations.\n verify(game, atLeast(6)).render(canvas)\n verify(game, atMost(14)).render(canvas)\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/sprite/BUILD.bazel",
"content": "# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Unit tests for Sprites in Catch.\nload(\"@rules_kotlin//kotlin:android.bzl\", \"kt_android_local_test\")\nload(\"@rules_kotlin//kotlin:core.bzl\", \"kt_kotlinc_options\")\n\nkt_kotlinc_options(\n name = \"kt_kotlinc_options\",\n jvm_target = \"11\", # Need to override default 1.8.\n x_no_param_assertions = True,\n)\n\nkt_android_local_test(\n name = \"BackgroundTest\",\n srcs = [\"BackgroundTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n deps = [\n \"//java/com/google/androidenv/catch/sprite:Background\",\n \"@maven//:com_google_guava_guava\",\n \"@maven//:com_google_testparameterinjector_test_parameter_injector\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_yaml_snakeyaml\",\n ],\n)\n\nkt_android_local_test(\n name = \"BallTest\",\n srcs = [\"BallTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n tags = [\"robolectric\"],\n deps = [\n \"//java/com/google/androidenv/catch/sprite:Ball\",\n \"//java/com/google/androidenv/catch/sprite:LineSegment\",\n \"//java/com/google/androidenv/catch/sprite:Point\",\n \"@maven//:androidx_test_ext_junit\",\n \"@maven//:com_google_guava_guava\",\n \"@maven//:com_google_truth_truth\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_robolectric_robolectric\",\n \"@robolectric//bazel:android-all\",\n ],\n)\n\nkt_android_local_test(\n name = \"PaddleTest\",\n srcs = [\"PaddleTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n tags = [\"robolectric\"],\n deps = [\n \"//java/com/google/androidenv/catch/sprite:Paddle\",\n \"//java/com/google/androidenv/catch/sprite:Point\",\n \"@maven//:androidx_test_ext_junit\",\n \"@maven//:com_google_guava_guava\",\n \"@maven//:com_google_truth_truth\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_robolectric_robolectric\",\n \"@robolectric//bazel:android-all\",\n ],\n)\n\nkt_android_local_test(\n name = \"SpriteTest\",\n srcs = [\"SpriteTest.kt\"],\n kotlinc_opts = \":kt_kotlinc_options\",\n deps = [\n \"//java/com/google/androidenv/catch/sprite:Sprite\",\n \"@maven//:org_mockito_kotlin_mockito_kotlin\",\n \"@maven//:org_mockito_mockito_core\",\n ],\n)\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/sprite/BackgroundTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport android.graphics.Color\nimport com.google.testing.junit.testparameterinjector.KotlinTestParameters.testValues\nimport com.google.testing.junit.testparameterinjector.TestParameter\nimport com.google.testing.junit.testparameterinjector.TestParameterInjector\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.times\nimport org.mockito.kotlin.verify\n\n@RunWith(TestParameterInjector::class)\nclass BackgroundTest {\n\n @Test\n fun draw_defaultConstructorIsBlack() {\n // Arrange.\n val mockCanvas: Canvas = mock()\n val background: Background = Background()\n\n // Act.\n background.draw(mockCanvas)\n\n // Assert.\n verify(mockCanvas, times(1)).drawColor(Color.BLACK)\n }\n\n @Test\n fun draw_customColors(\n @TestParameter colorInt: Int = testValues(0, 255, 13_579, 2_468, 12_384_173)\n ) {\n // Arrange.\n val mockCanvas: Canvas = mock()\n val background: Background = Background(color = colorInt)\n\n // Act.\n background.draw(mockCanvas)\n\n // Assert.\n verify(mockCanvas, times(1)).drawColor(colorInt)\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/sprite/BallTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport android.graphics.Color\nimport android.graphics.Paint\nimport androidx.test.ext.junit.runners.AndroidJUnit4\nimport com.google.common.truth.Truth.assertThat\nimport java.time.Duration\nimport kotlin.random.Random\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.junit.runners.Suite\nimport org.mockito.kotlin.any\nimport org.mockito.kotlin.argumentCaptor\nimport org.mockito.kotlin.doReturn\nimport org.mockito.kotlin.eq\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.verify\nimport org.robolectric.ParameterizedRobolectricTestRunner\n\n@RunWith(Suite::class)\n@Suite.SuiteClasses(\n BallTest.UpdateAndResetTests::class,\n BallTest.ColorIntTest::class,\n BallTest.CheckBoundsTest::class,\n BallTest.IntersectsTest::class,\n)\nclass BallTest {\n\n @RunWith(AndroidJUnit4::class)\n class UpdateAndResetTests() {\n @Test\n fun isOutOfBounds_initialState_isFalse() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n assertThat(isOutOfBounds()).isEqualTo(false)\n }\n }\n\n @Test\n fun isOutOfBounds_initialState_isTrueIfRadiusExceedsMaxY() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 100, maxY = 10, radius = 11.0f, speed = 1.0f, rand = mockRandom)) {\n assertThat(isOutOfBounds()).isEqualTo(true)\n }\n }\n\n @Test\n fun isOutOfBounds_initialState_isFalseIfRadiusExceedsOnlyMaxX() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 10, maxY = 100, radius = 11.0f, speed = 1.0f, rand = mockRandom)) {\n assertThat(isOutOfBounds()).isEqualTo(false)\n }\n }\n\n @Test\n fun update_zeroDurationDoesNotMove_withinBounds() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n // Act.\n update(Duration.ofMillis(0)) // The ball should not move.\n\n // Assert.\n assertThat(isOutOfBounds()).isEqualTo(false) // It should still be within the bounds.\n }\n }\n\n @Test\n fun update_zeroDurationDoesNotMove_outOfBounds() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n update(Duration.ofMillis(110)) // Place the ball out of bounds.\n assertThat(isOutOfBounds()).isEqualTo(true)\n\n // Act.\n update(Duration.ofMillis(0)) // The ball should not move.\n\n // Assert.\n assertThat(isOutOfBounds()).isEqualTo(true) // It should still be out of bounds.\n }\n }\n\n @Test\n fun update_negativeDurationsMovesUp() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n update(Duration.ofMillis(30)) // Move the ball down 30 pixels.\n assertThat(isOutOfBounds()).isEqualTo(false)\n\n // Act.\n update(Duration.ofMillis(-50)) // Move the ball _up_ 50 pixels.\n\n // Assert.\n assertThat(isOutOfBounds()).isEqualTo(true) // Now it should be out-of-bounds.\n }\n }\n\n @Test\n fun update_singleThrow() {\n // Ensures that a complete throw of a ball with radius==3.0f and maxY=100 behaves as expected.\n // [isOutOfBounds()] should return [false] for the first (100-3.0f-3.0f)=94 [update()] calls,\n // but [true] afterwards.\n\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n // Act.\n repeat(94) {\n update(Duration.ofMillis(1))\n assertThat(isOutOfBounds()).isEqualTo(false)\n }\n update(Duration.ofMillis(1))\n\n // Assert.\n assertThat(isOutOfBounds()).isEqualTo(true)\n }\n }\n\n @Test\n fun intersects_afterUpdate() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n\n // Act & Assert.\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n assertThat(intersects(LineSegment(Point(40, 0), Point(60, 0)))).isEqualTo(true)\n update(Duration.ofMillis(1))\n assertThat(intersects(LineSegment(Point(40, 0), Point(60, 0)))).isEqualTo(false)\n }\n }\n\n @Test\n fun reset_intersectsInitialPositionShouldBeTrue() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n\n with(Ball(maxX = 100, maxY = 100, radius = 3.0f, speed = 1.0f, rand = mockRandom)) {\n // Act.\n assertThat(intersects(LineSegment(Point(40, 0), Point(60, 0)))).isEqualTo(true)\n\n update(Duration.ofMillis(1)) // Move the ball 1 pixels down.\n assertThat(intersects(LineSegment(Point(40, 0), Point(60, 0))))\n .isEqualTo(false) // Segment is now outside of the ball.\n\n reset() // Resetting should move the ball up again.\n\n // Assert.\n assertThat(intersects(LineSegment(Point(40, 0), Point(60, 0))))\n .isEqualTo(true) // Segment is now inside of the ball.\n }\n }\n\n @Test\n fun reset_differentInitialXCoordinates() {\n // Arrange.\n val ball: Ball = Ball(maxX = 100, maxY = 100, radius = 3.0f)\n\n // Act.\n var pointInside: Boolean = false\n var pointOutside: Boolean = false\n while (!pointInside || !pointOutside) {\n if (ball.intersects(LineSegment(Point(45, 0), Point(55, 0)))) {\n pointInside = true\n } else {\n pointOutside = true\n }\n ball.reset() // Sample a new initial position for the ball.\n }\n\n // Assert.\n // Eventually after many initial positions the ball should satisfy both conditions.\n assertThat(pointInside).isEqualTo(true)\n assertThat(pointOutside).isEqualTo(true)\n }\n }\n\n @RunWith(ParameterizedRobolectricTestRunner::class)\n class ColorIntTest(private val c: Int) {\n\n @Test\n fun draw_customBallColors() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 37 }\n val mockCanvas: Canvas = mock()\n val paintCaptor = argumentCaptor()\n val ball: Ball = Ball(maxX = 50, maxY = 80, radius = 1.23f, color = c, rand = mockRandom)\n\n // Act.\n ball.draw(mockCanvas)\n\n // Assert.\n verify(mockCanvas).drawCircle(eq(37.0f), eq(2.0f), eq(1.23f), paintCaptor.capture())\n with(paintCaptor.lastValue) {\n assertThat(color).isEqualTo(c)\n assertThat(style).isEqualTo(Paint.Style.FILL)\n }\n }\n\n companion object {\n @JvmStatic\n @ParameterizedRobolectricTestRunner.Parameters(name = \"color = {0}\")\n fun parameters() = listOf(0, 255, -1, 13579, 2468, 12384173, Color.WHITE, Color.BLUE)\n }\n }\n\n @RunWith(ParameterizedRobolectricTestRunner::class)\n class CheckBoundsTest(private val p: ParamPack) {\n\n @Test\n fun intersects_checkBounds() {\n // Arrange.\n val mockRandom: Random =\n mock() { on { nextInt(any()) } doReturn p.maxX / 2 } // Horizontal middle.\n\n // Act.\n val ball: Ball = Ball(maxX = p.maxX, maxY = p.maxY, radius = p.radius, rand = mockRandom)\n\n // Assert.\n assertThat(ball.intersects(LineSegment(Point(p.x - 1, p.y), Point(p.x + 1, p.y))))\n .isEqualTo(p.expected)\n }\n\n data class ParamPack(\n val maxX: Int,\n val maxY: Int,\n val radius: Float,\n val x: Int,\n val y: Int,\n val expected: Boolean,\n )\n\n companion object {\n @JvmStatic\n @ParameterizedRobolectricTestRunner.Parameters(name = \"param = {0}\")\n fun parameters() =\n listOf(\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 0,\n y = 0,\n expected = false,\n ), // Ball to the right of `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 39,\n y = 0,\n expected = false,\n ), // Ball to the right of `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 40,\n y = 10,\n expected = true,\n ), // Ball contains `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 50,\n y = 0,\n expected = true,\n ), // Ball contains `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 60,\n y = 10,\n expected = true,\n ), // Ball contains `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 61,\n y = 0,\n expected = false,\n ), // Ball to the left of `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 100,\n y = 0,\n expected = false,\n ), // Ball to the left of `x`.\n ParamPack(\n maxX = 100,\n maxY = 100,\n radius = 10.0f,\n x = 50,\n y = 21,\n expected = false,\n ), // Ball above `y`.\n )\n }\n }\n\n @RunWith(ParameterizedRobolectricTestRunner::class)\n class IntersectsTest(private val p: ParamPack) {\n\n @Test\n fun intersects_ballAtx50y10radius10() {\n // Arrange.\n val mockRandom: Random = mock() { on { nextInt(any()) } doReturn 50 } // Horizontal middle.\n\n // Act.\n val ball: Ball = Ball(maxX = 100, maxY = 100, radius = 10.0f, rand = mockRandom)\n\n // Assert.\n assertThat(ball.intersects(p.segment)).isEqualTo(p.expected)\n }\n\n data class ParamPack(val segment: LineSegment, val expected: Boolean)\n\n companion object {\n @JvmStatic\n @ParameterizedRobolectricTestRunner.Parameters(name = \"param = {0}\")\n fun parameters() =\n listOf(\n ParamPack(\n segment = LineSegment(Point(50, 10), Point(80, 40)),\n expected = true,\n ), // Segment that starts at the center of the ball so it should always intersect.\n ParamPack(\n segment = LineSegment(Point(49, 0), Point(51, 0)),\n expected = true,\n ), // Tangential segment that touches the bottom of the ball.\n ParamPack(\n segment = LineSegment(Point(40, 5), Point(65, 7)),\n expected = true,\n ), // Segment longer than diameter, touching the circumference twice.\n ParamPack(\n segment = LineSegment(Point(42, 2), Point(58, 1)),\n expected = true,\n ), // Segment shorter than diameter, touching the circumference twice.\n ParamPack(\n segment = LineSegment(Point(44, 4), Point(54, 3)),\n expected = true,\n ), // Segment shorter than diameter, fully inside the circle, not touching the\n // circumference.\n ParamPack(\n segment = LineSegment(Point(35, 4), Point(54, 3)),\n expected = true,\n ), // Segment that touches the circumference once \"from the left\".\n ParamPack(\n segment = LineSegment(Point(54, 7), Point(67, 13)),\n expected = true,\n ), // Segment that touches the circumference once \"from the right\".\n ParamPack(\n segment = LineSegment(Point(36, 7), Point(45, 0)),\n expected = false,\n ), // Segment \"to the left of the ball\". No intersection.\n ParamPack(\n segment = LineSegment(Point(58, -3), Point(60, 3)),\n expected = false,\n ), // Segment \"to the right of the ball\". No intersection.\n )\n }\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/sprite/PaddleTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport android.graphics.Color\nimport android.graphics.Paint\nimport android.graphics.Rect\nimport androidx.test.ext.junit.runners.AndroidJUnit4\nimport com.google.common.truth.Truth.assertThat\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.junit.runners.Suite\nimport org.mockito.kotlin.argumentCaptor\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.verify\nimport org.robolectric.ParameterizedRobolectricTestRunner\n\n@RunWith(Suite::class)\n@Suite.SuiteClasses(\n PaddleTest.ConstructorTests::class,\n PaddleTest.MoveTests::class,\n PaddleTest.XSetterTests::class,\n PaddleTest.DrawTests::class,\n)\nclass PaddleTest {\n\n @RunWith(AndroidJUnit4::class)\n class ConstructorTests() {\n\n @Test\n fun x_initialValueShouldBeAtCenter() {\n with(Paddle(maxX = 30)) { assertThat(x).isEqualTo(15) }\n with(Paddle(maxX = 31)) { assertThat(x).isEqualTo(15) }\n }\n\n @Test\n fun topLeft_correspondsToGivenValues() {\n with(Paddle(width = 10, height = 6, maxX = 40, y = 33)) {\n assertThat(topLeft()).isEqualTo(Point(x = 15, y = 30))\n }\n }\n\n @Test\n fun topRight_correspondsToGivenValues() {\n with(Paddle(width = 10, height = 6, maxX = 40, y = 33)) {\n assertThat(topRight()).isEqualTo(Point(x = 25, y = 30))\n }\n }\n }\n\n @RunWith(ParameterizedRobolectricTestRunner::class)\n class MoveTests(private val p: ParamPack) {\n\n @Test\n fun move_expectedDestination() {\n // Arrange.\n with(Paddle(maxX = 50)) {\n // Act.\n move(deltaX = p.displacement)\n\n // Assert.\n assertThat(x).isEqualTo(p.expectedX)\n }\n }\n\n data class ParamPack(val displacement: Int, val expectedX: Int)\n\n companion object {\n @JvmStatic\n @ParameterizedRobolectricTestRunner.Parameters(name = \"param = {0}\")\n fun parameters() =\n listOf(\n // Initial position is x==25.\n ParamPack(displacement = 10, expectedX = 35),\n ParamPack(displacement = -10, expectedX = 15),\n ParamPack(displacement = 0, expectedX = 25),\n // Going beyond the left and right walls should clamp the values to 0 and 50.\n ParamPack(displacement = -26, expectedX = 0),\n ParamPack(displacement = 26, expectedX = 50),\n )\n }\n }\n\n @RunWith(ParameterizedRobolectricTestRunner::class)\n class XSetterTests(private val p: ParamPack) {\n\n @Test\n fun xSetter_expectedDestination() {\n // Arrange.\n with(Paddle(maxX = 50)) {\n // Act.\n x = p.target\n\n // Assert.\n assertThat(x).isEqualTo(p.expectedX)\n }\n }\n\n data class ParamPack(val target: Int, val expectedX: Int)\n\n companion object {\n @JvmStatic\n @ParameterizedRobolectricTestRunner.Parameters(name = \"param = {0}\")\n fun parameters() =\n listOf(\n // Initial position is x==25.\n ParamPack(target = 0, expectedX = 0),\n ParamPack(target = 15, expectedX = 15),\n ParamPack(target = 25, expectedX = 25),\n ParamPack(target = 35, expectedX = 35),\n ParamPack(target = 50, expectedX = 50),\n // Going beyond the left and right walls should clamp the values to 0 and 50.\n ParamPack(target = -1, expectedX = 0),\n ParamPack(target = 51, expectedX = 50),\n )\n }\n }\n\n @RunWith(AndroidJUnit4::class)\n class DrawTests() {\n\n @Test\n fun draw_initialPosition() {\n // Arrange.\n val mockCanvas: Canvas = mock()\n val rectCaptor = argumentCaptor()\n val paintCaptor = argumentCaptor()\n with(Paddle(color = Color.RED, width = 100, height = 20, maxX = 300, y = 400)) {\n // Act.\n draw(mockCanvas)\n\n // Assert.\n assertThat(x).isEqualTo(150)\n verify(mockCanvas).drawRect(rectCaptor.capture(), paintCaptor.capture())\n with(rectCaptor.lastValue) {\n assertThat(bottom).isEqualTo(400 + 10)\n assertThat(top).isEqualTo(400 - 10)\n assertThat(left).isEqualTo(150 - 50)\n assertThat(right).isEqualTo(150 + 50)\n }\n }\n }\n\n @Test\n fun draw_afterMove() {\n // Arrange.\n val mockCanvas: Canvas = mock()\n val rectCaptor = argumentCaptor()\n val paintCaptor = argumentCaptor()\n with(Paddle(color = Color.RED, width = 100, height = 20, maxX = 300, y = 400)) {\n // Act.\n move(50)\n draw(mockCanvas)\n\n // Assert.\n assertThat(x).isEqualTo(200)\n verify(mockCanvas).drawRect(rectCaptor.capture(), paintCaptor.capture())\n with(rectCaptor.lastValue) {\n assertThat(bottom).isEqualTo(400 + 10)\n assertThat(top).isEqualTo(400 - 10)\n assertThat(left).isEqualTo(200 - 50)\n assertThat(right).isEqualTo(200 + 50)\n }\n with(paintCaptor.lastValue) {\n assertThat(color).isEqualTo(Color.RED)\n assertThat(style).isEqualTo(Paint.Style.FILL)\n }\n }\n }\n }\n}\n"
},
{
"path": "android_env/apps/javatests/com/google/androidenv/catch/sprite/SpriteTest.kt",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\npackage com.google.androidenv.catch.sprite\n\nimport android.graphics.Canvas\nimport org.junit.Test\nimport org.junit.runner.RunWith\nimport org.junit.runners.JUnit4\nimport org.mockito.Mockito.verifyNoInteractions\nimport org.mockito.kotlin.mock\nimport org.mockito.kotlin.times\nimport org.mockito.kotlin.verify\n\n/** Trivial tests to ensure the types in the API are correct. */\n@RunWith(JUnit4::class)\nclass SpriteTest {\n\n @Test\n fun defaultImplementationDoesNothing() {\n // Arrange.\n val mockCanvas: Canvas = mock()\n val sprite = Sprite()\n\n // Act.\n sprite.draw(mockCanvas)\n\n // Assert.\n verifyNoInteractions(mockCanvas) // No methods should be called on the canvas.\n }\n\n @Test\n fun draw_argumentsAreForwarded() {\n // Arrange.\n val mockSprite: Sprite = mock()\n val mockCanvas: Canvas = mock()\n\n // Act.\n mockSprite.draw(mockCanvas)\n\n // Assert.\n verify(mockSprite, times(1)).draw(mockCanvas)\n }\n}\n"
},
{
"path": "android_env/components/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/components/action_fns.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Functions to convert actions between different components' formats.\"\"\"\n\nfrom absl import logging\nfrom android_env.components import action_type as action_type_lib\nfrom android_env.components import errors\nfrom android_env.components import pixel_fns\nfrom android_env.components.simulators import base_simulator\nimport numpy as np\n\n\ndef send_action_to_simulator(\n action: dict[str, np.ndarray],\n simulator: base_simulator.BaseSimulator,\n screen_width: int,\n screen_height: int,\n num_fingers: int,\n) -> bool:\n \"\"\"Sends the selected action to the given simulator.\n\n The simulator will interpret the action according to `action[\"action_type\"]`.\n The effect this action triggers in the Android OS will be determined by the\n currently running application.\n\n Args:\n action: action which will get interpreted as a touchscreen event.\n simulator: The simulator that will receive the action.\n screen_width: The width of the touchscreen in pixels.\n screen_height: The height of the touchscreen in pixels.\n num_fingers: The number of fingers used in this simulator.\n \"\"\"\n\n try:\n match action['action_type']:\n # If the action is a TOUCH or LIFT, send a touch event to the simulator.\n case action_type_lib.ActionType.TOUCH | action_type_lib.ActionType.LIFT:\n prepared_action = _prepare_touch_action(\n action, screen_width, screen_height, num_fingers\n )\n simulator.send_touch(prepared_action)\n # If the action is a key event, send a key event to the simulator.\n case action_type_lib.ActionType.KEYDOWN:\n simulator.send_key(action['keycode'].item(0), event_type='keydown')\n case action_type_lib.ActionType.KEYUP:\n simulator.send_key(action['keycode'].item(0), event_type='keyup')\n case action_type_lib.ActionType.KEYPRESS:\n simulator.send_key(action['keycode'].item(0), event_type='keypress')\n except errors.SendActionError:\n logging.exception('Unable to execute action: %r', action)\n return False\n\n return True\n\n\ndef _prepare_touch_action(\n action: dict[str, np.ndarray],\n screen_width: int,\n screen_height: int,\n num_fingers: int,\n) -> list[tuple[int, int, bool, int]]:\n \"\"\"Turns an AndroidEnv action into values that the simulator can interpret.\n\n Converts float-valued 'touch_position' to integer coordinates corresponding\n to specific pixels, and 'action_type' to booleans indicating whether the\n screen is touched at said location or not. The result of this function can\n be sent directly to the underlying simulator (e.g. the Android Emulator,\n virtual machine, or a phone).\n\n Args:\n action: An action containing 'action_type' and 'touch_position'.\n\n Returns:\n A tuple with the format (x: int, y: int, down/up: bool, finger_index: int).\n \"\"\"\n\n touch_events = []\n for i, finger_action in enumerate(_split_touch_action(action, num_fingers)):\n is_touch = finger_action['action_type'] == action_type_lib.ActionType.TOUCH\n touch_position = finger_action['touch_position']\n touch_pixels = pixel_fns.touch_position_to_pixel_position(\n touch_position, width_height=(screen_width, screen_height)\n )\n touch_events.append((touch_pixels[0], touch_pixels[1], is_touch, i))\n return touch_events\n\n\ndef _split_touch_action(\n action: dict[str, np.ndarray], num_fingers: int\n) -> list[dict[str, np.ndarray]]:\n \"\"\"Splits a multitouch action into a list of single-touch actions.\"\"\"\n\n single_touch_actions = [{\n 'action_type': action['action_type'],\n 'touch_position': action['touch_position'],\n }]\n for i in range(2, num_fingers + 1):\n single_touch_actions.append({\n 'action_type': action[f'action_type_{i}'],\n 'touch_position': action[f'touch_position_{i}'],\n })\n return single_touch_actions\n\n\ndef lift_all_fingers_action(num_fingers: int) -> dict[str, np.ndarray]:\n \"\"\"A lift action with each finger.\"\"\"\n\n # There's always at least one finger.\n lift_action = {\n 'action_type': np.array(action_type_lib.ActionType.LIFT),\n 'touch_position': np.array([0, 0]),\n }\n # Subsequent fingers have separate dict entries.\n for i in range(2, num_fingers + 1):\n lift_action |= {\n f'action_type_{i}': np.array(action_type_lib.ActionType.LIFT),\n f'touch_position_{i}': np.array([0, 0]),\n }\n return lift_action\n"
},
{
"path": "android_env/components/action_fns_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import action_fns\nfrom android_env.components import action_type as action_type_lib\nfrom android_env.components import errors\nfrom android_env.components.simulators import base_simulator\nimport numpy as np\n\n\nclass ActionFnsTest(parameterized.TestCase):\n\n def test_send_action_to_simulator_missing_action_type(self):\n \"\"\"A `KeyError` should be raised if the action is missing \"action_type\".\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n action = {'some_key': np.array(123, np.int32)}\n\n # Act & Assert.\n self.assertRaises(\n KeyError,\n action_fns.send_action_to_simulator,\n action,\n simulator,\n 800,\n 600,\n 1,\n )\n\n def test_send_action_to_simulator_sendactionerror(self):\n \"\"\"Returns `False` if the simulator raises a SendActionError.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n simulator.send_touch.side_effect = errors.SendActionError('oops!')\n action = {\n 'action_type': action_type_lib.ActionType.TOUCH,\n 'touch_position': np.array([0.3, 0.5], np.float32),\n }\n\n # Act.\n output = action_fns.send_action_to_simulator(\n action,\n simulator,\n 800,\n 600,\n 1,\n )\n\n # Assert.\n self.assertFalse(output)\n simulator.send_touch.assert_called_once()\n\n def test_send_action_to_simulator_touch_success_one_finger(self):\n \"\"\"Returns `True` with a proper 1-finger touch action.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n action = {\n 'action_type': action_type_lib.ActionType.TOUCH,\n 'touch_position': np.array([0.2, 0.5], np.float32),\n }\n\n # Act.\n output = action_fns.send_action_to_simulator(\n action,\n simulator,\n 800,\n 600,\n 1,\n )\n\n # Assert.\n self.assertTrue(output)\n simulator.send_touch.assert_called_once_with(\n [(np.int32(160), np.int32(300), True, 0)]\n )\n\n def test_send_action_to_simulator_touch_success_multiple_finger(self):\n \"\"\"Returns `True` with a proper 3-finger touch action.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n action = {\n 'action_type': action_type_lib.ActionType.TOUCH,\n 'touch_position': np.array([0.2, 0.5], np.float32),\n 'action_type_2': action_type_lib.ActionType.LIFT,\n 'touch_position_2': np.array([0.1, 0.2], np.float32),\n 'action_type_3': action_type_lib.ActionType.TOUCH,\n 'touch_position_3': np.array([0.5, 0.2], np.float32),\n }\n\n # Act.\n output = action_fns.send_action_to_simulator(\n action,\n simulator,\n 800,\n 600,\n 3,\n )\n\n # Assert.\n self.assertTrue(output)\n simulator.send_touch.assert_called_once_with([\n (np.int32(160), np.int32(300), True, 0),\n (np.int32(80), np.int32(120), False, 1),\n (np.int32(400), np.int32(120), True, 2),\n ])\n\n def test_send_action_to_simulator_keydown_success(self):\n \"\"\"Returns `True` with a proper keydown action.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n action = {\n 'action_type': action_type_lib.ActionType.KEYDOWN,\n 'keycode': np.array([21], np.int32),\n }\n\n # Act.\n output = action_fns.send_action_to_simulator(\n action,\n simulator,\n 800,\n 600,\n 1,\n )\n\n # Assert.\n self.assertTrue(output)\n simulator.send_key.assert_called_once_with(21, event_type='keydown')\n\n def test_send_action_to_simulator_keyup_success(self):\n \"\"\"Returns `True` with a proper keyup action.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n action = {\n 'action_type': action_type_lib.ActionType.KEYUP,\n 'keycode': np.array([42], np.int32),\n }\n\n # Act.\n output = action_fns.send_action_to_simulator(\n action,\n simulator,\n 800,\n 600,\n 1,\n )\n\n # Assert.\n self.assertTrue(output)\n simulator.send_key.assert_called_once_with(42, event_type='keyup')\n\n def test_send_action_to_simulator_keypress_success(self):\n \"\"\"Returns `True` with a proper keypress action.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n action = {\n 'action_type': action_type_lib.ActionType.KEYPRESS,\n 'keycode': np.array([96], np.int32),\n }\n\n # Act.\n output = action_fns.send_action_to_simulator(\n action,\n simulator,\n 800,\n 600,\n 1,\n )\n\n # Assert.\n self.assertTrue(output)\n simulator.send_key.assert_called_once_with(96, event_type='keypress')\n\n @parameterized.named_parameters(\n (\n 'one_finger',\n 1,\n {\n 'action_type': np.array(action_type_lib.ActionType.LIFT),\n 'touch_position': np.array([0, 0]),\n },\n ),\n (\n 'two_fingers',\n 2,\n {\n 'action_type': np.array(action_type_lib.ActionType.LIFT),\n 'touch_position': np.array([0, 0]),\n 'action_type_2': np.array(action_type_lib.ActionType.LIFT),\n 'touch_position_2': np.array([0, 0]),\n },\n ),\n )\n def test_lift_all_fingers_action(\n self, num_fingers: int, expected_action: dict[str, np.ndarray]\n ):\n \"\"\"Returns the expected action.\"\"\"\n\n output = action_fns.lift_all_fingers_action(num_fingers)\n for k, v in expected_action.items():\n np.testing.assert_array_equal(v, output[k])\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/action_type.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"The different kinds of actions that AndroidEnv supports.\n\nThe native action space of AndroidEnv consists of a tuple consisting of\n- A position (x, y) ∈ [0, 1] x [0, 1], determining the location of the action on\n the screen, and\n- A discrete value, indicating the action type, which is in this file.\n\nSee https://arxiv.org/abs/2105.13231, section 2.2 for details.\n\"\"\"\n\nimport enum\n\n\n@enum.unique\nclass ActionType(enum.IntEnum):\n \"\"\"Integer values to describe each supported action in AndroidEnv.\n\n Note for KEY* types:\n - Only meaningful if connected to a _physical_ keyboard, _not_ virtual\n keyboard.\n - Added afterwards so they did not appear in the paper.\n\n Attributes:\n TOUCH: Touching the screen at a location.\n LIFE: Lifting the (imaginary) pointer from the screen at a location.\n REPEAT: Repeating the last chosen action.\n KEYDOWN: Sending a key down event.\n KEYUP: Sending a key up event.\n KEYPRESS: Sending a key down event, immediately followed by a key up event.\n \"\"\"\n\n TOUCH = 0\n LIFT = 1\n REPEAT = 2\n KEYDOWN = 3\n KEYUP = 4\n KEYPRESS = 5\n"
},
{
"path": "android_env/components/adb_call_parser.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Processes adb_pb2.AdbRequest commands.\"\"\"\n\nimport os\nimport re\nimport subprocess\nimport sys\nimport tempfile\n\nfrom absl import logging\nfrom android_env.components import adb_controller as adb_control\nfrom android_env.proto import adb_pb2\n\n# A mapping from a Button enum to keycode strings.\n#\n# Please see https://developer.android.com/reference/android/view/KeyEvent\n#\n# We currently only accept the following entries:\n_BUTTON_TO_KEYCODE = {\n adb_pb2.AdbRequest.PressButton.Button.HOME: 'KEYCODE_HOME',\n adb_pb2.AdbRequest.PressButton.Button.BACK: 'KEYCODE_BACK',\n adb_pb2.AdbRequest.PressButton.Button.ENTER: 'KEYCODE_ENTER',\n}\n\n\nclass AdbCallParser:\n \"\"\"Parses AdbRequest messages and executes corresponding adb commands.\"\"\"\n\n def __init__(self, adb_controller: adb_control.AdbController):\n self._adb_controller = adb_controller\n self._handlers = {\n 'install_apk': self._install_apk,\n 'start_activity': self._start_activity,\n 'force_stop': self._force_stop,\n 'tap': self._tap,\n 'press_button': self._press_button,\n 'start_screen_pinning': self._start_screen_pinning,\n 'send_broadcast': self._send_broadcast,\n 'uninstall_package': self._handle_uninstall_package,\n 'get_current_activity': self._get_current_activity,\n 'get_orientation': self._get_orientation,\n 'push': self._push,\n 'pull': self._pull,\n 'input_text': self._input_text,\n 'settings': self._handle_settings,\n 'generic': self._handle_generic,\n 'package_manager': self._handle_package_manager,\n 'dumpsys': self._handle_dumpsys,\n }\n\n def _execute_command(\n self, command_args: list[str], timeout: float | None\n ) -> tuple[adb_pb2.AdbResponse, bytes]:\n \"\"\"Executes the command, catches errors and populates the response status.\n\n Args:\n command_args: a list of arguments for the ADB request.\n timeout: Timeout in seconds.\n\n Returns:\n A tuple of the AdbResponse with the status populated, and the output\n bytes from the command.\n \"\"\"\n response = adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)\n command_output = b''\n try:\n command_output = self._adb_controller.execute_command(\n command_args, timeout=timeout)\n except subprocess.CalledProcessError as adb_error:\n if adb_error.stdout is not None:\n response.status = adb_pb2.AdbResponse.Status.ADB_ERROR\n response.error_message = adb_error.stdout\n except subprocess.TimeoutExpired:\n response.status = adb_pb2.AdbResponse.Status.TIMEOUT\n response.error_message = 'Timeout'\n\n return response, command_output\n\n def parse(self, request: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:\n \"\"\"Executes `request` and returns an appropriate response.\"\"\"\n\n response = adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)\n command_type = request.WhichOneof('command')\n logging.debug('AdbRequest command type: %s', command_type)\n if command_type is None:\n response.status = adb_pb2.AdbResponse.Status.UNKNOWN_COMMAND\n response.error_message = 'AdbRequest.command is None.'\n return response\n\n if request.timeout_sec < 0:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = ('AdbRequest.timeout_sec cannot be negative. '\n f'Got: {request.timeout_sec}')\n return response\n\n timeout: float | None = request.timeout_sec or None\n return self._handlers[command_type](request, timeout)\n\n def _force_stop(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Stops an application.\n\n Args:\n request: The external request containing the package to force stop.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n force_stop = request.force_stop\n response = adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)\n if not force_stop.package_name:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = '`force_stop.package_name` cannot be empty.'\n return response\n\n response, _ = self._execute_command(\n ['shell', 'am', 'force-stop', force_stop.package_name], timeout)\n\n return response\n\n def _fetch_current_task_id(\n self, full_activity_name: str, timeout: float | None = None\n ) -> int:\n \"\"\"Returns the task ID of the given `full_activity_name`.\n\n Args:\n full_activity_name: The full name of the activity whose corresponding\n task id we are looking for.\n timeout: Optional time limit in seconds.\n Returns:\n task_id: An integer corresponding to the specified activity.\n \"\"\"\n\n stack = self._adb_controller.execute_command(\n ['shell', 'am', 'stack', 'list'], timeout=timeout)\n lines = stack.decode('utf-8').splitlines()\n\n regex = re.compile(\n r'^\\ *taskId=(?P[0-9]*): (?P[^\\s]*) .*visible=true'\n r'.*topActivity=ComponentInfo{(?P[^\\s]*)}$')\n\n for line in lines:\n match = regex.search(line)\n if match is None:\n continue\n\n current_task_id_str = match.group('id')\n base_activity = match.group('base_activity')\n top_activity = match.group('top_activity')\n\n # If neither of the matched activities equals the activity we are\n # looking for, we discard their task id and continue the search.\n if full_activity_name not in {base_activity, top_activity}:\n logging.info('Full activity %s was not found in current line %s',\n full_activity_name, line)\n continue\n\n # Otherwise return the integer task id.\n try:\n return int(current_task_id_str)\n except ValueError:\n logging.info('Failed to parse task ID [%r].', current_task_id_str)\n\n # At this point if we could not find a task ID, there's nothing we can do.\n logging.error('Could not find current activity in stack list: %r', lines)\n return -1\n\n def _start_screen_pinning(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Pins an application.\n\n Args:\n request: The request containing the activity to pin.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n full_activity = request.start_screen_pinning.full_activity\n response = adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)\n if not full_activity:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n '`start_screen_pinning.full_activity` cannot be empty.')\n return response\n\n current_task_id = self._fetch_current_task_id(full_activity, timeout)\n if current_task_id == -1:\n response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR\n response.error_message = ('Could not find task ID for activity '\n f'[{full_activity}]')\n return response\n\n response, _ = self._execute_command(\n ['shell', 'am', 'task', 'lock',\n str(current_task_id)], timeout=timeout)\n\n return response\n\n def _send_broadcast(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Sends a broadcast.\n\n Args:\n request: The request with the information for the broadcast event.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n send_broadcast = request.send_broadcast\n response = adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)\n if not send_broadcast.action:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = ('`send_broadcast.{action}` cannot be empty.')\n return response\n\n if send_broadcast.component:\n component_args = ['-n', send_broadcast.component]\n else:\n component_args = []\n\n response, _ = self._execute_command(\n ['shell', 'am', 'broadcast', '-a', send_broadcast.action]\n + component_args,\n timeout=timeout,\n )\n\n return response\n\n def _install_apk(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Installs an app given its local path in the filesystem.\n\n Args:\n request: The external request with an install_apk field.\n Contains information for the .apk installation.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n install_apk = request.install_apk\n response = adb_pb2.AdbResponse()\n location_type = install_apk.WhichOneof('location')\n logging.info('location_type: %s', location_type)\n\n match location_type:\n case 'filesystem':\n fpath = install_apk.filesystem.path\n if not os.path.exists(fpath):\n response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR\n response.error_message = f'Could not find local_apk_path: {fpath}'\n return response\n\n response, _ = self._execute_command(\n ['install', '-r', '-t', '-g', fpath], timeout=timeout\n )\n case 'blob':\n\n # `delete_on_close` was only added in Python 3.12 so we add a switch\n # here to still support previous Python versions.\n if sys.version_info >= (3, 12):\n kwargs = {'suffix': '.apk', 'delete_on_close': False}\n else:\n kwargs = {'suffix': '.apk'}\n\n with tempfile.NamedTemporaryFile(**kwargs) as f:\n fpath = f.name\n f.write(install_apk.blob.contents)\n\n response, _ = self._execute_command(\n ['install', '-r', '-t', '-g', fpath], timeout=timeout\n )\n case _:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Unsupported `install_apk.location` type: {location_type}'\n )\n return response\n\n return response\n\n def _start_activity(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Starts a given activity.\n\n Options for `start_activity`:\n `am start` command options:\n -D: enable debugging\n -W: wait for launch to complete\n --start-profiler : start profiler and send results to \n -P : like above, but profiling stops when app goes idle\n -R: repeat the activity launch times. Prior to each repeat,\n the top activity will be finished.\n -S: force stop the target app before starting the activity\n --opengl-trace: enable tracing of OpenGL functions\n\n Args:\n request: The request with information on what activity to start.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse. If successful, StartActivityResponse will contain the\n activity name and adb command output.\n \"\"\"\n\n activity = request.start_activity.full_activity\n if not activity:\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.FAILED_PRECONDITION,\n error_message='`start_activity.full_activity` cannot be empty.')\n\n force_stop = '-S' if request.start_activity.force_stop else ''\n response, command_output = self._execute_command(\n ['shell', 'am', 'start', force_stop, '-W', '-n', activity] +\n list(request.start_activity.extra_args or []),\n timeout=timeout)\n\n # Check command output for potential errors.\n expected_error = re.compile(r\"\"\".*Error.*\"\"\", re.VERBOSE)\n if expected_error.match(str(command_output)):\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.INTERNAL_ERROR,\n error_message=f'start_activity failed with error: {command_output}')\n\n response.start_activity.full_activity = activity\n response.start_activity.output = command_output\n return response\n\n def _press_button(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Presses a keyboard key.\n\n Args:\n request: The request with information on what button to press.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n button = request.press_button.button\n if button not in _BUTTON_TO_KEYCODE:\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.FAILED_PRECONDITION,\n error_message=('PressButton.button must be one of '\n f'[{_BUTTON_TO_KEYCODE.keys()}]. '\n f'Got: {button}. Please see `adb.proto`.'))\n\n keycode = _BUTTON_TO_KEYCODE[button]\n response, command_output = self._execute_command(\n ['shell', 'input', 'keyevent', keycode], timeout=timeout)\n response.press_button.output = command_output\n return response\n\n def _handle_uninstall_package(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Handles UninstallPackage messages.\n\n Args:\n request: The specification of what to uninstall.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse\n \"\"\"\n\n package_name = request.uninstall_package.package_name\n response = adb_pb2.AdbResponse()\n # Every UninstallPackage should have a package_name.\n if not package_name:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n '`uninstall_package.package_name` cannot be empty.')\n return response\n\n # Get list of installed packages and issue an uninstall only if it's\n # already installed.\n package_response = self._handle_package_manager(\n adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n packages=adb_pb2.AdbRequest.PackageManagerRequest.List\n .Packages()))))\n if package_name in package_response.package_manager.list.items:\n response, _ = self._execute_command(['uninstall', package_name], timeout)\n else:\n msg = (f'Cannot uninstall {package_name} since it is not installed.')\n logging.warning(msg)\n response.error_message = msg\n\n return response\n\n def _get_current_activity(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Fetches current activity.\n\n Args:\n request: The request with the `.get_current_activity` field set. This is\n unused, but it's in the signature so that all calls are uniform.\n timeout: Optional time limit in seconds.\n\n Returns:\n AdbResponse containing the current activity.\n \"\"\"\n\n del request # Unused.\n\n response, visible_task = self._execute_command(\n ['shell', 'am', 'stack', 'list', '|', 'grep', '-E', 'visible=true'],\n timeout=timeout)\n\n if response.status != adb_pb2.AdbResponse.Status.OK:\n return response\n\n if not visible_task:\n _, am_stack_list = self._execute_command(['shell', 'am', 'stack', 'list'],\n timeout=timeout)\n response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR\n response.error_message = ('Empty visible_task. `am stack list`: '\n f'{am_stack_list}')\n return response\n\n visible_task = visible_task.decode('utf-8')\n if sys.platform == 'win32':\n visible_task_list = re.findall(\n r'visible=true topActivity=ComponentInfo{(.+?)}', visible_task)\n if not visible_task_list:\n visible_task = ''\n else:\n visible_task = 'ComponentInfo{' + visible_task_list[0] + '}'\n\n p = re.compile(r'.*\\{(.*)\\}')\n matches = p.search(visible_task)\n if matches is None:\n _, am_stack_list = self._execute_command(['shell', 'am', 'stack', 'list'],\n timeout=timeout)\n response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR\n response.error_message = (\n 'Could not extract current activity. Will return nothing. '\n f'`am stack list`: {am_stack_list}')\n return response\n\n response.get_current_activity.full_activity = matches.group(1)\n return response\n\n def _get_orientation(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Fetches current device orientation.\n\n Args:\n request: The request with the `.get_orientation` field set.\n timeout: Optional time limit in seconds.\n\n Returns:\n AdbResponse containing the current device orientation. This is\n unused, but it's in the signature so that all calls are uniform.\n \"\"\"\n\n del request # Unused.\n\n logging.info('Getting orientation...')\n response = self._handle_dumpsys(\n adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(service='input')),\n timeout=timeout)\n output = response.dumpsys.output\n if not output:\n logging.error('Empty dumpsys output.')\n response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR\n response.error_message = 'Failed to execute `dumpsys input`'\n return response\n\n output = output.decode('utf-8')\n lines = output.split('\\n') # Split by lines.\n skip_next = False\n for line in lines:\n # There may be multiple devices in output. An invalid device can be\n # identified by negative PhysicalWidth.\n physical_width = re.match(r'\\s+PhysicalWidth:\\s+(-?\\d+)px', line)\n if physical_width:\n skip_next = int(physical_width.group(1)) < 0\n # Depending on the device type, the orientation could take these forms:\n # SurfaceOrientation: 0\n # InputDeviceOrientation: Rotation0\n surface_orientation = re.match(\n r'\\s+(SurfaceOrientation|InputDeviceOrientation):\\s+.*(\\d)', line\n )\n\n if surface_orientation is not None:\n if skip_next:\n continue\n if surface_orientation.re.groups < 2:\n continue\n orientation = surface_orientation.group(2)\n logging.info('Done getting orientation: %r', orientation)\n response.get_orientation.orientation = int(orientation)\n return response\n\n response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR\n response.error_message = (\n 'Could not find SurfaceOrientation/InputDeviceOrientation in dumpsys '\n 'output'\n )\n return response\n\n def _push(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Uploads contents to the device.\n\n Args:\n request: The request with the contents to push to the device.\n timeout: Optional time limit in seconds.\n\n Returns:\n An empty AdbResponse.\n \"\"\"\n\n path = request.push.path\n if not path:\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.FAILED_PRECONDITION,\n error_message='Push.path is empty.')\n\n # Create temporary file with `push` contents.\n with tempfile.NamedTemporaryFile(delete=False) as f:\n fname = f.name\n f.write(request.push.content)\n # Issue `adb push` command to upload file.\n logging.info('Uploading %r to %r.', fname, path)\n response, _ = self._execute_command(['push', fname, path], timeout=timeout)\n # Delete it.\n os.remove(fname)\n\n return response\n\n def _pull(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Downloads file content from the device.\n\n Args:\n request: The request with the information on what to get from the device.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse with the contents of the specified file.\n \"\"\"\n\n path = request.pull.path\n if not path:\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.FAILED_PRECONDITION,\n error_message='Pull.path is empty.')\n\n # Issue `adb pull` command to copy it to a temporary file.\n with tempfile.NamedTemporaryFile(delete=False) as f:\n fname = f.name\n logging.debug('Downloading %r to %r.', path, fname)\n response, _ = self._execute_command(['pull', path, fname],\n timeout=timeout)\n # Read the content of the file.\n with open(fname, 'rb') as f:\n response.pull.content = f.read()\n # Delete it.\n os.remove(fname)\n\n return response\n\n def _input_text(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Inserts text as keyboard events.\n\n Args:\n request: The external request.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse\n \"\"\"\n\n text = request.input_text.text\n if not text:\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.FAILED_PRECONDITION,\n error_message='InputText.text is empty.')\n\n response, _ = self._execute_command(['shell', 'input', 'text', text],\n timeout=timeout)\n return response\n\n def _tap(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Taps the device screen.\n\n Args:\n request: The request with information on where to tap the screen.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse\n \"\"\"\n\n x = request.tap.x\n y = request.tap.y\n # Check for negative coordinates.\n # Notice that zero coordinates are valid coordinates (i.e. the first\n # column/row of the screen).\n if x < 0 or y < 0:\n return adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.FAILED_PRECONDITION,\n error_message=(\n f'Tap coordinates must be non-negative. Got: {request.tap}.'))\n\n response, _ = self._execute_command(\n ['shell', 'input', 'tap', str(x),\n str(y)], timeout=timeout)\n\n return response\n\n def _handle_settings(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Handles SettingsRequest messages.\n\n Args:\n request: The specification of what to do with settings.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse\n \"\"\"\n\n request = request.settings\n response = adb_pb2.AdbResponse()\n # Every SettingsRequest should have a namespace.\n if (\n request.name_space\n == adb_pb2.AdbRequest.SettingsRequest.Namespace.UNKNOWN\n ):\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Unknown SettingsRequest.name_space. Got: {request}.')\n return response\n\n namespace = adb_pb2.AdbRequest.SettingsRequest.Namespace.Name(\n request.name_space).lower()\n\n match request.WhichOneof('verb'):\n case 'get':\n get = request.get\n if not get.key:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Empty SettingsRequest.get.key. Got: {request}.'\n )\n return response\n response, command_output = self._execute_command(\n ['shell', 'settings', 'get', namespace, get.key], timeout=timeout\n )\n response.settings.output = command_output\n case 'put':\n put = request.put\n if not put.key or not put.value:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Empty SettingsRequest.put key or value. Got: {request}.'\n )\n return response\n response, command_output = self._execute_command(\n ['shell', 'settings', 'put', namespace, put.key, put.value],\n timeout=timeout,\n )\n response.settings.output = command_output\n case 'delete_key':\n delete = request.delete_key\n if not delete.key:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Empty SettingsRequest.delete_key.key. Got: {request}.'\n )\n return response\n response, command_output = self._execute_command(\n ['shell', 'settings', 'delete', namespace, delete.key],\n timeout=timeout,\n )\n response.settings.output = command_output\n case 'reset':\n reset = request.reset\n # At least one of `package_name` or `mode` should be given.\n if (\n not reset.package_name\n and reset.mode\n == adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNKNOWN\n ):\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n 'At least one of SettingsRequest.reset package_name or mode'\n f' should be given. Got: {request}.'\n )\n return response\n\n mode = adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.Name(\n reset.mode\n ).lower()\n arg = reset.package_name or mode\n response, command_output = self._execute_command(\n ['shell', 'settings', 'reset', namespace, arg], timeout=timeout\n )\n response.settings.output = command_output\n case 'list':\n response, command_output = self._execute_command(\n ['shell', 'settings', 'list', namespace], timeout=timeout\n )\n response.settings.output = command_output\n case _:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Unknown SettingsRequest.verb. Got: {request}.'\n )\n\n return response\n\n def _handle_generic(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Handles GenericRequest messages.\n\n Args:\n request: The request with the `.generic` field set indicating what `adb`\n shell command to issue\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse\n \"\"\"\n\n response, command_output = self._execute_command(\n list(request.generic.args), timeout)\n response.generic.output = command_output\n return response\n\n def _handle_package_manager(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Handles PackageManagerRequest messages.\n\n Args:\n request: The request with the `.package_manager` field set containing the\n sub-commands to issue to `adb pm`.\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n request = request.package_manager\n response = adb_pb2.AdbResponse()\n\n match request.WhichOneof('verb'):\n case 'list':\n what = request.list.WhichOneof('what')\n response, output = self._execute_command(\n ['shell', 'pm', 'list', what], timeout=timeout\n )\n\n if output:\n items = output.decode('utf-8').split()\n # Remove prefix for each item.\n prefix = {\n 'features': 'feature:',\n 'libraries': 'library:',\n 'packages': 'package:',\n }[what]\n items = [x[len(prefix) :] for x in items if x.startswith(prefix)]\n response.package_manager.list.items.extend(items)\n response.package_manager.output = output\n case 'clear':\n package_name = request.clear.package_name\n if not package_name:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n f'Empty PackageManagerRequest.clear.package_name. Got: {request}.'\n )\n return response\n\n args = ['shell', 'pm', 'clear', package_name]\n if request.clear.user_id:\n args.insert(3, '-f')\n args.insert(4, request.clear.user_id)\n response, response.package_manager.output = self._execute_command(\n args, timeout=timeout\n )\n case 'grant':\n grant = request.grant\n if not grant.package_name:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = '`grant.package_name` cannot be empty.'\n return response\n\n if not grant.permissions:\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = '`grant.permissions` cannot be empty.'\n return response\n\n for permission in grant.permissions:\n logging.info('Granting permission: %r', permission)\n response, response.package_manager.output = self._execute_command(\n ['shell', 'pm', 'grant', grant.package_name, permission],\n timeout=timeout,\n )\n\n return response\n\n def _handle_dumpsys(\n self, request: adb_pb2.AdbRequest, timeout: float | None = None\n ) -> adb_pb2.AdbResponse:\n \"\"\"Handles DumpsysRequest messages.\n\n Args:\n request: The request with the `.dumpsys` field set containing\n sub-commands to `adb dumpsys` shell command..\n timeout: Optional time limit in seconds.\n\n Returns:\n An AdbResponse.\n \"\"\"\n\n request = request.dumpsys\n cmd = ['shell', 'dumpsys']\n\n if request.timeout_sec < 0 or request.timeout_ms < 0:\n response = adb_pb2.AdbResponse()\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n 'DumpsysRequest.timeout_{sec, ms} should be non-negative. '\n f'Got: {request}.')\n return response\n\n if request.list_only:\n # `-l` cannot be combined with the following options.\n if request.service or request.args or request.skip_services:\n response = adb_pb2.AdbResponse()\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n 'DumpsysRequest.list_only cannot be combined with other options. '\n f'Got: {request}.')\n return response\n\n cmd.append('-l')\n\n if request.timeout_sec > 0:\n cmd.append('-t')\n cmd.append(str(request.timeout_sec))\n elif request.timeout_ms > 0:\n cmd.append('-T')\n cmd.append(str(request.timeout_ms))\n\n if (\n request.priority\n != adb_pb2.AdbRequest.DumpsysRequest.PriorityLevel.UNSET\n ):\n cmd.append('--priority')\n cmd.append(adb_pb2.AdbRequest.DumpsysRequest.PriorityLevel.Name(\n request.priority))\n\n if request.skip_services:\n if request.service:\n response = adb_pb2.AdbResponse()\n response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n response.error_message = (\n 'DumpsysRequest.skip_services cannot be combined with `service`. '\n f'Got: {request}.')\n return response\n\n cmd.append('--skip')\n cmd.append(','.join(request.skip_services))\n\n if request.service:\n cmd.append(request.service)\n\n if request.args:\n cmd += list(request.args)\n\n if request.proto:\n cmd.append('--proto')\n\n response, response.dumpsys.output = self._execute_command(\n cmd, timeout=timeout)\n\n return response\n"
},
{
"path": "android_env/components/adb_call_parser_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport builtins\nimport os\nimport subprocess\nimport sys\nimport tempfile\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import adb_call_parser\nfrom android_env.components import adb_controller\nfrom android_env.proto import adb_pb2\n\n\nclass AdbCallParserTest(parameterized.TestCase):\n\n def test_unknown_command(self):\n \"\"\"Gets UNKNOWN_COMMAND for an empty request.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n response = parser.parse(request)\n self.assertEqual(\n response.status, adb_pb2.AdbResponse.Status.UNKNOWN_COMMAND\n )\n\n def test_invalid_timeout(self):\n \"\"\"AdbRequest.timeout_sec must be positive.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.tap.x = 123\n request.timeout_sec = -5\n response = parser.parse(request)\n self.assertEqual(\n response.status, adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n )\n\n @mock.patch.object(os.path, 'exists', autospec=True)\n def test_install_apk_file_not_found(self, mock_exists):\n \"\"\"Should fail installing APK when it is not found.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.install_apk.filesystem.path = '/my/home/game.apk'\n mock_exists.return_value = False\n\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n @mock.patch.object(os.path, 'exists', autospec=True)\n def test_install_apk_successful(self, mock_exists):\n \"\"\"Should succeed installing an arbitrary APK.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.install_apk.filesystem.path = '/my/home/game.apk'\n mock_exists.return_value = True\n\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['install', '-r', '-t', '-g', '/my/home/game.apk'], None)\n\n @mock.patch.object(tempfile, 'NamedTemporaryFile', autospec=True)\n def test_install_apk_from_blob(self, mock_tempfile):\n \"\"\"Should succeed installing APK from blob.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n blob_content = b'A fake blob content'\n request.install_apk.blob.contents = blob_content\n mock_tempfile.return_value.__enter__.return_value.name = '/my/home/test.apk'\n mock_tempfile.return_value.__enter__.return_value.write.return_value = None\n\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['install', '-r', '-t', '-g', '/my/home/test.apk'], None\n )\n # pytype: disable=attribute-error\n expected_tempfile_kwargs = (\n {'suffix': '.apk', 'delete_on_close': False}\n if sys.version_info > (3, 12)\n else {'suffix': '.apk'}\n )\n mock_tempfile.assert_has_calls([\n mock.call(**expected_tempfile_kwargs), # Constructor\n mock.call().__enter__(), # Enter context\n mock.call().__enter__().write(blob_content), # Call write function\n mock.call().__exit__(None, None, None), # Exit context\n ])\n # pytype: enable=attribute-error\n\n def test_start_activity_empty_full_activity(self):\n \"\"\"A start_activity command should always have a nonempty activity.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_activity.extra_args.extend(['blah'])\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n\n def test_start_activity_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n command_output = (b'Stopping: my.project.SplashActivity\\n'\n b'Starting: Intent { cmp=my.project.SplashActivity }\\n')\n adb.execute_command.return_value = command_output\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_activity.full_activity = 'my.project.SplashActivity'\n request.start_activity.extra_args.extend(['blah'])\n request.start_activity.force_stop = True\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call([\n 'shell', 'am', 'start', '-S', '-W', '-n',\n 'my.project.SplashActivity', 'blah'\n ],\n timeout=None),\n ])\n\n def test_start_activity_successful_no_force_stop(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n command_output = (b'Stopping: my.project.SplashActivity\\n'\n b'Starting: Intent { cmp=my.project.SplashActivity }\\n')\n adb.execute_command.return_value = command_output\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_activity.full_activity = 'my.project.SplashActivity'\n request.start_activity.extra_args.extend(['blah'])\n request.start_activity.force_stop = False\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call([\n 'shell', 'am', 'start', '', '-W', '-n', 'my.project.SplashActivity',\n 'blah'\n ],\n timeout=None),\n ])\n\n def test_start_activity_error(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n command_output = (b'Stopping: my.project.SplashActivity\\n'\n b'Starting: Intent { cmp=my.project.SplashActivity }\\n'\n b'Error: Activity not started, unknown error code 101\\n')\n adb.execute_command.return_value = command_output\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_activity.full_activity = 'my.project.SplashActivity'\n request.start_activity.extra_args.extend(['blah'])\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertEqual(\n response.error_message,\n f'start_activity failed with error: {str(command_output)}')\n\n def test_force_stop(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.force_stop.package_name = 'my.project'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'am', 'force-stop', 'my.project'], None)\n\n def test_grant_permissions_empty_package_name(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.package_manager.grant.permissions.extend(['perm1', 'perm2'])\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n\n def test_grant_permissions_empty_permissions(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.package_manager.grant.package_name = 'my.project'\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n\n def test_grant_permissions_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.package_manager.grant.package_name = 'my.project'\n request.package_manager.grant.permissions.extend(['perm1', 'perm2'])\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call(['shell', 'pm', 'grant', 'my.project', 'perm1'], None),\n mock.call(['shell', 'pm', 'grant', 'my.project', 'perm2'], None),\n ])\n\n def test_press_button_invalid_button(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.press_button.button = 99999\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n\n def test_press_button_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b''\n parser = adb_call_parser.AdbCallParser(adb)\n # HOME.\n request = adb_pb2.AdbRequest()\n request.press_button.button = adb_pb2.AdbRequest.PressButton.Button.HOME\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_with(\n ['shell', 'input', 'keyevent', 'KEYCODE_HOME'], None)\n # BACK.\n request = adb_pb2.AdbRequest()\n request.press_button.button = adb_pb2.AdbRequest.PressButton.Button.BACK\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_with(\n ['shell', 'input', 'keyevent', 'KEYCODE_BACK'], None)\n # ENTER.\n request = adb_pb2.AdbRequest()\n request.press_button.button = adb_pb2.AdbRequest.PressButton.Button.ENTER\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_with(\n ['shell', 'input', 'keyevent', 'KEYCODE_ENTER'], None)\n\n def test_start_screen_pinning_package_not_found(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = (\n b' taskId=12345: my.project.AnotherActivity visible=true'\n b' topActivity=ComponentInfo{my.project.AnotherActivity}')\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_screen_pinning.full_activity = 'my.project.AmazingActivity'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'am', 'stack', 'list'], None)\n\n def test_start_screen_pinning_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = (\n b' taskId=12345: my.project.AmazingActivity visible=true'\n b' topActivity=ComponentInfo{my.project.AmazingActivity}')\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_screen_pinning.full_activity = 'my.project.AmazingActivity'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call(['shell', 'am', 'stack', 'list'], None),\n mock.call(['shell', 'am', 'task', 'lock', '12345'], None),\n ])\n\n def test_start_screen_pinning_base_activity(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = (\n b' taskId=12345: my.project.MainActivity visible=true'\n b' topActivity=ComponentInfo{my.project.TopActivity}')\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_screen_pinning.full_activity = 'my.project.MainActivity'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call(['shell', 'am', 'stack', 'list'], None),\n mock.call(['shell', 'am', 'task', 'lock', '12345'], None),\n ])\n\n def test_start_screen_pinning_top_activity(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = (\n b' taskId=12345: my.project.MainActivity visible=true'\n b' topActivity=ComponentInfo{my.project.TopActivity}')\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.start_screen_pinning.full_activity = 'my.project.TopActivity'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call(['shell', 'am', 'stack', 'list'], None),\n mock.call(['shell', 'am', 'task', 'lock', '12345'], None),\n ])\n\n def test_send_broadcast_empty_action(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n send_broadcast=adb_pb2.AdbRequest.SendBroadcast())\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n\n def test_send_broadcast_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.send_broadcast.action = 'SOME-ACTION'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n\n def test_send_broadcast_with_component_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.send_broadcast.action = 'SOME-ACTION'\n request.send_broadcast.component = 'SOME-COMPONENT'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n\n def test_uninstall_package_empty_package_name(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.uninstall_package.package_name = ''\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n\n def test_uninstall_package_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'package:my.package'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest()\n request.uninstall_package.package_name = 'my.package'\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n\n def test_get_current_activity_no_visible_task(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = None\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n get_current_activity=adb_pb2.AdbRequest.GetCurrentActivity())\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call(\n ['shell', 'am', 'stack', 'list', '|', 'grep', '-E', 'visible=true'],\n None),\n mock.call(['shell', 'am', 'stack', 'list'], None),\n ])\n\n def test_get_orientation_empty_dumpsys(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b''\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n get_orientation=adb_pb2.AdbRequest.GetOrientationRequest())\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(['shell', 'dumpsys', 'input'],\n None)\n\n def test_get_orientation_invalid_device_no_surface_orientation(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b' PhysicalWidth: -123px'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n get_orientation=adb_pb2.AdbRequest.GetOrientationRequest())\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(['shell', 'dumpsys', 'input'],\n None)\n\n @parameterized.named_parameters(\n ('rotation_0', b\"\"\" SurfaceOrientation: 0\"\"\", 0),\n ('rotation_90', b\"\"\" SurfaceOrientation: 1\"\"\", 1),\n ('rotation_180', b\"\"\" SurfaceOrientation: 2\"\"\", 2),\n ('rotation_270', b\"\"\" SurfaceOrientation: 3\"\"\", 3),\n ('rotation_0_new', b\"\"\" InputDeviceOrientation: 0\"\"\", 0),\n ('rotation_90_new', b\"\"\" InputDeviceOrientation: 1\"\"\", 1),\n ('rotation_180_new', b\"\"\" InputDeviceOrientation: 2\"\"\", 2),\n ('rotation_270_new', b\"\"\" InputDeviceOrientation: 3\"\"\", 3),\n )\n def test_get_orientation_success(\n self, orientation: bytes, expected_orientation: int\n ):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = (\n b\"\"\"SomeRandomKey: 12345\\n\"\"\" + orientation + b\"\"\"\n MoreRandomStuff: awesome_value\n\"\"\"\n )\n\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n get_orientation=adb_pb2.AdbRequest.GetOrientationRequest())\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.get_orientation.orientation, expected_orientation)\n adb.execute_command.assert_called_once_with(['shell', 'dumpsys', 'input'],\n None)\n\n def test_get_current_activity_no_matches(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n get_current_activity=adb_pb2.AdbRequest.GetCurrentActivity())\n for platform in ['win32', 'linux']:\n with mock.patch.object(\n sys, 'platform', autospec=True, return_value=platform):\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.INTERNAL_ERROR)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_has_calls([\n mock.call([\n 'shell', 'am', 'stack', 'list', '|', 'grep', '-E',\n 'visible=true'\n ], None),\n mock.call(['shell', 'am', 'stack', 'list'], None),\n ])\n\n def test_get_current_activity_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'{MyAwesomeActivity}'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n get_current_activity=adb_pb2.AdbRequest.GetCurrentActivity())\n for platform in ['win32', 'linux']:\n with mock.patch.object(\n sys, 'platform', autospec=True, return_value=platform):\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n # `execute_command` will be called once for each platform.\n adb.execute_command.assert_called_with(\n ['shell', 'am', 'stack', 'list', '|', 'grep', '-E', 'visible=true'],\n None)\n self.assertEqual(response.get_current_activity.full_activity,\n 'MyAwesomeActivity')\n\n def test_push_no_path(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n push=adb_pb2.AdbRequest.Push(content=b'Has content but no path'))\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_push_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n push=adb_pb2.AdbRequest.Push(\n content=b'My text.', path='/sdcard/my_file.txt'))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once()\n args, kwargs = adb.execute_command.call_args\n self.assertLen(args, 1)\n cmd_args = args[0]\n self.assertLen(cmd_args, 3)\n self.assertEqual(cmd_args[0], 'push')\n self.assertEqual(cmd_args[2], '/sdcard/my_file.txt')\n self.assertIn('timeout', kwargs)\n self.assertIsNone(kwargs['timeout'])\n\n def test_pull_no_path(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(pull=adb_pb2.AdbRequest.Pull())\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n @mock.patch.object(builtins, 'open', autospec=True)\n def test_pull_successful(self, mock_open):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n mock_open.return_value.__enter__ = mock_open\n mock_open.return_value.read.return_value = b'S3cR3t. dO nOt TeLl ANYONE'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n pull=adb_pb2.AdbRequest.Pull(path='/sdcard/my_file.txt'))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.pull.content, b'S3cR3t. dO nOt TeLl ANYONE')\n adb.execute_command.assert_called_once()\n args, kwargs = adb.execute_command.call_args\n self.assertLen(args, 1)\n cmd_args = args[0]\n self.assertLen(cmd_args, 3)\n self.assertEqual(cmd_args[0], 'pull')\n self.assertEqual(cmd_args[1], '/sdcard/my_file.txt')\n self.assertIn('timeout', kwargs)\n self.assertIsNone(kwargs['timeout'])\n\n def test_input_text_no_text(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(input_text=adb_pb2.AdbRequest.InputText())\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_input_text_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n input_text=adb_pb2.AdbRequest.InputText(\n text='The Greatest Text of All Time'))\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'input', 'text', 'The Greatest Text of All Time'], None)\n\n @parameterized.named_parameters(\n ('negative_x_and_negative_y',\n adb_pb2.AdbRequest(tap=adb_pb2.AdbRequest.Tap(x=-1, y=-1))),\n ('negative_x',\n adb_pb2.AdbRequest(tap=adb_pb2.AdbRequest.Tap(x=-1, y=123))),\n ('negative_y',\n adb_pb2.AdbRequest(tap=adb_pb2.AdbRequest.Tap(x=456, y=-1))),\n )\n def test_tap_failed(self, request: adb_pb2.AdbRequest):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_tap_successful(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(tap=adb_pb2.AdbRequest.Tap(x=135, y=246))\n response = parser.parse(request)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'input', 'tap', '135', '246'], None)\n\n @parameterized.named_parameters(\n ('empty_request', adb_pb2.AdbRequest.SettingsRequest()),\n ('no_namespace',\n adb_pb2.AdbRequest.SettingsRequest(\n get=adb_pb2.AdbRequest.SettingsRequest.Get(key='my_key'))),\n ('get_no_key',\n adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n get=adb_pb2.AdbRequest.SettingsRequest.Get())),\n ('put_no_key',\n adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n put=adb_pb2.AdbRequest.SettingsRequest.Put())),\n ('put_no_value',\n adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(key='another_key'))),\n ('delete_no_key',\n adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n delete_key=adb_pb2.AdbRequest.SettingsRequest.Delete())),\n ('reset_no_package_name_and_no_mode',\n adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n reset=adb_pb2.AdbRequest.SettingsRequest.Reset())),\n )\n def test_settings_failures(self, request):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(settings=request)\n response = parser.parse(request)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_settings_success_get(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'here it is!'\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n get=adb_pb2.AdbRequest.SettingsRequest.Get(key='some_key'))\n request = adb_pb2.AdbRequest(settings=request)\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.settings.output, b'here it is!')\n adb.execute_command.assert_called_once_with(\n ['shell', 'settings', 'get', 'system', 'some_key'], None)\n\n def test_settings_success_put(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'Done for ya!'\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SECURE,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(key='key1', value='val2'))\n request = adb_pb2.AdbRequest(settings=request)\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.settings.output, b'Done for ya!')\n adb.execute_command.assert_called_once_with(\n ['shell', 'settings', 'put', 'secure', 'key1', 'val2'], None)\n\n def test_settings_success_delete(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'Key deleted.'\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,\n delete_key=adb_pb2.AdbRequest.SettingsRequest.Delete(key='useless_key'))\n request = adb_pb2.AdbRequest(settings=request)\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.settings.output, b'Key deleted.')\n adb.execute_command.assert_called_once_with(\n ['shell', 'settings', 'delete', 'global', 'useless_key'], None)\n\n @parameterized.named_parameters(\n ('mode_untrusted_defaults',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNTRUSTED_DEFAULTS, '',\n 'untrusted_defaults'),\n ('mode_untrusted_clear',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNTRUSTED_CLEAR, '',\n 'untrusted_clear'),\n ('mode_trusted_defaults',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.TRUSTED_DEFAULTS, '',\n 'trusted_defaults'),\n # If `package_name` is given, it takes precedence over `mode`.\n ('mode_unknown_package_given',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNKNOWN, 'great.package',\n 'great.package'),\n ('mode_untrusted_defaults_package_given',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNTRUSTED_DEFAULTS,\n 'great.package', 'great.package'),\n ('mode_untrusted_clear_package_given',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNTRUSTED_CLEAR,\n 'great.package', 'great.package'),\n ('mode_trusted_defaults_package_given',\n adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.TRUSTED_DEFAULTS,\n 'great.package', 'great.package'),\n )\n def test_settings_success_reset(self, mode, package_name, expected_arg):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'Pkg reset.'\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,\n reset=adb_pb2.AdbRequest.SettingsRequest.Reset(\n package_name=package_name, mode=mode))\n request = adb_pb2.AdbRequest(settings=request)\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.settings.output, b'Pkg reset.')\n adb.execute_command.assert_called_once_with(\n ['shell', 'settings', 'reset', 'global', expected_arg], None)\n\n def test_settings_success_list(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'volume_ring=5\\nvolume_system=7'\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n list=adb_pb2.AdbRequest.SettingsRequest.List())\n request = adb_pb2.AdbRequest(settings=request)\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.settings.output,\n b'volume_ring=5\\nvolume_system=7')\n adb.execute_command.assert_called_once_with(\n ['shell', 'settings', 'list', 'system'], None)\n\n def test_generic_command(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n expected_output = b'generic_output'\n args = ['shell', 'am', 'broadcast', '-n', 'receiver', '-a', 'action']\n adb.execute_command.return_value = expected_output\n parser = adb_call_parser.AdbCallParser(adb)\n\n generic_request = adb_pb2.AdbRequest.GenericRequest(args=args)\n request = adb_pb2.AdbRequest(generic=generic_request)\n response = parser.parse(request)\n\n self.assertEqual(adb_pb2.AdbResponse.Status.OK, response.status)\n self.assertEmpty(response.error_message)\n self.assertEqual(response.generic.output, expected_output)\n adb.execute_command.assert_called_once_with(args, None)\n\n def test_generic_command_adb_error(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n args = ['shell', 'am', 'broadcast', '-n', 'receiver', '-a', 'action']\n adb.execute_command.side_effect = subprocess.CalledProcessError(\n cmd='cmd', output='adb_error', returncode=-1)\n parser = adb_call_parser.AdbCallParser(adb)\n\n generic_request = adb_pb2.AdbRequest.GenericRequest(args=args)\n request = adb_pb2.AdbRequest(generic=generic_request)\n response = parser.parse(request)\n\n self.assertEqual(adb_pb2.AdbResponse.Status.ADB_ERROR, response.status)\n self.assertEqual('adb_error', response.error_message)\n self.assertEmpty(response.generic.output)\n adb.execute_command.assert_called_once_with(args, None)\n\n def test_generic_command_timeout(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n args = ['shell', 'am', 'broadcast', '-n', 'receiver', '-a', 'action']\n adb.execute_command.side_effect = subprocess.TimeoutExpired(\n cmd='cmd', timeout=10)\n parser = adb_call_parser.AdbCallParser(adb)\n\n generic_request = adb_pb2.AdbRequest.GenericRequest(args=args)\n request = adb_pb2.AdbRequest(generic=generic_request)\n response = parser.parse(request)\n\n self.assertEqual(adb_pb2.AdbResponse.Status.TIMEOUT, response.status)\n self.assertEqual('Timeout', response.error_message)\n self.assertEmpty(response.generic.output)\n adb.execute_command.assert_called_once_with(args, None)\n\n @parameterized.named_parameters(\n ('features',\n adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n features=adb_pb2.AdbRequest.PackageManagerRequest.List\n .Features())))),\n ('libraries',\n adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n libraries=adb_pb2.AdbRequest.PackageManagerRequest.List\n .Libraries())))),\n ('packages',\n adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n packages=adb_pb2.AdbRequest.PackageManagerRequest.List\n .Packages())))),\n )\n def test_package_manager_list_bad_output(self, request):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b\"\"\"Something irrelevant.\"\"\"\n parser = adb_call_parser.AdbCallParser(adb)\n response = parser.parse(request)\n response.package_manager.output = b\"\"\"Something irrelevant.\"\"\"\n self.assertEmpty(response.package_manager.list.items)\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once()\n\n def test_package_manager_list_features(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n output = b\"\"\"\nfeature:android.hardware.audio.output\nfeature:android.hardware.bluetooth\nfeature:android.hardware.camera\nfeature:android.hardware.fingerprint\nfeature:android.software.autofill\nfeature:android.software.backup\nfeature:android.software.webview\n\"\"\"\n adb.execute_command.return_value = output\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n features=adb_pb2.AdbRequest.PackageManagerRequest.List.Features(\n ))))\n response = parser.parse(request)\n self.assertEqual(response.package_manager.output, output)\n self.assertEqual(response.package_manager.list.items, [\n 'android.hardware.audio.output',\n 'android.hardware.bluetooth',\n 'android.hardware.camera',\n 'android.hardware.fingerprint',\n 'android.software.autofill',\n 'android.software.backup',\n 'android.software.webview',\n ])\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'pm', 'list', 'features'], None)\n\n def test_package_manager_list_libraries(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n output = b\"\"\"\nlibrary:android.ext.shared\nlibrary:android.hidl.base-V1.0-java\nlibrary:android.hidl.manager-V1.0-java\nlibrary:android.net.ipsec.ike\nlibrary:android.test.base\nlibrary:android.test.mock\nlibrary:android.test.runner\nlibrary:androidx.window.sidecar\nlibrary:com.android.future.usb.accessory\nlibrary:com.android.location.provider\nlibrary:com.android.media.remotedisplay\nlibrary:com.android.mediadrm.signer\nlibrary:com.android.nfc_extras\nlibrary:com.google.android.gms\nlibrary:com.google.android.trichromelibrary\nlibrary:javax.obex\nlibrary:org.apache.http.legacy\n\"\"\"\n adb.execute_command.return_value = output\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n libraries=adb_pb2.AdbRequest.PackageManagerRequest.List\n .Libraries())))\n response = parser.parse(request)\n self.assertEqual(response.package_manager.output, output)\n self.assertEqual(response.package_manager.list.items, [\n 'android.ext.shared',\n 'android.hidl.base-V1.0-java',\n 'android.hidl.manager-V1.0-java',\n 'android.net.ipsec.ike',\n 'android.test.base',\n 'android.test.mock',\n 'android.test.runner',\n 'androidx.window.sidecar',\n 'com.android.future.usb.accessory',\n 'com.android.location.provider',\n 'com.android.media.remotedisplay',\n 'com.android.mediadrm.signer',\n 'com.android.nfc_extras',\n 'com.google.android.gms',\n 'com.google.android.trichromelibrary',\n 'javax.obex',\n 'org.apache.http.legacy',\n ])\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'pm', 'list', 'libraries'], None)\n\n def test_package_manager_list_packages(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n output = b\"\"\"\npackage:com.android.phone\npackage:com.awesome.company\npackage:com.another.great.thingie\n\"\"\"\n adb.execute_command.return_value = output\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n packages=adb_pb2.AdbRequest.PackageManagerRequest.List.Packages(\n ))))\n response = parser.parse(request)\n self.assertEqual(response.package_manager.output, output)\n self.assertEqual(response.package_manager.list.items, [\n 'com.android.phone',\n 'com.awesome.company',\n 'com.another.great.thingie',\n ])\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'pm', 'list', 'packages'], None)\n\n def test_package_manager_clear_no_package_name(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b\"\"\"Something irrelevant.\"\"\"\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n clear=adb_pb2.AdbRequest.PackageManagerRequest.Clear(\n package_name='')))\n response = parser.parse(request)\n\n self.assertEmpty(response.package_manager.output)\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_package_manager_clear_successful_no_user_id(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b\"\"\"Some successful message.\"\"\"\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n clear=adb_pb2.AdbRequest.PackageManagerRequest.Clear(\n package_name='my.package')))\n response = parser.parse(request)\n\n self.assertEqual(response.package_manager.output,\n b\"\"\"Some successful message.\"\"\")\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'pm', 'clear', 'my.package'], None)\n\n def test_package_manager_clear_successful_with_user_id(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b\"\"\"Some successful message.\"\"\"\n parser = adb_call_parser.AdbCallParser(adb)\n\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n clear=adb_pb2.AdbRequest.PackageManagerRequest.Clear(\n package_name='my.package', user_id='mrawesome')))\n response = parser.parse(request)\n\n self.assertEqual(response.package_manager.output,\n b\"\"\"Some successful message.\"\"\")\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'pm', 'clear', '-f', 'mrawesome', 'my.package'], None)\n\n def test_dumpsys_empty_request(self):\n \"\"\"An empty `DumpsysRequest` is a valid request.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(dumpsys=adb_pb2.AdbRequest.DumpsysRequest())\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(['shell', 'dumpsys'],\n timeout=None)\n\n @parameterized.named_parameters(\n ('negative_timeout_sec',\n adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(timeout_sec=-1))),\n ('negative_timeout_ms',\n adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(timeout_ms=-2))),\n )\n def test_dumpsys_negative_timeouts(self, request):\n \"\"\"`DumpsysRequest.timeout_{sec, ms}` if passed, should be positive.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n parser = adb_call_parser.AdbCallParser(adb)\n\n response = parser.parse(request)\n\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n @parameterized.named_parameters(\n ('both_timeouts_zero', 0, 0, ['shell', 'dumpsys']),\n ('sec_takes_precedence_zero', 123, 0, ['shell', 'dumpsys', '-t', '123']),\n ('sec_takes_precedence', 123, 456, ['shell', 'dumpsys', '-t', '123']),\n ('ms_if_no_sec', 0, 456, ['shell', 'dumpsys', '-T', '456']),\n )\n def test_dumpsys_timeout_successful(self, timeout_sec, timeout_ms, expected):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(\n timeout_sec=timeout_sec, timeout_ms=timeout_ms))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(expected, timeout=None)\n\n @parameterized.named_parameters(\n ('priority_undefined',\n adb_pb2.AdbRequest.DumpsysRequest.PriorityLevel.UNSET,\n ['shell', 'dumpsys']),\n ('priority_normal',\n adb_pb2.AdbRequest.DumpsysRequest.PriorityLevel.NORMAL,\n ['shell', 'dumpsys', '--priority', 'NORMAL']),\n ('priority_high', adb_pb2.AdbRequest.DumpsysRequest.PriorityLevel.HIGH,\n ['shell', 'dumpsys', '--priority', 'HIGH']),\n ('priority_critical',\n adb_pb2.AdbRequest.DumpsysRequest.PriorityLevel.CRITICAL,\n ['shell', 'dumpsys', '--priority', 'CRITICAL']),\n )\n def test_dumpsys_priority_timeout_successful(self, priority, expected):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(priority=priority))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(expected, timeout=None)\n\n @parameterized.named_parameters(\n (\n 'window_service',\n adb_pb2.AdbRequest.DumpsysRequest(list_only=True, service='window'),\n ),\n (\n 'arbitrary_args',\n adb_pb2.AdbRequest.DumpsysRequest(\n list_only=True, args=['myoption', 'anotheroption']\n ),\n ),\n (\n 'skip_usb',\n adb_pb2.AdbRequest.DumpsysRequest(\n list_only=True, skip_services=['usb']\n ),\n ),\n )\n def test_dumpsys_list_only_cannot_be_combined(\n self, dumpsys_request: adb_pb2.AdbRequest.DumpsysRequest\n ):\n \"\"\"When `list_only==True`, the request cannot contain a few fields.\"\"\"\n\n # Arrange.\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(dumpsys=dumpsys_request)\n\n # Act.\n response = parser.parse(request)\n\n # Assert.\n self.assertEqual(\n response.status, adb_pb2.AdbResponse.Status.FAILED_PRECONDITION\n )\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_dumpsys_list_only_success(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(list_only=True))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(['shell', 'dumpsys', '-l'],\n timeout=None)\n\n def test_dumpsys_skip_services_cannot_combine_with_service(self):\n \"\"\"When using `DumpsysRequest.skip_service`, it cannot contain `.service`.\"\"\"\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(\n service='wifi', skip_services=['window', 'usb']))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status,\n adb_pb2.AdbResponse.Status.FAILED_PRECONDITION)\n self.assertNotEmpty(response.error_message)\n adb.execute_command.assert_not_called()\n\n def test_dumpsys_skip_services(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(\n skip_services=['window', 'usb']))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'dumpsys', '--skip', 'window,usb'], timeout=None)\n\n def test_dumpsys_single_service(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(service='window'))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(['shell', 'dumpsys', 'window'],\n timeout=None)\n\n def test_dumpsys_single_service_with_args(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'whatever'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(\n service='window', args=['arg1', 'arg2']))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'dumpsys', 'window', 'arg1', 'arg2'], timeout=None)\n\n def test_dumpsys_single_service_with_proto(self):\n adb = mock.create_autospec(adb_controller.AdbController)\n adb.execute_command.return_value = b'some binary output'\n parser = adb_call_parser.AdbCallParser(adb)\n request = adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(service='window', proto=True))\n\n response = parser.parse(request)\n\n self.assertEqual(response.status, adb_pb2.AdbResponse.Status.OK)\n self.assertEmpty(response.error_message)\n adb.execute_command.assert_called_once_with(\n ['shell', 'dumpsys', 'window', '--proto'], timeout=None)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/adb_controller.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A class to manage and control an external ADB process.\"\"\"\n\nimport os\nimport subprocess\nimport time\n\nfrom absl import logging\nfrom android_env.components import config_classes\nfrom android_env.components import errors\n\n\nclass AdbController:\n \"\"\"Manages communication with adb.\"\"\"\n\n def __init__(self, config: config_classes.AdbControllerConfig):\n \"\"\"Instantiates an AdbController object.\"\"\"\n\n self._config = config\n logging.info('config: %r', self._config)\n\n if not self._config.use_adb_server_port_from_os_env:\n # Unset problematic environment variables. ADB commands will fail if these\n # are set. They are normally exported by AndroidStudio.\n if 'ANDROID_HOME' in os.environ:\n logging.info('Removing ANDROID_HOME from os.environ')\n del os.environ['ANDROID_HOME']\n if 'ANDROID_ADB_SERVER_PORT' in os.environ:\n logging.info('Removing ANDROID_ADB_SERVER_PORT from os.environ')\n del os.environ['ANDROID_ADB_SERVER_PORT']\n\n # Explicitly expand the $HOME environment variable.\n self._os_env_vars = dict(os.environ).copy()\n self._os_env_vars.update(\n {'HOME': os.path.expandvars(self._os_env_vars.get('HOME', ''))}\n )\n logging.info('self._os_env_vars: %r', self._os_env_vars)\n\n def command_prefix(self, include_device_name: bool = True) -> list[str]:\n \"\"\"The command for instantiating an adb client to this server.\"\"\"\n if self._config.use_adb_server_port_from_os_env:\n # When using the adb server port set from the OS environment, we don't\n # need to pass the port explicitly.\n adb_port_args = []\n else:\n # When using the adb server port set from the config, we need to pass the\n # port explicitly.\n adb_port_args = ['-P', str(self._config.adb_server_port)]\n command_prefix = [\n self._config.adb_path,\n *adb_port_args,\n ]\n if include_device_name:\n command_prefix.extend(['-s', self._config.device_name])\n return command_prefix\n\n def init_server(self, timeout: float | None = None):\n \"\"\"Initialize the ADB server deamon on the given port.\n\n This function should be called immediately after initializing the first\n adb_controller, and before launching the simulator.\n\n Args:\n timeout: A timeout to use for this operation. If not set the default\n timeout set on the constructor will be used.\n \"\"\"\n # Make an initial device-independent call to ADB to start the deamon.\n self.execute_command(['devices'], timeout, device_specific=False)\n time.sleep(0.2)\n\n def _restart_server(self, timeout: float | None = None):\n \"\"\"Kills and restarts the adb server.\n\n Args:\n timeout: A timeout to use for this operation. If not set the default\n timeout set on the constructor will be used.\n \"\"\"\n logging.info('Restarting adb server.')\n self.execute_command(\n ['kill-server'], timeout=timeout, device_specific=False\n )\n time.sleep(0.2)\n cmd_output = self.execute_command(\n ['start-server'], timeout=timeout, device_specific=False\n )\n logging.info('start-server output: %r', cmd_output.decode('utf-8'))\n time.sleep(2.0)\n self.execute_command(['devices'], timeout=timeout, device_specific=False)\n time.sleep(0.2)\n\n def execute_command(\n self,\n args: list[str],\n timeout: float | None = None,\n device_specific: bool = True,\n ) -> bytes:\n \"\"\"Executes an adb command.\n\n Args:\n args: A list of strings representing each adb argument. For example:\n ['install', '/my/app.apk']\n timeout: A timeout to use for this operation. If not set the default\n timeout set on the constructor will be used.\n device_specific: Whether the call is device-specific or independent.\n\n Returns:\n The output of running such command as a binary string.\n \"\"\"\n timeout = self._config.default_timeout if timeout is None else timeout\n command = self.command_prefix(include_device_name=device_specific) + args\n command_str = 'adb ' + ' '.join(command[1:])\n\n n_tries = 2\n latest_error = None\n for i in range(n_tries):\n try:\n logging.info('Executing ADB command: [%s]', command_str)\n cmd_output = subprocess.check_output(\n command,\n stderr=subprocess.STDOUT,\n timeout=timeout,\n env=self._os_env_vars,\n )\n logging.debug('ADB command output: %s', cmd_output)\n return cmd_output\n except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:\n logging.exception(\n 'Failed to execute ADB command (try %d of %d): [%s]',\n i + 1,\n n_tries,\n command_str,\n )\n if e.stdout is not None:\n logging.error('**stdout**:')\n for line in e.stdout.splitlines():\n logging.error(' %s', line)\n if e.stderr is not None:\n logging.error('**stderr**:')\n for line in e.stderr.splitlines():\n logging.error(' %s', line)\n latest_error = e\n if device_specific and i < n_tries - 1:\n self._restart_server(timeout=timeout)\n\n raise errors.AdbControllerError(\n f'Error executing adb command: [{command_str}]\\n'\n f'Caused by: {latest_error}'\n ) from latest_error\n"
},
{
"path": "android_env/components/adb_controller_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport subprocess\nimport time\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import adb_controller as adb_controller_lib\nfrom android_env.components import config_classes\nfrom android_env.components import errors\n\n# Timeout to be used by default in tests below. Set to a small value to avoid\n# hanging on a failed test.\n_TIMEOUT = 2\n\n\nclass AdbControllerTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n # Set env vars.\n os.environ['MY_ENV_VAR'] = '/some/path/'\n os.environ['HOME'] = '$MY_ENV_VAR'\n self._env_before = os.environ.copy()\n\n def tearDown(self):\n super().tearDown()\n if 'ANDROID_HOME' in os.environ:\n del os.environ['ANDROID_HOME']\n if 'ANDROID_ADB_SERVER_PORT' in os.environ:\n del os.environ['ANDROID_ADB_SERVER_PORT']\n\n @mock.patch.object(subprocess, 'check_output', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_init_server(self, mock_sleep, mock_check_output):\n \"\"\"We expect an `adb devices` call when initializing the server.\"\"\"\n\n # Arrange.\n adb_controller = adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n use_adb_server_port_from_os_env=True,\n )\n )\n\n # Act.\n adb_controller.init_server(timeout=_TIMEOUT)\n\n # Assert.\n expected_env = self._env_before\n expected_env['HOME'] = '/some/path/'\n mock_check_output.assert_called_once_with(\n ['my_adb', 'devices'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n )\n mock_sleep.assert_called_once()\n\n @mock.patch.object(subprocess, 'check_output', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_init_server_with_adb_server_port_from_os_env(\n self, mock_sleep, mock_check_output\n ):\n \"\"\"Us OS env vars if `use_adb_server_port_from_os_env` is True.\"\"\"\n\n # Arrange.\n # Set the ADB server port to 1234 in the OS environment.\n os.environ['ANDROID_ADB_SERVER_PORT'] = '1234'\n os.environ['ANDROID_HOME'] = '/some/path/to/android'\n adb_controller = adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n adb_server_port=9999,\n use_adb_server_port_from_os_env=True,\n )\n )\n\n # Act.\n adb_controller.init_server(timeout=_TIMEOUT)\n\n # Assert.\n expected_env = self._env_before\n expected_env['HOME'] = '/some/path/'\n expected_env['ANDROID_HOME'] = '/some/path/to/android'\n expected_env['ANDROID_ADB_SERVER_PORT'] = '1234'\n\n mock_check_output.assert_called_once_with(\n ['my_adb', 'devices'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n )\n mock_sleep.assert_called_once()\n\n @mock.patch.object(subprocess, 'check_output', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_restart_server(self, mock_sleep, mock_check_output):\n \"\"\"When an adb command fails, we expect the server to be restarted.\"\"\"\n\n # Arrange.\n mock_check_output.side_effect = [\n subprocess.CalledProcessError(returncode=1, cmd='blah'),\n ] + ['fake_output'.encode('utf-8')] * 4\n adb_controller = adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n use_adb_server_port_from_os_env=True,\n )\n )\n\n # Act.\n adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)\n\n # Assert.\n expected_env = self._env_before\n expected_env['HOME'] = '/some/path/'\n mock_check_output.assert_has_calls([\n mock.call(\n ['my_adb', '-s', 'awesome_device', 'my_command'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', 'kill-server'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', 'start-server'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', 'devices'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', '-s', 'awesome_device', 'my_command'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n ])\n mock_sleep.assert_has_calls(\n [mock.call(0.2), mock.call(2.0), mock.call(0.2)]\n )\n\n @mock.patch.object(subprocess, 'check_output', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_invalid_command(self, mock_sleep, mock_check_output):\n \"\"\"Restart the server when given an invalid command.\"\"\"\n\n # Arrange.\n restart_sequence = ['fake_output'.encode('utf-8')] * 3\n mock_check_output.side_effect = (\n [\n subprocess.CalledProcessError(returncode=1, cmd='blah'),\n ]\n + restart_sequence\n + [subprocess.CalledProcessError(returncode=1, cmd='blah')]\n # Don't restart if last call fails.\n )\n adb_controller = adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n use_adb_server_port_from_os_env=True,\n )\n )\n\n # Act.\n with self.assertRaises(errors.AdbControllerError):\n adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)\n\n # Assert.\n expected_env = self._env_before\n expected_env['HOME'] = '/some/path/'\n mock_check_output.assert_has_calls(\n [\n mock.call(\n ['my_adb', '-s', 'awesome_device', 'my_command'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', 'kill-server'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', 'start-server'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', 'devices'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n mock.call(\n ['my_adb', '-s', 'awesome_device', 'my_command'],\n stderr=subprocess.STDOUT,\n timeout=_TIMEOUT,\n env=expected_env,\n ),\n ],\n any_order=False,\n )\n mock_sleep.assert_has_calls(\n [mock.call(0.2), mock.call(2.0), mock.call(0.2)]\n )\n\n @mock.patch.object(subprocess, 'check_output', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_avoid_infinite_recursion(self, mock_sleep, mock_check_output):\n \"\"\"Raise an error if the command fails even after restarts.\"\"\"\n\n del mock_sleep\n mock_check_output.side_effect = subprocess.CalledProcessError(\n returncode=1, cmd='blah'\n )\n adb_controller = adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n use_adb_server_port_from_os_env=True,\n )\n )\n self.assertRaises(\n errors.AdbControllerError,\n adb_controller.execute_command,\n ['my_command'],\n timeout=_TIMEOUT,\n )\n\n\nclass AdbControllerInitTest(absltest.TestCase):\n\n def test_deletes_problem_env_vars(self):\n os.environ['ANDROID_HOME'] = '/usr/local/Android/Sdk'\n os.environ['ANDROID_ADB_SERVER_PORT'] = '1337'\n adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n adb_server_port=9999,\n default_timeout=_TIMEOUT,\n )\n )\n self.assertNotIn('ANDROID_HOME', os.environ)\n self.assertNotIn('ANDROID_ADB_SERVER_PORT', os.environ)\n\n def test_use_adb_server_port_from_os_env_retains_os_env_vars(self):\n os.environ['ANDROID_HOME'] = '/usr/local/Android/Sdk'\n os.environ['ANDROID_ADB_SERVER_PORT'] = '1337'\n adb_controller_lib.AdbController(\n config_classes.AdbControllerConfig(\n adb_path='my_adb',\n device_name='awesome_device',\n adb_server_port=9999,\n default_timeout=_TIMEOUT,\n use_adb_server_port_from_os_env=True,\n )\n )\n self.assertIn('ANDROID_ADB_SERVER_PORT', os.environ)\n self.assertEqual(os.environ['ANDROID_ADB_SERVER_PORT'], '1337')\n self.assertIn('ANDROID_HOME', os.environ)\n self.assertEqual(os.environ['ANDROID_HOME'], '/usr/local/Android/Sdk')\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/adb_log_stream.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Class for a stream of logs output by a locally running emulator.\"\"\"\n\nimport subprocess\n\nfrom absl import logging\nfrom android_env.components import log_stream\n\n\n_LOGCAT_COMMAND = ['logcat', '-v', 'epoch']\n\n\nclass AdbLogStream(log_stream.LogStream):\n \"\"\"Manages adb logcat process for a locally running emulator.\"\"\"\n\n def __init__(self, adb_command_prefix: list[str], verbose: bool = False):\n super().__init__(verbose=verbose)\n self._adb_command_prefix = adb_command_prefix\n\n def _get_stream_output(self):\n\n # Before spawning a long-lived process, we issue `logcat -b all -c` to clear\n # all buffers to avoid interference from previous runs.\n clear_buffer_output = subprocess.check_output(\n self._adb_command_prefix + ['logcat', '-b', 'all', '-c'],\n stderr=subprocess.STDOUT,\n timeout=100)\n logging.info('clear_buffer_output: %r', clear_buffer_output)\n cmd = self._adb_command_prefix + _LOGCAT_COMMAND + self._filters\n self._adb_subprocess = subprocess.Popen(\n cmd,\n stdout=subprocess.PIPE,\n stderr=subprocess.STDOUT,\n bufsize=1,\n universal_newlines=True)\n return self._adb_subprocess.stdout\n\n def stop_stream(self):\n if not hasattr(self, '_adb_subprocess') or self._adb_subprocess is None:\n logging.error('`stop_stream()` called before `get_stream_output()`. '\n 'This violates the `LogStream` API.')\n else:\n self._adb_subprocess.kill()\n"
},
{
"path": "android_env/components/adb_log_stream_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for adb_log_stream.\"\"\"\n\nimport subprocess\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import adb_log_stream\n\n\nclass FakeAdbSubprocess:\n\n @property\n def stdout(self):\n return [f'line_{i}' for i in range(100)]\n\n def kill(self):\n pass\n\n\nclass AdbLogStreamTest(absltest.TestCase):\n\n @mock.patch.object(subprocess, 'check_output', return_value=b'')\n @mock.patch.object(subprocess, 'Popen', return_value=FakeAdbSubprocess())\n def test_get_stream_output(self, mock_popen, unused_mock_check_output):\n stream = adb_log_stream.AdbLogStream(adb_command_prefix=['foo'])\n stream.set_log_filters(['bar'])\n stream_output = stream.get_stream_output()\n\n for i, line in enumerate(stream_output):\n self.assertEqual(line, f'line_{i}')\n\n mock_popen.assert_called_with(\n ['foo', 'logcat', '-v', 'epoch', 'bar', '*:S'],\n stderr=subprocess.STDOUT,\n stdout=subprocess.PIPE,\n bufsize=1,\n universal_newlines=True)\n\n def test_stop_stream_before_get_stream_output(self):\n \"\"\"Calling `stop_stream()` before `get_stream_output()` should not crash.\"\"\"\n\n # Arrange.\n stream = adb_log_stream.AdbLogStream(adb_command_prefix=['foo'])\n\n # Act.\n stream.stop_stream()\n\n # Assert.\n # Nothing to assert. The test should just finish without raising an\n # exception.\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/app_screen_checker.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Determines if the current app screen matches an expected app screen.\"\"\"\n\nfrom collections.abc import Callable, Sequence\nimport enum\nimport re\nimport time\nfrom typing import Self\n\nfrom absl import logging\nfrom android_env.components import adb_call_parser as adb_call_parser_lib\nfrom android_env.components import errors\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import task_pb2\n\n\nclass _DumpsysNode:\n \"\"\"A node in a dumpsys tree.\"\"\"\n\n def __init__(self, data: str):\n self._children = []\n self._data = data\n\n @property\n def data(self) -> str:\n return self._data\n\n @property\n def children(self) -> list[Self]:\n return self._children\n\n def find_child(\n self, predicate: Callable[[Self], bool], max_levels: int = 0\n ) -> Self | None:\n \"\"\"Returns the first direct child that matches `predicate`, None otherwise.\n\n Args:\n predicate: Function-like that accepts a _DumpsysNode and returns boolean.\n max_levels: Maximum number of levels down the tree to search for a child.\n If non-positive, only direct children will be searched for.\n\n Returns:\n A _DumpsysNode or None.\n \"\"\"\n if not self.children:\n return None\n\n try:\n return next(x for x in self.children if predicate(x))\n except StopIteration:\n logging.info('Failed to find child. max_levels: %i.', max_levels)\n # Search children.\n if max_levels:\n for child in self.children:\n child_result = child.find_child(predicate, max_levels - 1)\n if child_result is not None:\n return child_result\n\n return None\n\n def __repr__(self):\n return self._data\n\n def print_tree(self, indent: int = 2):\n \"\"\"Prints this tree in logging.info().\"\"\"\n logging.info(' ' * indent + self.data)\n for c in self.children:\n c.print_tree(indent + 2)\n\n\ndef build_tree_from_dumpsys_output(dumpsys_output: str) -> _DumpsysNode:\n \"\"\"Constructs a tree from a dumpsys string output.\n\n Args:\n dumpsys_output: string Verbatim output from adb dumpsys. The expected format\n is a list where each line is a node and the indentation marks the\n relationship with its parent or sibling.\n\n Returns:\n _DumpsysNode The root of the tree.\n \"\"\"\n lines = dumpsys_output.split('\\n') # Split by lines.\n lines = [x.rstrip(' \\r') for x in lines]\n lines = [x for x in lines if len(x)] # Remove empty lines.\n\n root = _DumpsysNode('___root___') # The root of all nodes.\n parents_stack = [root]\n for line in lines:\n stripped_line = line.lstrip(' ')\n indent = len(line) - len(stripped_line) # Number of indent spaces.\n new_node = _DumpsysNode(stripped_line) # Create a node without indentation.\n\n parent = parents_stack.pop()\n if parent.data == '___root___': # The root is an exception for indentation.\n parent_indent = -2\n else:\n parent_indent = (len(parents_stack) - 1) * 2\n\n if indent == parent_indent: # `new_node` is a sibiling.\n parent = parents_stack.pop()\n elif indent < parent_indent: # Indentation reduced (i.e. a block finished)\n num_levels = (indent // 2) + 1\n parents_stack = parents_stack[:num_levels]\n parent = parents_stack.pop()\n elif indent > parent_indent: # `new_node` is a child.\n pass # No need to change the current parent.\n\n parent.children.append(new_node)\n parents_stack.append(parent)\n parents_stack.append(new_node)\n\n return root\n\n\ndef matches_path(\n dumpsys_activity_output: str,\n expected_view_hierarchy_path: Sequence[re.Pattern[str]],\n max_levels: int = 0,\n) -> bool:\n \"\"\"Returns True if the current dumpsys output matches the expected path.\n\n Args:\n dumpsys_activity_output: The output of running `dumpsys activity ...`.\n expected_view_hierarchy_path: [regex] A list of regular expressions to be\n tested at each level of the tree.\n max_levels: How many levels to search from root for View Hierarchy.\n\n Returns:\n True if the dumpsys tree contains one path that matches all regexes.\n \"\"\"\n root = build_tree_from_dumpsys_output(dumpsys_activity_output)\n\n # Find the View Hierarchy.\n view_hierarchy = root.find_child(\n lambda x: x.data.startswith('View Hierarchy'), max_levels)\n if view_hierarchy is None:\n logging.error(\n 'view_hierarchy is None. Dumpsys activity output: %s. tree: %r',\n str(dumpsys_activity_output), root.print_tree())\n logging.error('Tree root: %s', str(root))\n return False\n\n current_node = view_hierarchy\n for i, regex in enumerate(expected_view_hierarchy_path):\n\n def regex_predicate(node, expr=regex):\n matches = expr.match(node.data)\n return matches is not None\n\n child = current_node.find_child(regex_predicate)\n if child is None:\n logging.error('Mismatched regex (%i, %s). current_node: %s', i,\n regex.pattern, current_node)\n logging.error('Dumpsys activity output: %s', str(dumpsys_activity_output))\n logging.error('Tree root: %s', str(root))\n return False\n else:\n current_node = child\n return True\n\n\nclass AppScreenChecker:\n \"\"\"Checks that the current app screen matches an expected screen.\"\"\"\n\n class Outcome(enum.IntEnum):\n \"\"\"Possible return vales from checking the current app screen.\"\"\"\n # The current app screen matches the expected app screen.\n SUCCESS = 0\n # There's no activity to check.\n EMPTY_EXPECTED_ACTIVITY = 1\n # We were unable to determine the current activity.\n FAILED_ACTIVITY_EXTRACTION = 2\n # The current activity does not match the expected activity.\n UNEXPECTED_ACTIVITY = 3\n # The current view hierarchy does not match the expected view hierarchy.\n UNEXPECTED_VIEW_HIERARCHY = 4\n\n def __init__(self, adb_call_parser: adb_call_parser_lib.AdbCallParser,\n expected_app_screen: task_pb2.AppScreen):\n self._adb_call_parser = adb_call_parser\n self._expected_app_screen = expected_app_screen\n self._expected_activity = expected_app_screen.activity\n self._expected_view_hierarchy_path = [\n re.compile(regex) for regex in expected_app_screen.view_hierarchy_path\n ]\n\n # Return type is AppScreenChecker.Outcome, but pytype doesn't understand that.\n def matches_current_app_screen(self) -> enum.IntEnum:\n \"\"\"Determines whether the current app screen matches `expected_app_screen`.\"\"\"\n if not self._expected_activity:\n return AppScreenChecker.Outcome.EMPTY_EXPECTED_ACTIVITY\n\n # Check if we are still on the expected Activity.\n response = self._adb_call_parser.parse(\n adb_pb2.AdbRequest(\n get_current_activity=adb_pb2.AdbRequest.GetCurrentActivity()))\n if response.status != adb_pb2.AdbResponse.OK:\n return AppScreenChecker.Outcome.FAILED_ACTIVITY_EXTRACTION\n\n current_activity = response.get_current_activity.full_activity\n if current_activity != self._expected_activity:\n logging.error('current_activity: %s, expected_activity: %s',\n current_activity, self._expected_activity)\n return AppScreenChecker.Outcome.UNEXPECTED_ACTIVITY\n\n # Extract just the package name from the full activity name.\n package_name = self._expected_activity.split('/')[0]\n\n # Check if we are in the expected view hierarchy path.\n if self._expected_view_hierarchy_path:\n dumpsys_response = self._adb_call_parser.parse(\n adb_pb2.AdbRequest(\n dumpsys=adb_pb2.AdbRequest.DumpsysRequest(\n service='activity', args=[package_name, package_name])))\n if dumpsys_response.status != adb_pb2.AdbResponse.OK:\n return AppScreenChecker.Outcome.FAILED_ACTIVITY_EXTRACTION\n\n if dumpsys_response.dumpsys.output:\n if not matches_path(\n dumpsys_response.dumpsys.output.decode('utf-8'),\n self._expected_view_hierarchy_path,\n max_levels=3):\n return AppScreenChecker.Outcome.UNEXPECTED_VIEW_HIERARCHY\n\n return AppScreenChecker.Outcome.SUCCESS\n\n def wait_for_app_screen(self, timeout_sec: float) -> float:\n \"\"\"Waits for `self._expected_app_screen` to be the current screen.\n\n Args:\n timeout_sec: Maximum total time to wait for the screen to pop up.\n\n Returns:\n The total amount of time in seconds spent waiting for the screen to pop\n up.\n Raises:\n errors.WaitForAppScreenError if the screen does not pop up within\n `timeout_sec`.\n \"\"\"\n\n logging.info('Waiting for app screen...')\n start_time = time.time()\n while time.time() - start_time < timeout_sec:\n if self.matches_current_app_screen() == AppScreenChecker.Outcome.SUCCESS:\n wait_time = time.time() - start_time\n logging.info('Successfully waited for app screen in %r seconds: [%r]',\n wait_time, self._expected_app_screen)\n return wait_time\n time.sleep(0.1)\n\n wait_time = time.time() - start_time\n logging.error('Failed to wait for app screen in %r seconds: [%r].',\n wait_time, self._expected_app_screen)\n\n raise errors.WaitForAppScreenError()\n"
},
{
"path": "android_env/components/app_screen_checker_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.app_screen_checker.\"\"\"\n\nimport re\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import adb_call_parser\nfrom android_env.components import app_screen_checker\nfrom android_env.components import errors\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import task_pb2\n\n\ndef _flatten_tree(\n tree: app_screen_checker._DumpsysNode, flat_tree: list[str], indent: int = 2\n):\n \"\"\"Appends a list of strings to `flat_tree` from `tree`.\"\"\"\n flat_tree.append(' ' * indent + tree.data)\n for c in tree.children:\n _flatten_tree(c, flat_tree, indent + 2)\n\n\nclass AppScreenCheckerTest(absltest.TestCase):\n\n # Ensures that build_tree_from_dumpsys_output produces a node whose flat\n # representation matches our expectation from an arbitrary hierarchy.\n def test_build_tree_from_dumpsys_output(self):\n dumpsys_output = \"\"\"\nQueen Elizabeth II\n Charles\n William\n George\n Charlotte\n Louis\n Harry\n Archie\n Anne\n Peter\n Savannah\n Isla\n Zara\n Mia\n Lena\n Andrew\n Beatrice\n Eugenie\n Edward\n Louise\n James\n\"\"\"\n tree = app_screen_checker.build_tree_from_dumpsys_output(dumpsys_output)\n flat_tree = []\n _flatten_tree(tree, flat_tree, indent=2)\n self.assertEqual(flat_tree, [\n ' ___root___',\n ' Queen Elizabeth II',\n ' Charles',\n ' William',\n ' George',\n ' Charlotte',\n ' Louis',\n ' Harry',\n ' Archie',\n ' Anne',\n ' Peter',\n ' Savannah',\n ' Isla',\n ' Zara',\n ' Mia',\n ' Lena',\n ' Andrew',\n ' Beatrice',\n ' Eugenie',\n ' Edward',\n ' Louise',\n ' James',\n ])\n\n # Ensures that build_tree_from_dumpsys_output produces a node whose flat\n # representation matches our expectation from an arbitrary hierarchy.\n def test_build_forest_from_dumpsys_output(self):\n dumpsys_output = \"\"\"\nTree1\n Branch1\n Leaf1\n Leaf2\n Branch2\n Leaf3\n Leaf4\n Leaf5\nTree2\n Branch3\n Leaf6\n Leaf7\n Branch4\n Leaf8\n Leaf9\n Leaf10\n Leaf11\n\"\"\"\n tree = app_screen_checker.build_tree_from_dumpsys_output(dumpsys_output)\n flat_tree = []\n _flatten_tree(tree, flat_tree, indent=2)\n self.assertEqual(flat_tree, [\n ' ___root___',\n ' Tree1',\n ' Branch1',\n ' Leaf1',\n ' Leaf2',\n ' Branch2',\n ' Leaf3',\n ' Leaf4',\n ' Leaf5',\n ' Tree2',\n ' Branch3',\n ' Leaf6',\n ' Leaf7',\n ' Branch4',\n ' Leaf8',\n ' Leaf9',\n ' Leaf10',\n ' Leaf11',\n ])\n\n def test_no_view_hierarchy_matches_path(self):\n dumpsys_output = \"\"\"\nTASK\n ACTIVITY\n Missing View Hierarchy\n A\n B\n C\n D\n E\n F\n\"\"\"\n expected_path = ['^A$', 'B$']\n expected_view_hierarchy_path = [\n re.compile(regex) for regex in expected_path\n ]\n self.assertFalse(\n app_screen_checker.matches_path(dumpsys_output,\n expected_view_hierarchy_path))\n\n def test_matches_path(self):\n dumpsys_output = \"\"\"\nTASK\n ACTIVITY\n Some node we don't care\n Blah\n\n View Hierarchy\n Hirohito\n Akihito\n Naruhito\n Aiko\n Fumihito\n Mako\n Kako\n Hisahito\n Masahito\n\"\"\"\n expected_path = ['^Hirohito$', 'Akihito$', 'Fumihito$', 'Kako$']\n expected_view_hierarchy_path = [\n re.compile(regex) for regex in expected_path\n ]\n self.assertTrue(\n app_screen_checker.matches_path(\n dumpsys_output, expected_view_hierarchy_path, max_levels=2))\n\n # Also check that the following path does not match anything in the tree.\n expected_path = ['^Hirohito$', 'Akihito$', 'Fumihito$', 'Kenji$']\n expected_view_hierarchy_path = [\n re.compile(regex) for regex in expected_path\n ]\n self.assertFalse(\n app_screen_checker.matches_path(dumpsys_output,\n expected_view_hierarchy_path))\n\n def test_matches_path_one_level_deep(self):\n dumpsys_output = \"\"\"\nTASK\n ACTIVITY\n Some node we don't care\n Blah\n\n Some intermediate node\n View Hierarchy\n Hirohito\n Akihito\n Naruhito\n Aiko\n Fumihito\n Mako\n Kako\n Hisahito\n Masahito\n\"\"\"\n expected_path = ['^Hirohito$', 'Akihito$', 'Fumihito$', 'Kako$']\n expected_view_hierarchy_path = [\n re.compile(regex) for regex in expected_path\n ]\n self.assertTrue(\n app_screen_checker.matches_path(\n dumpsys_output, expected_view_hierarchy_path, max_levels=3))\n\n # Also check that the view hierarchy is not found when searching only grand\n # children of TASK.\n expected_path = ['^Hirohito$', 'Akihito$', 'Fumihito$', 'Kako$']\n expected_view_hierarchy_path = [\n re.compile(regex) for regex in expected_path\n ]\n self.assertFalse(\n app_screen_checker.matches_path(\n dumpsys_output, expected_view_hierarchy_path, max_levels=2))\n\n def test_wait_for_app_screen_zero_timeout(self):\n \"\"\"Ensures that an exception is raised if the timeout is passed.\"\"\"\n app_screen = task_pb2.AppScreen(activity='whatever.MyActivity')\n call_parser = mock.create_autospec(adb_call_parser.AdbCallParser)\n screen_checker = app_screen_checker.AppScreenChecker(\n adb_call_parser=call_parser,\n expected_app_screen=app_screen)\n # With a zero timeout, the method should never be able to wait for the\n # screen to pop up and an exception should be raised.\n self.assertRaises(\n errors.WaitForAppScreenError,\n screen_checker.wait_for_app_screen,\n timeout_sec=0.0)\n\n def test_wait_for_app_screen_successful(self):\n \"\"\"Ensures that with the right conditions, the app screen should pop up.\"\"\"\n app_screen = task_pb2.AppScreen(activity='my.favorite.AwesomeActivity')\n call_parser = mock.create_autospec(adb_call_parser.AdbCallParser)\n call_parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK,\n get_current_activity=adb_pb2.AdbResponse.GetCurrentActivityResponse(\n full_activity='my.favorite.AwesomeActivity'))\n\n screen_checker = app_screen_checker.AppScreenChecker(\n call_parser, app_screen)\n timeout = 1.0\n wait_time = screen_checker.wait_for_app_screen(timeout_sec=timeout)\n\n # The call should not generate an exception and the return value should be\n # less than the timeout given.\n self.assertLess(wait_time, timeout)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/config_classes.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Dataclass definitions used for instantiating AndroidEnv components.\"\"\"\n\nimport dataclasses\nimport enum\n\n\n@dataclasses.dataclass\nclass AdbControllerConfig:\n \"\"\"Settings for instatiating an `AdbController` instance.\"\"\"\n\n # Filesystem path to the `adb` binary.\n # NOTE: This must be a full path and must not contain environment variables\n # or user folder shorthands (e.g. `~/some/path/to/adb`) since they will not be\n # expanded internally by AndroidEnv.\n adb_path: str = '~/Android/Sdk/platform-tools/adb'\n # Port for adb server.\n adb_server_port: int = 5037\n # Default timeout in seconds for internal commands.\n default_timeout: float = 120.0\n # Name of the device to communicate with.\n device_name: str = ''\n # Whether to use adb server port set in OS Environment variables.\n # When True, adb will use the ANDROID_ADB_SERVER_PORT OS environment variable\n # for selecting its server port if available (or the default 5037 if not set).\n # When False, the ANDROID_ADB_SERVER_PORT OS environment var will be unset\n # and adb_server_port will be used and supplied as an argument to all adb\n # commands.\n use_adb_server_port_from_os_env: bool = False\n\n\n@dataclasses.dataclass\nclass DeviceSettingsConfig:\n \"\"\"Config class for DeviceSettings.\"\"\"\n\n # Whether to show circles on the screen indicating touch position.\n show_touches: bool = True\n # Whether to show blue lines on the screen indicating touch position.\n show_pointer_location: bool = True\n # Whether or not to show the status (top) bar.\n show_status_bar: bool = False\n # Whether or not to show the navigation (bottom) bar.\n show_navigation_bar: bool = False\n\n\n@dataclasses.dataclass\nclass CoordinatorConfig:\n \"\"\"Config class for Coordinator.\"\"\"\n\n # Number of virtual \"fingers\" of the agent.\n num_fingers: int = 1\n # Whether to enable keyboard key events.\n enable_key_events: bool = False\n # Time between periodic restarts in minutes. If > 0, will trigger\n # a simulator restart at the beginning of the next episode once the time has\n # been reached.\n periodic_restart_time_min: float = 0.0\n # General Android settings.\n device_settings: DeviceSettingsConfig = dataclasses.field(\n default_factory=DeviceSettingsConfig\n )\n\n\n@dataclasses.dataclass\nclass SimulatorConfig:\n \"\"\"Base class for all simulator configs.\"\"\"\n\n # If true, the log stream of the simulator will be verbose.\n verbose_logs: bool = False\n # How often to (asynchronously) grab the screenshot from the simulator.\n # If <= 0, stepping the environment blocks on fetching the screenshot (the\n # environment is synchronous).\n interaction_rate_sec: float = 0.0\n\n\n@enum.unique\nclass GPUMode(enum.Enum):\n \"\"\"Emulator GPU Mode.\"\"\"\n\n HOST = 'host'\n SWANGLE_INDIRECT = 'swangle_indirect'\n SWIFTSHADER_INDIRECT = 'swiftshader_indirect'\n\n\n@dataclasses.dataclass\nclass EmulatorLauncherConfig:\n \"\"\"Config class for EmulatorLauncher.\"\"\"\n\n # NOTE: If `adb_port`, `emulator_console_port` and `grpc_port` are defined\n # (i.e. not all equal to 0), it is assumed that the emulator they point to\n # exists already and EmulatorLauncher will be skipped.\n\n # Filesystem path to the `emulator` binary.\n emulator_path: str = '~/Android/Sdk/emulator/emulator'\n # Filesystem path to the Android SDK root.\n android_sdk_root: str = '~/Android/Sdk'\n # Name of the AVD.\n avd_name: str = ''\n # Local directory for AVDs.\n android_avd_home: str = '~/.android/avd'\n # Name of the snapshot to load.\n snapshot_name: str = ''\n # Path to the KVM device.\n kvm_device: str = '/dev/kvm'\n # Path to directory which will hold temporary files.\n tmp_dir: str = '/tmp/android_env/simulator/'\n # GPU mode override.\n # Please see\n # https://developer.android.com/studio/run/emulator-acceleration#accel-graphics.\n gpu_mode: str = GPUMode.SWANGLE_INDIRECT.value\n # Whether to run in headless mode (i.e. without a graphical window).\n run_headless: bool = True\n # Whether to restrict network access.\n # If True, will disable networking on the device. This option is only\n # available for emulator version > 31.3.9 (June 2022).\n restrict_network: bool = False\n # Whether to set `SHOW_PERF_STATS=1` when launching the emulator to display\n # performance and memory statistics.\n show_perf_stats: bool = False\n\n # ADB port for the Android device.\n adb_port: int = 0\n # Port for telnet communication with the emulator.\n emulator_console_port: int = 0\n # Port for gRPC communication with the emulator.\n grpc_port: int = 0\n\n\n@dataclasses.dataclass\nclass EmulatorConfig(SimulatorConfig):\n \"\"\"Config class for EmulatorSimulator.\"\"\"\n\n # Configuration for launching the Android Emulator.\n emulator_launcher: EmulatorLauncherConfig = dataclasses.field(\n default_factory=EmulatorLauncherConfig\n )\n # Configuration for talking to adb.\n adb_controller: AdbControllerConfig = dataclasses.field(\n default_factory=AdbControllerConfig\n )\n # Path to file which holds emulator logs. If not provided, it will be\n # determined by the EmulatorLauncher.\n logfile_path: str = ''\n # The number of times to try launching the emulator before rebooting (reboot\n # on the n+1-st try).\n launch_n_times_without_reboot: int = 1\n # The number of times to try launching the emulator before reinstalling\n # (reinstall on the n+1-st try).\n launch_n_times_without_reinstall: int = 2\n\n\n@dataclasses.dataclass\nclass FakeSimulatorConfig(SimulatorConfig):\n \"\"\"Config class for FakeSimulator.\"\"\"\n\n # The dimensions in pixels of the device screen (HxW).\n screen_dimensions: tuple[int, int] = (0, 0)\n\n\n@dataclasses.dataclass\nclass TaskManagerConfig:\n \"\"\"Config class for TaskManager.\"\"\"\n\n # If max_bad_states episodes finish in a bad state in a row, restart\n # the simulation.\n max_bad_states: int = 3\n # The frequency to check for the current activity and view hierarchy.\n # The unit is raw observation (i.e. each call to AndroidEnv.step()).\n dumpsys_check_frequency: int = 150\n # The maximum number of tries for extracting the current activity before\n # forcing the episode to restart.\n max_failed_current_activity: int = 10\n # The maximum number of extras elements to store. If this number is exceeded,\n # elements are dropped in the order they were received.\n extras_max_buffer_size: int = 100\n\n\n@dataclasses.dataclass\nclass TaskConfig:\n \"\"\"Base config class for loading tasks.\"\"\"\n\n # The directory for temporary task-related resources.\n tmp_dir: str = ''\n\n\n@dataclasses.dataclass\nclass FilesystemTaskConfig(TaskConfig):\n \"\"\"Config for protobuf files stored in the local filesystem.\"\"\"\n\n # Filesystem path to `.binarypb` or `.textproto` protobuf Task.\n path: str = ''\n\n\n@dataclasses.dataclass\nclass AndroidEnvConfig:\n \"\"\"Config class for AndroidEnv.\"\"\"\n\n # Configs for main components.\n task: TaskConfig = dataclasses.field(default_factory=TaskConfig)\n task_manager: TaskManagerConfig = dataclasses.field(\n default_factory=TaskManagerConfig\n )\n coordinator: CoordinatorConfig = dataclasses.field(\n default_factory=CoordinatorConfig\n )\n simulator: SimulatorConfig = dataclasses.field(default_factory=EmulatorConfig)\n"
},
{
"path": "android_env/components/coordinator.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Coordinator handles interaction between internal components of AndroidEnv.\"\"\"\n\nimport copy\nimport time\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env.components import action_fns\nfrom android_env.components import adb_call_parser\nfrom android_env.components import config_classes\nfrom android_env.components import device_settings as device_settings_lib\nfrom android_env.components import errors\nfrom android_env.components import specs\nfrom android_env.components import task_manager as task_manager_lib\nfrom android_env.components.simulators import base_simulator\nfrom android_env.proto import adb_pb2\nimport dm_env\nimport numpy as np\n\n\nclass Coordinator:\n \"\"\"Handles interaction between internal components of AndroidEnv.\"\"\"\n\n def __init__(\n self,\n simulator: base_simulator.BaseSimulator,\n task_manager: task_manager_lib.TaskManager,\n device_settings: device_settings_lib.DeviceSettings,\n config: config_classes.CoordinatorConfig | None = None,\n ):\n \"\"\"Handles communication between AndroidEnv and its components.\n\n Args:\n simulator: A BaseSimulator instance.\n task_manager: The TaskManager, responsible for coordinating RL tasks.\n config: Settings to customize this Coordinator.\n \"\"\"\n self._simulator = simulator\n self._task_manager = task_manager\n self._config = config or config_classes.CoordinatorConfig()\n self._device_settings = device_settings\n self._adb_call_parser: adb_call_parser.AdbCallParser = None\n\n # Initialize stats.\n self._stats = {\n 'relaunch_count': 0,\n 'relaunch_count_periodic': 0,\n 'relaunch_count_setup_steps': 0,\n 'relaunch_count_reset_steps': 0,\n 'relaunch_count_simulator_launch': 0,\n 'relaunch_count_simulator_reset': 0,\n 'relaunch_count_execute_action': 0,\n 'relaunch_count_fetch_observation': 0,\n 'relaunch_count_update_settings': 0,\n 'failed_task_updates': 0,\n }\n\n # Initialize counters.\n self._simulator_healthy = False\n self._latest_observation_time = 0\n self._simulator_start_time = None\n\n logging.info('Starting the simulator...')\n self._launch_simulator()\n\n def action_spec(self) -> dict[str, dm_env.specs.Array]:\n return specs.base_action_spec(\n num_fingers=self._config.num_fingers,\n enable_key_events=self._config.enable_key_events,\n )\n\n def observation_spec(self) -> dict[str, dm_env.specs.Array]:\n return specs.base_observation_spec(\n height=self._device_settings.screen_height(),\n width=self._device_settings.screen_width(),\n )\n\n def _should_periodic_relaunch(self) -> bool:\n \"\"\"Checks if it is time to restart the simulator.\n\n If a periodic restart time was specified, the Coordinator will re-launch\n the simulator at regular time intervals. This helps to make sure that the\n simulator is not in a stale state even if the environment has been running\n for a significant amount of time.\n\n Returns:\n Boolean indicating if it is time to restart the simulator.\n \"\"\"\n\n if self._config.periodic_restart_time_min and self._simulator_start_time:\n sim_alive_time = (time.time() - self._simulator_start_time) / 60.0\n logging.info('Simulator has been running for %f mins', sim_alive_time)\n if sim_alive_time > self._config.periodic_restart_time_min:\n logging.info('Maximum alive time reached. Restarting simulator.')\n self._stats['relaunch_count_periodic'] += 1\n return True\n return False\n\n def _launch_simulator(self, max_retries: int = 3):\n \"\"\"Launches the simulator.\n\n Sets up the simulator and other task-related settings.\n\n Args:\n max_retries: Number of times to attempt a restart before raising an error.\n \"\"\"\n\n self._simulator_healthy = False\n\n # Attempt to restart the system a given number of times.\n num_tries = 1\n latest_error = None\n while True:\n if num_tries > max_retries:\n raise errors.TooManyRestartsError(\n 'Maximum number of restart attempts reached.'\n ) from latest_error\n logging.info('Simulator launch attempt %d of %d', num_tries, max_retries)\n\n self._task_manager.stop()\n\n # Launch the simulator.\n self._simulator.launch()\n self._simulator_start_time = time.time()\n\n # From here on, the simulator is assumed to be up and running.\n self._adb_call_parser = self._create_adb_call_parser()\n try:\n self._device_settings.update(self._config.device_settings)\n except errors.AdbControllerError as e:\n logging.exception('device_settings.update() failed.')\n self._stats['relaunch_count_update_settings'] += 1\n self._latest_error = e\n num_tries += 1\n continue\n\n # Start the task.\n self._task_manager.start(\n adb_call_parser_factory=self._create_adb_call_parser,\n log_stream=self._simulator.create_log_stream(),\n )\n try:\n self._task_manager.setup_task()\n except errors.StepCommandError as error:\n logging.exception('Failed to set up the task. Restarting simulator.')\n self._stats['relaunch_count_setup_steps'] += 1\n latest_error = error\n num_tries += 1\n continue\n\n # Restart was successful.\n self._simulator_healthy = True\n self._stats['relaunch_count'] += 1\n break\n\n def _create_adb_call_parser(self):\n \"\"\"Creates a new AdbCallParser instance.\"\"\"\n return adb_call_parser.AdbCallParser(\n adb_controller=self._simulator.create_adb_controller()\n )\n\n def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:\n return self._adb_call_parser.parse(call)\n\n def rl_reset(self) -> dm_env.TimeStep:\n \"\"\"Resets the RL episode.\"\"\"\n\n # Relaunch the simulator if necessary.\n if not self._simulator_healthy or self._should_periodic_relaunch():\n self._launch_simulator()\n\n # Reset counters.\n self._latest_observation_time = 0\n for key in self._stats:\n if key.startswith('episode'):\n self._stats[key] = 0.0\n\n # Execute a lift action before resetting the task.\n if not action_fns.send_action_to_simulator(\n action_fns.lift_all_fingers_action(self._config.num_fingers),\n self._simulator,\n self._device_settings.screen_width(),\n self._device_settings.screen_height(),\n self._config.num_fingers,\n ):\n self._stats['relaunch_count_execute_action'] += 1\n self._simulator_healthy = False\n\n # Reset the task.\n self._task_manager.reset_task()\n self._device_settings.get_orientation()\n\n # Get data from the simulator.\n simulator_signals = self._gather_simulator_signals()\n\n return self._task_manager.rl_reset(simulator_signals)\n\n def rl_step(self, agent_action: dict[str, np.ndarray]) -> dm_env.TimeStep:\n \"\"\"Executes the selected action and returns a timestep.\n\n Args:\n agent_action: Selected action to perform on the simulated Android device.\n If `agent_action` is `None` it means that this is an RL reset (to start\n a new episode).\n\n Returns:\n An RL timestep.\n \"\"\"\n\n if not action_fns.send_action_to_simulator(\n agent_action,\n self._simulator,\n self._device_settings.screen_width(),\n self._device_settings.screen_height(),\n self._config.num_fingers,\n ):\n self._stats['relaunch_count_execute_action'] += 1\n self._simulator_healthy = False\n\n # Get data from the simulator.\n try:\n simulator_signals = self._gather_simulator_signals()\n except errors.ReadObservationError:\n logging.exception('Unable to fetch observation. Restarting simulator.')\n self._stats['relaunch_count_fetch_observation'] += 1\n self._simulator_healthy = False\n\n if not self._simulator_healthy:\n return dm_env.truncation(reward=0.0, observation=None)\n\n return self._task_manager.rl_step(simulator_signals)\n\n def _gather_simulator_signals(self) -> dict[str, np.ndarray]:\n \"\"\"Gathers data from various sources to assemble the RL observation.\"\"\"\n\n # Get current timestamp and update the delta.\n now = time.time()\n timestamp_delta = (\n 0\n if self._latest_observation_time == 0\n else (now - self._latest_observation_time) * 1e6\n )\n self._latest_observation_time = now\n\n return {\n 'pixels': self._simulator.get_screenshot(),\n 'orientation': self._device_settings.get_orientation(),\n 'timedelta': np.array(timestamp_delta, dtype=np.int64),\n }\n\n def __del__(self):\n self.close()\n\n def stats(self) -> dict[str, Any]:\n \"\"\"Returns various statistics.\"\"\"\n\n return copy.deepcopy(self._stats)\n\n def close(self):\n \"\"\"Cleans up the state of this Coordinator.\"\"\"\n\n if hasattr(self, '_task_manager'):\n try:\n self._task_manager.stop()\n except: # pylint: disable=bare-except\n logging.exception('Failed to stop task manager. Continuing.')\n if hasattr(self, '_simulator'):\n try:\n self._simulator.close()\n except: # pylint: disable=bare-except\n logging.exception('Failed to close simulator. Continuing.')\n"
},
{
"path": "android_env/components/coordinator_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.coordinator.\"\"\"\n\nimport tempfile\nimport time\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import action_type\nfrom android_env.components import adb_call_parser\nfrom android_env.components import config_classes\nfrom android_env.components import coordinator as coordinator_lib\nfrom android_env.components import device_settings as device_settings_lib\nfrom android_env.components import errors\nfrom android_env.components import task_manager\nfrom android_env.components.simulators import base_simulator\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import state_pb2\nfrom android_env.proto import task_pb2\nimport dm_env\nimport numpy as np\n\n\nclass CoordinatorTest(parameterized.TestCase):\n\n def setUp(self):\n super().setUp()\n self.addCleanup(mock.patch.stopall) # Disable previous patches.\n\n self._simulator = mock.create_autospec(base_simulator.BaseSimulator)\n self._random_screenshot = np.random.randint(\n low=0, high=255, size=(800, 600, 3), dtype=np.uint8)\n self._simulator.get_screenshot.return_value = self._random_screenshot\n self._task_manager = mock.create_autospec(task_manager.TaskManager)\n self._adb_call_parser = mock.create_autospec(adb_call_parser.AdbCallParser)\n self.enter_context(\n mock.patch.object(\n adb_call_parser,\n 'AdbCallParser',\n autospec=True,\n return_value=self._adb_call_parser))\n self._coordinator = coordinator_lib.Coordinator(\n simulator=self._simulator,\n task_manager=self._task_manager,\n device_settings=device_settings_lib.DeviceSettings(self._simulator),\n )\n\n def tearDown(self):\n super().tearDown()\n self._coordinator.close()\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_relaunch_simulator(self, unused_mock_sleep):\n relaunch_count = self._coordinator.stats()['relaunch_count']\n self._coordinator._launch_simulator()\n self.assertEqual(self._coordinator.stats()['relaunch_count'],\n relaunch_count + 1)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_reset(self, unused_mock_sleep):\n \"\"\"'relaunch_count_execute_action' should be zero if there's no error.\"\"\"\n\n self._coordinator.rl_reset()\n stats = self._coordinator.stats()\n self.assertIn('relaunch_count_execute_action', stats)\n self.assertEqual(stats['relaunch_count_execute_action'], 0)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_reset_error_sending_action(self, unused_mock_sleep):\n \"\"\"'relaunch_count_execute_action' should be positive if there's an error.\"\"\"\n\n self._simulator.send_touch.side_effect = errors.SendActionError()\n self._coordinator.rl_reset()\n stats = self._coordinator.stats()\n self.assertIn('relaunch_count_execute_action', stats)\n self.assertEqual(stats['relaunch_count_execute_action'], 1)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_lift_all_fingers(self, unused_mock_sleep):\n self._coordinator = coordinator_lib.Coordinator(\n simulator=self._simulator,\n task_manager=self._task_manager,\n device_settings=device_settings_lib.DeviceSettings(self._simulator),\n config=config_classes.CoordinatorConfig(num_fingers=3),\n )\n self._coordinator.rl_reset()\n expected_actions = [\n # (x, y, is_down, identifier).\n (0, 0, False, 0),\n (0, 0, False, 1),\n (0, 0, False, 2),\n ]\n actual_actions = self._simulator.send_touch.call_args[0][0]\n for actual, expected in zip(actual_actions, expected_actions):\n np.testing.assert_array_equal(actual, expected)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_process_action(self, unused_mock_sleep):\n\n def fake_rl_step(simulator_signals):\n return dm_env.transition(\n reward=10.0,\n observation={\n 'pixels': simulator_signals['pixels'],\n 'orientation': simulator_signals['orientation'],\n 'timedelta': simulator_signals['timedelta'],\n 'extras': {\n 'extra': [0.0]\n }\n })\n\n self._task_manager.rl_step.side_effect = fake_rl_step\n timestep = self._coordinator.rl_step(\n agent_action={\n 'action_type': np.array(action_type.ActionType.LIFT),\n 'touch_position': np.array([0.5, 0.5]),\n })\n obs = timestep.observation\n self.assertEqual(obs['pixels'].shape, (800, 600, 3))\n np.testing.assert_equal(obs['orientation'],\n np.array([0, 0, 0, 0], dtype=np.uint8))\n self.assertEqual(timestep.reward, 10.0)\n self.assertEqual(obs['extras'], {'extra': [0.0]})\n self.assertFalse(timestep.last())\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_process_action_error(self, unused_mock_sleep):\n\n def fake_rl_step(simulator_signals):\n self.assertFalse(simulator_signals['simulator_healthy'])\n return dm_env.truncation(reward=0.0, observation=None)\n\n self._task_manager.rl_step.side_effect = fake_rl_step\n self._simulator.get_screenshot.side_effect = errors.ReadObservationError()\n timestep = self._coordinator.rl_step(\n agent_action={\n 'action_type': np.array(action_type.ActionType.LIFT),\n 'touch_position': np.array([0.5, 0.5]),\n })\n self.assertIsNone(timestep.observation)\n self.assertEqual(timestep.reward, 0.0)\n self.assertTrue(timestep.last())\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_execute_action_touch(self, unused_mock_sleep):\n\n def fake_rl_step(simulator_signals):\n return dm_env.transition(\n reward=123.0,\n observation={\n 'pixels': simulator_signals['pixels'],\n 'orientation': simulator_signals['orientation'],\n 'timedelta': simulator_signals['timedelta'],\n 'extras': {\n 'extra': [0.0]\n }\n })\n\n self._task_manager.rl_step.side_effect = fake_rl_step\n timestep = self._coordinator.rl_step(\n agent_action={\n 'action_type': np.array(action_type.ActionType.TOUCH),\n 'touch_position': np.array([0.5, 0.5])\n })\n self.assertEqual(timestep.reward, 123.0)\n np.testing.assert_equal(timestep.observation['pixels'],\n self._random_screenshot)\n self._simulator.send_touch.assert_called_once_with([(300, 400, True, 0)])\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_execute_multitouch_action(self, unused_mock_sleep):\n self._coordinator = coordinator_lib.Coordinator(\n simulator=self._simulator,\n task_manager=self._task_manager,\n device_settings=device_settings_lib.DeviceSettings(self._simulator),\n config=config_classes.CoordinatorConfig(num_fingers=3),\n )\n\n def fake_rl_step(simulator_signals):\n return dm_env.transition(\n reward=456.0,\n observation={\n 'pixels': simulator_signals['pixels'],\n 'orientation': simulator_signals['orientation'],\n 'timedelta': simulator_signals['timedelta'],\n 'extras': {\n 'extra': [0.0]\n }\n })\n\n self._task_manager.rl_step.side_effect = fake_rl_step\n action = {\n 'action_type': np.array([action_type.ActionType.TOUCH]),\n 'touch_position': np.array([0.25, 0.75]),\n 'action_type_2': np.array([action_type.ActionType.TOUCH]),\n 'touch_position_2': np.array([0.75, 0.25]),\n 'action_type_3': np.array([action_type.ActionType.LIFT]),\n 'touch_position_3': np.array([0.5, 0.5]),\n }\n timestep = self._coordinator.rl_step(action)\n self._simulator.send_touch.assert_called_once_with([(150, 600, True, 0),\n (450, 200, True, 1),\n (300, 400, False, 2)])\n self.assertEqual(timestep.reward, 456.0)\n np.testing.assert_equal(timestep.observation['pixels'],\n self._random_screenshot)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_execute_action_repeat(self, unused_mock_sleep):\n def fake_rl_step(simulator_signals):\n return dm_env.transition(\n reward=10.0,\n observation={\n 'pixels': simulator_signals['pixels'],\n 'orientation': simulator_signals['orientation'],\n 'timedelta': simulator_signals['timedelta'],\n 'extras': {\n 'extra': [0.0]\n }\n })\n\n self._task_manager.rl_step.side_effect = fake_rl_step\n timestep = self._coordinator.rl_step(\n {'action_type': np.array(action_type.ActionType.REPEAT)})\n self._simulator.send_touch.assert_not_called()\n np.testing.assert_equal(timestep.observation['pixels'],\n self._random_screenshot)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_execute_action_error(self, unused_mock_sleep):\n def fake_rl_step(simulator_signals):\n self.assertFalse(simulator_signals['simulator_healthy'])\n return dm_env.truncation(reward=0.0, observation=None)\n\n self._task_manager.rl_step.side_effect = fake_rl_step\n self._simulator.send_touch.side_effect = errors.SendActionError\n timestep = self._coordinator.rl_step({\n 'action_type': np.array(action_type.ActionType.TOUCH),\n 'touch_position': np.array([0.3, 0.8])\n })\n self.assertIsNone(timestep.observation)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_max_restarts_setup_steps(self, unused_mock_sleep):\n init_fn_call = self._task_manager.setup_task.call_count\n self._task_manager.setup_task.side_effect = errors.StepCommandError\n self.assertRaises(errors.TooManyRestartsError,\n self._coordinator._launch_simulator)\n # The method was called three more times when attempting to relaunch.\n self.assertEqual(init_fn_call + 3,\n self._task_manager.setup_task.call_count)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_execute_adb_call(self, unused_mock_sleep):\n call = adb_pb2.AdbRequest(\n force_stop=adb_pb2.AdbRequest.ForceStop(package_name='blah'))\n expected_response = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n self._adb_call_parser.parse.side_effect = [expected_response]\n\n response = self._coordinator.execute_adb_call(call)\n\n self.assertEqual(response, expected_response)\n self._adb_call_parser.parse.assert_called_with(call)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/device_settings.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sets and gets some global settings on an Android device.\"\"\"\n\nfrom typing import Final\nfrom unittest import mock\n\nfrom absl import logging\nfrom android_env.components import adb_call_parser\nfrom android_env.components import config_classes\nfrom android_env.components.simulators import base_simulator\nfrom android_env.proto import adb_pb2\nimport numpy as np\n\n\n# The internal `AdbCallParser` instance is lazily instantiated within\n# `DeviceSettings`. If we make it optional (i.e. `| None`), pytype will think\n# that it could be `None`, requiring either explicit runtime checks or escape\n# hatches in every actual call, even if it's never actually `None` if reached\n# via the public API.\n# The trick here is to create this dummy instance of the right type that's used\n# as a sentinel to indicate that it hasn't been initialized yet.\n_PLACEHOLDER_ADB_CALL_PARSER: Final[adb_call_parser.AdbCallParser] = (\n mock.create_autospec(adb_call_parser.AdbCallParser)\n)\n\n\nclass DeviceSettings:\n \"\"\"An abstraction for general properties and settings of an Android device.\"\"\"\n\n def __init__(self, simulator: base_simulator.BaseSimulator):\n self._simulator = simulator\n self._adb_call_parser = _PLACEHOLDER_ADB_CALL_PARSER\n\n # The size of the device screen in pixels.\n self._screen_width: int = 0\n self._screen_height: int = 0\n # The device orientation.\n self._orientation = np.zeros(4, dtype=np.uint8)\n\n def update(self, config: config_classes.DeviceSettingsConfig) -> None:\n \"\"\"Sets the configuration of the device according to `config`.\"\"\"\n\n if self._adb_call_parser is _PLACEHOLDER_ADB_CALL_PARSER:\n self._adb_call_parser = adb_call_parser.AdbCallParser(\n adb_controller=self._simulator.create_adb_controller()\n )\n\n self._update_screen_size()\n self._set_show_touches(config.show_touches)\n self._set_show_pointer_location(config.show_pointer_location)\n self._set_status_navigation_bars(\n config.show_navigation_bar, config.show_status_bar\n )\n\n def screen_width(self) -> int:\n \"\"\"The screen width in pixels. Only valid after `update()` is called.\"\"\"\n\n return self._screen_width\n\n def screen_height(self) -> int:\n \"\"\"The screen height in pixels. Only valid after `update()` is called.\"\"\"\n\n return self._screen_height\n\n def get_orientation(self) -> np.ndarray:\n \"\"\"Returns the device orientation. Please see specs.py for details.\"\"\"\n\n if self._adb_call_parser is _PLACEHOLDER_ADB_CALL_PARSER:\n self._adb_call_parser = adb_call_parser.AdbCallParser(\n adb_controller=self._simulator.create_adb_controller()\n )\n\n self._update_orientation()\n return self._orientation\n\n def _update_screen_size(self) -> None:\n \"\"\"Sets the screen size from a screenshot ignoring the color channel.\"\"\"\n\n screenshot = self._simulator.get_screenshot()\n self._screen_height = screenshot.shape[0]\n self._screen_width = screenshot.shape[1]\n\n def _set_show_touches(self, show: bool) -> None:\n \"\"\"Whether to display circles indicating the touch position.\"\"\"\n\n self._adb_call_parser.parse(\n adb_pb2.AdbRequest(\n settings=adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(\n key='show_touches', value='1' if show else '0'\n ),\n )\n )\n )\n\n def _set_show_pointer_location(self, show: bool) -> None:\n \"\"\"Whether to display blue lines on the screen indicating touch position.\"\"\"\n\n self._adb_call_parser.parse(\n adb_pb2.AdbRequest(\n settings=adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(\n key='pointer_location', value='1' if show else '0'\n ),\n )\n )\n )\n\n def _set_status_navigation_bars(\n self, show_navigation: bool, show_status: bool\n ) -> None:\n \"\"\"Whether to display the status (top) and navigation (bottom) bars.\"\"\"\n\n if show_navigation and show_status:\n policy_control_value = 'null*'\n elif show_navigation and not show_status:\n policy_control_value = 'immersive.status=*'\n elif not show_navigation and show_status:\n policy_control_value = 'immersive.navigation=*'\n else:\n policy_control_value = 'immersive.full=*'\n\n self._adb_call_parser.parse(\n adb_pb2.AdbRequest(\n settings=adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(\n key='policy_control', value=policy_control_value\n ),\n )\n )\n )\n\n def _update_orientation(self) -> None:\n \"\"\"Updates the current device orientation.\"\"\"\n\n # Skip fetching the orientation if we already have it.\n if not np.all(self._orientation == np.zeros(4)):\n return\n\n orientation_response = self._adb_call_parser.parse(\n adb_pb2.AdbRequest(\n get_orientation=adb_pb2.AdbRequest.GetOrientationRequest()\n )\n )\n if orientation_response.status != adb_pb2.AdbResponse.Status.OK:\n logging.error('Got bad orientation: %r', orientation_response)\n return\n\n orientation = orientation_response.get_orientation.orientation\n if orientation not in {0, 1, 2, 3}:\n logging.error('Got bad orientation: %r', orientation)\n return\n\n # Transform into one-hot format.\n orientation_onehot = np.zeros([4], dtype=np.uint8)\n orientation_onehot[orientation] = 1\n self._orientation = orientation_onehot\n"
},
{
"path": "android_env/components/device_settings_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import config_classes\nfrom android_env.components import device_settings as device_settings_lib\nfrom android_env.components.simulators import base_simulator\nimport numpy as np\n\n\nclass DeviceSettingsTest(parameterized.TestCase):\n\n def test_screen_size_before_update(self):\n \"\"\"The screen size should be 0x0 without calling `update()`.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n device_settings = device_settings_lib.DeviceSettings(simulator)\n\n # Act.\n height = device_settings.screen_height()\n width = device_settings.screen_width()\n\n # Assert.\n self.assertEqual(height, 0)\n self.assertEqual(width, 0)\n\n def test_screen_size_after_update(self):\n \"\"\"The screen size should be set after calling `update()`.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n simulator.get_screenshot.return_value = np.random.randint(\n low=0, high=255, size=(123, 456, 3), dtype=np.uint8\n )\n adb_controller = simulator.create_adb_controller.return_value\n adb_controller.execute_command.return_value = b''\n device_settings = device_settings_lib.DeviceSettings(simulator)\n\n # Act.\n device_settings.update(config_classes.DeviceSettingsConfig())\n height = device_settings.screen_height()\n width = device_settings.screen_width()\n\n # Assert.\n self.assertEqual(height, 123)\n self.assertEqual(width, 456)\n\n @parameterized.named_parameters(\n (\n 'show_touches',\n config_classes.DeviceSettingsConfig(show_touches=True),\n mock.call(\n ['shell', 'settings', 'put', 'system', 'show_touches', '1'],\n timeout=None,\n ),\n ),\n (\n 'show_touches_false',\n config_classes.DeviceSettingsConfig(show_touches=False),\n mock.call(\n ['shell', 'settings', 'put', 'system', 'show_touches', '0'],\n timeout=None,\n ),\n ),\n (\n 'show_pointer_location',\n config_classes.DeviceSettingsConfig(show_pointer_location=True),\n mock.call(\n ['shell', 'settings', 'put', 'system', 'pointer_location', '1'],\n timeout=None,\n ),\n ),\n (\n 'show_pointer_location_false',\n config_classes.DeviceSettingsConfig(show_pointer_location=False),\n mock.call(\n ['shell', 'settings', 'put', 'system', 'pointer_location', '0'],\n timeout=None,\n ),\n ),\n (\n 'show_navigation_and_status',\n config_classes.DeviceSettingsConfig(\n show_navigation_bar=True, show_status_bar=True\n ),\n mock.call(\n ['shell', 'settings', 'put', 'global', 'policy_control', 'null*'],\n timeout=None,\n ),\n ),\n (\n 'show_navigation_and_no_status',\n config_classes.DeviceSettingsConfig(\n show_navigation_bar=True, show_status_bar=False\n ),\n mock.call(\n [\n 'shell',\n 'settings',\n 'put',\n 'global',\n 'policy_control',\n 'immersive.status=*',\n ],\n timeout=None,\n ),\n ),\n (\n 'show_no_navigation_and_status',\n config_classes.DeviceSettingsConfig(\n show_navigation_bar=False, show_status_bar=True\n ),\n mock.call(\n [\n 'shell',\n 'settings',\n 'put',\n 'global',\n 'policy_control',\n 'immersive.navigation=*',\n ],\n timeout=None,\n ),\n ),\n (\n 'show_no_navigation_and_no_status',\n config_classes.DeviceSettingsConfig(\n show_navigation_bar=False, show_status_bar=False\n ),\n mock.call(\n [\n 'shell',\n 'settings',\n 'put',\n 'global',\n 'policy_control',\n 'immersive.full=*',\n ],\n timeout=None,\n ),\n ),\n )\n def test_update(\n self, settings: config_classes.DeviceSettingsConfig, expected_call\n ):\n \"\"\"We expect the right call for each setting.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n adb_controller = simulator.create_adb_controller.return_value\n adb_controller.execute_command.return_value = b''\n device_settings = device_settings_lib.DeviceSettings(simulator)\n\n # Act.\n device_settings.update(settings)\n\n # Assert.\n adb_controller.execute_command.assert_has_calls(\n [expected_call], any_order=True\n )\n\n def test_get_orientation_bad_response(self):\n \"\"\"The orientation should be unset if the underlying response is bad.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n adb_controller = simulator.create_adb_controller.return_value\n adb_controller.execute_command.return_value = b''\n device_settings = device_settings_lib.DeviceSettings(simulator)\n\n # Act.\n orientation = device_settings.get_orientation()\n\n # Assert.\n np.testing.assert_array_equal(orientation, np.zeros(4))\n\n def test_get_orientation_bad_orientation(self):\n \"\"\"The orientation should be unset if the underlying orientation is bad.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n adb_controller = simulator.create_adb_controller.return_value\n adb_controller.execute_command.return_value = b' InputDeviceOrientation: 9'\n device_settings = device_settings_lib.DeviceSettings(simulator)\n\n # Act.\n orientation = device_settings.get_orientation()\n\n # Assert.\n np.testing.assert_array_equal(orientation, np.zeros(4))\n\n def test_get_orientation_success(self):\n \"\"\"Checks that the orientation comes back as expected.\"\"\"\n\n # Arrange.\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n adb_controller = simulator.create_adb_controller.return_value\n adb_controller.execute_command.return_value = b' InputDeviceOrientation: 3'\n device_settings = device_settings_lib.DeviceSettings(simulator)\n\n # Act.\n orientation = device_settings.get_orientation()\n # The output should be idempotent if the underlying system did not change.\n orientation_again = device_settings.get_orientation()\n\n # Assert.\n np.testing.assert_array_equal(orientation, np.array([0, 0, 0, 1]))\n np.testing.assert_array_equal(orientation, orientation_again)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/dumpsys_thread.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A ThreadFunction that runs and parses adb dumpsys.\"\"\"\n\nimport concurrent.futures\n\nfrom absl import logging\nfrom android_env.components import app_screen_checker as app_screen_checker_lib\n\n_Outcome = app_screen_checker_lib.AppScreenChecker.Outcome\n\n\nclass DumpsysThread:\n \"\"\"A thread that checks if the user is in the expected app screen.\"\"\"\n\n def __init__(\n self,\n app_screen_checker: app_screen_checker_lib.AppScreenChecker,\n check_frequency: int = 10,\n max_failed_current_activity: int = 10,\n ):\n \"\"\"Initializes the dumpsys reader thread.\n\n This loops forever checking if the user is in the expected screen dictated\n by `app_screen_checker`. These analyses are too expensive to be in the\n critical path of AndroidEnv::step() so we consume them async from this\n separate thread.\n\n Args:\n app_screen_checker: The class that actually determines if the current\n screen matches the expected screen.\n check_frequency: Integer. We only call dumpsys 1/check_frequency times in\n each iteration of the while loop below.\n max_failed_current_activity: Integer. We try to fetch the current activity\n but sometimes it fails. If it fails more than\n `max_failed_current_activity` consecutive times, we declare that the\n user has exited `expected_activity`.\n \"\"\"\n\n self._app_screen_checker = app_screen_checker\n self._main_loop_counter = 0\n self._check_frequency = check_frequency\n self._max_failed_activity_extraction = max_failed_current_activity\n self._num_failed_activity_extraction = 0\n self._latest_check: concurrent.futures.Future | None = None\n\n def check_user_exited(self, timeout: float | None = None) -> bool:\n \"\"\"Returns True if the user is not in the expected screen.\n\n Args:\n timeout: An optional time in seconds to block waiting for the result of\n the (expensive) checking operation. If None, the function will return\n immediately with `False`.\n\n Returns:\n Whether the user of the Android device has exited the expected screen\n determined by `AppScreenChecker` given at __init__().\n \"\"\"\n\n # Update and check loop_counter against check_frequency.\n self._main_loop_counter += 1\n if (self._check_frequency <= 0 or\n self._main_loop_counter < self._check_frequency):\n return False\n self._main_loop_counter = 0\n\n # If the latest check is None, perform a check and return.\n if self._latest_check is None:\n with concurrent.futures.ThreadPoolExecutor() as executor:\n self._latest_check = executor.submit(self._check_impl)\n return False\n\n # If there's a check in flight, continue only if it's finished.\n if not timeout and not self._latest_check.done():\n return False\n\n v = self._latest_check.result(timeout=timeout)\n self._latest_check = None # Reset the check.\n return v\n\n def _check_impl(self) -> bool:\n \"\"\"The synchronous implementation of Dumpsys.\"\"\"\n\n outcome = self._app_screen_checker.matches_current_app_screen()\n\n # We were unable to determine the current activity.\n if outcome == _Outcome.FAILED_ACTIVITY_EXTRACTION:\n self._num_failed_activity_extraction += 1\n logging.info('self._num_failed_activity_extraction: %s',\n self._num_failed_activity_extraction)\n if (self._num_failed_activity_extraction >=\n self._max_failed_activity_extraction):\n logging.error('Maximum number of failed activity extraction reached.')\n self._num_failed_activity_extraction = 0\n return True\n else:\n self._num_failed_activity_extraction = 0\n\n # The current app screen matches all expectations.\n if (outcome == _Outcome.SUCCESS or\n outcome == _Outcome.EMPTY_EXPECTED_ACTIVITY):\n return False\n\n # Player has exited the app. Terminate the episode.\n elif outcome == _Outcome.UNEXPECTED_ACTIVITY:\n return True\n\n # Player has exited the main game. Terminate the episode.\n elif outcome == _Outcome.UNEXPECTED_VIEW_HIERARCHY:\n return True\n\n return False\n"
},
{
"path": "android_env/components/dumpsys_thread_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.dumpsys_thread.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import app_screen_checker as screen_checker\nfrom android_env.components import dumpsys_thread\n\n\nclass DumpsysThreadTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self._app_screen_checker = mock.create_autospec(\n screen_checker.AppScreenChecker)\n\n def test_unexpected_activity(self):\n dumpsys = dumpsys_thread.DumpsysThread(\n app_screen_checker=self._app_screen_checker, check_frequency=1)\n outcome = screen_checker.AppScreenChecker.Outcome.UNEXPECTED_ACTIVITY\n self._app_screen_checker.matches_current_app_screen.return_value = outcome\n # The first time that `check_user_exited()` is called, it'll only trigger\n # the processing, but it should return immediately.\n self.assertFalse(dumpsys.check_user_exited(timeout=1.0))\n # The second time it should then wait for the result.\n self.assertTrue(dumpsys.check_user_exited(timeout=1.0))\n\n def test_unexpected_view_hierarchy(self):\n dumpsys = dumpsys_thread.DumpsysThread(\n app_screen_checker=self._app_screen_checker, check_frequency=1)\n outcome = screen_checker.AppScreenChecker.Outcome.UNEXPECTED_VIEW_HIERARCHY\n self._app_screen_checker.matches_current_app_screen.return_value = outcome\n self.assertFalse(dumpsys.check_user_exited(timeout=1.0))\n self.assertTrue(dumpsys.check_user_exited(timeout=1.0))\n\n def test_success(self):\n dumpsys = dumpsys_thread.DumpsysThread(\n app_screen_checker=self._app_screen_checker, check_frequency=1)\n outcome = screen_checker.AppScreenChecker.Outcome.SUCCESS\n self._app_screen_checker.matches_current_app_screen.return_value = outcome\n self.assertFalse(dumpsys.check_user_exited(timeout=1.0))\n self.assertFalse(dumpsys.check_user_exited(timeout=1.0))\n\n def test_skipped(self):\n dumpsys = dumpsys_thread.DumpsysThread(\n app_screen_checker=self._app_screen_checker, check_frequency=5)\n self._app_screen_checker.matches_current_app_screen.side_effect = [\n screen_checker.AppScreenChecker.Outcome.SUCCESS,\n screen_checker.AppScreenChecker.Outcome.FAILED_ACTIVITY_EXTRACTION\n ]\n\n for _ in range(17):\n self.assertFalse(dumpsys.check_user_exited(timeout=1.0))\n\n # The first 4 calls will hit the early exit from `check_frequency`.\n # The 5th call will trigger the processing (increasing the call count to\n # matches_current_app_screen() by 1), but it should return early.\n # The 10th call will find a result of the previous processing, and it should\n # be SUCCESS.\n # The next 4 calls (11, 12, 13, 14) will hit the early exit from\n # `check_frequency`.\n # The 15th call should trigger the processing again (increasing the call\n # count to matches_current_app_screen() by 1), but it should return early.\n # The next 2 calls (16, 17) will hit the early exit from `check_frequency`.\n # In total there should be only two calls to `matches_current_app_screen()`.\n expected_call_count = 2\n self.assertEqual(\n self._app_screen_checker.matches_current_app_screen.call_count,\n expected_call_count)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/errors.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Definitions of exceptions used by AndroidEnv.\"\"\"\n\n\nclass AndroidEnvError(Exception):\n \"\"\"Base class for all known errors generated by AndroidEnv.\"\"\"\n\n # An integer that identifies this class of error.\n # Subclasses should use a different value.\n ERROR_CODE: int = 0\n\n\nclass ReadObservationError(AndroidEnvError):\n \"\"\"When the environment is unable to obtain an observation from a simulator.\"\"\"\n\n ERROR_CODE = 1\n\n\nclass CoordinatorError(AndroidEnvError):\n \"\"\"Error raised by the Coordinator.\"\"\"\n\n ERROR_CODE = 2\n\n\nclass TooManyRestartsError(CoordinatorError):\n \"\"\"The number of restarts has exceeded _MAX_RESTART_TRIES.\"\"\"\n\n ERROR_CODE = 3\n\n\nclass AdbControllerError(AndroidEnvError):\n \"\"\"Errors that can be raised by ADBController.\"\"\"\n\n ERROR_CODE = 4\n\n\nclass SimulatorError(AndroidEnvError):\n \"\"\"Errors that can be raised by a simulator.\"\"\"\n\n ERROR_CODE = 5\n\n\nclass SendActionError(AndroidEnvError):\n \"\"\"Raised when action couldn't be sent successfully.\"\"\"\n\n ERROR_CODE = 6\n\n\nclass StepCommandError(AndroidEnvError):\n \"\"\"Raised when setup step interpreter cannot process a command.\"\"\"\n\n ERROR_CODE = 7\n\n\nclass WaitForAppScreenError(StepCommandError):\n \"\"\"Raised when the wait_for_app_screen success check is not met.\"\"\"\n\n ERROR_CODE = 8\n\n\nclass CheckInstallError(StepCommandError):\n \"\"\"Raised when the check_install success check is not met.\"\"\"\n\n ERROR_CODE = 9\n\n\ndef from_code(code: int, msg: str = '') -> AndroidEnvError | None:\n \"\"\"Returns an AndroidEnvError instance from the given arguments.\"\"\"\n\n code_to_error = {\n 0: AndroidEnvError,\n 1: ReadObservationError,\n 2: CoordinatorError,\n 3: TooManyRestartsError,\n 4: AdbControllerError,\n 5: SimulatorError,\n 6: SendActionError,\n 7: StepCommandError,\n 8: WaitForAppScreenError,\n 9: CheckInstallError,\n }\n\n if code in code_to_error:\n return code_to_error[code](msg)\n"
},
{
"path": "android_env/components/errors_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for errors.py.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import errors\n\n\nclass ErrorsTest(parameterized.TestCase):\n\n @parameterized.parameters(\n (errors.ReadObservationError, 1),\n (errors.CoordinatorError, 2),\n (errors.TooManyRestartsError, 3),\n (errors.AdbControllerError, 4),\n (errors.SimulatorError, 5),\n (errors.SendActionError, 6),\n (errors.StepCommandError, 7),\n (errors.WaitForAppScreenError, 8),\n (errors.CheckInstallError, 9),\n )\n def test_error_codes(self, error, expected_error_code):\n with self.assertRaises(error) as context:\n raise error()\n self.assertEqual(context.exception.ERROR_CODE, expected_error_code)\n\n def test_error_codes_unique(self):\n error_codes = set()\n errors_list = (\n errors.ReadObservationError,\n errors.CoordinatorError,\n errors.TooManyRestartsError,\n errors.AdbControllerError,\n errors.SimulatorError,\n errors.SendActionError,\n errors.StepCommandError,\n errors.WaitForAppScreenError,\n errors.CheckInstallError,\n )\n for error in errors_list:\n self.assertNotIn(error.ERROR_CODE, error_codes)\n error_codes.add(error.ERROR_CODE)\n\n @parameterized.parameters([\n errors.ReadObservationError(),\n errors.CoordinatorError(),\n errors.TooManyRestartsError(),\n errors.AdbControllerError(),\n errors.SimulatorError(),\n errors.SendActionError(),\n errors.StepCommandError(),\n errors.WaitForAppScreenError(),\n errors.CheckInstallError(),\n ])\n def test_all_errors_are_androidenv_errors(self, error):\n self.assertIsInstance(error, errors.AndroidEnvError)\n\n @parameterized.named_parameters([\n ('less_than_zero', -1),\n # The largest `ERROR_CODE` is currently `CheckInstallError == 10`.\n ('greater_than_all_errors', 10 + 1),\n ('less_than_zero_float', -3.14159265),\n ('greater_than_all_errors_float', 123.456),\n ])\n def test_from_code_unsupported_code(self, code: int):\n \"\"\"Unsupported errors should raise `RuntimeError`.\"\"\"\n\n self.assertIsNone(errors.from_code(code))\n\n @parameterized.parameters([\n (-1, None, 'No such error code.'),\n (0, errors.AndroidEnvError, 'hello'),\n (0, errors.AndroidEnvError, ''),\n (1, errors.ReadObservationError, 'Could not read obs.'),\n (2, errors.CoordinatorError, 'Some error'),\n (3, errors.TooManyRestartsError, 'Too many already...'),\n (4, errors.AdbControllerError, 'Some adb error...'),\n (5, errors.SimulatorError, 'Simulator is not coping.'),\n (6, errors.SendActionError, 'Could not send action.'),\n (7, errors.StepCommandError, 'Some issue setting up the task.'),\n (8, errors.WaitForAppScreenError, 'Waited for too long!'),\n (9, errors.CheckInstallError, 'App did not install correctly.'),\n ])\n def test_from_code(self, code: int, expected_class: errors.AndroidEnvError,\n msg: str):\n \"\"\"`from_code` should produce consistent outputs for known errors.\"\"\"\n\n error = errors.from_code(code, msg)\n if error is not None:\n self.assertIsInstance(error, expected_class)\n self.assertEqual(error.ERROR_CODE, code)\n self.assertEqual(str(error), msg)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/log_stream.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Abstract class for handling a stream of logs from a simulator.\"\"\"\n\nimport abc\nfrom collections.abc import Generator, Sequence\nimport threading\nfrom absl import logging\n\n\nclass LogStream(metaclass=abc.ABCMeta):\n \"\"\"Manages the stream of logs output by a simulator.\"\"\"\n\n def __init__(self, verbose: bool = False):\n self._verbose = verbose\n self._filters = []\n self._should_stream = threading.Event()\n\n def get_stream_output(self) -> Generator[str, None, None]:\n \"\"\"Starts log process and returns the stream of logs.\"\"\"\n for line in self._get_stream_output():\n if self._verbose:\n logging.info('line: %r', line)\n if self._should_stream.is_set():\n yield line\n\n @abc.abstractmethod\n def _get_stream_output(self):\n \"\"\"Starts log process and returns the stream of logs.\"\"\"\n pass\n\n @abc.abstractmethod\n def stop_stream(self) -> None:\n \"\"\"Terminates the log stream.\n\n NOTE: This should only be called _after_ `get_stream_output()`.\n \"\"\"\n\n def pause_stream(self) -> None:\n \"\"\"No lines are yielded while the event is not set.\"\"\"\n logging.info('Pausing LogStream.')\n self._should_stream.clear()\n\n def resume_stream(self) -> None:\n \"\"\"The stream will continue yielding lines if the event is set.\"\"\"\n logging.info('Resuming LogStream.')\n self._should_stream.set()\n\n def set_log_filters(self, log_filters: Sequence[str]):\n \"\"\"Sets the filters for the log stream.\"\"\"\n self._filters = list(log_filters) + ['*:S']\n"
},
{
"path": "android_env/components/log_stream_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for log_stream.\"\"\"\n\nfrom absl.testing import absltest\nfrom android_env.components import log_stream\n\n\nclass FakeLogStream(log_stream.LogStream):\n\n def __init__(self, filter_name: str):\n super().__init__()\n self._filter_name = filter_name\n\n def _get_stream_output(self):\n \"\"\"Starts a log process and returns the stream of logs.\"\"\"\n lines = [\n f'{self._filter_name} fake_line_1',\n 'fake_line_2',\n f'{self._filter_name} fake_line_3',\n f'{self._filter_name} fake_line_4',\n 'fake_line_5',\n 'fake_line_6',\n ]\n for line in lines:\n if f'{self._filter_name}:V' in self._filters:\n if self._filter_name in line:\n yield line\n else:\n yield line\n\n def stop_stream(self):\n \"\"\"Stops the log stream from the simulator.\"\"\"\n pass\n\n\nclass LogStreamTest(absltest.TestCase):\n\n def test_get_stream_output(self):\n filter_name = 'AndroidRLTask'\n stream = FakeLogStream(filter_name=filter_name)\n stream.resume_stream()\n stream_output = stream.get_stream_output()\n expected_lines = [\n f'{filter_name} fake_line_1',\n 'fake_line_2',\n f'{filter_name} fake_line_3',\n f'{filter_name} fake_line_4',\n 'fake_line_5',\n 'fake_line_6',\n ]\n for line, expected_line in zip(stream_output, expected_lines):\n self.assertEqual(line, expected_line)\n\n def test_set_log_filters(self):\n filter_name = 'AndroidRLTask'\n stream = FakeLogStream(filter_name=filter_name)\n stream.set_log_filters([f'{filter_name}:V'])\n stream.resume_stream()\n stream_output = stream.get_stream_output()\n expected_lines = [\n f'{filter_name} fake_line_1',\n f'{filter_name} fake_line_3',\n f'{filter_name} fake_line_4',\n ]\n for line, expected_line in zip(stream_output, expected_lines):\n self.assertEqual(line, expected_line)\n\n def test_pause_resume_stream(self):\n filter_name = 'AndroidRLTask'\n stream = FakeLogStream(filter_name=filter_name)\n stream.resume_stream()\n stream_output = stream.get_stream_output()\n expected_lines = [\n f'{filter_name} fake_line_1',\n 'fake_line_2',\n f'{filter_name} fake_line_3',\n f'{filter_name} fake_line_4',\n 'fake_line_5',\n 'fake_line_6',\n ]\n for line, expected_line in zip(stream_output, expected_lines):\n self.assertEqual(line, expected_line)\n # If the stream is paused, we expect no lines to be yielded.\n stream.pause_stream()\n stream_output = list(stream.get_stream_output())\n self.assertEmpty(stream_output)\n # If the stream is resumed, we expect to see all lines yielded.\n stream.resume_stream()\n stream_output = stream.get_stream_output()\n for line, expected_line in zip(stream_output, expected_lines):\n self.assertEqual(line, expected_line)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/logcat_thread.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A class that launches a thread to read Android log outputs.\"\"\"\n\nfrom collections.abc import Callable\nimport dataclasses\nimport re\nimport threading\n\nfrom absl import logging\nfrom android_env.components import log_stream as log_stream_lib\n\n\n@dataclasses.dataclass\nclass EventListener:\n \"\"\"A function that's called when an event is triggered.\"\"\"\n\n regexp: re.Pattern[str]\n handler_fn: Callable[[re.Pattern[str], re.Match[str]], None]\n\n\nclass LogcatThread:\n \"\"\"Reads log entries in a separate thread.\"\"\"\n\n def __init__(self, log_stream: log_stream_lib.LogStream):\n \"\"\"Initializes this LogcatThread with optional filters.\n\n Please see https://developer.android.com/studio/command-line/logcat for more\n info on `logcat`.\n\n Args:\n log_stream: Stream of logs from simulator.\n \"\"\"\n\n self._log_stream = log_stream\n self._listeners = {}\n self._line_ready = threading.Event()\n self._line_ready.set()\n self._should_stop = threading.Event()\n self._thread = threading.Thread(target=self._process_logs)\n self._thread.daemon = True\n self._thread.start()\n\n def add_event_listener(self, event_listener: EventListener) -> None:\n \"\"\"Adds `fn` to the list of handlers to call when `event` occurs.\"\"\"\n event_regexp = event_listener.regexp\n if event_regexp not in self._listeners:\n self._listeners[event_regexp] = []\n self._listeners[event_regexp].append(event_listener.handler_fn)\n\n def remove_event_listener(self, event_listener: EventListener) -> None:\n \"\"\"Removes `fn` from the list of handlers to call when `event` occurs.\"\"\"\n event_regexp = event_listener.regexp\n if event_regexp not in self._listeners:\n logging.error('Event: %r is not registered.', event_regexp)\n return\n self._listeners[event_regexp].remove(event_listener.handler_fn)\n\n def line_ready(self) -> threading.Event:\n \"\"\"Indicates whether all listeners have been notified for a given line.\"\"\"\n return self._line_ready\n\n def pause(self):\n self._log_stream.pause_stream()\n\n def resume(self):\n \"\"\"Resume or restart the thread if it's dead after resetting environment.\"\"\"\n if not self._thread.is_alive():\n self._should_stop.clear()\n self._thread = threading.Thread(target=self._process_logs)\n self._thread.daemon = True\n self._thread.start()\n self._log_stream.resume_stream()\n\n def kill(self):\n self._should_stop.set()\n self._log_stream.stop_stream()\n self._thread.join(timeout=3.0)\n\n def _process_logs(self) -> None:\n \"\"\"A loop that runs until `self._should_stop` is set().\"\"\"\n\n # pylint: disable=g-line-too-long\n # Format is: \"TIME_SEC PID TID PRIORITY TAG: MESSAGE\"\n #\n # Example:\n # ' 1553110400.424 5583 5658 D NostalgicRacer: com.google.example.games.nostalgicracer.views.renderers.OpenGLRenderDriver@912fb8.onSurfaceChanged 480x320' #\n # pylint: enable=g-line-too-long\n\n logline_regexp = r\"\"\"\n ^ # Beginning of the line.\n [ ]+(?P[0-9]+\\.[0-9]+) # Spaces and a float.\n [ ]+(?P[0-9]+) # Spaces and an int.\n [ ]+(?P[0-9]+) # Spaces and an int.\n [ ]+(?P.) # Spaces and any single character.\n [ ]+(?P[^:]*): # Spaces and any char that's not ':'.\n [ ](?P.*)$ # The actual log message.\n \"\"\"\n logline_re = re.compile(logline_regexp, re.VERBOSE)\n\n for line in self._log_stream.get_stream_output():\n if self._should_stop.is_set():\n break\n\n if not line: # Skip empty lines.\n continue\n\n matches = logline_re.match(line)\n if not matches or len(matches.groups()) != 6:\n continue\n\n # Make sure that values are not read until all listeners are notified.\n self._line_ready.clear()\n\n # We're currently only consuming `message`, but we may use the other\n # fields in the future.\n content = matches.group('message')\n for ev, listeners in self._listeners.items():\n ev_matches = ev.match(content)\n if ev_matches:\n for listener in listeners: # Notify listeners.\n listener(ev, ev_matches)\n\n self._line_ready.set() # Allow consumers to read values.\n"
},
{
"path": "android_env/components/logcat_thread_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.logcat_thread.\"\"\"\n\nimport re\nimport threading\n\nfrom absl.testing import absltest\nfrom android_env.components import log_stream\nfrom android_env.components import logcat_thread\nfrom android_env.proto import task_pb2\n\n\nclass FakeStream:\n \"\"\"This class simulates the logs coming from ADB.\"\"\"\n\n def __init__(self):\n self._values = []\n self._kill = False\n self._lock = threading.Lock()\n\n def send_value(self, value):\n with self._lock:\n self._values.append(value)\n\n def has_next_value(self):\n return bool(self._values)\n\n def kill(self):\n self._kill = True\n\n def reset(self):\n self._kill = False\n\n def __iter__(self):\n while True:\n if self._kill:\n return\n if not self._values:\n continue\n else:\n with self._lock:\n next_value = self._values.pop(0)\n yield next_value\n\n\ndef make_stdout(data):\n \"\"\"Returns a valid log output with given data as message.\"\"\"\n return ' 1553110400.424 5583 5658 D Tag: %s' % data\n\n\nclass FakeLogStream(log_stream.LogStream):\n \"\"\"FakeLogStream class that wraps a FakeStream.\"\"\"\n\n def __init__(self):\n super().__init__(verbose=False)\n self.logs = FakeStream()\n self.stream_is_alive = True\n\n def _get_stream_output(self):\n return self.logs\n\n def stop_stream(self):\n self.stream_is_alive = False\n self.logs.kill()\n\n def reset(self):\n self.stream_is_alive = True\n self.logs.reset()\n\n\nclass LogcatThreadTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self.fake_log_stream = FakeLogStream()\n\n def tearDown(self):\n self.fake_log_stream.stop_stream()\n super().tearDown()\n\n def test_set_filters(self):\n log_parsing_config = task_pb2.LogParsingConfig(filters=['AndroidRLTask:V'])\n self.fake_log_stream.set_log_filters(log_parsing_config.filters)\n _ = logcat_thread.LogcatThread(log_stream=self.fake_log_stream)\n expected_filters = ['AndroidRLTask:V', '*:S']\n self.assertEqual(expected_filters, self.fake_log_stream._filters)\n\n def test_kill(self):\n logcat = logcat_thread.LogcatThread(log_stream=self.fake_log_stream)\n self.assertTrue(self.fake_log_stream.stream_is_alive)\n logcat.kill()\n self.assertFalse(self.fake_log_stream.stream_is_alive)\n\n def test_listeners(self):\n \"\"\"Ensures that we can wait for a specific message without polling.\"\"\"\n logcat = logcat_thread.LogcatThread(log_stream=self.fake_log_stream)\n # Start yielding lines from LogStream.\n logcat.resume()\n\n # Set up a listener that modifies an arbitrary state.\n some_state = threading.Event()\n\n def my_handler(event: re.Pattern[str], match: re.Match[str]):\n del event, match\n nonlocal some_state\n some_state.set()\n\n # Create a desired event and hook up the listener.\n my_event = re.compile('Hello world')\n listener = logcat_thread.EventListener(my_event, my_handler)\n logcat.add_event_listener(listener)\n self.fake_log_stream.logs.send_value('Hi there!') # This should not match.\n self.assertFalse(some_state.is_set())\n self.fake_log_stream.logs.send_value(make_stdout('Hello world'))\n some_state.wait(timeout=1.0)\n self.assertTrue(some_state.is_set())\n\n # Waiting for any events should also trigger the listener.\n some_state.clear()\n self.fake_log_stream.logs.send_value(make_stdout('Hello world'))\n some_state.wait(timeout=1.0)\n self.assertTrue(some_state.is_set())\n\n # After removing the listener, it should not be called anymore.\n some_state.clear()\n logcat.remove_event_listener(listener)\n self.fake_log_stream.logs.send_value(make_stdout('Hello world'))\n some_state.wait(timeout=1.0)\n self.assertFalse(some_state.is_set())\n\n def test_resume_does_not_recreate_alive_thread(self):\n logcat = logcat_thread.LogcatThread(log_stream=self.fake_log_stream)\n thread_before = logcat._thread\n self.assertTrue(thread_before.is_alive())\n logcat.resume()\n thread_after = logcat._thread\n self.assertTrue(thread_after.is_alive())\n self.assertIs(thread_before, thread_after)\n\n def test_resume_recreates_thread(self):\n logcat = logcat_thread.LogcatThread(log_stream=self.fake_log_stream)\n self.assertTrue(logcat._thread.is_alive())\n logcat.kill()\n self.assertFalse(logcat._thread.is_alive())\n self.assertTrue(logcat._should_stop.is_set())\n self.fake_log_stream.reset()\n logcat.resume()\n self.assertTrue(logcat._thread.is_alive())\n self.assertFalse(logcat._should_stop.is_set())\n self.assertTrue(logcat._thread.daemon)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/pixel_fns.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utils for AndroidEnv.\"\"\"\n\nfrom collections.abc import Sequence\n\nfrom dm_env import specs\nimport numpy as np\n\n\ndef touch_position_to_pixel_position(\n touch_position: np.ndarray,\n width_height: Sequence[int],\n) -> tuple[int, int]:\n \"\"\"Maps touch position in [0,1] to the corresponding pixel on the screen.\"\"\"\n touch_pixels = (touch_position * width_height).astype(np.int32)\n cap_idx = lambda v, idx_len: min(v, idx_len - 1)\n return tuple(map(cap_idx, touch_pixels, width_height))\n\n\ndef transpose_pixels(frame: np.ndarray) -> np.ndarray:\n \"\"\"Converts image from shape (H, W, C) to (W, H, C) and vice-versa.\"\"\"\n return np.transpose(frame, axes=(1, 0, 2))\n\n\ndef orient_pixels(frame: np.ndarray, orientation: int) -> np.ndarray:\n \"\"\"Rotates screen pixels according to the given orientation.\"\"\"\n\n match orientation:\n case 0: # PORTRAIT_90\n return frame\n case 1: # LANDSCAPE_90\n return np.rot90(frame, k=3, axes=(0, 1))\n case 2: # PORTRAIT_180\n return np.rot90(frame, k=2, axes=(0, 1))\n case 3: # LANDSCAPE_270\n return np.rot90(frame, k=1, axes=(0, 1))\n case _:\n raise ValueError(\n 'Orientation must be an integer in [0, 3] but is %r' % orientation\n )\n\n\ndef convert_int_to_float(data: np.ndarray, data_spec: specs.Array):\n \"\"\"Converts an array of int values to floats between 0 and 1.\"\"\"\n\n if not np.issubdtype(data.dtype, np.integer):\n raise TypeError(f'{data.dtype} is not an integer type')\n if isinstance(data_spec, specs.BoundedArray):\n value_min = data_spec.minimum\n value_max = data_spec.maximum\n else:\n # We use the int type to figure out the boundaries.\n iinfo = np.iinfo(data_spec.dtype)\n value_min = iinfo.min\n value_max = iinfo.max\n return np.float32(1.0 * (data - value_min) / (value_max - value_min))\n"
},
{
"path": "android_env/components/pixel_fns_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for pixel_fns.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import pixel_fns\nfrom dm_env import specs\nimport numpy as np\n\n\nclass UtilsTest(parameterized.TestCase):\n\n @parameterized.parameters(\n ([0.5, 0.5], [320, 480], (160, 240)),\n ([0.25, 0.75], [320, 480], (80, 360)),\n ([0.0, 0.0], [320, 480], (0, 0)),\n ([1.0, 1.0], [320, 480], (319, 479)),\n )\n def test_touch_position_to_pixel_position(\n self, touch_pos, width_height, pixel_pos):\n self.assertEqual(\n pixel_fns.touch_position_to_pixel_position(\n np.array(touch_pos), width_height\n ),\n pixel_pos,\n )\n\n def test_transpose_pixels(self):\n image = np.reshape(np.array(range(12)), (3, 2, 2))\n expected = [[[0, 1], [4, 5], [8, 9]], [[2, 3], [6, 7], [10, 11]]]\n self.assertEqual(pixel_fns.transpose_pixels(image).shape, (2, 3, 2))\n self.assertTrue((pixel_fns.transpose_pixels(image) == expected).all())\n\n def test_orient_pixels(self):\n image = np.reshape(np.array(range(12)), (3, 2, 2))\n\n expected_90 = [[[8, 9], [4, 5], [0, 1]], [[10, 11], [6, 7], [2, 3]]]\n rot_90 = 1 # LANDSCAPE_90\n rotated = pixel_fns.orient_pixels(image, rot_90)\n self.assertEqual(rotated.shape, (2, 3, 2))\n self.assertTrue((rotated == expected_90).all())\n\n expected_180 = [[[10, 11], [8, 9]], [[6, 7], [4, 5]], [[2, 3], [0, 1]]]\n rot_180 = 2 # PORTRAIT_180\n rotated = pixel_fns.orient_pixels(image, rot_180)\n self.assertEqual(rotated.shape, (3, 2, 2))\n self.assertTrue((rotated == expected_180).all())\n\n expected_270 = [[[2, 3], [6, 7], [10, 11]], [[0, 1], [4, 5], [8, 9]]]\n rot_270 = 3 # LANDSCAPE_270\n rotated = pixel_fns.orient_pixels(image, rot_270)\n self.assertEqual(rotated.shape, (2, 3, 2))\n self.assertTrue((rotated == expected_270).all())\n\n rot_0 = 0 # PORTRAIT_0\n rotated = pixel_fns.orient_pixels(image, rot_0)\n self.assertEqual(rotated.shape, (3, 2, 2))\n self.assertTrue((rotated == image).all())\n\n def test_convert_int_to_float_bounded_array(self):\n spec = specs.BoundedArray(\n shape=(4,),\n dtype=np.int32,\n minimum=[0, 1, 10, -2],\n maximum=[5, 5, 20, 2],\n name='bounded_array')\n data = np.array([2, 2, 10, 0], dtype=np.int32)\n float_data = pixel_fns.convert_int_to_float(data, spec)\n np.testing.assert_equal(\n np.array([2.0 / 5.0, 1.0 / 4.0, 0.0, 0.5], dtype=np.float32), float_data\n )\n\n def test_convert_int_to_float_bounded_array_broadcast(self):\n spec = specs.BoundedArray(\n shape=(3,), dtype=np.int16, minimum=2, maximum=4, name='bounded_array')\n data = np.array([2, 3, 4], dtype=np.int16)\n float_data = pixel_fns.convert_int_to_float(data, spec)\n np.testing.assert_equal(\n np.array([0.0, 0.5, 1.0], dtype=np.float32), float_data)\n\n def test_convert_int_to_float_no_bounds(self):\n spec = specs.Array(\n shape=(3,),\n dtype=np.int8, # int8 implies min=-128, max=127\n name='bounded_array')\n data = np.array([-128, 0, 127], dtype=np.int16)\n float_data = pixel_fns.convert_int_to_float(data, spec)\n np.testing.assert_equal(\n np.array([0.0, 128. / 255., 1.0], dtype=np.float32), float_data)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/setup_step_interpreter.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A component that parses and processes SetupSteps.\"\"\"\n\nfrom collections.abc import Sequence\nimport copy\nimport time\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env.components import adb_call_parser as adb_call_parser_lib\nfrom android_env.components import app_screen_checker\nfrom android_env.components import errors\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import task_pb2\n\n\nclass SetupStepInterpreter:\n \"\"\"An interpreter for SetupSteps.\"\"\"\n\n def __init__(self, adb_call_parser: adb_call_parser_lib.AdbCallParser):\n \"\"\"Initializes this interpreter.\n\n Args:\n adb_call_parser: An object to communicate with Android via ADB.\n \"\"\"\n self._adb_call_parser = adb_call_parser\n self._stats = {\n 'error_count_adb_request': 0,\n 'error_count_wait_for_app_screen': 0,\n 'error_count_check_install': 0,\n 'error_count_wait_for_message': 0,\n 'total_time_waiting_for_app_screen': 0\n }\n\n def stats(self) -> dict[str, Any]:\n return copy.deepcopy(self._stats)\n\n def interpret(self, setup_steps: Sequence[task_pb2.SetupStep]) -> None:\n \"\"\"Returns True if parsing and processing `setup_steps` is successful.\"\"\"\n if setup_steps:\n logging.info('Executing setup steps: %s', setup_steps)\n for step in setup_steps:\n self._process_step_command(step)\n logging.info('Done executing setup steps.')\n\n def _process_step_command(self, step_cmd: task_pb2.SetupStep) -> None:\n \"\"\"Processes a single step command from a reset or extra setup.\"\"\"\n\n if not step_cmd:\n logging.info('Empty step_cmd')\n return\n\n logging.info('Executing step_cmd: %r', step_cmd)\n step_type = step_cmd.WhichOneof('step')\n success_condition = step_cmd.success_condition\n success_check = success_condition.WhichOneof('check')\n assert step_type or success_check, (\n 'At least one of step and success_condition must be defined.')\n\n num_tries = 0\n max_retries = max(success_condition.num_retries, 3)\n latest_error = None\n while num_tries < max_retries:\n\n num_tries += 1\n\n try:\n unused_adb_response = self._execute_step_cmd(step_cmd, step_type)\n time.sleep(0.5)\n self._check_success(success_check, success_condition)\n return\n\n except NotImplementedError:\n logging.exception('Not implemented error! Skipping this step command.')\n return\n\n except errors.AdbControllerError as error:\n latest_error = error\n self._stats['error_count_adb_request'] += 1\n logging.exception('ADB call [%r] has failed. Try %d of %d.',\n step_cmd.adb_request, num_tries, max_retries)\n\n except errors.WaitForAppScreenError as error:\n latest_error = error\n self._stats['error_count_wait_for_app_screen'] += 1\n logging.exception('Failed to wait for app screen. Try %d of %d.',\n num_tries, max_retries)\n\n except errors.CheckInstallError as error:\n latest_error = error\n self._stats['error_count_check_install'] += 1\n logging.exception('Package [%r] not installed. Try %d of %d.',\n success_condition.check_install.package_name,\n num_tries, max_retries)\n\n raise errors.StepCommandError(\n f'Step failed: [{step_cmd}]') from latest_error\n\n def _execute_step_cmd(\n self, step_cmd: task_pb2.SetupStep, step_type: str | None\n ) -> adb_pb2.AdbResponse | None:\n \"\"\"Executes a step command of given type.\"\"\"\n\n match step_type:\n case None:\n return None\n case 'sleep':\n time.sleep(step_cmd.sleep.time_sec)\n return None\n case 'adb_request':\n response = self._adb_call_parser.parse(step_cmd.adb_request)\n if response.status != adb_pb2.AdbResponse.Status.OK:\n raise errors.AdbControllerError(\n f'Failed to execute AdbRequest [{step_cmd.adb_request}].\\n'\n f'Status: {response.status}\\n'\n f'Error: {response.error_message}'\n )\n return response\n case _:\n raise NotImplementedError(f'No step command of type [{step_type}].')\n\n def _check_success(\n self,\n success_check: str | None,\n success_condition: task_pb2.SuccessCondition,\n ) -> None:\n \"\"\"Checks whether the given success condition was met.\"\"\"\n\n match success_check:\n case None:\n return None\n case 'wait_for_app_screen':\n wait_for_app_screen = success_condition.wait_for_app_screen\n screen_checker = app_screen_checker.AppScreenChecker(\n adb_call_parser=self._adb_call_parser,\n expected_app_screen=wait_for_app_screen.app_screen,\n )\n wait_time = screen_checker.wait_for_app_screen(\n timeout_sec=wait_for_app_screen.timeout_sec\n )\n self._stats['total_time_waiting_for_app_screen'] += wait_time\n case 'check_install':\n self._check_install(success_condition.check_install)\n case _:\n raise NotImplementedError(f'No success check called [{success_check}].')\n\n def _check_install(self, check_install: task_pb2.CheckInstall) -> None:\n \"\"\"Checks that the given package is installed.\"\"\"\n\n package = check_install.package_name\n logging.info('Checking if package is installed: [%r]', package)\n\n request = adb_pb2.AdbRequest(\n package_manager=adb_pb2.AdbRequest.PackageManagerRequest(\n list=adb_pb2.AdbRequest.PackageManagerRequest.List(\n packages=adb_pb2.AdbRequest.PackageManagerRequest.List.Packages(\n ))))\n\n start_time = time.time()\n while time.time() - start_time < check_install.timeout_sec:\n response = self._adb_call_parser.parse(request)\n if package in response.package_manager.list.items:\n logging.info('Done confirming that package is installed.')\n return\n time.sleep(0.1)\n\n logging.error('Package not found.')\n raise errors.CheckInstallError()\n"
},
{
"path": "android_env/components/setup_step_interpreter_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.setup_step_interpreter.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import adb_call_parser\nfrom android_env.components import errors\nfrom android_env.components import setup_step_interpreter\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import task_pb2\n\nfrom google.protobuf import text_format\n\n\ndef _to_proto(proto_class, text):\n proto = proto_class()\n text_format.Parse(text, proto)\n return proto\n\n\nclass SetupStepInterpreterTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self._parser = mock.create_autospec(\n adb_call_parser.AdbCallParser, instance=True)\n\n def test_empty_setup_steps(self):\n \"\"\"Simple test where nothing should break, and nothing should be done.\n\n The test simply expects this test to not crash.\n \"\"\"\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([])\n\n def test_none_setup_steps(self):\n \"\"\"Simple test where nothing should break, and nothing should be done.\n\n The test simply expects this test to not crash.\n \"\"\"\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n # Empty setup steps should be ignored.\n interpreter.interpret([])\n\n def test_invalid_setup_step(self):\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n # Empty setup steps should be ignored.\n self.assertRaises(AssertionError, interpreter.interpret,\n [task_pb2.SetupStep()])\n\n def test_adb_install_apk_filesystem(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nadb_request: {\n install_apk: {\n filesystem: {\n path: \"/my/favorite/dir/my_apk.apk\"\n }\n }\n}\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(\n install_apk=adb_pb2.AdbRequest.InstallApk(\n filesystem=adb_pb2.AdbRequest.InstallApk.Filesystem(\n path='/my/favorite/dir/my_apk.apk'))))\n\n def test_adb_force_stop(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nadb_request: { force_stop: { package_name: \"my.app.Activity\" } }\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(\n force_stop=adb_pb2.AdbRequest.ForceStop(\n package_name='my.app.Activity')))\n\n def test_adb_start_activity(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nadb_request: {\n start_activity: {\n full_activity: \"my.app.Activity\"\n extra_args: \"arg1\"\n extra_args: \"arg2\"\n }\n}\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(\n start_activity=adb_pb2.AdbRequest.StartActivity(\n full_activity='my.app.Activity', extra_args=['arg1', 'arg2'])))\n\n def test_adb_single_tap(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(task_pb2.SetupStep, \"\"\"\nadb_request: {\n tap: {\n x: 321\n y: 654\n }\n}\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(tap=adb_pb2.AdbRequest.Tap(x=321, y=654)))\n\n def test_adb_press_button(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(task_pb2.SetupStep,\n \"\"\" adb_request: { press_button: { button: HOME } }\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(\n press_button=adb_pb2.AdbRequest.PressButton(\n button=adb_pb2.AdbRequest.PressButton.Button.HOME)))\n\n self._parser.reset_mock()\n interpreter.interpret([\n _to_proto(task_pb2.SetupStep,\n \"\"\" adb_request: { press_button: { button: BACK } }\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(\n press_button=adb_pb2.AdbRequest.PressButton(\n button=adb_pb2.AdbRequest.PressButton.Button.BACK)))\n\n def test_adb_start_screen_pinning(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nadb_request: {\n start_screen_pinning: {\n full_activity: \"my.app.HighlanderApp\" # \"There can be only one\".\n }\n}\"\"\")\n ])\n self._parser.parse.assert_called_once_with(\n adb_pb2.AdbRequest(\n start_screen_pinning=adb_pb2.AdbRequest.StartScreenPinning(\n full_activity='my.app.HighlanderApp')))\n\n @mock.patch('time.sleep')\n def test_time_sleep(self, mock_sleep):\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret(\n [_to_proto(task_pb2.SetupStep, \"\"\"sleep: { time_sec: 0.875 }\"\"\")])\n assert mock_sleep.call_count == 2\n mock_sleep.assert_has_calls([mock.call(0.875), mock.call(0.5)])\n\n @mock.patch('time.sleep')\n def test_wait_for_app_screen_empty_activity(self, unused_mock_sleep):\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n with self.assertRaises(errors.StepCommandError):\n interpreter.interpret([\n _to_proto(task_pb2.SetupStep,\n \"\"\"success_condition: {wait_for_app_screen: { }}\"\"\")\n ])\n\n @mock.patch('time.sleep')\n def test_check_install_not_installed(self, unused_mock_sleep):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(items=[\n 'com.some.package',\n 'not.what.you.are.looking.for',\n ])))\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n with self.assertRaises(errors.StepCommandError):\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nsuccess_condition: {\n check_install: {\n package_name: \"faz\"\n timeout_sec: 0.0001\n }\n}\n\"\"\")\n ])\n\n def test_check_install_installed(self):\n self._parser.parse.return_value = adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(items=[\n 'com.some.package',\n 'baz',\n ])))\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n # The test checks that this command raises no AssertionError.\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nsuccess_condition: {\n check_install: {\n package_name: \"baz\"\n timeout_sec: 0.0001\n }\n}\"\"\")\n ])\n\n def test_num_retries_failure(self):\n self._parser.parse.side_effect = [\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(\n items=[]))),\n ] * 3\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n with self.assertRaises(errors.StepCommandError):\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nsuccess_condition: {\n check_install: {\n package_name: \"faz\"\n timeout_sec: 0.0001\n }\n num_retries: 3\n}\"\"\")\n ])\n # We retried 3 times after the first call, so we expect 3+1 calls.\n self.assertEqual(self._parser.parse.call_count, 3)\n\n @mock.patch('time.sleep')\n def test_num_retries_success(self, unused_mock_sleep):\n self._parser.parse.side_effect = [\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(\n items=[]))),\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(\n items=[]))),\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(items=[\n 'com.some.package',\n 'bar',\n ]))),\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(items=[])))\n ]\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nsuccess_condition: {\n check_install: {\n package_name: \"bar\"\n timeout_sec: 0.0001\n }\n num_retries: 5\n}\"\"\")\n ])\n # The check should succeed on the third try.\n self.assertEqual(self._parser.parse.call_count, 3)\n\n def test_retry_step(self):\n self._parser.parse.side_effect = [\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(\n items=[]))),\n adb_pb2.AdbResponse(\n package_manager=adb_pb2.AdbResponse.PackageManagerResponse(\n list=adb_pb2.AdbResponse.PackageManagerResponse.List(items=[\n 'com.some.package',\n 'bar',\n ]))),\n ]\n interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=self._parser)\n interpreter.interpret([\n _to_proto(\n task_pb2.SetupStep, \"\"\"\nsuccess_condition: {\n check_install: {\n package_name: \"bar\"\n timeout_sec: 0.0001\n }\n num_retries: 2\n}\"\"\")\n ])\n # We expect the check to fail once and succeed on the second pass.\n self.assertEqual(self._parser.parse.call_count, 2)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/simulators/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/components/simulators/base_simulator.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A base class for talking to different types of Android simulators.\"\"\"\n\nimport abc\nfrom collections.abc import Callable\nimport threading\nimport time\n\nfrom absl import logging\nfrom android_env.components import adb_controller\nfrom android_env.components import config_classes\nfrom android_env.components import errors\nfrom android_env.components import log_stream\nfrom android_env.proto import state_pb2\nimport numpy as np\n\n\nclass BaseSimulator(metaclass=abc.ABCMeta):\n \"\"\"An interface for communicating with an Android simulator.\"\"\"\n\n def __init__(self, config: config_classes.SimulatorConfig):\n \"\"\"Instantiates a BaseSimulator object.\n\n The simulator may be an emulator, virtual machine or even a physical device.\n Each simulator has its own AdbController that is used for internal\n bookkeeping.\n\n Args:\n config: Settings for this simulator.\n \"\"\"\n\n self._config = config\n self._interaction_thread: InteractionThread | None = None\n\n # An increasing number that tracks the attempt at launching the simulator.\n self._num_launch_attempts: int = 0\n\n def get_logs(self) -> str:\n \"\"\"Returns logs recorded by the simulator (if provided).\"\"\"\n return 'No simulator logs provided.'\n\n @abc.abstractmethod\n def adb_device_name(self) -> str:\n \"\"\"Returns the device name that the adb client will connect to.\"\"\"\n\n @abc.abstractmethod\n def create_adb_controller(self) -> adb_controller.AdbController:\n \"\"\"Returns an ADB controller which can communicate with this simulator.\"\"\"\n\n @abc.abstractmethod\n def create_log_stream(self) -> log_stream.LogStream:\n \"\"\"Creates a stream of logs from the simulator.\"\"\"\n\n def launch(self) -> None:\n \"\"\"Starts the simulator.\"\"\"\n\n # Stop screenshot thread if it's enabled.\n if self._interaction_thread is not None:\n self._interaction_thread.stop()\n self._interaction_thread.join()\n\n self._num_launch_attempts += 1\n try:\n self._launch_impl()\n except Exception as error:\n for line in self.get_logs().splitlines():\n logging.error(line)\n raise errors.SimulatorError(\n 'Exception caught in simulator. Please see the simulator logs '\n 'above for more details.'\n ) from error\n\n # Start interaction thread.\n if self._config.interaction_rate_sec > 0:\n self._interaction_thread = InteractionThread(\n self._get_screenshot_impl, self._config.interaction_rate_sec\n )\n self._interaction_thread.start()\n\n @abc.abstractmethod\n def _launch_impl(self) -> None:\n \"\"\"Platform specific launch implementation.\"\"\"\n\n @abc.abstractmethod\n def send_touch(self, touches: list[tuple[int, int, bool, int]]) -> None:\n \"\"\"Sends a touch event to be executed on the simulator.\n\n Args:\n touches: A list of touch events. Each element in the list corresponds to a\n single touch event. Each touch event tuple should have:\n 0 x: The horizontal coordinate of this event.\n 1 y: The vertical coordinate of this event.\n 2 is_down: Whether the finger is touching or not the screen.\n 3 identifier: Identifies a particular finger in a multitouch event.\n \"\"\"\n\n @abc.abstractmethod\n def send_key(self, keycode: np.int32, event_type: str) -> None:\n \"\"\"Sends a keyboard event.\n\n Args:\n keycode: Represents a specific keyboard key. This is platform and\n simulator-specific.\n event_type: Type of key event to be sent.\n \"\"\"\n\n def load_state(\n self, request: state_pb2.LoadStateRequest\n ) -> state_pb2.LoadStateResponse:\n \"\"\"Loads a state.\n\n Args:\n request: A `LoadStateRequest` containing any parameters necessary to\n specify how/what state to load.\n\n Returns:\n A `LoadStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n raise NotImplementedError('This simulator does not support load_state()')\n\n def save_state(\n self, request: state_pb2.SaveStateRequest\n ) -> state_pb2.SaveStateResponse:\n \"\"\"Saves a state.\n\n Args:\n request: A `SaveStateRequest` containing any parameters necessary to\n specify how/what state to save.\n\n Returns:\n A `SaveStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n raise NotImplementedError('This simulator does not support save_state()')\n\n def get_screenshot(self) -> np.ndarray:\n \"\"\"Returns pixels representing the current screenshot of the simulator.\"\"\"\n\n if self._config.interaction_rate_sec > 0:\n assert self._interaction_thread is not None\n return self._interaction_thread.screenshot() # Async mode.\n else:\n return self._get_screenshot_impl() # Sync mode.\n\n @abc.abstractmethod\n def _get_screenshot_impl(self) -> np.ndarray:\n \"\"\"Actual implementation of `get_screenshot()`.\n\n The output numpy array should have shape [height, width, num_channels] and\n can be loaded into PIL using Image.fromarray(img, mode='RGB') and be saved\n as a PNG file using my_pil.save('/tmp/my_screenshot.png', 'PNG').\n \"\"\"\n\n def close(self):\n \"\"\"Frees up resources allocated by this object.\"\"\"\n\n if self._interaction_thread is not None:\n self._interaction_thread.stop()\n self._interaction_thread.join()\n\n\nclass InteractionThread(threading.Thread):\n \"\"\"A thread that gets screenshot in the background.\"\"\"\n\n def __init__(\n self,\n get_screenshot_fn: Callable[[], np.ndarray],\n interaction_rate_sec: float,\n ):\n super().__init__()\n self._get_screenshot_fn = get_screenshot_fn\n self._interaction_rate_sec = interaction_rate_sec\n self._should_stop = threading.Event()\n self._screenshot = self._get_screenshot_fn()\n\n def run(self):\n last_read = time.time()\n while not self._should_stop.is_set():\n self._screenshot = self._get_screenshot_fn()\n now = time.time()\n elapsed = now - last_read\n last_read = now\n sleep_time = self._interaction_rate_sec - elapsed\n if sleep_time > 0.0:\n time.sleep(sleep_time)\n logging.info('InteractionThread.run() finished.')\n\n def stop(self):\n logging.info('Stopping InteractionThread.')\n self._should_stop.set()\n\n def screenshot(self) -> np.ndarray:\n return self._screenshot\n"
},
{
"path": "android_env/components/simulators/base_simulator_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport itertools\nimport time\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import config_classes\nfrom android_env.components import errors\n# fake_simulator.FakeSimulator inherits from BaseSimulator, so there's no need\n# to import it here explicitly.\nfrom android_env.components.simulators import base_simulator\nfrom android_env.components.simulators.fake import fake_simulator\nimport numpy as np\n\n\nclass BaseSimulatorTest(absltest.TestCase):\n\n def test_launch(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(640, 480))\n )\n # The simulator should launch and not crash.\n simulator.launch()\n\n def test_launch_close(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig()\n )\n # The simulator should launch and not crash.\n simulator.launch()\n # Closing the simulator should also not crash.\n simulator.close()\n\n def test_get_screenshot(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(640, 480))\n )\n # The simulator should launch and not crash.\n simulator.launch()\n\n screenshot = simulator.get_screenshot()\n np.testing.assert_equal(screenshot.shape, [640, 480, 3])\n\n def test_print_logs_on_exception(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig()\n )\n with mock.patch.object(\n simulator, 'get_logs'\n ) as mock_get_logs, mock.patch.object(\n simulator, '_launch_impl', autospec=True\n ) as mock_launch:\n mock_launch.side_effect = ValueError('Oh no!')\n self.assertRaises(errors.SimulatorError, simulator.launch)\n mock_get_logs.assert_called_once()\n\n def test_get_screenshot_error_async(self):\n \"\"\"An exception in the underlying interaction thread should bubble up.\"\"\"\n\n # Arrange.\n mock_interaction_thread = mock.create_autospec(\n base_simulator.InteractionThread\n )\n mock_interaction_thread.screenshot.side_effect = (\n errors.ReadObservationError()\n )\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(interaction_rate_sec=0.5)\n )\n with mock.patch.object(\n base_simulator,\n 'InteractionThread',\n autospec=True,\n return_value=mock_interaction_thread,\n ):\n simulator.launch()\n\n # Act & Assert.\n self.assertRaises(errors.ReadObservationError, simulator.get_screenshot)\n\n # Cleanup.\n simulator.close()\n\n def test_get_screenshot_faster_than_screenshot_impl(self):\n \"\"\"Return same screenshot when step is faster than the interaction rate.\"\"\"\n\n # Arrange.\n slow_rate = 0.5\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(interaction_rate_sec=slow_rate)\n )\n\n # Act.\n with mock.patch.object(\n simulator, '_get_screenshot_impl', autospec=True\n ) as mock_get_screenshot_impl:\n mock_get_screenshot_impl.side_effect = (\n np.array(i, ndmin=3) for i in itertools.count(0, 1)\n )\n simulator.launch()\n # Get two screenshots one after the other without pausing.\n screenshot1 = simulator.get_screenshot()\n screenshot2 = simulator.get_screenshot()\n\n # Assert.\n self.assertAlmostEqual(screenshot1[0][0][0], screenshot2[0][0][0])\n\n # Cleanup.\n simulator.close()\n\n def test_get_screenshot_slower_than_screenshot_impl(self):\n \"\"\"Return different screenshots when step slower than the interaction rate.\"\"\"\n\n # Arrange.\n fast_rate = 0.01\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(interaction_rate_sec=fast_rate)\n )\n\n # Act.\n with mock.patch.object(\n simulator, '_get_screenshot_impl', autospec=True\n ) as mock_get_screenshot_impl:\n mock_get_screenshot_impl.side_effect = (\n np.array(i, ndmin=3) for i in itertools.count(0, 1)\n )\n simulator.launch()\n # Sleep for 500ms between two screenshots.\n screenshot1 = simulator.get_screenshot()\n time.sleep(0.5)\n screenshot2 = simulator.get_screenshot()\n\n # Assert.\n self.assertNotEqual(screenshot1[0][0][0], screenshot2[0][0][0])\n\n # Cleanup.\n simulator.close()\n\n def test_interaction_thread_closes_upon_relaunch(self):\n \"\"\"Async interaction should kill the InteractionThread when relaunching.\"\"\"\n\n # Arrange.\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(interaction_rate_sec=0.01)\n )\n mock_interaction_thread = mock.create_autospec(\n base_simulator.InteractionThread\n )\n\n # Act & Assert.\n with mock.patch.object(\n base_simulator,\n 'InteractionThread',\n autospec=True,\n return_value=mock_interaction_thread,\n ):\n simulator.launch()\n mock_interaction_thread.stop.assert_not_called()\n mock_interaction_thread.join.assert_not_called()\n simulator.launch()\n mock_interaction_thread.stop.assert_called_once()\n mock_interaction_thread.join.assert_called_once()\n simulator.close()\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/simulators/emulator/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/components/simulators/emulator/emulator_launcher.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Prepares and launches an emulator process.\"\"\"\n\nimport glob\nimport os\nimport subprocess\nimport tempfile\n\nfrom absl import logging\nfrom android_env.components import config_classes\n\n\nclass EmulatorLauncher:\n \"\"\"Handles launching an emulator.\"\"\"\n\n def __init__(\n self,\n config: config_classes.EmulatorLauncherConfig,\n adb_controller_config: config_classes.AdbControllerConfig,\n ):\n \"\"\"Launches an emulator.\"\"\"\n\n self._config = config\n self._adb_controller_config = adb_controller_config\n\n self._emulator = None\n self._emulator_output = None\n self._is_closed = False\n\n # Create directory for tmp files.\n # Note: this will be deleted once EmulatorLauncher instance is cleaned up.\n os.makedirs(config.tmp_dir, exist_ok=True)\n self._local_tmp_dir_handle = tempfile.TemporaryDirectory(\n dir=config.tmp_dir, prefix='simulator_instance_'\n )\n self._local_tmp_dir = self._local_tmp_dir_handle.name\n self._logfile_path = os.path.join(self._local_tmp_dir, 'emulator_output')\n logging.info('Simulator local_tmp_dir: %s', self._local_tmp_dir)\n\n def logfile_path(self) -> str:\n return self._logfile_path\n\n def launch_emulator_process(self) -> None:\n \"\"\"Launches the emulator.\"\"\"\n\n logging.info('Booting new emulator: %s', self._config.emulator_path)\n\n # Set necessary environment variables.\n base_lib_dir = self._config.emulator_path[:-8] + 'lib64/'\n ld_library_path = ':'.join([\n base_lib_dir + 'x11/', base_lib_dir + 'qt/lib/',\n base_lib_dir + 'gles_swiftshader/', base_lib_dir\n ])\n extra_env_vars = {\n 'ANDROID_HOME': '',\n 'ANDROID_SDK_ROOT': self._config.android_sdk_root,\n 'ANDROID_AVD_HOME': self._config.android_avd_home,\n 'ANDROID_EMULATOR_KVM_DEVICE': self._config.kvm_device,\n 'ANDROID_ADB_SERVER_PORT': str(\n self._adb_controller_config.adb_server_port\n ),\n 'LD_LIBRARY_PATH': ld_library_path,\n 'QT_XKB_CONFIG_ROOT': str(\n self._config.emulator_path[:-8] + 'qt_config/'\n ),\n 'ANDROID_EMU_ENABLE_CRASH_REPORTING': '1',\n 'SHOW_PERF_STATS': str(1 if self._config.show_perf_stats else 0),\n }\n logging.info('extra_env_vars: %s',\n ' '.join(f'{k}={v}' for k, v in extra_env_vars.items()))\n env_vars = dict(os.environ).copy()\n env_vars.update(extra_env_vars)\n\n # Compile command.\n grpc_port = (\n ['-grpc', str(self._config.grpc_port)]\n if self._config.grpc_port >= 0\n else []\n )\n run_headless = (\n ['-no-skin', '-no-window'] if self._config.run_headless else []\n )\n ports = [\n '-ports',\n '%s,%s' % (self._config.emulator_console_port, self._config.adb_port),\n ]\n snapshot = [\n '-snapshot',\n self._config.snapshot_name,\n '-feature',\n 'AllowSnapshotMigration,MigratableSnapshotSave',\n ]\n snapshot = snapshot if self._config.snapshot_name else ['-no-snapshot']\n restrict_network_args = [\n '-network-user-mode-options', 'restrict=y', '-wifi-user-mode-options',\n 'restrict=y'\n ]\n network_args = (\n restrict_network_args if self._config.restrict_network else []\n )\n command = (\n [\n self._config.emulator_path,\n '-adb-path',\n self._adb_controller_config.adb_path,\n '-gpu',\n self._config.gpu_mode,\n '-no-audio',\n '-show-kernel',\n '-verbose',\n '-avd',\n self._config.avd_name,\n ]\n + grpc_port\n + run_headless\n + ports\n + snapshot\n + network_args\n )\n logging.info('Emulator launch command: %s', ' '.join(command))\n # Prepare logfile.\n self._emulator_output = open(self._logfile_path, 'wb')\n\n # Spawn the emulator process.\n self._emulator = subprocess.Popen(\n command,\n env=env_vars,\n stdout=self._emulator_output,\n stderr=self._emulator_output)\n\n def confirm_shutdown(self) -> None:\n \"\"\"Shuts down the emulator process.\"\"\"\n if self._emulator is not None:\n logging.info('Checking if emulator process has finished...')\n try:\n self._emulator.wait(timeout=30.0)\n except subprocess.TimeoutExpired:\n logging.exception(\n 'The emulator process did not finish after 30s. '\n 'returncode: %s. Will now try to kill() it.',\n self._emulator.returncode)\n self._emulator.kill()\n self._emulator = None\n self._emulator_output.close()\n logging.info('The emulator process has finished.')\n\n def close(self):\n \"\"\"Clean up launcher files and processes.\"\"\"\n if not self._is_closed:\n self._local_tmp_dir_handle.cleanup()\n self.confirm_shutdown()\n self._is_closed = True\n\n def __del__(self):\n self.close()\n"
},
{
"path": "android_env/components/simulators/emulator/emulator_launcher_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.emulator_launcher.\"\"\"\n\nimport builtins\nimport os\nimport subprocess\nimport tempfile\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import config_classes\nfrom android_env.components.simulators.emulator import emulator_launcher\n\n\nclass EmulatorLauncherTest(parameterized.TestCase):\n\n def setUp(self):\n super().setUp()\n\n self._emulator_path = 'fake/path/emulator'\n self._adb_path = 'fake/path/adb'\n self._adb_port = 5554\n self._adb_server_port = 1234\n self._emulator_console_port = 5555\n self._avd_name = 'my_avd_name'\n\n self._expected_command = [\n self._emulator_path,\n '-adb-path',\n 'fake/path/adb',\n '-gpu',\n 'swangle_indirect',\n '-no-audio',\n '-show-kernel',\n '-verbose',\n '-avd',\n self._avd_name,\n ]\n self._headless = ['-no-skin', '-no-window']\n self._ports = ['-ports', f'{self._emulator_console_port},{self._adb_port}']\n self._snapshot = ['-no-snapshot']\n\n base_lib_dir = self._emulator_path[:-8] + 'lib64/'\n ld_library_path = ':'.join([\n base_lib_dir + 'x11/', base_lib_dir + 'qt/lib/',\n base_lib_dir + 'gles_swiftshader/', base_lib_dir\n ])\n\n # Instantiate the config to extract default values.\n config = config_classes.EmulatorLauncherConfig()\n self._expected_env_vars = {\n 'ANDROID_HOME': '',\n 'ANDROID_SDK_ROOT': config.android_sdk_root,\n 'ANDROID_AVD_HOME': config.android_avd_home,\n 'ANDROID_EMULATOR_KVM_DEVICE': '/dev/kvm',\n 'ANDROID_ADB_SERVER_PORT': '1234',\n 'LD_LIBRARY_PATH': ld_library_path,\n 'QT_XKB_CONFIG_ROOT': str(self._emulator_path[:-8] + 'qt_config/'),\n 'ANDROID_EMU_ENABLE_CRASH_REPORTING': '1',\n }\n\n @parameterized.named_parameters([\n ('hide_perf_stats', False),\n ('show_perf_stats', True),\n ])\n @mock.patch.object(os, 'makedirs')\n @mock.patch.object(os, 'environ', autospec=True, return_value=dict())\n @mock.patch.object(tempfile, 'TemporaryDirectory', instance=True)\n def test_launch(\n self,\n show_perf_stats: bool,\n mock_tmp_dir,\n unused_os_environ,\n unused_os_makedirs,\n ):\n mock_tmp_dir.return_value.name.return_value = 'local_tmp_dir'\n\n config = config_classes.EmulatorLauncherConfig(\n adb_port=self._adb_port,\n emulator_console_port=self._emulator_console_port,\n emulator_path=self._emulator_path,\n avd_name=self._avd_name,\n grpc_port=-1,\n show_perf_stats=show_perf_stats,\n )\n adb_controller_config = config_classes.AdbControllerConfig(\n adb_path=self._adb_path,\n adb_server_port=self._adb_server_port,\n )\n launcher = emulator_launcher.EmulatorLauncher(\n config=config, adb_controller_config=adb_controller_config\n )\n\n expected_env_vars = self._expected_env_vars\n expected_env_vars['SHOW_PERF_STATS'] = '1' if show_perf_stats else '0'\n\n with mock.patch.object(\n subprocess, 'Popen', autospec=True\n ) as emulator_init, mock.patch.object(builtins, 'open', autospec=True) as f:\n f.return_value.__enter__ = f()\n launcher.launch_emulator_process()\n emulator_init.assert_called_once_with(\n args=self._expected_command\n + self._headless\n + self._ports\n + self._snapshot,\n env=expected_env_vars,\n stdout=f(),\n stderr=f(),\n )\n\n @parameterized.named_parameters([\n ('hide_perf_stats', False),\n ('show_perf_stats', True),\n ])\n @mock.patch.object(os, 'makedirs')\n @mock.patch.object(os, 'environ', autospec=True, return_value=dict())\n @mock.patch.object(tempfile, 'TemporaryDirectory', instance=True)\n def test_grpc_port(\n self,\n show_perf_stats: bool,\n mock_tmp_dir,\n unused_os_environ,\n unused_os_makedirs,\n ):\n mock_tmp_dir.return_value.name.return_value = 'local_tmp_dir'\n\n config = config_classes.EmulatorLauncherConfig(\n adb_port=self._adb_port,\n emulator_console_port=self._emulator_console_port,\n emulator_path=self._emulator_path,\n avd_name=self._avd_name,\n grpc_port=8554,\n show_perf_stats=show_perf_stats,\n )\n adb_controller_config = config_classes.AdbControllerConfig(\n adb_path=self._adb_path,\n adb_server_port=self._adb_server_port,\n )\n launcher = emulator_launcher.EmulatorLauncher(\n config=config, adb_controller_config=adb_controller_config\n )\n\n expected_env_vars = self._expected_env_vars\n expected_env_vars['SHOW_PERF_STATS'] = '1' if show_perf_stats else '0'\n\n with mock.patch.object(\n subprocess, 'Popen', autospec=True\n ) as emulator_init, mock.patch.object(builtins, 'open', autospec=True) as f:\n f.return_value.__enter__ = f()\n launcher.launch_emulator_process()\n emulator_init.assert_called_once_with(\n args=self._expected_command\n + ['-grpc', '8554']\n + self._headless\n + self._ports\n + self._snapshot,\n env=expected_env_vars,\n stdout=f(),\n stderr=f(),\n )\n\n @parameterized.named_parameters([\n ('hide_perf_stats', False),\n ('show_perf_stats', True),\n ])\n @mock.patch.object(os, 'makedirs')\n @mock.patch.object(os, 'environ', autospec=True, return_value=dict())\n @mock.patch.object(tempfile, 'TemporaryDirectory', instance=True)\n def test_snapshot(\n self,\n show_perf_stats: bool,\n mock_tmp_dir,\n unused_os_environ,\n unused_os_makedirs,\n ):\n mock_tmp_dir.return_value.name.return_value = 'local_tmp_dir'\n\n config = config_classes.EmulatorLauncherConfig(\n adb_port=self._adb_port,\n emulator_console_port=self._emulator_console_port,\n emulator_path=self._emulator_path,\n avd_name=self._avd_name,\n grpc_port=-1,\n snapshot_name='my_snapshot',\n show_perf_stats=show_perf_stats,\n )\n adb_controller_config = config_classes.AdbControllerConfig(\n adb_path=self._adb_path,\n adb_server_port=self._adb_server_port,\n )\n launcher = emulator_launcher.EmulatorLauncher(\n config=config, adb_controller_config=adb_controller_config\n )\n\n expected_snapshot = [\n '-snapshot', 'my_snapshot', '-feature',\n 'AllowSnapshotMigration,MigratableSnapshotSave'\n ]\n\n expected_env_vars = self._expected_env_vars\n expected_env_vars['SHOW_PERF_STATS'] = '1' if show_perf_stats else '0'\n\n with mock.patch.object(\n subprocess, 'Popen', autospec=True) as emulator_init, \\\n mock.patch.object(builtins, 'open', autospec=True) as f:\n f.return_value.__enter__ = f()\n launcher.launch_emulator_process()\n emulator_init.assert_called_once_with(\n args=self._expected_command\n + self._headless\n + self._ports\n + expected_snapshot,\n env=expected_env_vars,\n stdout=f(),\n stderr=f(),\n )\n\n @parameterized.named_parameters([\n ('hide_perf_stats', False),\n ('show_perf_stats', True),\n ])\n @mock.patch.object(os, 'makedirs')\n @mock.patch.object(os, 'environ', autospec=True, return_value=dict())\n @mock.patch.object(tempfile, 'TemporaryDirectory', instance=True)\n def test_network_restrict(\n self,\n show_perf_stats: bool,\n mock_tmp_dir,\n unused_os_environ,\n unused_os_makedirs,\n ):\n mock_tmp_dir.return_value.name.return_value = 'local_tmp_dir'\n\n config = config_classes.EmulatorLauncherConfig(\n adb_port=self._adb_port,\n emulator_console_port=self._emulator_console_port,\n emulator_path=self._emulator_path,\n avd_name=self._avd_name,\n grpc_port=-1,\n restrict_network=True,\n show_perf_stats=show_perf_stats,\n )\n adb_controller_config = config_classes.AdbControllerConfig(\n adb_path=self._adb_path,\n adb_server_port=self._adb_server_port,\n )\n launcher = emulator_launcher.EmulatorLauncher(\n config=config, adb_controller_config=adb_controller_config\n )\n\n expected_snapshot = ['-no-snapshot']\n expected_network_restrict = [\n '-network-user-mode-options', 'restrict=y', '-wifi-user-mode-options',\n 'restrict=y'\n ]\n\n expected_env_vars = self._expected_env_vars\n expected_env_vars['SHOW_PERF_STATS'] = '1' if show_perf_stats else '0'\n\n with mock.patch.object(\n subprocess, 'Popen', autospec=True) as emulator_init, \\\n mock.patch.object(builtins, 'open', autospec=True) as f:\n f.return_value.__enter__ = f()\n launcher.launch_emulator_process()\n emulator_init.assert_called_once_with(\n self._expected_command\n + self._headless\n + self._ports\n + expected_snapshot\n + expected_network_restrict,\n env=expected_env_vars,\n stdout=f(),\n stderr=f(),\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/simulators/emulator/emulator_simulator.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A class that manages an Android Emulator.\"\"\"\n\nimport os\nimport time\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env.components import adb_controller\nfrom android_env.components import adb_log_stream\nfrom android_env.components import config_classes\nfrom android_env.components import errors\nfrom android_env.components import log_stream\nfrom android_env.components.simulators import base_simulator\nfrom android_env.components.simulators.emulator import emulator_launcher\nfrom android_env.proto import state_pb2\nimport grpc\nimport numpy as np\nimport portpicker\n\nfrom android_env.proto import emulator_controller_pb2\nfrom android_env.proto import emulator_controller_pb2_grpc\nfrom android_env.proto import snapshot_service_pb2\nfrom android_env.proto import snapshot_service_pb2_grpc\nfrom google.protobuf import empty_pb2\n\n\n_DEFAULT_SNAPSHOT_NAME: str = 'default_snapshot'\n\n\ndef _is_existing_emulator_provided(\n launcher_config: config_classes.EmulatorLauncherConfig,\n) -> bool:\n \"\"\"Returns true if all necessary args were provided.\"\"\"\n\n return bool(\n launcher_config.adb_port\n and launcher_config.emulator_console_port\n and launcher_config.grpc_port\n )\n\n\ndef _pick_adb_port() -> int:\n \"\"\"Tries to pick a port in the recommended range 5555-5585.\n\n If no such port can be found, will return a random unused port. More info:\n https://developer.android.com/studio/command-line/adb#howadbworks.\n\n Returns:\n port: an available port for adb.\n \"\"\"\n\n for p in range(5555, 5587, 2):\n if portpicker.is_port_free(p):\n return p\n return portpicker.pick_unused_port()\n\n\ndef _pick_emulator_grpc_port() -> int:\n \"\"\"Tries to pick the recommended port for grpc.\n\n If no such port can be found, will return a random unused port. More info:\n https://android.googlesource.com/platform/external/qemu/+/emu-master-dev/android/android-grpc/docs/.\n\n Returns:\n port: an available port for emulator grpc.\n \"\"\"\n\n if portpicker.is_port_free(8554):\n return 8554\n else:\n return portpicker.pick_unused_port()\n\n\nclass EmulatorBootError(errors.SimulatorError):\n \"\"\"Raised when an emulator failed to boot.\"\"\"\n\n\nclass EmulatorCrashError(errors.SimulatorError):\n \"\"\"Raised when a simulator crashed.\"\"\"\n\n\nclass EmulatorSimulator(base_simulator.BaseSimulator):\n \"\"\"Controls an Android Emulator.\"\"\"\n\n def __init__(self, config: config_classes.EmulatorConfig):\n \"\"\"Instantiates an EmulatorSimulator.\"\"\"\n\n super().__init__(config)\n self._config = config\n\n # If adb_port, console_port and grpc_port are all already provided,\n # we assume the emulator already exists and there's no need to launch.\n if _is_existing_emulator_provided(self._config.emulator_launcher):\n self._existing_emulator_provided = True\n logging.info('Connecting to existing emulator \"%r\"',\n self.adb_device_name())\n else:\n self._existing_emulator_provided = False\n self._config.emulator_launcher.adb_port = _pick_adb_port()\n self._config.emulator_launcher.emulator_console_port = (\n portpicker.pick_unused_port()\n )\n self._config.emulator_launcher.grpc_port = _pick_emulator_grpc_port()\n\n self._channel = None\n self._emulator_stub: emulator_controller_pb2_grpc.EmulatorControllerStub | None = (\n None\n )\n self._snapshot_stub = None\n # Set the image format to RGBA. The width and height of the returned\n # screenshots will use the device's width and height.\n self._image_format = emulator_controller_pb2.ImageFormat(\n format=emulator_controller_pb2.ImageFormat.ImgFormat.RGBA8888)\n\n if (\n self._config.launch_n_times_without_reboot\n > self._config.launch_n_times_without_reinstall\n ):\n raise ValueError(\n 'Number of launch attempts before reboot'\n f' ({self._config.launch_n_times_without_reboot}) should not be'\n ' greater than number of launch attempts before reinstall'\n f' ({self._config.launch_n_times_without_reinstall})'\n )\n\n # Initialize own ADB controller.\n self._config.adb_controller.device_name = self.adb_device_name()\n self._adb_controller = self.create_adb_controller()\n self._adb_controller.init_server()\n logging.info(\n 'Initialized simulator with ADB server port %r.',\n self._config.adb_controller.adb_server_port,\n )\n\n # If necessary, create EmulatorLauncher.\n if self._existing_emulator_provided:\n self._logfile_path = self._config.logfile_path or None\n self._launcher = None\n else:\n logging.info(\n 'emulator_launcher config: %r', self._config.emulator_launcher\n )\n self._launcher = emulator_launcher.EmulatorLauncher(\n config=self._config.emulator_launcher,\n adb_controller_config=self._config.adb_controller,\n )\n self._logfile_path = (\n self._config.logfile_path or self._launcher.logfile_path()\n )\n\n def _reconnect_on_grpc_error(func):\n \"\"\"Decorator function for reconnecting to emulator upon grpc errors.\"\"\"\n\n def wrapper(self, *args, **kwargs):\n try:\n return func(self, *args, **kwargs)\n except grpc.RpcError:\n logging.exception('RpcError caught. Reconnecting to emulator...')\n self._emulator_stub, self._snapshot_stub = self._connect_to_emulator(\n self._config.emulator_launcher.grpc_port\n )\n return func(self, *args, **kwargs)\n\n return wrapper\n\n def get_logs(self) -> str:\n \"\"\"Returns logs recorded by the emulator.\"\"\"\n if self._logfile_path and os.path.exists(self._logfile_path):\n with open(self._logfile_path, 'rb') as f:\n return f.read().decode('utf-8')\n else:\n return f'Logfile does not exist: {self._logfile_path}.'\n\n def adb_device_name(self) -> str:\n return 'emulator-%s' % (self._config.emulator_launcher.adb_port - 1)\n\n def create_adb_controller(self):\n \"\"\"Returns an ADB controller which can communicate with this simulator.\"\"\"\n return adb_controller.AdbController(self._config.adb_controller)\n\n def create_log_stream(self) -> log_stream.LogStream:\n return adb_log_stream.AdbLogStream(\n adb_command_prefix=self._adb_controller.command_prefix(),\n verbose=self._config.verbose_logs,\n )\n\n def _launch_impl(self) -> None:\n \"\"\"Prepares an Android Emulator for RL interaction.\n\n The behavior depends on `self._num_launch_attempts`'s value:\n * <= self._config.launch_n_times_without_reboot -> Normal boot behavior.\n * > self._config.launch_n_times_without_reboot but <=\n self._config.launch_n_times_without_reinstall -> reboot (i.e. process\n is killed and started again).\n * > self._config.launch_n_times_without_reinstall -> reinstall (i.e.\n process is killed, emulator files are deleted and the process started\n again).\n \"\"\"\n\n logging.info('Attempt %r at launching the Android Emulator (%r)',\n self._num_launch_attempts, self.adb_device_name())\n\n if self._launcher is not None:\n # If not the first time, then shutdown the emulator first.\n if (\n self._emulator_stub is not None\n and self._num_launch_attempts\n > self._config.launch_n_times_without_reboot\n ):\n self._shutdown_emulator()\n # Subsequent attempts cause the emulator files to be reinstalled.\n if (\n self._num_launch_attempts\n > self._config.launch_n_times_without_reinstall\n ):\n logging.info('Closing emulator (%r)', self.adb_device_name())\n self._launcher.close()\n self._launcher = emulator_launcher.EmulatorLauncher(\n config=self._config.emulator_launcher,\n adb_controller_config=self._config.adb_controller,\n )\n self._launcher.launch_emulator_process()\n # Establish grpc connection to emulator process.\n self._emulator_stub, self._snapshot_stub = self._connect_to_emulator(\n self._config.emulator_launcher.grpc_port\n )\n\n # Confirm booted status.\n try:\n self._confirm_booted()\n except EmulatorCrashError:\n logging.exception('Failed to confirm booted status of emulator.')\n\n logging.info('Done booting the Android Emulator.')\n\n def load_state(\n self, request: state_pb2.LoadStateRequest\n ) -> state_pb2.LoadStateResponse:\n \"\"\"Loads a state using the emulator's snapshotting mechanism.\n\n Args:\n request: The `LoadStateRequest`. In this case, `args` should be a dict\n containing the key 'snapshot_name', representing the name of the\n snapshot to load. If `request.args.snapshot_name` is `None`, a default\n snapshot name is used.\n\n Returns:\n A response indicating whether the snapshot was successfully loaded.\n * If the snapshot was loaded successfully, the status will be `OK`.\n * If no snapshot of the given name was found, the status will be\n `NOT_FOUND`.\n * If an error occurred during the snapshot loading process, the status\n will be `ERROR` and the `error_message` field will be filled.\n \"\"\"\n assert self._snapshot_stub is not None\n snapshot_name = request.args.get('snapshot_name', _DEFAULT_SNAPSHOT_NAME)\n snapshot_list = self._snapshot_stub.ListSnapshots(\n snapshot_service_pb2.SnapshotFilter(\n statusFilter=snapshot_service_pb2.SnapshotFilter.LoadStatus.All\n )\n )\n if any(\n snapshot.snapshot_id == snapshot_name\n for snapshot in snapshot_list.snapshots\n ):\n snapshot_result = self._snapshot_stub.LoadSnapshot(\n snapshot_service_pb2.SnapshotPackage(snapshot_id=snapshot_name)\n )\n if snapshot_result.success:\n return state_pb2.LoadStateResponse(\n status=state_pb2.LoadStateResponse.Status.OK\n )\n else:\n return state_pb2.LoadStateResponse(\n status=state_pb2.LoadStateResponse.Status.ERROR,\n error_message=snapshot_result.err.decode('utf-8'),\n )\n\n else:\n return state_pb2.LoadStateResponse(\n status=state_pb2.LoadStateResponse.Status.NOT_FOUND\n )\n\n def save_state(\n self, request: state_pb2.SaveStateRequest\n ) -> state_pb2.SaveStateResponse:\n \"\"\"Saves a state using the emulator's snapshotting mechanism.\n\n Args:\n request: The `SaveStateRequest`. In this case, `args` should be a dict\n containing the key 'snapshot_name', representing the name of the\n snapshot to save. If `request.args.snapshot_name` is `None`, a default\n snapshot name is used.\n\n Returns:\n A response indicating whether the snapshot was successfully saved.\n * If the snapshot was saved successfully, the status will be `OK`.\n * If an error occurred during the snapshot saving process, the status\n will be `ERROR` and the `error_message` field will be filled.\n \"\"\"\n assert self._snapshot_stub is not None\n snapshot_name = request.args.get('snapshot_name', _DEFAULT_SNAPSHOT_NAME)\n snapshot_result = self._snapshot_stub.SaveSnapshot(\n snapshot_service_pb2.SnapshotPackage(snapshot_id=snapshot_name)\n )\n if snapshot_result.success:\n return state_pb2.SaveStateResponse(\n status=state_pb2.SaveStateResponse.Status.OK\n )\n else:\n return state_pb2.SaveStateResponse(\n status=state_pb2.SaveStateResponse.Status.ERROR,\n error_message=snapshot_result.err.decode('utf-8'),\n )\n\n def _connect_to_emulator(\n self,\n grpc_port: int,\n timeout_sec: int = 100,\n ) -> tuple[\n emulator_controller_pb2_grpc.EmulatorControllerStub,\n snapshot_service_pb2_grpc.SnapshotServiceStub,\n ]:\n \"\"\"Connects to an emulator and returns a corresponsing stub.\"\"\"\n\n logging.info('Creating gRPC channel to the emulator on port %r', grpc_port)\n port = f'localhost:{grpc_port}'\n options = [('grpc.max_send_message_length', -1),\n ('grpc.max_receive_message_length', -1)]\n creds = grpc.local_channel_credentials()\n\n try:\n self._channel = grpc.secure_channel(port, creds, options=options)\n grpc.channel_ready_future(self._channel).result(timeout=timeout_sec)\n except (grpc.RpcError, grpc.FutureTimeoutError) as grpc_error:\n logging.exception('Failed to connect to the emulator.')\n raise EmulatorBootError(\n 'Failed to connect to the emulator.') from grpc_error\n\n logging.info('Added gRPC channel for the Emulator on port %s', port)\n emulator_controller_stub = (\n emulator_controller_pb2_grpc.EmulatorControllerStub(self._channel)\n )\n snapshot_stub = snapshot_service_pb2_grpc.SnapshotServiceStub(self._channel)\n return emulator_controller_stub, snapshot_stub\n\n @_reconnect_on_grpc_error\n def _confirm_booted(self, startup_wait_time_sec: int = 300):\n \"\"\"Waits until the emulator is fully booted.\"\"\"\n\n assert (\n self._emulator_stub is not None\n ), 'Emulator stub has not been initialized yet.'\n start_time = time.time()\n deadline = start_time + startup_wait_time_sec\n success = False\n while time.time() < deadline:\n emu_status = self._emulator_stub.getStatus(empty_pb2.Empty())\n logging.info('Waiting for emulator (%r) to start... (%rms)',\n self.adb_device_name(), emu_status.uptime)\n if emu_status.booted:\n success = True\n break\n time.sleep(5.0)\n\n elapsed_time = time.time() - start_time\n if not success:\n raise EmulatorCrashError(\n f'The emulator failed to boot after {startup_wait_time_sec} seconds')\n\n logging.info('Done booting the emulator (in %f seconds).', elapsed_time)\n logging.info('********** Emulator logs **********')\n for line in self.get_logs().splitlines():\n logging.info(line)\n logging.info('******* End of emulator logs *******')\n logging.info('See the full emulator logs at %r', self._logfile_path)\n\n @_reconnect_on_grpc_error\n def send_touch(self, touches: list[tuple[int, int, bool, int]]) -> None:\n \"\"\"Sends a touch event to be executed on the simulator.\n\n Args:\n touches: A list of touch events. Each element in the list corresponds to a\n single touch event. Each touch event tuple should have:\n 0 x: The horizontal coordinate of this event.\n 1 y: The vertical coordinate of this event.\n 2 is_down: Whether the finger is touching or not the screen.\n 3 identifier: Identifies a particular finger in a multitouch event.\n \"\"\"\n\n assert (\n self._emulator_stub is not None\n ), 'Emulator stub has not been initialized yet.'\n touch_events = [\n emulator_controller_pb2.Touch(\n x=t[0], y=t[1], pressure=int(t[2]), identifier=t[3])\n for t in touches\n ]\n self._emulator_stub.sendTouch(\n emulator_controller_pb2.TouchEvent(touches=touch_events))\n\n @_reconnect_on_grpc_error\n def send_key(self, keycode: np.int32, event_type: str) -> None:\n \"\"\"Sends a key event to the emulator.\n\n Args:\n keycode: Code representing the desired key press in XKB format.\n See the emulator_controller_pb2 for details.\n event_type: Type of key event to be sent.\n \"\"\"\n\n event_types = emulator_controller_pb2.KeyboardEvent.KeyEventType.keys()\n if event_type not in event_types:\n raise ValueError(\n f'Event type must be one of {event_types} but is {event_type}.')\n\n assert (\n self._emulator_stub is not None\n ), 'Emulator stub has not been initialized yet.'\n self._emulator_stub.sendKey(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.Value(\n event_type\n ),\n keyCode=int(keycode),\n )\n )\n\n @_reconnect_on_grpc_error\n def _get_screenshot_impl(self) -> np.ndarray:\n \"\"\"Fetches the latest screenshot from the emulator.\"\"\"\n\n assert (\n self._emulator_stub is not None\n ), 'Emulator stub has not been initialized yet.'\n assert self._image_format, 'ImageFormat has not been initialized yet.'\n image_proto = self._emulator_stub.getScreenshot(self._image_format)\n h, w = image_proto.format.height, image_proto.format.width\n image = np.frombuffer(image_proto.image, dtype='uint8', count=h * w * 4)\n image.shape = (h, w, 4)\n return image[:, :, :3]\n\n @_reconnect_on_grpc_error\n def _shutdown_emulator(self):\n \"\"\"Sends a signal to trigger emulator shutdown.\"\"\"\n\n if self._emulator_stub is None:\n logging.info('Emulator (%r) is not up.', self.adb_device_name())\n return\n\n assert self._launcher is not None, 'Launcher is already down.'\n\n logging.info('Shutting down the emulator (%r)...', self.adb_device_name())\n self._emulator_stub.setVmState(\n emulator_controller_pb2.VmRunState(\n state=emulator_controller_pb2.VmRunState.RunState.SHUTDOWN))\n self._launcher.confirm_shutdown()\n\n def close(self):\n super().close()\n\n if self._launcher is not None:\n self._shutdown_emulator()\n logging.info('Closing emulator (%r)', self.adb_device_name())\n self._launcher.close()\n self._emulator_stub = None\n self._snapshot_stub = None\n if self._channel is not None:\n self._channel.close()\n super().close()\n"
},
{
"path": "android_env/components/simulators/emulator/emulator_simulator_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.emulator_simulator.\"\"\"\n\nimport builtins\nimport os\nimport time\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import adb_call_parser\nfrom android_env.components import adb_controller\nfrom android_env.components import config_classes\nfrom android_env.components.simulators.emulator import emulator_launcher\nfrom android_env.components.simulators.emulator import emulator_simulator\nfrom android_env.proto import state_pb2\nimport grpc\nfrom PIL import Image\nimport portpicker\n\nfrom android_env.proto import emulator_controller_pb2\nfrom android_env.proto import emulator_controller_pb2_grpc\nfrom android_env.proto import snapshot_service_pb2\n\n\nclass EmulatorSimulatorTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self.addCleanup(mock.patch.stopall) # Disable previous patches.\n\n self._adb_controller = mock.create_autospec(adb_controller.AdbController)\n self._adb_call_parser = mock.create_autospec(adb_call_parser.AdbCallParser)\n self._launcher = mock.create_autospec(emulator_launcher.EmulatorLauncher)\n self._launcher.logfile_path.return_value = 'logfile_path'\n self._emulator_stub = mock.create_autospec(\n emulator_controller_pb2_grpc.EmulatorControllerStub)\n\n self._grpc_channel = mock.create_autospec(grpc.Channel)\n mock.patch.object(\n grpc.aio, 'secure_channel', return_value=self._grpc_channel).start()\n mock.patch.object(\n grpc, 'secure_channel', return_value=self._grpc_channel).start()\n mock.patch.object(\n grpc, 'local_channel_credentials',\n return_value=self._grpc_channel).start()\n self._mock_future = mock.create_autospec(grpc.Future)\n mock.patch.object(\n grpc, 'channel_ready_future', return_value=self._mock_future).start()\n mock.patch.object(time, 'time', return_value=12345).start()\n\n mock.patch.object(\n adb_controller, 'AdbController',\n return_value=self._adb_controller).start()\n mock.patch.object(\n adb_call_parser,\n 'AdbCallParser',\n autospec=True,\n return_value=self._adb_call_parser).start()\n mock.patch.object(\n emulator_launcher, 'EmulatorLauncher',\n return_value=self._launcher).start()\n\n def test_adb_device_name_not_empty(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n self.assertNotEmpty(simulator.adb_device_name())\n\n def test_logfile_path(self):\n \"\"\"The log file's path should correspond to the one from the config.\"\"\"\n\n config = config_classes.EmulatorConfig(\n logfile_path='fake/logfile/path',\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n with mock.patch.object(\n os.path, 'exists', autospec=True, return_value=True\n ), mock.patch.object(builtins, 'open', autospec=True) as mock_open:\n mock_file = mock_open.return_value.__enter__.return_value\n mock_file.read.return_value = b'fake_logs'\n logs = simulator.get_logs()\n mock_open.assert_called_once_with('fake/logfile/path', 'rb')\n self.assertEqual(logs, 'fake_logs')\n\n @mock.patch.object(portpicker, 'is_port_free', return_value=True)\n def test_grpc_port(self, unused_mock_portpicker):\n\n launcher_config = config_classes.EmulatorLauncherConfig(\n tmp_dir=self.create_tempdir().full_path\n )\n config = config_classes.EmulatorConfig(\n emulator_launcher=launcher_config,\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n self.assertEqual(launcher_config.grpc_port, 8554)\n\n @mock.patch.object(portpicker, 'is_port_free', return_value=False)\n def test_grpc_port_unavailable(self, unused_mock_portpicker):\n\n launcher_config = config_classes.EmulatorLauncherConfig(\n tmp_dir=self.create_tempdir().full_path\n )\n config = config_classes.EmulatorConfig(\n emulator_launcher=launcher_config,\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n self.assertNotEqual(launcher_config.grpc_port, 8554)\n\n def test_launch_operation_order(self):\n \"\"\"Makes sure that adb_controller is started before Emulator is launched.\"\"\"\n\n # Arrange.\n call_order = []\n self._adb_controller.init_server.side_effect = lambda: call_order.append(\n 'init_server'\n )\n self._launcher.launch_emulator_process.side_effect = (\n lambda: call_order.append('launch_emulator_process')\n )\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # Act.\n simulator.launch() # The simulator should launch and not crash.\n\n # Assert.\n # The adb server should be initialized before launching the emulator.\n self.assertEqual(call_order, ['init_server', 'launch_emulator_process'])\n\n def test_close(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n # For whatever reason clients may want to close the EmulatorSimulator.\n # We just want to check that the simulator does not crash and/or leak\n # resources.\n simulator.close()\n\n def test_value_error_if_launch_attempt_params_incorrect(self):\n self.assertRaises(\n ValueError,\n emulator_simulator.EmulatorSimulator,\n config=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n launch_n_times_without_reboot=2,\n launch_n_times_without_reinstall=1,\n ),\n )\n\n def test_launch_attempt_reboot(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n launch_n_times_without_reboot=1,\n launch_n_times_without_reinstall=2,\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n self._launcher.launch_emulator_process.assert_called_once()\n self._launcher.reset_mock()\n\n # Launch attempt 2.\n simulator.launch()\n self._launcher.confirm_shutdown.assert_called_once()\n self._launcher.close.assert_not_called()\n self._launcher.launch_emulator_process.assert_called_once()\n\n def test_launch_attempt_reinstall_after_zero_attempts(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n launch_n_times_without_reboot=0,\n launch_n_times_without_reinstall=0,\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should not reboot or reinstall on its very first launch.\n simulator.launch()\n self._launcher.launch_emulator_process.assert_called_once()\n self._launcher.confirm_shutdown.assert_not_called()\n self._launcher.close.assert_not_called()\n\n # Every subsequent attempt should reboot and reinstall.\n self._launcher.reset_mock()\n simulator.launch()\n self._launcher.confirm_shutdown.assert_called_once()\n self._launcher.close.assert_called_once() # Now this should `close()`.\n self._launcher.launch_emulator_process.assert_called_once()\n\n def test_launch_attempt_reinstall(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n launch_n_times_without_reboot=1,\n launch_n_times_without_reinstall=2,\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n self._launcher.launch_emulator_process.assert_called_once()\n\n # Launch attempt 2.\n self._launcher.reset_mock()\n simulator.launch()\n self._launcher.confirm_shutdown.assert_called_once()\n self._launcher.close.assert_not_called() # Reboots don't `close()`.\n self._launcher.launch_emulator_process.assert_called_once()\n\n # Launch attempt 3.\n self._launcher.reset_mock()\n simulator.launch()\n self._launcher.confirm_shutdown.assert_called_once()\n self._launcher.close.assert_called_once() # Now this should `close()`.\n self._launcher.launch_emulator_process.assert_called_once()\n\n def test_get_screenshot(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n simulator._emulator_stub.getScreenshot = mock.MagicMock(\n return_value=emulator_controller_pb2.Image(\n format=emulator_controller_pb2.ImageFormat(width=5678, height=1234),\n image=Image.new('RGBA', (1234, 5678)).tobytes(),\n timestampUs=123))\n\n screenshot = simulator.get_screenshot()\n # The screenshot should have the same screen dimensions as reported by ADB\n # and it should have 3 channels (RGB).\n self.assertEqual(screenshot.shape, (1234, 5678, 3))\n\n def test_load_state(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n with mock.patch.object(\n simulator, '_snapshot_stub', create_autospec=True\n ) as mock_snapshot_stub:\n snapshot_list = snapshot_service_pb2.SnapshotList()\n snapshot_list.snapshots.add(snapshot_id='snapshot_name_foo')\n snapshot_list.snapshots.add(snapshot_id='snapshot_name_bar')\n mock_snapshot_stub.ListSnapshots.return_value = snapshot_list\n mock_snapshot_stub.LoadSnapshot.return_value = (\n snapshot_service_pb2.SnapshotPackage(success=True)\n )\n load_response = simulator.load_state(\n request=state_pb2.LoadStateRequest(\n args={'snapshot_name': 'snapshot_name_foo'}\n )\n )\n self.assertEqual(\n load_response.status, state_pb2.LoadStateResponse.Status.OK\n )\n load_response = simulator.load_state(\n request=state_pb2.LoadStateRequest(\n args={'snapshot_name': 'snapshot_name_baz'}\n )\n )\n self.assertEqual(\n load_response.status, state_pb2.LoadStateResponse.Status.NOT_FOUND\n )\n mock_snapshot_stub.LoadSnapshot.return_value = (\n snapshot_service_pb2.SnapshotPackage(success=False, err=b'error')\n )\n load_response = simulator.load_state(\n request=state_pb2.LoadStateRequest(\n args={'snapshot_name': 'snapshot_name_bar'}\n )\n )\n self.assertEqual(\n load_response.status, state_pb2.LoadStateResponse.Status.ERROR\n )\n self.assertEqual(load_response.error_message, 'error')\n\n def test_save_state(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n with mock.patch.object(\n simulator, '_snapshot_stub', create_autospec=True\n ) as mock_snapshot_stub:\n mock_snapshot_stub.SaveSnapshot.return_value = (\n snapshot_service_pb2.SnapshotPackage(success=True)\n )\n save_response = simulator.save_state(\n request=state_pb2.SaveStateRequest(\n args={'snapshot_name': 'snapshot_name_foo'}\n )\n )\n self.assertEqual(\n save_response.status, state_pb2.SaveStateResponse.Status.OK\n )\n mock_snapshot_stub.SaveSnapshot.return_value = (\n snapshot_service_pb2.SnapshotPackage(success=False, err=b'error')\n )\n save_response = simulator.save_state(\n request=state_pb2.SaveStateRequest(\n args={'snapshot_name': 'snapshot_name_bar'}\n )\n )\n self.assertEqual(\n save_response.status, state_pb2.SaveStateResponse.Status.ERROR\n )\n self.assertEqual(save_response.error_message, 'error')\n\n def test_send_touch(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n simulator._emulator_stub.sendTouch = mock.MagicMock(return_value=None)\n\n simulator.send_touch([(123, 456, True, 0), (135, 246, True, 1)])\n simulator.send_touch([(1, 2, True, 0), (3, 4, True, 1)])\n simulator.send_touch([(321, 654, False, 0), (531, 642, False, 1)])\n\n simulator._emulator_stub.sendTouch.assert_has_calls([\n mock.call(\n emulator_controller_pb2.TouchEvent(touches=[{\n 'x': 123,\n 'y': 456,\n 'pressure': 1\n }, {\n 'x': 135,\n 'y': 246,\n 'pressure': 1,\n 'identifier': 1\n }])),\n mock.call(\n emulator_controller_pb2.TouchEvent(touches=[{\n 'x': 1,\n 'y': 2,\n 'pressure': 1\n }, {\n 'x': 3,\n 'y': 4,\n 'pressure': 1,\n 'identifier': 1\n }])),\n mock.call(\n emulator_controller_pb2.TouchEvent(touches=[{\n 'x': 321,\n 'y': 654,\n 'pressure': 0\n }, {\n 'x': 531,\n 'y': 642,\n 'pressure': 0,\n 'identifier': 1\n }])),\n ])\n\n def test_send_key(self):\n config = config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n grpc_port=1234, tmp_dir=self.create_tempdir().full_path\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='/my/adb',\n adb_server_port=5037,\n ),\n )\n simulator = emulator_simulator.EmulatorSimulator(config)\n\n # The simulator should launch and not crash.\n simulator.launch()\n\n simulator._emulator_stub.sendTouch = mock.MagicMock(return_value=None)\n\n simulator.send_key(123, 'keydown')\n simulator.send_key(321, 'keydown')\n simulator.send_key(321, 'keyup')\n simulator.send_key(123, 'keyup')\n simulator.send_key(321, 'keypress')\n simulator.send_key(123, 'keypress')\n\n simulator._emulator_stub.sendKey.assert_has_calls([\n mock.call(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType\n .keydown,\n keyCode=123,\n )),\n mock.call(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType\n .keydown,\n keyCode=321,\n )),\n mock.call(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType\n .keyup,\n keyCode=321,\n )),\n mock.call(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType\n .keyup,\n keyCode=123,\n )),\n mock.call(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType\n .keypress,\n keyCode=321,\n )),\n mock.call(\n emulator_controller_pb2.KeyboardEvent(\n codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,\n eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType\n .keypress,\n keyCode=123,\n ))\n ])\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/simulators/fake/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/components/simulators/fake/fake_simulator.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Fake Simulator for testing AndroidEnv infrastructure.\"\"\"\n\nimport random\nimport threading\nimport time\n\nfrom absl import logging\nfrom android_env.components import adb_controller\nfrom android_env.components import config_classes\nfrom android_env.components import log_stream\nfrom android_env.components.simulators import base_simulator\nimport numpy as np\n\n\nclass FakeStream:\n \"\"\"This class simulates the logs coming from ADB.\"\"\"\n\n def __init__(self):\n self._values = [\n '',\n self._make_stdout('reward: 0.5'),\n self._make_stdout('reward: 1.0'),\n self._make_stdout('extra: my_extra [1.0]'),\n self._make_stdout('episode end'),\n ]\n self._kill = False\n self._lock = threading.Lock()\n\n def _make_stdout(self, data):\n \"\"\"Returns a valid log output with given data as message.\"\"\"\n return f' 1553110400.424 5583 5658 D Tag: {data}'\n\n def kill(self):\n self._kill = True\n\n def __iter__(self):\n while True:\n if self._kill:\n return\n else:\n with self._lock:\n next_value = random.choices(\n self._values, weights=[0.49, 0.15, 0.15, 0.15, 0.01], k=1)[0]\n time.sleep(0.1)\n yield next_value\n\n\nclass FakeLogStream(log_stream.LogStream):\n \"\"\"FakeLogStream class that wraps a FakeStream.\"\"\"\n\n def __init__(self):\n super().__init__(verbose=False)\n self.stream = FakeStream()\n\n def _get_stream_output(self):\n return self.stream\n\n def stop_stream(self):\n self.stream.kill()\n\n\nclass FakeAdbController(adb_controller.AdbController):\n \"\"\"Fake adb controller for FakeSimulator.\"\"\"\n\n def execute_command(\n self,\n args: list[str],\n timeout: float | None = None,\n device_specific: bool = True,\n ) -> bytes:\n \"\"\"Returns fake output for adb commands.\"\"\"\n\n del timeout, device_specific\n\n # Fake \"service is ready\" output.\n if args[:3] == ['shell', 'service', 'check']:\n return f'Service {args[-1]}: found'.encode('utf-8')\n\n # Fake dumpsys output for getting orientation.\n if args == ['shell', 'dumpsys', 'input']:\n return b' SurfaceOrientation: 0'\n\n # app_screen_checker: fake_task expects 'fake_activity'.\n if args[:4] == ['shell', 'am', 'stack', 'list']:\n return (b'taskId=0 fake_activity visible=true '\n b'topActivity=ComponentInfo{fake_activity}')\n\n return b'fake output'\n\n\nclass FakeSimulator(base_simulator.BaseSimulator):\n \"\"\"FakeSimulator class.\"\"\"\n\n def __init__(self, config: config_classes.FakeSimulatorConfig):\n \"\"\"FakeSimulator class that can replace EmulatorSimulator in AndroidEnv.\"\"\"\n super().__init__(config)\n self._screen_dimensions = np.array(config.screen_dimensions)\n logging.info('Created FakeSimulator.')\n\n def get_logs(self) -> str:\n return 'FakeSimulator: fake logs'\n\n def adb_device_name(self) -> str:\n return 'fake_simulator'\n\n def create_adb_controller(self):\n return FakeAdbController(config_classes.AdbControllerConfig())\n\n def create_log_stream(self) -> log_stream.LogStream:\n return FakeLogStream()\n\n def _launch_impl(self) -> None:\n pass\n\n def send_touch(self, touches: list[tuple[int, int, bool, int]]) -> None:\n del touches\n\n def send_key(self, keycode: np.int32, event_type: str) -> None:\n del keycode, event_type\n\n def _get_screenshot_impl(self) -> np.ndarray:\n return np.random.randint(\n low=0, high=255, size=(*self._screen_dimensions, 3), dtype=np.uint8)\n"
},
{
"path": "android_env/components/simulators/fake/fake_simulator_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for fake_simulator.\"\"\"\n\nimport re\nfrom absl.testing import absltest\nfrom android_env.components import config_classes\nfrom android_env.components.simulators.fake import fake_simulator\nimport numpy as np\n\n\nclass FakeSimulatorTest(absltest.TestCase):\n\n def test_device_name(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n self.assertEqual(simulator.adb_device_name(), 'fake_simulator')\n\n def test_launch_close(self):\n # The simulator should launch and not crash.\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n simulator.launch()\n # Closing the simulator should also not crash.\n simulator.close()\n\n def test_get_screenshot(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n simulator.launch()\n\n screenshot = simulator.get_screenshot()\n np.testing.assert_equal(screenshot.shape, [320, 480, 3])\n np.testing.assert_equal(screenshot.dtype, np.uint8)\n\n def test_log_stream(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n simulator.launch()\n log_stream = simulator.create_log_stream()\n # Start yielding lines from LogStream.\n log_stream.resume_stream()\n lines = [\n '',\n ' 1553110400.424 5583 5658 D Tag: reward: 0.5',\n ' 1553110400.424 5583 5658 D Tag: reward: 1.0',\n ' 1553110400.424 5583 5658 D Tag: extra: my_extra [1.0]',\n ' 1553110400.424 5583 5658 D Tag: episode end',\n ]\n for i, line in enumerate(log_stream.get_stream_output()):\n self.assertIn(line, lines)\n if i > 10:\n break\n\n def test_adb_output(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n simulator.launch()\n adb_controller = simulator.create_adb_controller()\n line = adb_controller.execute_command(['shell', 'dumpsys', 'input'])\n line = line.decode('utf-8')\n matches = re.match(r'\\s+SurfaceOrientation:\\s+(\\d)', line)\n self.assertIsNotNone(matches)\n orientation = matches.group(1)\n self.assertEqual(orientation, '0')\n line = adb_controller.execute_command(['shell', 'service', 'check', 'foo'])\n line = line.decode('utf-8')\n self.assertEqual(line, 'Service foo: found')\n line = adb_controller.execute_command(['shell', 'am', 'stack', 'list'])\n line = line.decode('utf-8')\n self.assertEqual(line, 'taskId=0 fake_activity visible=true '\n 'topActivity=ComponentInfo{fake_activity}')\n\n def test_send_touch(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n simulator.launch()\n simulator.send_touch([(0, 1, True, 0)])\n simulator.send_touch([(0, 1, False, 0)])\n # No assertions, we just want to ensure that `send_touch()` can be called\n # without crashing anything.\n\n def test_send_key(self):\n simulator = fake_simulator.FakeSimulator(\n config_classes.FakeSimulatorConfig(screen_dimensions=(320, 480))\n )\n simulator.launch()\n simulator.send_key(np.int32(123), 'keydown')\n simulator.send_key(np.int32(123), 'keyup')\n simulator.send_key(np.int32(123), 'keypress')\n # No assertions, we just want to ensure that `send_key()` can be called\n # without crashing anything.\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/specs.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Base specs for AndroidEnv.\"\"\"\n\nfrom android_env.components import action_type\nfrom android_env.proto import task_pb2\nfrom dm_env import specs\nimport numpy as np\n\n\n_PROTO_DTYPE_TO_NUMPY_DTYPE = {\n task_pb2.ArraySpec.DataType.FLOAT: np.float32,\n task_pb2.ArraySpec.DataType.DOUBLE: np.float64,\n task_pb2.ArraySpec.DataType.INT8: np.int8,\n task_pb2.ArraySpec.DataType.INT16: np.int16,\n task_pb2.ArraySpec.DataType.INT32: np.int32,\n task_pb2.ArraySpec.DataType.INT64: np.int64,\n task_pb2.ArraySpec.DataType.UINT8: np.uint8,\n task_pb2.ArraySpec.DataType.UINT16: np.uint16,\n task_pb2.ArraySpec.DataType.UINT32: np.uint32,\n task_pb2.ArraySpec.DataType.UINT64: np.uint64,\n task_pb2.ArraySpec.DataType.BOOL: np.bool_,\n task_pb2.ArraySpec.DataType.STRING_U1: np.dtype(('U1')),\n task_pb2.ArraySpec.DataType.STRING_U16: np.dtype((' dict[str, specs.Array]:\n \"\"\"Default action spec for AndroidEnv.\n\n Args:\n num_fingers: Number of virtual fingers of the agent.\n enable_key_events: Whether keyboard key events are enabled.\n\n Returns:\n A dict of action specs, each item corresponding to a virtual finger.\n action_type: An integer of type ActionType: TOUCH=0, LIFT=1, REPEAT=2\n touch_position: Position [x, y] of the touch action, where x, y are float\n values between 0.0 and 1.0 corresponding to the relative position on the\n screen. IGNORED when (action_type != ActionType.TOUCH).\n keycode: code representing the desired key press in XKB format. See the\n emulator_controller_pb2 for details.\n action_type_i: Action type for additional fingers (i>1).\n touch_position_i: Touch position for additional fingers (i>1).\n \"\"\"\n\n num_actions = len(action_type.ActionType) if enable_key_events else 3\n\n action_spec = {\n 'action_type':\n specs.DiscreteArray(num_values=num_actions, name='action_type'),\n 'touch_position':\n specs.BoundedArray(\n shape=(2,),\n dtype=np.float32,\n minimum=[0.0, 0.0],\n maximum=[1.0, 1.0],\n name='touch_position'),\n }\n\n for i in range(2, num_fingers + 1):\n action_spec.update({\n f'action_type_{i}':\n specs.DiscreteArray(\n num_values=len(action_type.ActionType),\n name=f'action_type_{i}'),\n f'touch_position_{i}':\n specs.BoundedArray(\n shape=(2,),\n dtype=np.float32,\n minimum=[0.0, 0.0],\n maximum=[1.0, 1.0],\n name=f'touch_position_{i}'),\n })\n\n if enable_key_events:\n action_spec['keycode'] = specs.DiscreteArray(\n num_values=(1 << 16) - 1, name='keycode')\n\n return action_spec\n\n\ndef base_observation_spec(height: int, width: int) -> dict[str, specs.Array]:\n \"\"\"Default observation spec for AndroidEnv.\n\n Args:\n height: Height of the device screen in pixels.\n width: Width of the device screen in pixels.\n\n Returns:\n pixels: Spec for the RGB screenshot of the device. Has shape (H, W, 3)\n timedelta: Spec for time delta since the last observation (in microseconds).\n The first timestep immediately after reset() will have this value set to\n 0.\n orientation: Spec for the latest orientation in a one-hot representation:\n [1, 0, 0, 0]: PORTRAIT (0 degrees)\n [0, 1, 0, 0]: LANDSCAPE (90 degrees clockwise)\n [0, 0, 1, 0]: PORTRAIT (180 degrees) (\"upside down\")\n [0, 0, 0, 1]: LANDSCAPE (270 degrees clockwise)\n \"\"\"\n\n return {\n 'pixels':\n specs.BoundedArray(\n shape=(height, width, 3),\n dtype=np.uint8,\n name='pixels',\n minimum=0,\n maximum=255),\n 'timedelta':\n specs.Array(shape=(), dtype=np.int64, name='timedelta'),\n 'orientation':\n specs.BoundedArray(\n shape=np.array([4]),\n dtype=np.uint8,\n name='orientation',\n minimum=0,\n maximum=1),\n }\n"
},
{
"path": "android_env/components/specs_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for specs.py.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.components import specs\nfrom android_env.proto import task_pb2\nfrom dm_env import specs as dm_env_specs\nimport numpy as np\n\n\nclass SpecsTest(parameterized.TestCase):\n\n def test_base_action_spec(self):\n action_spec = specs.base_action_spec(num_fingers=1)\n for spec in action_spec.values():\n self.assertIsInstance(spec, dm_env_specs.Array)\n self.assertEqual(action_spec['action_type'].shape, ())\n self.assertEqual(action_spec['action_type'].dtype, np.int32)\n self.assertEqual(action_spec['touch_position'].shape, (2,))\n self.assertEqual(action_spec['touch_position'].dtype, np.float32)\n\n def test_base_action_spec_with_key_events(self):\n action_spec = specs.base_action_spec(num_fingers=1, enable_key_events=True)\n for spec in action_spec.values():\n self.assertIsInstance(spec, dm_env_specs.Array)\n self.assertEqual(action_spec['action_type'].shape, ())\n self.assertEqual(action_spec['action_type'].dtype, np.int32)\n self.assertEqual(action_spec['touch_position'].shape, (2,))\n self.assertEqual(action_spec['touch_position'].dtype, np.float32)\n self.assertEqual(action_spec['keycode'].shape, ())\n self.assertEqual(action_spec['keycode'].dtype, np.int32)\n\n def test_base_action_spec_multitouch(self):\n action_spec = specs.base_action_spec(num_fingers=3)\n self.assertLen(action_spec.keys(), 6)\n for spec in action_spec.values():\n self.assertIsInstance(spec, dm_env_specs.Array)\n self.assertEqual(action_spec['action_type'].shape, ())\n self.assertEqual(action_spec['action_type'].dtype, np.int32)\n self.assertEqual(action_spec['touch_position'].shape, (2,))\n self.assertEqual(action_spec['touch_position'].dtype, np.float32)\n self.assertEqual(action_spec['action_type_2'].shape, ())\n self.assertEqual(action_spec['action_type_2'].dtype, np.int32)\n self.assertEqual(action_spec['touch_position_2'].shape, (2,))\n self.assertEqual(action_spec['touch_position_2'].dtype, np.float32)\n self.assertEqual(action_spec['action_type_3'].shape, ())\n self.assertEqual(action_spec['action_type_3'].dtype, np.int32)\n self.assertEqual(action_spec['touch_position_3'].shape, (2,))\n self.assertEqual(action_spec['touch_position_3'].dtype, np.float32)\n\n @parameterized.parameters(\n (480, 320),\n (100, 100),\n (1440, 1960),\n )\n def test_base_observation_spec(self, height, width):\n observation_spec = specs.base_observation_spec(height, width)\n for spec in observation_spec.values():\n self.assertIsInstance(spec, dm_env_specs.Array)\n self.assertEqual(observation_spec['pixels'].shape, (height, width, 3))\n self.assertEqual(observation_spec['pixels'].dtype, np.uint8)\n self.assertEqual(observation_spec['timedelta'].shape, ())\n self.assertEqual(observation_spec['timedelta'].dtype, np.int64)\n self.assertEqual(observation_spec['orientation'].shape, (4,))\n self.assertEqual(observation_spec['orientation'].dtype, np.uint8)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/components/task_manager.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"TaskManager handles all events and information related to the task.\"\"\"\n\nimport ast\nfrom collections.abc import Callable, Iterable, Sequence\nimport copy\nimport datetime\nimport itertools\nimport json\nimport re\nimport threading\nimport time\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env.components import adb_call_parser as adb_call_parser_lib\nfrom android_env.components import app_screen_checker\nfrom android_env.components import config_classes\nfrom android_env.components import dumpsys_thread\nfrom android_env.components import log_stream as log_stream_lib\nfrom android_env.components import logcat_thread\nfrom android_env.components import setup_step_interpreter\nfrom android_env.proto import task_pb2\nimport dm_env\nimport numpy as np\n\n\nclass TaskManager:\n \"\"\"Handles all events and information related to the task.\"\"\"\n\n def __init__(\n self,\n task: task_pb2.Task,\n config: config_classes.TaskManagerConfig | None = None,\n ):\n \"\"\"Controls task-relevant events and information.\n\n Args:\n task: A task proto defining the RL task.\n config: Configuration for this instance.\n \"\"\"\n\n self._task = task\n self._config = config or config_classes.TaskManagerConfig()\n self._lock = threading.Lock()\n self._logcat_thread = None\n self._dumpsys_thread = None\n self._setup_step_interpreter = None\n\n # Initialize stats.\n self._stats = {\n 'episode_steps': 0,\n 'reset_count_step_timeout': 0,\n 'reset_count_user_exited': 0,\n 'reset_count_episode_end': 0,\n 'reset_count_max_duration_reached': 0,\n 'restart_count_max_bad_states': 0,\n 'task_updates': 0,\n }\n\n # Initialize internal state\n self._task_start_time = None\n self._bad_state_counter = 0\n self._is_bad_episode = False\n\n self._latest_values = {\n 'reward': 0.0,\n 'score': 0.0,\n 'extra': {},\n 'episode_end': False,\n }\n\n logging.info('Task config: %s', self._task)\n\n def stats(self) -> dict[str, Any]:\n \"\"\"Returns a dictionary of stats.\n\n This method is expected to be called after setup_task() has been called.\n \"\"\"\n output = copy.deepcopy(self._stats)\n if self._setup_step_interpreter is not None:\n output.update(self._setup_step_interpreter.stats())\n return output\n\n def setup_task(self) -> None:\n \"\"\"Performs one-off task setup..\"\"\"\n self._setup_step_interpreter.interpret(self._task.setup_steps)\n\n def stop(self) -> None:\n \"\"\"Suspends task processing.\"\"\"\n n_tries = 3\n for i in range(n_tries):\n try:\n self._stop_logcat_thread()\n break\n except: # pylint: disable=bare-except\n logging.exception(\n 'Failed to stop logcat thread [%d/%d]. Continuing.', i + 1, n_tries\n )\n time.sleep(1)\n\n def start(\n self,\n adb_call_parser_factory: Callable[[], adb_call_parser_lib.AdbCallParser],\n log_stream: log_stream_lib.LogStream,\n ) -> None:\n \"\"\"Starts task processing.\"\"\"\n\n self._start_logcat_thread(log_stream=log_stream)\n self._logcat_thread.resume()\n self._start_dumpsys_thread(adb_call_parser_factory())\n self._start_setup_step_interpreter(adb_call_parser_factory())\n\n def reset_task(self) -> None:\n \"\"\"Resets a task for a new run.\"\"\"\n\n self._logcat_thread.pause()\n self._setup_step_interpreter.interpret(self._task.reset_steps)\n self._logcat_thread.resume()\n\n # Reset some other variables.\n if not self._is_bad_episode:\n self._bad_state_counter = 0\n self._is_bad_episode = False\n\n self._task_start_time = datetime.datetime.now()\n with self._lock:\n self._latest_values = {\n 'reward': 0.0,\n 'score': 0.0,\n 'extra': {},\n 'episode_end': False,\n }\n\n def rl_reset(self, observation: dict[str, Any]) -> dm_env.TimeStep:\n \"\"\"Performs one RL step.\"\"\"\n\n self._stats['episode_steps'] = 0\n\n self._logcat_thread.line_ready().wait()\n with self._lock:\n extras = self._get_current_extras()\n\n observation['extras'] = extras\n\n return dm_env.TimeStep(\n step_type=dm_env.StepType.FIRST,\n reward=0.0,\n discount=0.0,\n observation=observation,\n )\n\n def rl_step(self, observation: dict[str, Any]) -> dm_env.TimeStep:\n \"\"\"Performs one RL step.\"\"\"\n\n self._stats['episode_steps'] += 1\n\n self._logcat_thread.line_ready().wait()\n with self._lock:\n reward = self._get_current_reward()\n extras = self._get_current_extras()\n transition_fn = self._determine_transition_fn()\n\n observation['extras'] = extras\n\n return transition_fn(reward=reward, observation=observation)\n\n def _get_current_reward(self) -> float:\n \"\"\"Returns total reward accumulated since the last step.\"\"\"\n reward = self._latest_values['reward']\n self._latest_values['reward'] = 0.0\n return reward\n\n def _get_current_extras(self) -> dict[str, Any]:\n \"\"\"Returns task extras accumulated since the last step.\"\"\"\n extras = {}\n for name, values in self._latest_values['extra'].items():\n extras[name] = np.stack(values)\n self._latest_values['extra'] = {}\n return extras\n\n def _determine_transition_fn(self) -> Callable[..., dm_env.TimeStep]:\n \"\"\"Determines the type of RL transition will be used.\"\"\"\n\n # Check if user existed the task\n if self._dumpsys_thread.check_user_exited():\n self._increment_bad_state()\n self._stats['reset_count_user_exited'] += 1\n logging.warning('User exited the task. Truncating the episode.')\n logging.info('************* END OF EPISODE *************')\n return dm_env.truncation\n\n # Check if episode has ended\n if self._latest_values['episode_end']:\n self._stats['reset_count_episode_end'] += 1\n logging.info('End of episode from logcat! Ending episode.')\n return dm_env.termination\n\n # Check if step limit or time limit has been reached\n if self._task.max_episode_steps > 0:\n if self._stats['episode_steps'] > self._task.max_episode_steps:\n self._stats['reset_count_max_duration_reached'] += 1\n logging.info(\n 'Maximum task duration (%r steps) reached. Truncating the episode.',\n self._task.max_episode_steps,\n )\n return dm_env.truncation\n\n if self._task.max_episode_sec > 0.0:\n task_duration = datetime.datetime.now() - self._task_start_time\n max_episode_sec = self._task.max_episode_sec\n if task_duration > datetime.timedelta(seconds=int(max_episode_sec)):\n self._stats['reset_count_max_duration_reached'] += 1\n logging.info(\n 'Maximum task duration (%r sec) reached. Truncating the episode.',\n max_episode_sec,\n )\n return dm_env.truncation\n\n return dm_env.transition\n\n def _start_setup_step_interpreter(\n self, adb_call_parser: adb_call_parser_lib.AdbCallParser\n ):\n self._setup_step_interpreter = setup_step_interpreter.SetupStepInterpreter(\n adb_call_parser=adb_call_parser\n )\n\n def _start_logcat_thread(self, log_stream: log_stream_lib.LogStream):\n log_stream.set_log_filters(list(self._task.log_parsing_config.filters))\n self._logcat_thread = logcat_thread.LogcatThread(log_stream=log_stream)\n\n for event_listener in self._logcat_listeners():\n self._logcat_thread.add_event_listener(event_listener)\n\n def _start_dumpsys_thread(\n self, adb_call_parser: adb_call_parser_lib.AdbCallParser\n ):\n self._dumpsys_thread = dumpsys_thread.DumpsysThread(\n app_screen_checker=app_screen_checker.AppScreenChecker(\n adb_call_parser=adb_call_parser,\n expected_app_screen=self._task.expected_app_screen,\n ),\n check_frequency=self._config.dumpsys_check_frequency,\n max_failed_current_activity=self._config.max_failed_current_activity,\n )\n\n def _stop_logcat_thread(self):\n if self._logcat_thread is not None:\n self._logcat_thread.kill()\n self._logcat_thread = None\n\n def _increment_bad_state(self) -> None:\n \"\"\"Increments the bad state counter.\n\n Bad states are errors that shouldn't happen and that trigger an\n episode reset. If enough bad states have been seen consecutively,\n we restart the simulation in the hope of returning the simulation\n to a good state.\n \"\"\"\n logging.warning('Bad state detected.')\n if self._config.max_bad_states:\n self._is_bad_episode = True\n self._bad_state_counter += 1\n logging.warning('Bad state counter: %d.', self._bad_state_counter)\n if self._bad_state_counter >= self._config.max_bad_states:\n logging.error('Too many consecutive bad states. Restarting simulator.')\n self._stats['restart_count_max_bad_states'] += 1\n self._should_restart = True\n else:\n logging.warning('Max bad states not set, bad states will be ignored.')\n\n def _logcat_listeners(self) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates list of EventListeners for logcat thread.\"\"\"\n\n # Defaults to 'a^' since that regex matches no string by definition.\n regexps = self._task.log_parsing_config.log_regexps\n return itertools.chain(\n self._reward_listeners(regexps),\n self._reward_event_listeners(regexps),\n self._score_listeners(regexps),\n self._episode_end_listeners(regexps),\n self._extras_listeners(regexps),\n self._json_extras_listeners(regexps),\n )\n\n def _reward_listeners(\n self, regexps: task_pb2.LogParsingConfig.LogRegexps\n ) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates an iterable of reward listeners.\"\"\"\n\n def _reward_handler(event: re.Pattern[str], match: re.Match[str]):\n del event\n reward = float(match.group(1))\n with self._lock:\n self._latest_values['reward'] += reward\n\n for regexp in regexps.reward:\n yield logcat_thread.EventListener(\n regexp=re.compile(regexp or 'a^'), handler_fn=_reward_handler\n )\n\n def _reward_event_listeners(\n self, regexps: task_pb2.LogParsingConfig.LogRegexps\n ) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates an iterable of reward event listeners.\"\"\"\n\n for reward_event in regexps.reward_event:\n\n def get_reward_event_handler(reward):\n\n def _reward_event_handler(event: re.Pattern[str], match: re.Match[str]):\n del event, match\n with self._lock:\n self._latest_values['reward'] += reward\n\n return _reward_event_handler\n\n yield logcat_thread.EventListener(\n regexp=re.compile(reward_event.event or 'a^'),\n handler_fn=get_reward_event_handler(reward_event.reward),\n )\n\n def _score_listeners(\n self, regexps: task_pb2.LogParsingConfig.LogRegexps\n ) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates an iterable of score listeners.\"\"\"\n\n def _score_handler(event: re.Pattern[str], match: re.Match[str]):\n del event\n current_score = float(match.group(1))\n with self._lock:\n current_reward = current_score - self._latest_values['score']\n self._latest_values['score'] = current_score\n self._latest_values['reward'] += current_reward\n\n yield logcat_thread.EventListener(\n regexp=re.compile(regexps.score or 'a^'), handler_fn=_score_handler\n )\n\n def _episode_end_listeners(\n self, regexps: task_pb2.LogParsingConfig.LogRegexps\n ) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates an iterable of episode end listeners.\"\"\"\n\n def _episode_end_handler(event: re.Pattern[str], match: re.Match[str]):\n del event, match\n with self._lock:\n self._latest_values['episode_end'] = True\n\n for regexp in regexps.episode_end:\n yield logcat_thread.EventListener(\n regexp=re.compile(regexp or 'a^'), handler_fn=_episode_end_handler\n )\n\n def _process_extra(self, extra_name: str, extra: Sequence[int | float]):\n extra = np.array(extra)\n with self._lock:\n latest_extras = self._latest_values['extra']\n if extra_name in latest_extras:\n # If latest extra is not flushed, append.\n if (\n len(latest_extras[extra_name])\n >= self._config.extras_max_buffer_size\n ):\n latest_extras[extra_name].pop(0)\n latest_extras[extra_name].append(extra)\n else:\n latest_extras[extra_name] = [extra]\n self._latest_values['extra'] = latest_extras\n\n def _extras_listeners(\n self, regexps: task_pb2.LogParsingConfig.LogRegexps\n ) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates an iterable of extras listeners.\"\"\"\n\n def _extras_handler(event: re.Pattern[str], match: re.Match[str]):\n del event\n extra_name = match.group('name')\n extra = match.group('extra')\n if extra:\n try:\n extra = ast.literal_eval(extra)\n except (\n ValueError,\n TypeError,\n SyntaxError,\n MemoryError,\n RecursionError,\n ):\n logging.exception('Could not parse extra: %s', extra)\n # Don't try to process the extra as text; that would probably crash.\n return\n else:\n # No extra value provided for boolean extra. Setting value to True.\n extra = 1\n self._process_extra(extra_name, extra)\n\n for regexp in regexps.extra:\n yield logcat_thread.EventListener(\n regexp=re.compile(regexp or 'a^'), handler_fn=_extras_handler\n )\n\n def _json_extras_listeners(\n self, regexps: task_pb2.LogParsingConfig.LogRegexps\n ) -> Iterable[logcat_thread.EventListener]:\n \"\"\"Creates an iterable of JSON extras listeners.\"\"\"\n\n def _json_extras_handler(event: re.Pattern[str], match: re.Match[str]):\n del event\n extra_data = match.group('json_extra')\n try:\n extra = dict(json.loads(extra_data))\n except ValueError:\n logging.error('JSON string could not be parsed: %s', extra_data)\n return\n for extra_name, extra_value in extra.items():\n self._process_extra(extra_name, extra_value)\n\n for regexp in regexps.json_extra:\n yield logcat_thread.EventListener(\n regexp=re.compile(regexp or 'a^'), handler_fn=_json_extras_handler\n )\n"
},
{
"path": "android_env/components/task_manager_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.components.task_manager.py.\"\"\"\n\nimport json\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.components import adb_call_parser as adb_call_parser_lib\nfrom android_env.components import dumpsys_thread\nfrom android_env.components import log_stream\nfrom android_env.components import logcat_thread\nfrom android_env.components import setup_step_interpreter\nfrom android_env.components import task_manager\nfrom android_env.proto import task_pb2\nimport numpy as np\n\n\nclass TaskManagerTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self.addCleanup(mock.patch.stopall) # Disable previous patches.\n\n self._setup_step_interpreter = mock.create_autospec(\n setup_step_interpreter.SetupStepInterpreter)\n self._dumpsys_thread = mock.create_autospec(dumpsys_thread.DumpsysThread)\n self._logcat_thread = mock.create_autospec(logcat_thread.LogcatThread)\n self._log_stream = mock.create_autospec(log_stream.LogStream)\n\n mock.patch.object(\n setup_step_interpreter,\n 'SetupStepInterpreter',\n return_value=self._setup_step_interpreter).start()\n mock.patch.object(\n dumpsys_thread, 'DumpsysThread',\n return_value=self._dumpsys_thread).start()\n mock.patch.object(\n logcat_thread, 'LogcatThread',\n return_value=self._logcat_thread).start()\n mock.patch.object(\n log_stream, 'LogStream',\n return_value=self._log_stream).start()\n\n def test_start(self):\n task_mgr = task_manager.TaskManager(task=task_pb2.Task())\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n self.assertIsNotNone(task_mgr._logcat_thread)\n self.assertIsNotNone(task_mgr._dumpsys_thread)\n self.assertIsNotNone(task_mgr._setup_step_interpreter)\n\n def test_setup_task(self):\n task_mgr = task_manager.TaskManager(task=task_pb2.Task())\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n self._setup_step_interpreter.interpret.assert_called_once()\n\n def test_step_count(self):\n task_mgr = task_manager.TaskManager(task=task_pb2.Task())\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n task_mgr.rl_reset(observation={})\n self.assertEqual(task_mgr.stats()['episode_steps'], 0)\n task_mgr.rl_step(observation={})\n self.assertEqual(task_mgr.stats()['episode_steps'], 1)\n task_mgr.rl_step(observation={})\n self.assertEqual(task_mgr.stats()['episode_steps'], 2)\n task_mgr.rl_reset(observation={})\n self.assertEqual(task_mgr.stats()['episode_steps'], 0)\n\n def test_get_current_reward(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n match = event_listener.regexp.match('Reward: 123.0')\n if match is None: # Ignore events that are not rewards.\n return\n\n event_listener.handler_fn(event_listener.regexp, match)\n\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.reward.extend([\n '^[Rr]eward: ([-+]?[0-9]*\\\\.?[0-9]*)$'\n ])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n self.assertEqual(timestep.reward, 123.0)\n np.testing.assert_equal(timestep.observation['pixels'], np.array([1, 2, 3]))\n\n def test_reward_event(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n match_1 = event_listener.regexp.match('foo_1')\n match_2 = event_listener.regexp.match('foo_2')\n match_3 = event_listener.regexp.match('Reward: 2.0')\n if match_1:\n event_listener.handler_fn(event_listener.regexp, match_1)\n if match_2:\n event_listener.handler_fn(event_listener.regexp, match_2)\n if match_3:\n event_listener.handler_fn(event_listener.regexp, match_3)\n\n task = task_pb2.Task()\n reward_event_1 = task_pb2.LogParsingConfig.LogRegexps.RewardEvent(\n event='foo_1', reward=5.0)\n reward_event_2 = task_pb2.LogParsingConfig.LogRegexps.RewardEvent(\n event='foo_2', reward=-1.0)\n task.log_parsing_config.log_regexps.reward_event.extend(\n [reward_event_1, reward_event_2])\n task.log_parsing_config.log_regexps.reward.extend(\n ['^[Rr]eward: ([-+]?[0-9]*\\\\.?[0-9]*)$'])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n self.assertEqual(timestep.reward, 6.0)\n\n def test_get_current_reward_via_score(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n event = event_listener.regexp\n match = event.match('score: 200.0')\n if match is None: # Ignore events that are not scores.\n return\n\n event_listener.handler_fn(event, match)\n\n # Scores are accumulated by their differences, so a subsequent lower score\n # means that the final reward decreases.\n match = event.match('score: 185')\n event_listener.handler_fn(event, match)\n\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.score = (\n '^score: ([-+]?[0-9]*\\\\.?[0-9]*)$')\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n self.assertEqual(timestep.reward, 185.0)\n\n def test_get_current_extras(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n event = event_listener.regexp\n match = event.match('extra: some_extra [1, 2]')\n if match is None: # Ignore events that are not extras.\n return\n\n # Emit events.\n fn = event_listener.handler_fn\n fn(event, event.match('extra: an_extra [1, 2, 3]'))\n fn(event, event.match('extra: an_extra [4, 5, 6]'))\n fn(event, event.match('extra: another_extra 0.5'))\n fn(event, event.match('extra: multi_dimension_extra [[9,8,7],[6,5,4]]'))\n fn(event, event.match('extra: boolean_extra'))\n\n # Setup the task and trigger the listener.\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.extra.extend([\n '^extra: (?P[^ ]*)[ ]?(?P.*)$'\n ])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n\n # Check expectations.\n self.assertIn('extras', timestep.observation)\n extras = timestep.observation['extras']\n np.testing.assert_almost_equal([[1, 2, 3], [4, 5, 6]],\n extras.get('an_extra'))\n np.testing.assert_almost_equal([0.5], extras.get('another_extra'))\n np.testing.assert_almost_equal([[[9, 8, 7], [6, 5, 4]]],\n extras.get('multi_dimension_extra'))\n np.testing.assert_equal([1], extras.get('boolean_extra'))\n\n def test_get_current_extras_json_format(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n event = event_listener.regexp\n match = event.match('json_extra: {}')\n if match is None: # Ignore events that are not extras.\n return\n\n # Emit events.\n extra = {\n 'extra_scalar': 0,\n 'extra_list': [1, 2, 3, 4],\n 'extra_dict': {\n 'foo': 'bar'\n },\n 'extra_string': 'a_string'\n }\n extra_update = {'extra_string': 'a_new_string', 'extra_float': 0.6}\n fn = event_listener.handler_fn\n fn(event, event.match(f'json_extra: {json.dumps(extra)}'))\n fn(event, event.match(f'json_extra: {json.dumps(extra_update)}'))\n\n # Setup the task and trigger the listener.\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.json_extra.extend([\n '^json_extra: (?P.*)$'\n ])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n\n # Check expectations.\n self.assertIn('extras', timestep.observation)\n extras = timestep.observation['extras']\n expected_extra = {\n 'extra_scalar': [0],\n 'extra_list': [[1, 2, 3, 4]],\n 'extra_dict': [{\n 'foo': 'bar'\n }],\n 'extra_string': ['a_string', 'a_new_string'],\n 'extra_float': [0.6]\n }\n np.testing.assert_almost_equal(\n expected_extra.get('extra_scalar'), extras.get('extra_scalar'))\n np.testing.assert_almost_equal(\n expected_extra.get('extra_list'), extras.get('extra_list'))\n np.testing.assert_equal(\n expected_extra.get('extra_string'), extras.get('extra_string'))\n np.testing.assert_almost_equal(\n expected_extra.get('extra_float'), extras.get('extra_float'))\n np.testing.assert_equal(\n expected_extra.get('extra_dict'), extras.get('extra_dict'))\n\n def test_get_current_extras_failed_to_parse(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n event = event_listener.regexp\n match = event.match('extra: some_extra [1, 2]')\n if match is None: # Ignore events that are not extras.\n return\n\n # Emit events.\n fn = event_listener.handler_fn\n fn(event, event.match('extra: extra_with_malformed_1 [1]'))\n fn(event, event.match('extra: extra_with_malformed_1 [\\'this is \\\\ bad]'))\n fn(event, event.match('extra: extra_with_malformed_1 [2]'))\n fn(event, event.match('extra: extra_with_malformed_2 [\\'this is bad]'))\n fn(event, event.match('extra: extra_with_malformed_2 [1]'))\n fn(event, event.match('extra: extra_malformed_only [_very_bad_news]'))\n\n # Setup the task and trigger the listener.\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.extra.extend([\n '^extra: (?P[^ ]*)[ ]?(?P.*)$'\n ])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n\n # Check expectations.\n self.assertIn('extras', timestep.observation)\n extras = timestep.observation['extras']\n np.testing.assert_almost_equal(extras.get('extra_with_malformed_1'),\n [[1], [2]])\n np.testing.assert_almost_equal(extras.get('extra_with_malformed_2'), [[1]])\n self.assertNotIn('extra_malformed_only', extras)\n\n def test_multi_log_regexp(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n match = event_listener.regexp.match('Reward_2: 123.0')\n if match is None: # Ignore events that are not rewards.\n return\n\n event_listener.handler_fn(event_listener.regexp, match)\n\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.reward.extend([\n '^[Rr]eward_1: ([-+]?[0-9]*\\\\.?[0-9]*)$',\n '^[Rr]eward_2: ([-+]?[0-9]*\\\\.?[0-9]*)$'\n ])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n self.assertEqual(timestep.reward, 123.0)\n\n def test_multi_reward_regexp(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.'\n\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n match_1 = event_listener.regexp.match('Reward_1: 5.0')\n match_2 = event_listener.regexp.match('Reward_2: 10.0')\n\n if match_1:\n event_listener.handler_fn(event_listener.regexp, match_1)\n\n if match_2:\n event_listener.handler_fn(event_listener.regexp, match_2)\n\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.reward.extend([\n '^[Rr]eward_1: ([-+]?[0-9]*\\\\.?[0-9]*)$',\n '^[Rr]eward_2: ([-+]?[0-9]*\\\\.?[0-9]*)$',\n ])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n self.assertEqual(timestep.reward, 15.0)\n\n def test_determine_transition_fn(self):\n # Replace `LogcatThread.add_event_listener` with one that simply calls `fn`\n # right away.\n def my_add_ev_listener(event_listener: logcat_thread.EventListener):\n # Check that the event matches what's expected.\n event = event_listener.regexp\n match = event.match('I am done!')\n if match is None: # Ignore events that are not episode end.\n return\n\n event_listener.handler_fn(event, match)\n\n task = task_pb2.Task()\n task.log_parsing_config.log_regexps.episode_end.extend(['I am done!'])\n task_mgr = task_manager.TaskManager(task=task)\n self._logcat_thread.add_event_listener.side_effect = my_add_ev_listener\n adb_call_parser = mock.create_autospec(adb_call_parser_lib.AdbCallParser)\n task_mgr.start(lambda: adb_call_parser, log_stream=self._log_stream)\n task_mgr.setup_task()\n timestep = task_mgr.rl_step(\n observation={\n 'pixels': np.array([1, 2, 3]),\n })\n self.assertTrue(timestep.last())\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/env_interface.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Abstract AndroidEnv interface.\n\nAndroidEnv is a standard dm_env.Environment instance, but it also offers a few\nextra methods that clients may use for extended functionality.\n\"\"\"\n\nimport abc\nfrom typing import Any\n\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import state_pb2\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\nclass AndroidEnvInterface(dm_env.Environment, metaclass=abc.ABCMeta):\n \"\"\"Pure virtual interface for AndroidEnv implementations.\"\"\"\n\n # Methods required by dm_env.Environment.\n\n @abc.abstractmethod\n def action_spec(self) -> dict[str, specs.Array]:\n \"\"\"Returns the action specification.\"\"\"\n\n @abc.abstractmethod\n def observation_spec(self) -> dict[str, specs.Array]:\n \"\"\"Returns the observation specification.\"\"\"\n\n @abc.abstractmethod\n def reset(self) -> dm_env.TimeStep:\n \"\"\"Resets the current episode.\"\"\"\n\n @abc.abstractmethod\n def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:\n \"\"\"Executes `action` and returns a `TimeStep`.\"\"\"\n\n @abc.abstractmethod\n def close(self) -> None:\n \"\"\"Frees up resources.\"\"\"\n\n # Extensions provided by AndroidEnv.\n\n def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:\n \"\"\"Returns extra info provided by tasks.\"\"\"\n\n return {}\n\n @property\n def raw_action(self) -> Any:\n \"\"\"Returns the latest action.\"\"\"\n\n @property\n def raw_observation(self) -> Any:\n \"\"\"Returns the latest observation.\"\"\"\n\n def stats(self) -> dict[str, Any]:\n \"\"\"Returns information generated inside the implementation.\"\"\"\n\n return {}\n\n def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:\n \"\"\"Executes `call` and returns its response.\"\"\"\n\n return adb_pb2.AdbResponse()\n\n def load_state(\n self, request: state_pb2.LoadStateRequest\n ) -> state_pb2.LoadStateResponse:\n \"\"\"Loads a state.\n\n Args:\n request: A `LoadStateRequest` containing any parameters necessary to\n specify how/what state to load.\n\n Returns:\n A `LoadStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n raise NotImplementedError('This environment does not support loading state')\n\n def save_state(\n self, request: state_pb2.SaveStateRequest\n ) -> state_pb2.SaveStateResponse:\n \"\"\"Saves a state.\n\n Args:\n request: A `SaveStateRequest` containing any parameters necessary to\n specify how/what state to save.\n\n Returns:\n A `SaveStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n raise NotImplementedError('This environment does not support saving state')\n"
},
{
"path": "android_env/environment.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Android environment implementation.\"\"\"\n\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env import env_interface\nfrom android_env.components import adb_call_parser\nfrom android_env.components import coordinator as coordinator_lib\nfrom android_env.components import task_manager as task_manager_lib\nfrom android_env.components.simulators import base_simulator\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import state_pb2\nimport dm_env\nimport numpy as np\n\n\nclass AndroidEnv(env_interface.AndroidEnvInterface):\n \"\"\"An RL environment that interacts with Android apps.\"\"\"\n\n def __init__(\n self,\n simulator: base_simulator.BaseSimulator,\n coordinator: coordinator_lib.Coordinator,\n task_manager: task_manager_lib.TaskManager,\n ):\n \"\"\"Initializes the state of this AndroidEnv object.\"\"\"\n\n self._simulator = simulator\n self._coordinator = coordinator\n self._task_manager = task_manager\n self._latest_action = {}\n self._latest_observation = {}\n self._latest_extras = {}\n self._reset_next_step = True\n self._is_closed = False\n\n logging.info('Action spec: %s', self.action_spec())\n logging.info('Observation spec: %s', self.observation_spec())\n\n def __del__(self) -> None:\n self.close()\n\n # Methods required by dm_env.Environment.\n\n def action_spec(self) -> dict[str, dm_env.specs.Array]:\n return self._coordinator.action_spec()\n\n def observation_spec(self) -> dict[str, dm_env.specs.Array]:\n return self._coordinator.observation_spec()\n\n def reset(self) -> dm_env.TimeStep:\n \"\"\"Resets the environment for a new RL episode.\"\"\"\n\n logging.info('Resetting AndroidEnv...')\n\n # Execute a reset. Timestep will be of type FIRST.\n timestep = self._coordinator.rl_reset()\n\n # Process relevant information.\n if timestep.observation is not None:\n self._latest_extras = timestep.observation.pop('extras')\n self._latest_observation = timestep.observation.copy()\n else:\n # If the observation is None, we return the latest observation again.\n timestep = timestep._replace(observation=self._latest_observation.copy())\n\n self._latest_action = {}\n self._reset_next_step = False\n\n logging.info('Done resetting AndroidEnv.')\n logging.info('************* NEW EPISODE *************')\n\n return timestep\n\n def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:\n \"\"\"Takes a step in the environment.\"\"\"\n\n # Check if it's time to reset the episode.\n if self._reset_next_step:\n return self.reset()\n\n # Execute selected action.\n timestep = self._coordinator.rl_step(action)\n\n # Process relevant information.\n if timestep.observation is not None:\n self._latest_extras = timestep.observation.pop('extras')\n self._latest_observation = timestep.observation.copy()\n else:\n # If the observation is None, we return the latest observation again.\n timestep = timestep._replace(observation=self._latest_observation.copy())\n\n self._latest_action = action.copy()\n\n if timestep.last():\n self._reset_next_step = True\n logging.info('************* END OF EPISODE *************')\n\n return timestep\n\n def close(self) -> None:\n \"\"\"Cleans up running processes, threads and local files.\"\"\"\n if not self._is_closed:\n logging.info('Cleaning up AndroidEnv...')\n if hasattr(self, '_coordinator'):\n self._coordinator.close()\n logging.info('Done cleaning up AndroidEnv.')\n self._is_closed = True\n\n # Extensions provided by AndroidEnv.\n\n def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:\n \"\"\"Returns latest task extras.\"\"\"\n\n task_extras = {} # Build a copy to avoid reusing objects.\n for k, spec in self._latest_extras.items():\n extra_values = spec.astype(spec.dtype)\n task_extras[k] = extra_values[-1] if latest_only else extra_values\n return task_extras\n\n @property\n def raw_action(self):\n return self._latest_action.copy()\n\n @property\n def raw_observation(self):\n return self._latest_observation.copy()\n\n def stats(self) -> dict[str, Any]:\n coordinator_stats = self._coordinator.stats()\n task_manager_stats = self._task_manager.stats()\n return coordinator_stats | task_manager_stats\n\n def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse:\n return self._coordinator.execute_adb_call(call)\n\n def load_state(\n self, request: state_pb2.LoadStateRequest\n ) -> state_pb2.LoadStateResponse:\n \"\"\"Loads a state.\n\n Args:\n request: A `LoadStateRequest` containing any parameters necessary to\n specify how/what state to load.\n\n Returns:\n A `LoadStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n\n self._task_manager.stop()\n response = self._simulator.load_state(request)\n self._task_manager.start(\n adb_call_parser_factory=lambda: adb_call_parser.AdbCallParser(\n self._simulator.create_adb_controller()\n ),\n log_stream=self._simulator.create_log_stream(),\n )\n return response\n\n def save_state(\n self, request: state_pb2.SaveStateRequest\n ) -> state_pb2.SaveStateResponse:\n \"\"\"Saves a state.\n\n Args:\n request: A `SaveStateRequest` containing any parameters necessary to\n specify how/what state to save.\n\n Returns:\n A `SaveStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n\n return self._simulator.save_state(request)\n"
},
{
"path": "android_env/environment_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Unit tests for AndroidEnv.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import environment\nfrom android_env.components import config_classes\nfrom android_env.components import coordinator as coordinator_lib\nfrom android_env.components import task_manager as task_manager_lib\nfrom android_env.components.simulators import base_simulator\nfrom android_env.components.simulators.fake import fake_simulator\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import state_pb2\nimport dm_env\nimport numpy as np\n\n\ndef _create_mock_coordinator() -> coordinator_lib.Coordinator:\n coordinator = mock.create_autospec(coordinator_lib.Coordinator)\n coordinator.action_spec.return_value = {\n 'action_type':\n dm_env.specs.DiscreteArray(num_values=3),\n 'touch_position':\n dm_env.specs.BoundedArray(\n shape=(2,), dtype=np.float32, minimum=0.0, maximum=1.0),\n }\n coordinator.observation_spec.return_value = {\n 'pixels': dm_env.specs.Array(shape=(123, 456, 3), dtype=np.uint8),\n 'timedelta': dm_env.specs.Array(shape=(), dtype=np.int64),\n 'orientation': dm_env.specs.Array(shape=(4,), dtype=np.uint8),\n }\n return coordinator\n\n\ndef _create_fake_simulator() -> fake_simulator.FakeSimulator:\n return fake_simulator.FakeSimulator(\n config=config_classes.FakeSimulatorConfig(screen_dimensions=(123, 456))\n )\n\n\nclass AndroidEnvTest(absltest.TestCase):\n\n def test_specs(self):\n simulator = _create_fake_simulator()\n coordinator = _create_mock_coordinator()\n task_manager = mock.create_autospec(task_manager_lib.TaskManager)\n env = environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n\n # Check action spec.\n self.assertNotEmpty(env.action_spec())\n self.assertIn('action_type', env.action_spec())\n self.assertIsInstance(env.action_spec()['action_type'],\n dm_env.specs.DiscreteArray)\n self.assertIn('touch_position', env.action_spec())\n self.assertIsInstance(env.action_spec()['touch_position'],\n dm_env.specs.BoundedArray)\n\n # Check observation spec.\n self.assertNotEmpty(env.observation_spec())\n self.assertIn('pixels', env.observation_spec())\n self.assertIsInstance(env.observation_spec()['pixels'], dm_env.specs.Array)\n # The `pixels` entry in the observation spec should match the screen size of\n # the simulator with three color channels (RGB).\n self.assertEqual(env.observation_spec()['pixels'].shape, (123, 456, 3))\n self.assertIn('timedelta', env.observation_spec())\n self.assertIsInstance(env.observation_spec()['timedelta'],\n dm_env.specs.Array)\n # The `timedelta` should be a scalar.\n self.assertEqual(env.observation_spec()['timedelta'].shape, ())\n self.assertIn('orientation', env.observation_spec())\n # The `orientation` should be a one-hot vector with four dimensions.\n self.assertIsInstance(env.observation_spec()['orientation'],\n dm_env.specs.Array)\n self.assertEqual(env.observation_spec()['orientation'].shape, (4,))\n\n def test_reset_and_step(self):\n simulator = _create_fake_simulator()\n coordinator = _create_mock_coordinator()\n task_manager = mock.create_autospec(task_manager_lib.TaskManager)\n coordinator.action_spec.return_value = {\n 'action_type':\n dm_env.specs.DiscreteArray(num_values=3),\n 'touch_position':\n dm_env.specs.BoundedArray(\n shape=(2,), dtype=np.float32, minimum=0.0, maximum=1.0),\n }\n coordinator.observation_spec.return_value = {\n 'pixels': dm_env.specs.Array(shape=(123, 456, 3), dtype=np.uint8),\n 'timedelta': dm_env.specs.Array(shape=(), dtype=np.int64),\n 'orientation': dm_env.specs.Array(shape=(4,), dtype=np.uint8),\n }\n env = environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n coordinator.rl_reset.return_value = dm_env.TimeStep(\n step_type=dm_env.StepType.FIRST,\n reward=0.0,\n discount=0.0,\n observation={\n 'pixels': np.random.rand(987, 654, 3),\n 'timedelta': 123456,\n 'orientation': np.array((1, 0, 0, 0)),\n 'extras': {\n 'click': np.array([[246]], dtype=np.int64)\n }\n },\n )\n\n ts = env.reset()\n self.assertIsInstance(ts, dm_env.TimeStep)\n # After a `reset()` the TimeStep should follow some expectations.\n self.assertTrue(ts.first())\n self.assertEqual(ts.reward, 0.0)\n self.assertEqual(ts.discount, 0.0)\n obs = ts.observation\n self.assertIn('pixels', obs)\n self.assertEqual(obs['pixels'].shape, (987, 654, 3))\n self.assertIn('timedelta', obs)\n self.assertEqual(obs['timedelta'], 123456)\n self.assertIn('orientation', obs)\n self.assertEqual(obs['orientation'].shape, (4,))\n np.testing.assert_equal(obs['orientation'], (1, 0, 0, 0))\n\n # Extras should also be provided.\n extras = env.task_extras()\n self.assertIn('click', extras)\n self.assertEqual(extras['click'], np.array([246], dtype=np.int64))\n\n coordinator.stats.return_value = {'my_measurement': 135}\n task_manager.stats.return_value = {'another_measurement': 79}\n\n # Step again in the environment and check expectations again.\n pixels = np.random.rand(987, 654, 3)\n latest_observation = {\n 'pixels': pixels,\n 'timedelta': 123456,\n 'orientation': np.array((1, 0, 0, 0)),\n 'extras': {\n 'click': np.array([[246]], dtype=np.int64)\n }\n }\n coordinator.rl_step.return_value = dm_env.transition(\n reward=0.0,\n discount=0.0,\n observation=latest_observation,\n )\n ts = env.step({'action_type': 1, 'touch_position': (10, 20)})\n self.assertIsInstance(ts, dm_env.TimeStep)\n # The StepType now should NOT be FIRST.\n self.assertFalse(ts.first())\n self.assertEqual(ts.reward, 0.0)\n self.assertEqual(ts.discount, 0.0)\n obs = ts.observation\n self.assertIn('pixels', obs)\n self.assertEqual(obs['pixels'].shape, (987, 654, 3))\n self.assertIn('timedelta', obs)\n self.assertEqual(obs['timedelta'], 123456)\n self.assertIn('orientation', obs)\n self.assertEqual(obs['orientation'].shape, (4,))\n np.testing.assert_equal(obs['orientation'], (1, 0, 0, 0))\n\n # Extras should still be provided.\n extras = env.task_extras()\n self.assertIn('click', extras)\n self.assertEqual(extras['click'], np.array([246], dtype=np.int64))\n\n # At this point these methods and properties should return something.\n self.assertNotEmpty(env.stats())\n self.assertNotEmpty(env.raw_observation)\n self.assertNotIn('extras', env.raw_observation)\n self.assertNotEmpty(env.raw_action)\n\n # If the observation is None, we want to return the latest observation.\n coordinator.rl_step.return_value = dm_env.truncation(\n reward=0.0,\n observation=None,\n )\n ts = env.step({'action_type': 1, 'touch_position': (10, 20)})\n self.assertIsInstance(ts, dm_env.TimeStep)\n # Assert the observation matches the latest observation.\n obs = ts.observation\n self.assertIn('pixels', obs)\n self.assertEqual(obs['pixels'].shape, (987, 654, 3))\n np.testing.assert_equal(obs['pixels'], pixels)\n self.assertIn('timedelta', obs)\n self.assertEqual(obs['timedelta'], 123456)\n self.assertIn('orientation', obs)\n self.assertEqual(obs['orientation'].shape, (4,))\n np.testing.assert_equal(obs['orientation'], (1, 0, 0, 0))\n\n def test_adb_call(self):\n simulator = _create_fake_simulator()\n coordinator = _create_mock_coordinator()\n task_manager = mock.create_autospec(task_manager_lib.TaskManager)\n env = environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n call = adb_pb2.AdbRequest(\n force_stop=adb_pb2.AdbRequest.ForceStop(package_name='blah'))\n expected_response = adb_pb2.AdbResponse(\n status=adb_pb2.AdbResponse.Status.OK)\n coordinator.execute_adb_call.return_value = expected_response\n\n response = env.execute_adb_call(call)\n\n self.assertEqual(response, expected_response)\n coordinator.execute_adb_call.assert_called_once_with(call)\n\n def test_load_state(self):\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n coordinator = _create_mock_coordinator()\n task_manager = mock.create_autospec(task_manager_lib.TaskManager)\n env = environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n expected_response = state_pb2.LoadStateResponse(\n status=state_pb2.LoadStateResponse.Status.OK\n )\n request = state_pb2.LoadStateRequest(args={'foo': 'bar'})\n simulator.load_state.return_value = expected_response\n response = env.load_state(request)\n self.assertEqual(response, expected_response)\n simulator.load_state.assert_called_once_with(request)\n task_manager.stop.assert_called_once()\n task_manager.start.assert_called_once()\n\n def test_save_state(self):\n simulator = mock.create_autospec(base_simulator.BaseSimulator)\n coordinator = _create_mock_coordinator()\n task_manager = mock.create_autospec(task_manager_lib.TaskManager)\n env = environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n expected_response = state_pb2.SaveStateResponse(\n status=state_pb2.SaveStateResponse.Status.OK\n )\n request = state_pb2.SaveStateRequest(args={'foo': 'bar'})\n simulator.save_state.return_value = expected_response\n response = env.save_state(request)\n self.assertEqual(response, expected_response)\n simulator.save_state.assert_called_once_with(request)\n\n def test_double_close(self):\n simulator = _create_fake_simulator()\n coordinator = _create_mock_coordinator()\n task_manager = mock.create_autospec(task_manager_lib.TaskManager)\n env = environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n env.close()\n env.close()\n coordinator.close.assert_called_once()\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/loader.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Function for loading AndroidEnv.\"\"\"\n\nimport os\n\nfrom absl import logging\nfrom android_env import environment\nfrom android_env.components import config_classes\nfrom android_env.components import coordinator as coordinator_lib\nfrom android_env.components import device_settings as device_settings_lib\nfrom android_env.components import task_manager as task_manager_lib\nfrom android_env.components.simulators.emulator import emulator_simulator\nfrom android_env.components.simulators.fake import fake_simulator\nfrom android_env.proto import task_pb2\n\nfrom google.protobuf import text_format\n\n\ndef _load_task(task_config: config_classes.TaskConfig) -> task_pb2.Task:\n \"\"\"Returns the task according to `task_config`.\"\"\"\n\n task = task_pb2.Task()\n match task_config:\n case config_classes.FilesystemTaskConfig():\n with open(task_config.path, 'r') as proto_file:\n text_format.Parse(proto_file.read(), task)\n case _:\n logging.error('Unsupported TaskConfig: %r', task_config)\n\n return task\n\n\ndef load(config: config_classes.AndroidEnvConfig) -> environment.AndroidEnv:\n \"\"\"Loads an AndroidEnv instance.\"\"\"\n\n task = _load_task(config.task)\n task_manager = task_manager_lib.TaskManager(task)\n\n match config.simulator:\n case config_classes.EmulatorConfig():\n _process_emulator_launcher_config(config.simulator)\n simulator = emulator_simulator.EmulatorSimulator(config=config.simulator)\n case config_classes.FakeSimulatorConfig():\n simulator = fake_simulator.FakeSimulator(config=config.simulator)\n case _:\n raise ValueError('Unsupported simulator config: {config.simulator}')\n\n device_settings = device_settings_lib.DeviceSettings(simulator)\n coordinator = coordinator_lib.Coordinator(\n simulator, task_manager, device_settings\n )\n return environment.AndroidEnv(\n simulator=simulator, coordinator=coordinator, task_manager=task_manager\n )\n\n\ndef _process_emulator_launcher_config(\n emulator_config: config_classes.EmulatorConfig,\n) -> None:\n \"\"\"Adjusts the configuration of the emulator depending on some conditions.\"\"\"\n\n # Expand the user directory if specified.\n launcher_config = emulator_config.emulator_launcher\n launcher_config.android_avd_home = os.path.expanduser(\n launcher_config.android_avd_home\n )\n launcher_config.android_sdk_root = os.path.expanduser(\n launcher_config.android_sdk_root\n )\n launcher_config.emulator_path = os.path.expanduser(\n launcher_config.emulator_path\n )\n emulator_config.adb_controller.adb_path = os.path.expanduser(\n emulator_config.adb_controller.adb_path\n )\n"
},
{
"path": "android_env/loader_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for loader.\"\"\"\n\nimport builtins\nimport os\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env import loader\nfrom android_env.components import config_classes\nfrom android_env.components import coordinator as coordinator_lib\nfrom android_env.components import device_settings as device_settings_lib\nfrom android_env.components import task_manager as task_manager_lib\nfrom android_env.components.simulators.emulator import emulator_simulator\nfrom android_env.components.simulators.fake import fake_simulator\nfrom android_env.proto import task_pb2\n\n\nclass LoaderTest(absltest.TestCase):\n\n @mock.patch.object(task_manager_lib, 'TaskManager', autospec=True)\n @mock.patch.object(emulator_simulator, 'EmulatorSimulator', autospec=True)\n @mock.patch.object(device_settings_lib, 'DeviceSettings', autospec=True)\n @mock.patch.object(coordinator_lib, 'Coordinator', autospec=True)\n @mock.patch.object(builtins, 'open', autospec=True)\n def test_load_emulator(\n self,\n mock_open,\n mock_coordinator,\n mock_device_settings,\n mock_simulator_class,\n mock_task_manager,\n ):\n\n # Arrange.\n mock_open.return_value.__enter__ = mock_open\n mock_open.return_value.read.return_value = ''\n config = config_classes.AndroidEnvConfig(\n task=config_classes.FilesystemTaskConfig(path='some/path/'),\n simulator=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n avd_name='my_avd',\n android_avd_home='~/.android/avd',\n android_sdk_root='~/Android/Sdk',\n emulator_path='~/Android/Sdk/emulator/emulator',\n run_headless=False,\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='~/Android/Sdk/platform-tools/adb',\n ),\n ),\n )\n\n # Act.\n env = loader.load(config)\n\n # Assert.\n self.assertIsInstance(env, env_interface.AndroidEnvInterface)\n mock_simulator_class.assert_called_with(\n config=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n avd_name='my_avd',\n android_avd_home=os.path.expanduser('~/.android/avd'),\n android_sdk_root=os.path.expanduser('~/Android/Sdk'),\n emulator_path=os.path.expanduser(\n '~/Android/Sdk/emulator/emulator'\n ),\n run_headless=False,\n gpu_mode='swangle_indirect',\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path=os.path.expanduser('~/Android/Sdk/platform-tools/adb'),\n adb_server_port=5037,\n ),\n )\n )\n mock_coordinator.assert_called_with(\n mock_simulator_class.return_value,\n mock_task_manager.return_value,\n mock_device_settings.return_value,\n )\n\n @mock.patch.object(task_manager_lib, 'TaskManager', autospec=True)\n @mock.patch.object(fake_simulator, 'FakeSimulator', autospec=True)\n @mock.patch.object(device_settings_lib, 'DeviceSettings', autospec=True)\n @mock.patch.object(coordinator_lib, 'Coordinator', autospec=True)\n @mock.patch.object(builtins, 'open', autospec=True)\n def test_load_fake_simulator(\n self,\n mock_open,\n mock_coordinator,\n mock_device_settings,\n mock_simulator_class,\n mock_task_manager,\n ):\n\n # Arrange.\n mock_open.return_value.__enter__ = mock_open\n mock_open.return_value.read.return_value = ''\n config = config_classes.AndroidEnvConfig(\n task=config_classes.FilesystemTaskConfig(path='some/path/'),\n simulator=config_classes.FakeSimulatorConfig(\n screen_dimensions=(1234, 5678)\n ),\n )\n\n # Act.\n env = loader.load(config)\n\n # Assert.\n self.assertIsInstance(env, env_interface.AndroidEnvInterface)\n mock_simulator_class.assert_called_with(\n config=config_classes.FakeSimulatorConfig(\n screen_dimensions=(1234, 5678)\n )\n )\n mock_coordinator.assert_called_with(\n mock_simulator_class.return_value,\n mock_task_manager.return_value,\n mock_device_settings.return_value,\n )\n\n @mock.patch.object(task_manager_lib, 'TaskManager', autospec=True)\n @mock.patch.object(emulator_simulator, 'EmulatorSimulator', autospec=True)\n @mock.patch.object(coordinator_lib, 'Coordinator', autospec=True)\n @mock.patch.object(builtins, 'open', autospec=True)\n def test_task(\n self, mock_open, mock_coordinator, mock_simulator, mock_task_manager\n ):\n\n # Arrange.\n del mock_coordinator, mock_simulator\n mock_open.return_value.__enter__ = mock_open\n mock_open.return_value.read.return_value = r'''\nid: \"fake_task\"\nname: \"Fake Task\"\ndescription: \"Task for testing loader.\"\nmax_episode_sec: 0\n'''\n config = config_classes.AndroidEnvConfig(\n task=config_classes.FilesystemTaskConfig(path='some/path/'),\n simulator=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n avd_name='my_avd'\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path='~/Android/Sdk/platform-tools/adb',\n ),\n ),\n )\n\n # Act.\n env = loader.load(config)\n\n # Assert.\n expected_task = task_pb2.Task()\n expected_task.id = 'fake_task'\n expected_task.name = 'Fake Task'\n expected_task.description = 'Task for testing loader.'\n expected_task.max_episode_sec = 0\n\n mock_task_manager.assert_called_with(expected_task)\n self.assertIsInstance(env, env_interface.AndroidEnvInterface)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/proto/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/proto/a11y/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/proto/a11y/a11y.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nimport \"android_env/proto/a11y/android_accessibility_forest.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// A service to send Accessibility information to a remote server.\n//\n// The client is assumed to be running inside an Android device (e.g. emulator\n// or real device) while the server is assumed to be running outside (e.g. in a\n// Python process).\nservice A11yService {\n // Sends a forest of Accessibility trees to a server.\n rpc SendForest(AndroidAccessibilityForest) returns (ForestResponse) {}\n // Sends an a11y event to a server.\n rpc SendEvent(EventRequest) returns (EventResponse) {}\n\n // Long-lived bidirection communication between the client and the server.\n rpc Bidi(stream ClientToServer) returns (stream ServerToClient) {}\n}\n\n// TODO(b/334952387): Remove `ForestResponse`, `EventRequest` and\n// `EventResponse` once bidi communication is in-place.\nmessage ForestResponse {\n // The error if anything.\n string error = 1;\n}\n\n// An Accessibility event.\nmessage EventRequest {\n // A single event as a dictionary.\n map event = 1;\n}\n\nmessage EventResponse {\n // The error if anything.\n string error = 1;\n}\n\n// The message sent from the Android device to the server running outside of the\n// device.\nmessage ClientToServer {\n oneof payload {\n EventRequest event = 1;\n AndroidAccessibilityForest forest = 2;\n }\n}\n\n// The message sent from the server running outside of the device to the Android\n// device.\nmessage ServerToClient {\n // A request to obtain the Accessibility forest.\n message GetA11yForest {}\n\n oneof payload {\n GetA11yForest get_forest = 1;\n }\n}\n"
},
{
"path": "android_env/proto/a11y/android_accessibility_action.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// An Android Accessibility Action.\n// Next index: 3\nmessage AndroidAccessibilityAction {\n // Required ID that uniquely identifies the action for this node.\n // Can be one of the standard action IDs listed in the documentation.\n // https://developer.android.com/reference/android/view/accessibility/AccessibilityNodeInfo.AccessibilityAction\n int32 id = 1;\n\n // Optional label describing what the action is.\n string label = 2;\n}\n"
},
{
"path": "android_env/proto/a11y/android_accessibility_forest.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nimport \"android_env/proto/a11y/android_accessibility_window_info.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// A forest of Android accessibility trees. Each tree belongs to a single\n// window. Next index: 2\nmessage AndroidAccessibilityForest {\n // All of the windows present on screen.\n repeated AndroidAccessibilityWindowInfo windows = 1;\n}\n"
},
{
"path": "android_env/proto/a11y/android_accessibility_node_info.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nimport \"android_env/proto/a11y/android_accessibility_action.proto\";\nimport \"android_env/proto/a11y/android_accessibility_node_info_clickable_span.proto\";\nimport \"android_env/proto/a11y/rect.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// An Android AccessibilityNodeInfo.\n// Next index: 32\nmessage AndroidAccessibilityNodeInfo {\n // Unique monotonically-increasing ID.\n int32 unique_id = 1;\n\n // The bounds of this node within the device's screen.\n ProtoRect bounds_in_screen = 2;\n\n // The name of the View class that created this node.\n string class_name = 3;\n\n // The content description of the node.\n string content_description = 4;\n\n // The hint text of the node.\n string hint_text = 5;\n\n // The name of the package this node comes from.\n string package_name = 6;\n\n // The text of this node.\n string text = 7;\n\n // The start index of the text selection.\n int64 text_selection_start = 8;\n\n // The end index of the text selection.\n int64 text_selection_end = 9;\n\n // The view ID resource name of the node.\n string view_id_resource_name = 10;\n\n // The ID of the window this node belongs to.\n int32 window_id = 11;\n\n // If true, this node can be checked.\n bool is_checkable = 12;\n\n // If true, this node is currently checked.\n bool is_checked = 13;\n\n // If true, this node (probably) responds to being clicked.\n bool is_clickable = 14;\n\n // If true, this node's text can be edited by the user.\n bool is_editable = 15;\n\n // If true, this node is enabled (e.g., if it is a button).\n bool is_enabled = 16;\n\n // If true, this node can be focused (e.g., a text input).\n bool is_focusable = 17;\n\n // If true, this node is currently focused.\n bool is_focused = 18;\n\n // If true, this node (probably) responds to being long pressed.\n bool is_long_clickable = 19;\n\n // If true, this node is a password input.\n bool is_password = 20;\n\n // If true, this node can be scrolled.\n bool is_scrollable = 21;\n\n // If true, this node is currently selected.\n bool is_selected = 22;\n\n // If true, this node is (probably) visible to the user.\n bool is_visible_to_user = 23;\n\n // List of actions that can be performed on this node.\n repeated AndroidAccessibilityAction actions = 24;\n\n // Ordered list of child IDs (i.e., unique_id).\n repeated int32 child_ids = 25 [packed = true];\n\n // List of clickable spans present in the node's text or content description.\n repeated AndroidAccessibilityNodeInfoClickableSpan clickable_spans = 26;\n\n // The depth of this node in the accessibility tree.\n int32 depth = 27;\n\n // Unique ID of the node that this node is declaring itself to be labeled by.\n int32 labeled_by_id = 28;\n\n // Unique ID of the node that this is node is declaring itself to be a label\n // for.\n int32 label_for_id = 29;\n\n // The drawing order for the node.\n int32 drawing_order = 30;\n\n // The tooltip text of the node.\n string tooltip_text = 31;\n}\n"
},
{
"path": "android_env/proto/a11y/android_accessibility_node_info_clickable_span.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// A single clickable span found in the accessibility node's text.\n// Next index: 6\nmessage AndroidAccessibilityNodeInfoClickableSpan {\n // The source of the span (so the client can find the correct spannable string\n // in the node).\n // Next index: 3\n enum SpanSource {\n UNKNOWN_TYPE = 0; // Catch all type for forward compatibility.\n TEXT = 1; // The span is from node#getText\n CONTENT_DESCRIPTION = 2; // The span is from node#getContentDescription.\n }\n\n // The text of the span (a substring of the spannable string).\n string text = 1;\n\n // The URL attached to the span if specified.\n string url = 2;\n\n // The source of the span.\n SpanSource source = 3;\n\n // The index of the first character of the span in the spannable string.\n // The end of the span would be a sum of span_start and text.length().\n int32 start = 4;\n\n // The unique_id from the corresponding AndroidAccessibilityNodeInfo.\n int32 node_id = 5;\n}\n"
},
{
"path": "android_env/proto/a11y/android_accessibility_tree.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nimport \"android_env/proto/a11y/android_accessibility_node_info.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// A tree (actually a graph) of Android accessibility nodes.\n// Next index: 3\nmessage AndroidAccessibilityTree {\n // All of the nodes in the graph. The root node is the node whose ID is 0.\n repeated AndroidAccessibilityNodeInfo nodes = 1;\n}\n"
},
{
"path": "android_env/proto/a11y/android_accessibility_window_info.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nimport \"android_env/proto/a11y/android_accessibility_tree.proto\";\nimport \"android_env/proto/a11y/rect.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// An Android AccessibilityWindowInfo.\n// Next index: 12\nmessage AndroidAccessibilityWindowInfo {\n // Type of the window.\n // Next index: 8\n enum WindowType {\n // The window type is an unknown value.\n UNKNOWN_TYPE = 0;\n\n // A standard application window.\n TYPE_APPLICATION = 1;\n\n // An IME window (e.g. GBoard).\n TYPE_INPUT_METHOD = 2;\n\n // A system window (e.g., a notification).\n TYPE_SYSTEM = 3;\n\n // An accessibility overlay.\n TYPE_ACCESSIBILITY_OVERLAY = 4;\n\n // A system window used to divide the screen in split-screen mode. This type\n // of window is present only in split-screen mode.\n TYPE_SPLIT_SCREEN_DIVIDER = 5;\n\n // Used to show the UI for window-based magnification.\n TYPE_MAGNIFICATION_OVERLAY = 6;\n\n // System window that has the function to control an associated window.\n TYPE_WINDOW_CONTROL = 7;\n }\n\n // Bounds of this window in the device's screen.\n ProtoRect bounds_in_screen = 1;\n\n // A unique ID identifying the display in which this window is shown.\n int32 display_id = 2;\n\n // Unique ID as defined by the Android platform.\n int32 id = 3;\n\n // Z-index of the window. Windows with a greater z-index appear in front of\n // those with a lesser z-index.\n int32 layer = 4;\n\n // The title of the window, if set.\n string title = 5;\n\n // The type of the window.\n WindowType window_type = 6;\n\n // If true, the window is currently accessibility-focused.\n bool is_accessibility_focused = 7;\n\n // If true, the window is currently active.\n bool is_active = 8;\n\n // If true, the window is currently focused.\n bool is_focused = 9;\n\n // If true, the window is in Picture in Picture mode.\n bool is_in_picture_in_picture_mode = 10;\n\n // The associated accessibility tree for this window.\n AndroidAccessibilityTree tree = 11;\n}\n"
},
{
"path": "android_env/proto/a11y/rect.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\noption java_multiple_files = true;\noption java_package = \"com.google.androidenv.accessibilityforwarder\";\n\n// Proto representation of Android Rect.\n// https://developer.android.com/reference/android/graphics/Rect\n// Next index: 5\nmessage ProtoRect {\n int32 left = 1;\n int32 top = 2;\n int32 right = 3;\n int32 bottom = 4;\n}\n"
},
{
"path": "android_env/proto/adb.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nmessage AdbRequest {\n // Installs an APK into the simulator.\n message InstallApk {\n\n // A location in the filesystem.\n message Filesystem {\n string path = 1;\n }\n\n // A byte sequence of a single APK file.\n message Blob {\n // The serialized file as bytes.\n bytes contents = 1;\n }\n\n oneof location {\n Filesystem filesystem = 2;\n Blob blob = 6;\n }\n }\n\n message StartActivity {\n string full_activity = 1;\n repeated string extra_args = 2;\n // Whether to stop the current app before starting the activity.\n // Notice that if this option is `true`, the activity probably needs the\n // `android:launchMode=\"singleTop\"` attribute in its `AndroidManifest.xml`,\n // otherwise intents may not be received by `onNewIntent()`. Please see more\n // info on `android:launchMode` at\n // https://developer.android.com/guide/topics/manifest/activity-element.\n bool force_stop = 3;\n }\n\n message SendBroadcast {\n // Action to send during the broadcast event.\n string action = 1;\n\n // Specify the component name with package name prefix to create an explicit\n // intent, such as com.example.app/.ExampleActivity (see -n specification at\n // https://developer.android.com/tools/adb#IntentSpec).\n string component = 2;\n }\n\n message UninstallPackage {\n string package_name = 1;\n }\n\n message ForceStop {\n string package_name = 1;\n }\n\n message Tap {\n // NOTE: These are absolute coordinates in the range of the screen\n // resolution. They are NOT floats in [0,1].\n // Precondition: `x` and `y` must be non-negative.\n int32 x = 1;\n int32 y = 2;\n }\n\n message PressButton {\n enum Button {\n HOME = 0;\n BACK = 1;\n ENTER = 2;\n }\n Button button = 1;\n }\n\n // Pins the given activity to the screen.\n // This essentially locks the user into a single app mode (aka \"Kiosk mode\").\n message StartScreenPinning {\n string full_activity = 1;\n }\n\n // Returns the full activity name that is currently opened to the user.\n // If successful, a GetCurrentActivityResponse is returned.\n message GetCurrentActivity {}\n\n // Returns the orientation of the device.\n message GetOrientationRequest {}\n\n // Performs `adb push`.\n // Please see https://developer.android.com/studio/command-line/adb#copyfiles.\n //\n // Notice that a source destination path for the file is not sent, but raw\n // bytes in `content` instead. Obviously, the `content` can be set from a real\n // file, but this is done to ensure Task definitions are as hermetic as\n // possible, without depending on the environment from where they're run.\n message Push {\n // The contents of the file.\n bytes content = 1;\n\n // Destination path _inside_ Android. E.g. /sdcard/my_file.txt.\n string path = 2;\n }\n\n // Performs `adb pull`.\n // Please see https://developer.android.com/studio/command-line/adb#copyfiles.\n //\n // Notice that a local destination for the copied file is not sent, as raw\n // bytes are returned instead (please see PullResponse). Obviously, these\n // bytes can be written to disk by the caller of this command.\n message Pull {\n // Path _inside_ Android. E.g. /sdcard/my_file.txt.\n string path = 1;\n }\n\n // Inserts text into the current text field (if any).\n // Essentially `adb shell input text `.\n message InputText {\n string text = 1;\n }\n\n // Issues an `adb shell settings` command.\n message SettingsRequest {\n // Each request has an associated namespace.\n enum Namespace {\n UNKNOWN = 0;\n SYSTEM = 1;\n SECURE = 2;\n GLOBAL = 3;\n }\n\n // Retrieves the current value for `key`.\n message Get {\n string key = 1;\n }\n\n // Changes the contents `key` to `value`.\n message Put {\n string key = 1;\n string value = 2;\n }\n\n // Deletes the entry for `key`.\n message Delete {\n string key = 1;\n }\n\n // Resets the global/secure table for a package with the given mode.\n message Reset {\n enum Mode {\n UNKNOWN = 0;\n UNTRUSTED_DEFAULTS = 1;\n UNTRUSTED_CLEAR = 2;\n TRUSTED_DEFAULTS = 3;\n }\n\n string package_name = 1;\n Mode mode = 2;\n }\n\n // Prints all defined keys in the given namespace.\n message List {}\n\n // The part of the system where this command will take place.\n // NOTE: We avoid the identifier `namespace` because it's a keyword in C++.\n Namespace name_space = 1;\n\n // The subcommand to issue to `adb settings`.\n // NOTE: We avoid the identifiers `delete` and `del` because they're\n // keywords in C++ and Python respectively.\n oneof verb {\n Get get = 2;\n Put put = 3;\n Delete delete_key = 4;\n Reset reset = 5;\n List list = 6;\n }\n }\n\n // Generic ADB command. Use this for commands that are not\n // explicitly implemented.\n // Calls `adb [args...]`.\n message GenericRequest {\n repeated string args = 1;\n }\n\n message PackageManagerRequest {\n message List {\n // Lists all features of the system.\n message Features {}\n\n // Lists all system libraries.\n message Libraries {}\n\n // Lists all packages; optionally only those whose name contains the text\n // in `filter`.\n message Packages {\n string filter = 1;\n\n // Extra options that control the output. Please see `pm help` for\n // details.\n repeated string options = 2;\n }\n\n oneof what {\n Features features = 1;\n Libraries libraries = 2;\n Packages packages = 3;\n }\n }\n\n // Deletes all data associated with a package.\n message Clear {\n // The package name to clear its cache.\n string package_name = 1;\n\n // Optional USER_ID.\n string user_id = 2;\n }\n\n message Grant {\n string package_name = 1;\n\n // Possible values listed at\n // https://developer.android.com/reference/android/Manifest.permission\n // To query an app's required permissions, use the following adb command:\n // > adb shell dumpsys package \n // The output will contain things like\n // android.permission.WRITE_SECURE_SETTINGS\n repeated string permissions = 2;\n }\n\n // The subcommand to issue to `pm`.\n oneof verb {\n List list = 1;\n Clear clear = 2;\n Grant grant = 3;\n }\n }\n\n // For executing `dumpsys` commands.\n message DumpsysRequest {\n enum PriorityLevel {\n UNSET = 0;\n NORMAL = 1;\n HIGH = 2;\n CRITICAL = 3;\n }\n\n // The service to dump. If empty, all services will be dumped.\n string service = 1;\n\n // Optional arguments to pass to the specific service dump.\n repeated string args = 2;\n\n // Lists services, does not dump them.\n // This effectively disables dumping information about any particular\n // service.\n bool list_only = 3;\n\n // Timeouts natively supported by `dumpsys`.\n int32 timeout_sec = 4;\n int32 timeout_ms = 5;\n\n // Whether to dump the process ID instead of the usual dump.\n bool pid = 6;\n\n // Whether dumps will be in proto format. Only works for services that\n // support dumping data in proto format.\n bool proto = 7;\n\n // Filters services based on specified priority.\n PriorityLevel priority = 8;\n\n // Excludes services from the dump.\n repeated string skip_services = 9;\n }\n\n oneof command {\n InstallApk install_apk = 1;\n StartActivity start_activity = 2;\n ForceStop force_stop = 3;\n Tap tap = 6;\n PressButton press_button = 7;\n StartScreenPinning start_screen_pinning = 10;\n UninstallPackage uninstall_package = 16;\n GetCurrentActivity get_current_activity = 17;\n GetOrientationRequest get_orientation = 24;\n Push push = 18;\n Pull pull = 19;\n InputText input_text = 20;\n SettingsRequest settings = 21;\n GenericRequest generic = 22;\n PackageManagerRequest package_manager = 23;\n DumpsysRequest dumpsys = 26;\n SendBroadcast send_broadcast = 25;\n }\n\n // Optional (soft) deadline in seconds for completing this command.\n // Expected to be >0. If ==0 (the default), it's ignored.\n // Notice that not all commands accept timeouts, but because it's such a\n // common parameter, we include it here instead of in each separate command.\n float timeout_sec = 100;\n}\n\nmessage AdbResponse {\n enum Status {\n // Reserved value for unset statuses.\n UNDEFINED = 0;\n // Returned when everything goes well.\n OK = 1;\n // Returned when handling unknown AdbRequest commands.\n UNKNOWN_COMMAND = 2;\n // Returned when an argument does not respect a precondition.\n FAILED_PRECONDITION = 3;\n // Returned when something internal did not work as expected.\n INTERNAL_ERROR = 4;\n // Returned when the adb command failed.\n ADB_ERROR = 5;\n // Returned when the adb command timed out.\n TIMEOUT = 6;\n }\n Status status = 1;\n\n // `error_message` is only populated in case of errors.\n string error_message = 2;\n\n // General stats that components may optionally report.\n map stats = 3;\n\n // Response for GetCurrentActivity requests.\n message GetCurrentActivityResponse {\n // The format of the output is `package/package.ActivityName', for example:\n // \"com.example.vokram/com.example.vokram.MainActivity\"\n string full_activity = 1;\n }\n\n // Response for GetOrientationRequests.\n message GetOrientationResponse {\n // Possible values are {0, 1, 2, 3} corresponding to {0, 90, 180, 270}\n // degrees respectively.\n // Please see https://developer.android.com/reference/android/view/Surface.\n int32 orientation = 1;\n }\n\n // Response for StartActivity requests.\n message StartActivityResponse {\n // The activity that was actually started. On a failed request, this will be\n // empty.\n string full_activity = 1;\n bytes output = 2;\n }\n\n // Response for PressButton requests.\n message PressButtonResponse {\n // The output, if any, by `adb` after sending a key press.\n // This is intentionally left as `bytes` instead of `string` so that content\n // other than `UTF-8` can be transmitted.\n bytes output = 1;\n }\n\n // Response for Push requests.\n message PushResponse {}\n\n // Response for Pull requests.\n message PullResponse {\n // The contents of the file.\n // This is intentionally left as `bytes` instead of `string` so that content\n // other than `UTF-8` can be transmitted.\n bytes content = 1;\n }\n\n // Response for InputText requests.\n message InputTextResponse {}\n\n // Response for SettingsRequests.\n message SettingsResponse {\n // The output, if any, of the `adb shell settings` command.\n bytes output = 1;\n }\n\n // Response for GenericRequests.\n message GenericResponse {\n // The output, if any, of the generic adb command.\n bytes output = 1;\n }\n\n // Response for PackageManagerRequests.\n message PackageManagerResponse {\n // The output, if any, of the `adb shell pm` command.\n bytes output = 1;\n\n message List {\n // A list of items. The actual content depends on the request, but it\n // could be things like features, libraries or package names.\n repeated string items = 1;\n }\n\n oneof verb {\n List list = 2;\n }\n }\n\n // Response for DumpsysRequests.\n message DumpsysResponse {\n // The output, if any, of the `dumpsys` command.\n bytes output = 1;\n }\n\n oneof payload {\n GetCurrentActivityResponse get_current_activity = 10;\n StartActivityResponse start_activity = 11;\n PressButtonResponse press_button = 12;\n PushResponse push = 13;\n PullResponse pull = 14;\n InputTextResponse input_text = 15;\n SettingsResponse settings = 16;\n GenericResponse generic = 17;\n PackageManagerResponse package_manager = 18;\n GetOrientationResponse get_orientation = 19;\n DumpsysResponse dumpsys = 21;\n }\n}\n"
},
{
"path": "android_env/proto/emulator_controller.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n// Copyright (C) 2018 The Android Open Source Project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n// Note that if you add/remove methods in this file you must update\n// the metrics sql as well ./android/scripts/gen-grpc-sql.py\n//\n// Please group deleted methods in a block including the date (MM/DD/YY)\n// it was removed. This enables us to easily keep metrics around after removal\n//\n// list of deleted methods\n// rpc iWasDeleted (03/12/12)\n// ...\n\n// LINT: LEGACY_NAMES\n\nsyntax = \"proto3\";\n\npackage android.emulation.control;\n\nimport \"google/protobuf/empty.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.android.emulator.control\";\noption objc_class_prefix = \"AEC\";\n\n// An EmulatorController service lets you control the emulator.\n// Note that this is currently an experimental feature, and that the\n// service definition might change without notice. Use at your own risk!\n//\n// We use the following rough conventions:\n//\n// streamXXX --> streams values XXX (usually for emulator lifetime). Values\n// are updated as soon as they become available.\n// getXXX --> gets a single value XXX\n// setXXX --> sets a single value XXX, does not returning state, these\n// usually have an observable lasting side effect.\n// sendXXX --> send a single event XXX, possibly returning state information.\n// android usually responds to these events.\nservice EmulatorController {\n // Set the sensor data\n rpc streamSensor(SensorValue) returns (stream SensorValue) {}\n // Get the sensor data\n rpc getSensor(SensorValue) returns (SensorValue) {}\n // Stream the sensor data\n rpc setSensor(SensorValue) returns (google.protobuf.Empty) {}\n\n // Set the physical model, this is likely the one you are\n // looking for when you wish to modify the device state.\n rpc setPhysicalModel(PhysicalModelValue) returns (google.protobuf.Empty) {}\n // Get the physical model\n rpc getPhysicalModel(PhysicalModelValue) returns (PhysicalModelValue) {}\n // Stream the physical model\n rpc streamPhysicalModel(PhysicalModelValue)\n returns (stream PhysicalModelValue) {}\n\n // Atomically set the current primary clipboard data.\n rpc setClipboard(ClipData) returns (google.protobuf.Empty) {}\n // Atomically get the current primary clipboard data.\n rpc getClipboard(google.protobuf.Empty) returns (ClipData) {}\n\n // Streams the current data on the clipboard. This will immediately produce\n // a result with the current state of the clipboard after which the stream\n // will block and wait until a new clip event is available from the guest.\n // Calling the setClipboard method above will not result in generating a clip\n // event. It is possible to lose clipboard events if the clipboard updates\n // very rapidly.\n rpc streamClipboard(google.protobuf.Empty) returns (stream ClipData) {}\n\n // Set the battery to the given state.\n rpc setBattery(BatteryState) returns (google.protobuf.Empty) {}\n // Get the battery to the given state.\n rpc getBattery(google.protobuf.Empty) returns (BatteryState) {}\n\n // Set the state of the gps, gps support will only work\n // properly if:\n //\n // - no location ui is active. That is the emulator\n // is launched in headless mode (-no-window) or the location\n // ui is disabled (-no-location-ui).\n // - the passiveUpdate is set to false. Setting this to false\n // will disable/break the LocationUI.\n //\n // Keep in mind that android usually only samples the gps at 1 hz.\n rpc setGps(GpsState) returns (google.protobuf.Empty) {}\n\n // Gets the latest gps state as delivered by the setGps call, or location ui\n // if active.\n //\n // Note: this is not necessarily the actual gps coordinate visible at the\n // time, due to gps sample frequency (usually 1hz).\n rpc getGps(google.protobuf.Empty) returns (GpsState) {}\n\n // Simulate a touch event on the finger print sensor.\n rpc sendFingerprint(Fingerprint) returns (google.protobuf.Empty) {}\n\n // Send a keyboard event. Translating the event.\n rpc sendKey(KeyboardEvent) returns (google.protobuf.Empty) {}\n\n // Send touch events. Note that mouse events can be simulated by touch events.\n rpc sendTouch(TouchEvent) returns (google.protobuf.Empty) {}\n // Send mouse events.\n rpc sendMouse(MouseEvent) returns (google.protobuf.Empty) {}\n\n // Make a phone call.\n rpc sendPhone(PhoneCall) returns (PhoneResponse) {}\n\n // Sends an sms message to the emulator.\n rpc sendSms(SmsMessage) returns (PhoneResponse) {}\n\n // Retrieve the status of the emulator. This will contain general\n // hardware information, and whether the device has booted or not.\n rpc getStatus(google.protobuf.Empty) returns (EmulatorStatus) {}\n\n // Gets an individual screenshot in the desired format.\n //\n // The image will be scaled to the desired ImageFormat, while maintaining\n // the aspect ratio. The returned image will never exceed the provided width\n // and height. Not setting the width or height (i.e. they are 0) will result\n // in using the device width and height.\n //\n // The resulting image will be properly oriented and can be displayed\n // directly without post processing. For example, if the device has a\n // 1080x1920 screen and is in landscape mode and called with no width or\n // height parameter, it will return an 1920x1080 image.\n //\n // This method will return an empty image if the display is not visible.\n rpc getScreenshot(ImageFormat) returns (Image) {}\n\n // Streams a series of screenshots in the desired format.\n // A new frame will be delivered whenever the device produces a new frame.\n // (Beware that this can produce a significant amount of data, and that\n // certain translations are (png transform) can be costly).\n //\n // If the requested display is not visible it will send a single empty image\n // and wait start producing images once the display becomes active, again\n // producing a single empty image when the display becomes inactive.\n rpc streamScreenshot(ImageFormat) returns (stream Image) {}\n\n // Streams a series of audio packets in the desired format.\n // A new frame will be delivered whenever the emulated device\n // produces a new audio frame. You can expect packets to be\n // delivered in intervals of 20-30ms.\n //\n // Be aware that this can block when the emulator does not\n // produce any audio whatsoever!\n rpc streamAudio(AudioFormat) returns (stream AudioPacket) {}\n\n // Injects a series of audio packets to the android microphone.\n // A new frame will be delivered whenever the emulated device\n // requests a new audio frame. Audio is usually delivered at a rate\n // that the emulator is requesting frames. Audio will be stored in a\n // temporary buffer that can hold 500ms of audio.\n //\n // Note: Currently the emulator will downsample to 16khz.\n //\n // - INVALID_ARGUMENT (code 3) The sampling rate was too high\n // - INVALID_ARGUMENT (code 3) The audio packet was too large to handle.\n // - FAILED_PRECONDITION (code 9) If there was a microphone registered\n // already.\n rpc injectAudio(stream AudioPacket) returns (google.protobuf.Empty) {}\n\n // Returns the last 128Kb of logcat output from the emulator\n // Note that parsed logcat messages are only available after L (Api >23).\n // it is possible that the logcat buffer gets overwritten, or falls behind.\n rpc getLogcat(LogMessage) returns (LogMessage) {}\n\n // Streams the logcat output from the emulator. The first call\n // can retrieve up to 128Kb. This call will not return.\n // Note that parsed logcat messages are only available after L (Api >23)\n // it is possible that the logcat buffer gets overwritten, or falls behind.\n rpc streamLogcat(LogMessage) returns (stream LogMessage) {}\n\n // Transition the virtual machine to the desired state. Note that\n // some states are only observable. For example you cannot transition\n // to the error state.\n rpc setVmState(VmRunState) returns (google.protobuf.Empty) {}\n\n // Gets the state of the virtual machine.\n rpc getVmState(google.protobuf.Empty) returns (VmRunState) {}\n\n // Atomically changes the current multi-display configuration.\n // After this call the given display configurations will be activated. You\n // can only update secondary displays. Displays with id 0 will be ignored.\n //\n // This call can result in the removal or addition of secondary displays, the\n // final display state can be observed by the returned configuration.\n //\n // The following gRPC error codes can be returned:\n // - FAILED_PRECONDITION (code 9) if the AVD does not support a configurable\n // secondary display.\n // - INVALID_ARGUMENT (code 3) if:\n // - The same display id is defined multiple times.\n // - The display configurations are outside valid ranges\n // (see DisplayConfiguration)\n // - INTERNAL (code 13) if there was an internal emulator failure.\n rpc setDisplayConfigurations(DisplayConfigurations)\n returns (DisplayConfigurations) {}\n\n // Returns all currently valid logical displays.\n // The gRPC error code FAILED_PRECONDITION (code 9) is returned if the AVD\n // does not support a configurable secondary display.\n rpc getDisplayConfigurations(google.protobuf.Empty)\n returns (DisplayConfigurations) {}\n\n // Notifies client of the following changes:\n //\n // - Virtual scene camera status change.\n // - Display configuration changes from extended ui. This will only be fired\n // if the user makes modifications the extended displays through the\n // extended control tab.\n //\n // Note that this method will send the initial virtual scene state\n // immediately.\n rpc streamNotification(google.protobuf.Empty) returns (stream Notification) {}\n\n // RotationRadian is relative to the camera's current orientation.\n rpc rotateVirtualSceneCamera(RotationRadian) returns (google.protobuf.Empty) {\n }\n // Velocity is absolute\n rpc setVirtualSceneCameraVelocity(Velocity) returns (google.protobuf.Empty) {}\n // Set foldable posture\n rpc setPosture(Posture) returns (google.protobuf.Empty) {}\n}\n\n// A Run State that describes the state of the Virtual Machine.\nmessage VmRunState {\n enum RunState {\n // The emulator is in an unknown state. You cannot transition to this state.\n UNKNOWN = 0;\n // Guest is actively running. You can transition to this state from the\n // paused state.\n RUNNING = 1;\n // Guest is paused to load a snapshot. You cannot transition to this state.\n RESTORE_VM = 2;\n // Guest has been paused. Transitioning to this state will pause the\n // emulator the guest will not be consuming any cpu cycles.\n PAUSED = 3;\n // Guest is paused to take or export a snapshot. You cannot\n // transition to this state.\n SAVE_VM = 4;\n // System shutdown, note that it is similar to power off. It tries to set\n // the system status and notify guest. The system is likely going to\n // disappear soon and do proper cleanup of resources, possibly taking\n // a snapshot. This is the same behavior as closing the emulator by clicking\n // the X (close) in the user interface.\n SHUTDOWN = 5;\n // Immediately terminate the emulator. No resource cleanup will take place.\n // There is a good change to corrupt the system.\n TERMINATE = 7;\n // Will cause the emulator to reset. This is not a state you can observe.\n RESET = 9;\n // Guest experienced some error state, you cannot transition to this state.\n INTERNAL_ERROR = 10;\n }\n\n RunState state = 1;\n}\n\nmessage ParameterValue {\n repeated float data = 1 [packed = true];\n}\n\nmessage PhysicalModelValue {\n enum State {\n OK = 0;\n NO_SERVICE = -3; // qemud service is not available/initiated.\n DISABLED = -2; // Sensor is disabled.\n UNKNOWN = -1; // Unknown sensor (should not happen)\n }\n\n // Details on the sensors documentation can be found here:\n // https://developer.android.com/reference/android/hardware/Sensor.html#TYPE_\n // The types must follow the order defined in\n // \"external/qemu/android/hw-sensors.h\"\n enum PhysicalType {\n POSITION = 0;\n\n // All values are angles in degrees.\n // values = [x,y,z]\n ROTATION = 1;\n\n MAGNETIC_FIELD = 2;\n\n // Temperature in °C\n TEMPERATURE = 3;\n\n // Proximity sensor distance measured in centimeters\n PROXIMITY = 4;\n\n // Ambient light level in SI lux units\n LIGHT = 5;\n\n // Atmospheric pressure in hPa (millibar)\n PRESSURE = 6;\n\n // Relative ambient air humidity in percent\n HUMIDITY = 7;\n\n VELOCITY = 8;\n AMBIENT_MOTION = 9;\n\n // Describing a hinge angle sensor in degrees.\n HINGE_ANGLE0 = 10;\n HINGE_ANGLE1 = 11;\n HINGE_ANGLE2 = 12;\n\n ROLLABLE0 = 13;\n ROLLABLE1 = 14;\n ROLLABLE2 = 15;\n }\n PhysicalType target = 1;\n\n // [Output Only]\n State status = 2;\n\n // Value interpretation depends on sensor, will contain at most 3 values.\n ParameterValue value = 3;\n}\n\n// A single sensor value.\nmessage SensorValue {\n enum State {\n OK = 0;\n NO_SERVICE = -3; // qemud service is not available/initiated.\n DISABLED = -2; // Sensor is disabled.\n UNKNOWN = -1; // Unknown sensor (should not happen)\n }\n\n // These are the various sensors that can be available in an emulated\n // devices.\n enum SensorType {\n // Measures the acceleration force in m/s2 that is applied to a device\n // on all three physical axes (x, y, and z), including the force of\n // gravity.\n ACCELERATION = 0;\n // Measures a device's rate of rotation in rad/s around each of the\n // three physical axes (x, y, and z).\n GYROSCOPE = 1;\n // Measures the ambient geomagnetic field for all three physical axes\n // (x, y, z) in μT.\n MAGNETIC_FIELD = 2;\n // Measures degrees of rotation that a device makes around all three\n // physical axes (x, y, z)\n ORIENTATION = 3;\n // Measures the temperature of the device in degrees Celsius (°C).\n TEMPERATURE = 4;\n // Measures the proximity of an object in cm relative to the view screen\n // of a device. This sensor is typically used to determine whether a\n // handset is being held up to a person's ear.\n PROXIMITY = 5;\n // Measures the ambient light level (illumination) in lx.\n LIGHT = 6;\n // Measures the ambient air pressure in hPa or mbar.\n PRESSURE = 7;\n // Measures the relative ambient humidity in percent (%).\n HUMIDITY = 8;\n MAGNETIC_FIELD_UNCALIBRATED = 9;\n GYROSCOPE_UNCALIBRATED = 10;\n }\n\n // Type of sensor\n SensorType target = 1;\n\n // [Output Only]\n State status = 2;\n\n // Value interpretation depends on sensor enum, will contain at most 3\n // values.\n ParameterValue value = 3;\n}\n\nmessage LogMessage {\n // [Output Only] The contents of the log output.\n string contents = 1;\n // The starting byte position of the output that was returned. This\n // should match the start parameter sent with the request. If the serial\n // console output exceeds the size of the buffer, older output will be\n // overwritten by newer content and the start values will be mismatched.\n int64 start = 2;\n //[Output Only] The position of the next byte of content from the serial\n // console output. Use this value in the next request as the start\n // parameter.\n int64 next = 3;\n\n // Set the sort of response you are interested it in.\n // It the type is \"Parsed\" the entries field will contain the parsed\n // results. otherwise the contents field will be set.\n LogType sort = 4;\n\n // [Output Only] The parsed logcat entries so far. Only set if sort is\n // set to Parsed\n repeated LogcatEntry entries = 5;\n\n enum LogType {\n Text = 0;\n Parsed = 1;\n }\n}\n\n// A parsed logcat entry.\nmessage LogcatEntry {\n // The possible log levels.\n enum LogLevel {\n UNKNOWN = 0;\n DEFAULT = 1;\n VERBOSE = 2;\n DEBUG = 3;\n INFO = 4;\n WARN = 5;\n ERR = 6;\n FATAL = 7;\n SILENT = 8;\n }\n\n // A Unix timestamps in milliseconds (The number of milliseconds that\n // have elapsed since January 1, 1970 (midnight UTC/GMT), not counting\n // leap seconds)\n uint64 timestamp = 1;\n\n // Process id.\n uint32 pid = 2;\n\n // Thread id.\n uint32 tid = 3;\n LogLevel level = 4;\n string tag = 5;\n string msg = 6;\n}\n\n// Information about the hypervisor that is currently in use.\nmessage VmConfiguration {\n enum VmHypervisorType {\n // An unknown hypervisor\n UNKNOWN = 0;\n\n // No hypervisor is in use. This usually means that the guest is\n // running on a different CPU than the host, or you are using a\n // platform where no hypervisor is available.\n NONE = 1;\n\n // The Kernel based Virtual Machine\n // (https://www.linux-kvm.org/page/Main_Page)\n KVM = 2;\n\n // Intel® Hardware Accelerated Execution Manager (Intel® HAXM)\n // https://github.com/intel/haxm\n HAXM = 3;\n\n // Hypervisor Framework.\n // https://developer.apple.com/documentation/hypervisor\n HVF = 4;\n\n // Window Hypervisor Platform\n // https://docs.microsoft.com/en-us/virtualization/api/\n WHPX = 5;\n\n GVM = 6;\n }\n\n VmHypervisorType hypervisorType = 1;\n int32 numberOfCpuCores = 2;\n int64 ramSizeBytes = 3;\n}\n\n// Representation of a clipped data object on the clipboard.\nmessage ClipData {\n // UTF-8 Encoded text.\n string text = 1;\n}\n\n// The Touch interface represents a single contact point on a\n// touch-sensitive device. The contact point is commonly a finger or stylus\n// and the device may be a touchscreen or trackpad.\nmessage Touch {\n // The horizontal coordinate. This is the physical location on the\n // screen For example 0 indicates the leftmost coordinate.\n int32 x = 1;\n\n // The vertical coordinate. This is the physical location on the screen\n // For example 0 indicates the top left coordinate.\n int32 y = 2;\n\n // The identifier is an arbitrary non-negative integer that is used to\n // identify and track each tool independently when multiple tools are\n // active. For example, when multiple fingers are touching the device,\n // each finger should be assigned a distinct tracking id that is used as\n // long as the finger remains in contact. Tracking ids may be reused\n // when their associated tools move out of range.\n //\n // The emulator currently supports up to 10 concurrent touch events. The\n // identifier can be any uninque value and will be mapped to the next\n // available internal identifier.\n int32 identifier = 3;\n\n // Reports the physical pressure applied to the tip of the tool or the\n // signal strength of the touch contact.\n //\n // The values reported must be non-zero when the tool is touching the\n // device and zero otherwise to indicate that the touch event is\n // completed.\n //\n // Make sure to deliver a pressure of 0 for the given identifier when\n // the touch event is completed, otherwise the touch identifier will not\n // be unregistered!\n int32 pressure = 4;\n\n // Optionally reports the cross-sectional area of the touch contact, or\n // the length of the longer dimension of the touch contact.\n int32 touch_major = 5;\n\n // Optionally reports the length of the shorter dimension of the touch\n // contact. This axis will be ignored if touch_major is reporting an\n // area measurement greater than 0.\n int32 touch_minor = 6;\n\n enum EventExpiration {\n // The system will use the default time of 120s to track\n // the touch event with the given identifier. If no update happens\n // within this timeframe the identifier is considered expired\n // and can be made available for re-use. This means that a touch event\n // with pressure 0 for this identifier will be send to the emulator.\n EVENT_EXPIRATION_UNSPECIFIED = 0;\n\n // Never expire the given slot. You must *ALWAYS* close the identifier\n // by sending a touch event with 0 pressure.\n NEVER_EXPIRE = 1;\n }\n\n EventExpiration expiration = 7;\n}\n\n// A TouchEvent contains a list of Touch objects that are in contact with\n// the touch surface.\n//\n// Touch events are delivered in sequence as specified in the touchList.\n//\n// TouchEvents are delivered to the emulated devices using [\"Protocol\n// B\"](https://www.kernel.org/doc/Documentation/input/multi-touch-protocol.txt)\nmessage TouchEvent {\n // The list of Touch objects, note that these do not need to be unique\n repeated Touch touches = 1;\n\n // The display device where the touch event occurred.\n // Omitting or using the value 0 indicates the main display.\n //\n // Touch events cannot be send to displays other than 0, due to\n // https://issuetracker.google.com/issues/150699691\n int32 display = 2;\n}\n\n// The MouseEvent interface represents events that occur due to the user\n// interacting with a pointing device (such as a mouse).\nmessage MouseEvent {\n // The horizontal coordinate. This is the physical location on the\n // screen For example 0 indicates the leftmost coordinate.\n int32 x = 1;\n\n // The vertical coordinate. This is the physical location on the screen\n // For example 0 indicates the top left coordinate.\n int32 y = 2;\n\n // Indicates which buttons are pressed.\n // 0: No button was pressed\n // 1: Primary button (left)\n // 2: Secondary button (right)\n int32 buttons = 3;\n\n // The display device where the mouse event occurred.\n // Omitting or using the value 0 indicates the main display.\n int32 display = 4;\n}\n\n// KeyboardEvent objects describe a user interaction with the keyboard; each\n// event describes a single interaction between the user and a key (or\n// combination of a key with modifier keys) on the keyboard.\n// This follows the pattern as set by\n// (javascript)[https://developer.mozilla.org/en-US/docs/Web/API/KeyboardEvent]\n//\n// Note: that only keyCode, key, or text can be set and that the semantics\n// will slightly vary.\nmessage KeyboardEvent {\n // Code types that the emulator can receive. Note that the emulator\n // will do its best to translate the code to an evdev value that\n // will be send to the emulator. This translation is based on\n // the chromium translation tables. See\n // (this)[https://android.googlesource.com/platform/external/qemu/+/refs/heads/emu-master-dev/android/android-grpc/android/emulation/control/keyboard/keycode_converter_data.inc]\n // for details on the translation.\n enum KeyCodeType {\n Usb = 0;\n Evdev = 1;\n XKB = 2;\n Win = 3;\n Mac = 4;\n }\n\n enum KeyEventType {\n // Indicates that this keyevent should be send to the emulator\n // as a key down event. Meaning that the key event will be\n // translated to an EvDev event type and bit 11 (0x400) will be\n // set before it is sent to the emulator.\n keydown = 0;\n\n // Indicates that the keyevent should be send to the emulator\n // as a key up event. Meaning that the key event will be\n // translated to an EvDev event type and\n // sent to the emulator.\n keyup = 1;\n\n // Indicates that the keyevent will be send to the emulator\n // as e key down event and immediately followed by a keyup event.\n keypress = 2;\n }\n\n // Type of keycode contained in the keyCode field.\n KeyCodeType codeType = 1;\n\n // The type of keyboard event that should be sent to the emulator\n KeyEventType eventType = 2;\n\n // This property represents a physical key on the keyboard (as opposed\n // to the character generated by pressing the key). In other words, this\n // property is a value which isn't altered by keyboard layout or the\n // state of the modifier keys. This value will be interpreted by the\n // emulator depending on the KeyCodeType. The incoming key code will be\n // translated to an evdev code type and send to the emulator.\n // The values in key and text will be ignored.\n int32 keyCode = 3;\n\n // The value of the key pressed by the user, taking into consideration\n // the state of modifier keys such as Shift as well as the keyboard\n // locale and layout. This follows the w3c standard used in browsers.\n // You can find an accurate description of valid values\n // [here](https://developer.mozilla.org/en-US/docs/Web/API/KeyboardEvent/key/Key_Values)\n //\n // Note that some keys can result in multiple evdev events that are\n // delivered to the emulator. for example the Key \"A\" will result in a\n // sequence:\n // [\"Shift\", \"a\"] -> [0x2a, 0x1e] whereas \"a\" results in [\"a\"] -> [0x1e].\n //\n // Not all documented keys are understood by android, and only printable\n // ASCII [32-127) characters are properly translated.\n //\n // Keep in mind that there are a set of key values that result in android\n // specific behavior\n // [see](https://developer.mozilla.org/en-US/docs/Web/API/KeyboardEvent/key/Key_Values#Phone_keys):\n //\n // - \"AppSwitch\": Behaves as the \"Overview\" button in android.\n // - \"GoBack\": The Back button.\n // - \"GoHome\": The Home button, which takes the user to the phone's main\n // screen (usually an application launcher).\n // - \"Power\": The Power button.\n string key = 4;\n\n // Series of utf8 encoded characters to send to the emulator. An attempt\n // will be made to translate every character will an EvDev event type and\n // send to the emulator as a keypress event. The values in keyCode,\n // eventType, codeType and key will be ignored.\n //\n // Note that most printable ASCII characters (range [32-127) can be send\n // individually with the \"key\" param. Do not expect arbitrary UTF symbols to\n // arrive in the emulator (most will be ignored).\n //\n // Note that it is possible to overrun the keyboard buffer by slamming this\n // endpoint with large quantities of text (>1kb). The clipboard api is better\n // suited for transferring large quantities of text.\n string text = 5;\n}\n\nmessage Fingerprint {\n // True when the fingprint is touched.\n bool isTouching = 1;\n\n // The identifier of the registered fingerprint.\n int32 touchId = 2;\n}\n\nmessage GpsState {\n // Setting this to false will disable auto updating from the LocationUI,\n // otherwise the location UI will override the location at a frequency of 1hz.\n //\n // - This is unused if the emulator is launched with -no-window, or when he\n // location ui is disabled.\n // - This will BREAK the location ui experience if it is set to false. For\n // example routing will no longer function.\n bool passiveUpdate = 1;\n\n // The latitude, in degrees.\n double latitude = 2;\n\n // The longitude, in degrees.\n double longitude = 3;\n\n // The speed if it is available, in meters/second over ground\n double speed = 4;\n\n // gets the horizontal direction of travel of this device, and is not\n // related to the device orientation. It is guaranteed to be in the\n // range [0.0, 360.0] if the device has a bearing. 0=North, 90=East,\n // 180=South, etc..\n double bearing = 5;\n\n // The altitude if available, in meters above the WGS 84 reference\n // ellipsoid.\n double altitude = 6;\n\n // The number of satellites used to derive the fix\n int32 satellites = 7;\n}\n\nmessage BatteryState {\n enum BatteryStatus {\n UNKNOWN = 0;\n CHARGING = 1;\n DISCHARGING = 2;\n NOT_CHARGING = 3;\n FULL = 4;\n }\n\n enum BatteryCharger {\n NONE = 0;\n AC = 1;\n USB = 2;\n WIRELESS = 3;\n }\n\n enum BatteryHealth {\n GOOD = 0;\n FAILED = 1;\n DEAD = 2;\n OVERVOLTAGE = 3;\n OVERHEATED = 4;\n }\n\n bool hasBattery = 1;\n bool isPresent = 2;\n BatteryCharger charger = 3;\n int32 chargeLevel = 4;\n BatteryHealth health = 5;\n BatteryStatus status = 6;\n}\n\n// An ImageTransport allows for specifying a side channel for\n// delivering image frames versus using the standard bytes array that is\n// returned with the gRPC request.\nmessage ImageTransport {\n enum TransportChannel {\n // Return full frames over the gRPC transport\n TRANSPORT_CHANNEL_UNSPECIFIED = 0;\n\n // Write images to the a file/shared memory handle.\n MMAP = 1;\n }\n\n // The desired transport channel used for delivering image frames. Only\n // relevant when streaming screenshots.\n TransportChannel channel = 1;\n\n // Handle used for writing image frames if transport is mmap. The client sets\n // and owns this handle. It can be either a shm region, or a mmap. A mmap\n // should be a url that starts with `file:///`\n // Note: the mmap can result in tearing.\n string handle = 2;\n}\n\n// The aspect ratio (width/height) will be different from the one\n// where the device is unfolded.\nmessage FoldedDisplay {\n uint32 width = 1;\n uint32 height = 2;\n // It is possible for the screen to be folded in different ways depending\n // on which surface is shown to the user. So xOffset and yOffset indicate\n // the top left corner of the folded screen within the original unfolded\n // screen.\n uint32 xOffset = 3;\n uint32 yOffset = 4;\n}\n\nmessage ImageFormat {\n enum ImgFormat {\n // Portable Network Graphics format\n // (https://en.wikipedia.org/wiki/Portable_Network_Graphics)\n PNG = 0;\n\n // Three-channel RGB color model supplemented with a fourth alpha\n // channel. https://en.wikipedia.org/wiki/RGBA_color_model\n // Each pixel consists of 4 bytes.\n RGBA8888 = 1;\n\n // Three-channel RGB color model, each pixel consists of 3 bytes\n RGB888 = 2;\n }\n\n // The (desired) format of the resulting bytes.\n ImgFormat format = 1;\n\n // [Output Only] The rotation of the image. The image will be rotated\n // based upon the coarse grained orientation of the device.\n Rotation rotation = 2;\n\n // The (desired) width of the image. When passed as input\n // the image will be scaled to match the given\n // width, while maintaining the aspect ratio of the device.\n // The returned image will never exceed the given width, but can be less.\n // Omitting this value (or passing in 0) will result in no scaling,\n // and the width of the actual device will be used.\n uint32 width = 3;\n\n // The (desired) height of the image. When passed as input\n // the image will be scaled to match the given\n // height, while maintaining the aspect ratio of the device.\n // The returned image will never exceed the given height, but can be less.\n // Omitting this value (or passing in 0) will result in no scaling,\n // and the height of the actual device will be used.\n uint32 height = 4;\n\n // The (desired) display id of the device. Setting this to 0 (or omitting)\n // indicates the main display.\n uint32 display = 5;\n\n // Set this if you wish to use a different transport channel to deliver image\n // frames.\n ImageTransport transport = 6;\n\n // [Output Only] Display configuration when screen is folded. The value is the\n // original configuration before scaling.\n FoldedDisplay foldedDisplay = 7;\n}\n\nmessage Image {\n ImageFormat format = 1;\n\n uint32 width = 2 [deprecated = true]; // width is contained in format.\n uint32 height = 3 [deprecated = true]; // height is contained in format.\n\n // The organization of the pixels in the image buffer is from left to\n // right and bottom up. This will be empty if an alternative image transport\n // is requested in the image format. In that case the side channel should\n // be used to obtain the image data.\n bytes image = 4;\n\n // [Output Only] Monotonically increasing sequence number in a stream of\n // screenshots. The first screenshot will have a sequence of 0. A single\n // screenshot will always have a sequence number of 0. The sequence is not\n // necessarily contiguous, and can be used to detect how many frames were\n // dropped. An example sequence could be: [0, 3, 5, 7, 9, 11].\n uint32 seq = 5;\n\n // [Output Only] Unix timestamp in microseconds when the emulator estimates\n // the frame was generated. The timestamp is before the actual frame is\n // copied and transformed. This can be used to calculate variance between\n // frame production time, and frame depiction time.\n uint64 timestampUs = 6;\n}\n\nmessage Rotation {\n enum SkinRotation {\n PORTRAIT = 0; // 0 degrees\n LANDSCAPE = 1; // 90 degrees\n REVERSE_PORTRAIT = 2; // -180 degrees\n REVERSE_LANDSCAPE = 3; // -90 degrees\n }\n\n // The rotation of the device, derived from the sensor state\n // of the emulator. The derivation reflects how android observes\n // the rotation state.\n SkinRotation rotation = 1;\n\n // Specifies the angle of rotation, in degrees [-180, 180]\n double xAxis = 2;\n double yAxis = 3;\n double zAxis = 4;\n}\n\nmessage PhoneCall {\n enum Operation {\n InitCall = 0;\n AcceptCall = 1;\n RejectCallExplicit = 2;\n RejectCallBusy = 3;\n DisconnectCall = 4;\n PlaceCallOnHold = 5;\n TakeCallOffHold = 6;\n }\n Operation operation = 1;\n string number = 2;\n}\n\nmessage PhoneResponse {\n enum Response {\n OK = 0;\n BadOperation = 1; // Enum out of range\n BadNumber = 2; // Mal-formed telephone number\n InvalidAction = 3; // E.g., disconnect when no call is in progress\n ActionFailed = 4; // Internal error\n RadioOff = 5; // Radio power off\n }\n Response response = 1;\n}\n\nmessage Entry {\n string key = 1;\n string value = 2;\n}\n\nmessage EntryList {\n repeated Entry entry = 1;\n}\n\nmessage EmulatorStatus {\n // The emulator version string.\n string version = 1;\n\n // The time the emulator has been active in .ms\n uint64 uptime = 2;\n\n // True if the device has completed booting.\n // For P and later this information will accurate,\n // for older images we rely on adb.\n bool booted = 3;\n\n // The current vm configuration\n VmConfiguration vmConfig = 4;\n\n // The hardware configuration of the running emulator as\n // key valure pairs.\n EntryList hardwareConfig = 5;\n}\n\nmessage AudioFormat {\n enum SampleFormat {\n AUD_FMT_U8 = 0; // Unsigned 8 bit\n AUD_FMT_S16 = 1; // Signed 16 bit (little endian)\n }\n\n enum Channels {\n Mono = 0;\n Stereo = 1;\n }\n\n // Sampling rate to use, defaulting to 44100 if this is not set.\n // Note, that android devices typically will not use a sampling\n // rate higher than 48kHz. See https://developer.android.com/ndk/guides/audio.\n uint64 samplingRate = 1;\n Channels channels = 2;\n SampleFormat format = 3;\n}\n\nmessage AudioPacket {\n AudioFormat format = 1;\n\n // Unix epoch in us when this frame was captured.\n uint64 timestamp = 2;\n\n // Contains a sample in the given audio format.\n bytes audio = 3;\n}\n\nmessage SmsMessage {\n // The source address where this message came from.\n //\n // The address should be a valid GSM-formatted address as specified by\n // 3GPP 23.040 Sec 9.1.2.5.\n //\n // For example: +3106225412 or (650) 555-1221\n string srcAddress = 1;\n\n // A utf8 encoded text message that should be delivered.\n string text = 2;\n}\n\n// A DisplayConfiguration describes a primary or secondary\n// display available to the emulator. The screen aspect ratio\n// cannot be longer (or wider) than 21:9 (or 9:21). Screen sizes\n// larger than 4k will be rejected.\n//\n// Common configurations (w x h) are:\n// - 480p (480x720) 142 dpi\n// - 720p (720x1280) 213 dpi\n// - 1080p (1080x1920) 320 dpi\n// - 4K (2160x3840) 320 dpi\n// - 4K (2160x3840) 640 dpi (upscaled)\n//\n// The behavior of the virtual display depends on the flags that are provided to\n// this method. By default, virtual displays are created to be private,\n// non-presentation and unsecure.\nmessage DisplayConfiguration {\n // These are the set of known android flags and their respective values.\n // you can combine the int values to (de)construct the flags field below.\n enum DisplayFlags {\n DISPLAYFLAGS_UNSPECIFIED = 0;\n\n // When this flag is set, the virtual display is public.\n // A public virtual display behaves just like most any other display\n // that is connected to the system such as an external or wireless\n // display. Applications can open windows on the display and the system\n // may mirror the contents of other displays onto it. see:\n // https://developer.android.com/reference/android/hardware/display/DisplayManager#VIRTUAL_DISPLAY_FLAG_PUBLIC\n VIRTUAL_DISPLAY_FLAG_PUBLIC = 1;\n\n // When this flag is set, the virtual display is registered as a\n // presentation display in the presentation display category.\n // Applications may automatically project their content to presentation\n // displays to provide richer second screen experiences.\n // https://developer.android.com/reference/android/hardware/display/DisplayManager#VIRTUAL_DISPLAY_FLAG_PRESENTATION\n VIRTUAL_DISPLAY_FLAG_PRESENTATION = 2;\n\n // When this flag is set, the virtual display is considered secure as\n // defined by the Display#FLAG_SECURE display flag. The caller promises\n // to take reasonable measures, such as over-the-air encryption, to\n // prevent the contents of the display from being intercepted or\n // recorded on a persistent medium.\n // see:\n // https://developer.android.com/reference/android/hardware/display/DisplayManager#VIRTUAL_DISPLAY_FLAG_SECURE\n VIRTUAL_DISPLAY_FLAG_SECURE = 4;\n\n // This flag is used in conjunction with VIRTUAL_DISPLAY_FLAG_PUBLIC.\n // Ordinarily public virtual displays will automatically mirror the\n // content of the default display if they have no windows of their own.\n // When this flag is specified, the virtual display will only ever show\n // its own content and will be blanked instead if it has no windows. See\n // https://developer.android.com/reference/android/hardware/display/DisplayManager#VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY\n VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY = 8;\n\n // Allows content to be mirrored on private displays when no content is\n // being shown.\n // This flag is mutually exclusive with\n // VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY. If both flags are specified\n // then the own-content only behavior will be applied.\n // see:\n // https://developer.android.com/reference/android/hardware/display/DisplayManager#VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR)\n VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR = 16;\n }\n\n // The width of the display, restricted to:\n // 320 * (dpi / 160) <= width\n uint32 width = 1;\n\n // The heigh of the display, restricted to:\n // * 320 * (dpi / 160) <= height\n uint32 height = 2;\n\n // The pixel density (dpi).\n // See https://developer.android.com/training/multiscreen/screendensities\n // for details. This value should be in the range [120, ..., 640]\n uint32 dpi = 3;\n\n // A combination of virtual display flags. These flags can be constructed\n // by combining the DisplayFlags enum described above.\n //\n // The behavior of the virtual display depends on the flags. By default\n // virtual displays are created to be private, non-presentation and\n // unsecure.\n uint32 flags = 4;\n\n // The id of the display.\n // The primary (default) display has the display ID of 0.\n // A secondary display has a display ID not 0.\n //\n // The id can be used to get or stream a screenshot.\n uint32 display = 5;\n}\n\nmessage DisplayConfigurations {\n repeated DisplayConfiguration displays = 1;\n}\n\nmessage Notification {\n enum EventType {\n VIRTUAL_SCENE_CAMERA_INACTIVE = 0;\n VIRTUAL_SCENE_CAMERA_ACTIVE = 1;\n\n // Fired when an update to a display event has been fired through\n // the extended ui. This does not fire events when the display\n // is changed through the console or gRPC endpoint.\n DISPLAY_CONFIGURATIONS_CHANGED_UI = 2;\n // Keep adding more for other event types\n }\n\n EventType event = 1;\n}\n\nmessage RotationRadian {\n float x = 1; // x axis is horizontal and orthogonal to the view direction.\n float y = 2; // y axis points up and is perpendicular to the floor.\n float z = 3; // z axis is the view direction and is set to 0.0 in\n // rotateVirtualSceneCamera call.\n}\n\nmessage Velocity {\n float x = 1; // x axis is horizontal and orthogonal to the view direction.\n float y = 2; // y axis points up and is perpendicular to the floor.\n float z = 3; // z axis is the view direction\n}\n\n// must follow the definition in \"external/qemu/android/hw-sensors.h\"\nmessage Posture {\n enum PostureValue {\n POSTURE_UNKNOWN = 0;\n POSTURE_CLOSED = 1;\n POSTURE_HALF_OPENED = 2;\n POSTURE_OPENED = 3;\n POSTURE_FLIPPED = 4;\n POSTURE_TENT = 5;\n POSTURE_MAX = 6;\n }\n PostureValue value = 3;\n}\n"
},
{
"path": "android_env/proto/snapshot.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n// Copyright (C) 2018 The Android Open Source Project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto2\";\n\n// This file must be synchronized between\n// Emulator (branch aosp/emu-master-dev):\n// external/qemu/android/android-emu/android/snapshot/proto/snapshot.proto\n//\n// Android Studio (branch goog/studio-master-dev):\n// tools/adt/idea/android/src/com/android/emulator/snapshot.proto\n//\n// If you modify one, please modify the other.\n\npackage emulator_snapshot;\n\noption java_package = \"com.android.emulator.snapshot\";\n\nmessage Image {\n enum Type {\n IMAGE_TYPE_UNKNOWN = 0;\n IMAGE_TYPE_KERNEL = 1;\n IMAGE_TYPE_KERNEL_RANCHU = 2;\n IMAGE_TYPE_SYSTEM = 3;\n IMAGE_TYPE_SYSTEM_COPY = 4;\n IMAGE_TYPE_DATA = 5;\n IMAGE_TYPE_DATA_COPY = 6;\n IMAGE_TYPE_RAMDISK = 7;\n IMAGE_TYPE_SDCARD = 8;\n IMAGE_TYPE_CACHE = 9;\n IMAGE_TYPE_VENDOR = 10;\n IMAGE_TYPE_ENCRYPTION_KEY = 11;\n }\n\n optional Type type = 1;\n optional string path = 2;\n optional bool present = 3;\n optional int64 size = 4;\n optional int64 modification_time = 5;\n}\n\nmessage Host {\n optional string gpu_driver = 4;\n optional int32 hypervisor = 5;\n}\n\nmessage Config {\n // Features are int32, not enums here to make sure we don't have to update\n // one more protobuf definition with every single new feature flag, even\n // when the code doesn't really care about the actual meaning for them,\n // only for the values.\n repeated int32 enabled_features = 1;\n\n // This holds the renderer; int32 for the same reason as |enabled_features|.\n optional int32 selected_renderer = 2;\n\n optional int32 cpu_core_count = 3;\n optional int64 ram_size_bytes = 4;\n}\n\nmessage SaveStats {\n // Type of save\n // 0: non-incremental\n // 1: incremental\n optional uint32 incremental = 1;\n // Time taken to save.\n optional uint64 duration = 2;\n // How many changed bytes in RAM.\n optional uint64 ram_changed_bytes = 3;\n}\n\nmessage Snapshot {\n // Update every time when introducing some breaking changes that make the\n // previous loading code break when trying to load the new snapshot.\n // NOTE: if the old code is fine with just skipping the new fields or not\n // getting the meaning of new values, |version| should remain\n // unchanged.\n optional int32 version = 1;\n\n // Purely informative: when this snapshot was created, Unix timestamp.\n optional int64 creation_time = 2;\n\n // list of mounted disk images used during the snapshot creation.\n repeated Image images = 3;\n\n // Description of the host machine properties needed to load this snapshot.\n optional Host host = 4;\n\n // Description of the emulator configuration needed for this snapshot.\n // NOTE: try not to duplicate the configuration that's already in\n // hardware-qemu.ini; only add what's either not there or what\n // could've been overridden during process initialization.\n optional Config config = 5;\n\n // Set if the snapshot failed to load during the last attempt.\n // Code is up to the application to define, with 0 meaning 'not failed' just\n // in case.\n optional int64 failed_to_load_reason_code = 7;\n\n // Set if data image is mounted.\n // User build and userdebug build mount data partition at different time.\n // But it should be done before boot finished, so this field is very likely\n // to be true.\n // We snapshot it here just in case someday we support snapshot during\n // booting.\n optional bool guest_data_partition_mounted = 8;\n\n // Emulator rotation angle, in right angles (e.g. 1 is 90 degrees, 2 is 180\n // etc).\n optional int32 rotation = 9;\n\n // Number of invalid loads / crashes that happened under this snapshot.\n optional int32 invalid_loads = 10;\n\n // Number of successful loads.\n optional int32 successful_loads = 11;\n\n // The name given to the snapshot by the user. Independent of the\n // file name.\n optional string logical_name = 12;\n\n // The file name of this snapshot's parent. The parent is the\n // snapshot that was loaded into the AVD prior to this snapshot\n // being taken\n optional string parent = 13;\n\n // Arbitrary description added by the user\n optional string description = 14;\n\n // Record of save stats.\n repeated SaveStats save_stats = 15;\n\n // Folded state.\n optional bool folded = 16;\n\n // Emulator boot parameters\n repeated string launch_parameters = 17;\n\n // Emulator build ID\n optional string emulator_build_id = 18;\n\n // System image build ID\n optional string system_image_build_id = 19;\n}\n"
},
{
"path": "android_env/proto/snapshot_service.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n// Copyright (C) 2018 The Android Open Source Project\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n// Note that if you add/remove methods in this file you must update\n// the metrics sql as well by running ./android/scripts/gen-grpc-sql.py\n//\n// Please group deleted methods in a block including the date (MM/DD/YY)\n// it was removed. This enables us to easily keep metrics around after removal\n//\n// list of deleted methods\n// rpc iWasDeleted (03/12/12)\n// ...\nsyntax = \"proto3\";\n\npackage android.emulation.control;\n\nimport \"android_env/proto/snapshot.proto\";\n\noption java_multiple_files = true;\noption java_package = \"com.android.emulator.control\";\noption objc_class_prefix = \"AEC\";\n\n// The SnapshotService enables you to list, insert, store, and retrieve\n// snapshots.\n//\n// Currently there are two types of snapshots:\n//\n// - Local (default): These are snapshots that are created locally. They are\n// stored internally inside qcow2 files and are very efficient. These are\n// the snapshots usually created by interacting with the UI.\n//\n// - Remote: These are snapshots that have been exported at a certain point.\n// an exported snapshot is normalized (completely self contained) and\n// can be imported into an emulator with a similar hardware configuration.\n//\n// Currently the emulator has limited support for importing snapshots:\n// - Once an imported snapshot has been loaded into an emulator it is no longer\n// possible to create new snapshots.\n// - The hardware configuration of the emulator your are pushing a snapshot to\n// must match (or be very similar) to the one you pulled the snapshot from.\n//\n// For example do not expect to be able to restore a snapshot on created on an\n// Intel cpu on an AMD cpu.\nservice SnapshotService {\n // Lists all the snapshots, filtered by the given query, that are stored\n // locally for the currently running avd. This includes all the snapshots that\n // were imported (pushed) into this emulator.\n //\n // Returns a list of snapshot_id's and associated details that describes\n // the hardware configuration, logical name, etc of the snapshot.\n rpc ListSnapshots(SnapshotFilter) returns (SnapshotList) {}\n\n // Pulls down the snapshot stored inside the AVD as a tar.gz/tar stream\n // This will normalize the snapshot, all relevant data to push a snapshot\n // into a similar emulator will be placed inside the tar file.\n //\n // Pulling down a snapshot will pause the emulator until the snapshots\n // are rebased and ready for exporting. Once the snapshot is rebased\n // the emulator will continue and downloading should commence.\n //\n // Note that pulling .gz stream is slow.\n //\n // You must provide the snapshot_id and (desired) format.\n //\n // If SnapshotPackage.path is set, the gRPC service will directly write the\n // exported snapshot to SnapshotPackage.path without streaming, which is\n // usually significantly faster. It would require emulator to have direct\n // access to SnapshotPackage.path, which usually means it can only be used\n // when pulling from a local emulator.\n rpc PullSnapshot(SnapshotPackage) returns (stream SnapshotPackage) {}\n\n // Push a tar.gz stream contain the snapshot. The tar file should\n // be a snapshot that was exported through the PullSnapshot in the past.\n // The emulator will try to import the snapshot. The hardware configuration\n // of the current emulator should match the one used for pulling.\n //\n // A detailed description of the snapshot (emulator_snapshot.Snapshot)\n // is stored in the snapshot.pb file inside the tar.\n //\n // You must provide the snapshot_id and format in the first message.\n // Will return success and a possible error message when a failure occurs.\n //\n // If SnapshotPackage.path is set, the gRPC service will directly unzip the\n // exported snapshot from SnapshotPackage.path without streaming, which is\n // usually significantly faster. It would require emulator to have direct\n // access to SnapshotPackage.path, which usually means it can only be used\n // when pushing to a local emulator.\n rpc PushSnapshot(stream SnapshotPackage) returns (SnapshotPackage) {}\n\n // Loads the given snapshot inside the emulator and activates it.\n // The device will be in the state as it was when the snapshot was created.\n //\n // You will no longer be able to call Save if this was an imported\n // snapshot that was pushed into this emulator.\n //\n // You must provide the snapshot_id to indicate which snapshot to load\n // Will return success and a possible error message when a failure occurs.\n rpc LoadSnapshot(SnapshotPackage) returns (SnapshotPackage) {}\n\n // Creates as a snapshot of the current state of the emulator.\n // You can only save a snapshot if you never activated (Load) an imported\n // snapshot (Push).\n //\n // For example:\n // - PushSnapshot(\"some_snap.tar.gz\");\n // - LoadSnapshot(\"some_snap\");\n // - SaveSnapshot(\"same_newer_snap\"); // <--- Will currently fail.\n //\n // You can provide the snapshot_id to indicate the name used for storing.\n // Will return success and a possible error message when a failure occurs.\n rpc SaveSnapshot(SnapshotPackage) returns (SnapshotPackage) {}\n\n // Deletes the snapshot with the given snapshot_id from the avd.\n //\n // You must provide the snapshot_id to indicate which snapshot to delete.\n // Will return success and a possible error message when a failure occurs.\n rpc DeleteSnapshot(SnapshotPackage) returns (SnapshotPackage) {}\n\n // Tracks the given process for automated snapshot creation in case of\n // assert failures.\n //\n // Will return success and a possible error message when a failure occurs.\n // The snapshot_id field will contain the name of the snapshot that\n // will be created. The pid field will contain the process id that is\n // being tracked.\n rpc TrackProcess(IceboxTarget) returns (IceboxTarget) {}\n}\n\n// Sets options for SnapshotService. Used for both request and response\n// messages.\nmessage SnapshotPackage {\n enum Format {\n TARGZ = 0;\n TAR = 1;\n DIRECTORY = 2;\n }\n // The identifier to the snapshot, only required for request messages. For\n // streaming service, only used in the first stream message of a gRPC call\n // (would be ignored in consequent stream messages of the same call).\n string snapshot_id = 1;\n\n // A stream of bytes. Encoded as a tar (possibly gzipped) file pendinf on the\n // value of format.\n bytes payload = 2;\n\n // [response only] status fields, usually indicates end of transmission.\n bool success = 3;\n bytes err = 4;\n\n // [request only] Format of the payload. Only used in request messages. For\n // streaming service, only used in the first stream message of a gRPC call\n // (would be ignored in consequent stream messages of the same call).\n Format format = 5;\n\n // [request only] Path to the snapshot package file. Only used in request\n // messages.\n //\n // When set in a request, the PullSnapshot/PushSnapshot operation will\n // directly write/read the exported snapshot in path without streaming, which\n // is usually significantly faster. It would require emulator to have direct\n // access to path, which usually means it can only be used with a local\n // emulator.\n string path = 6;\n}\n\n// A snapshot filter can be used to filter the results produced by ListSnapshots\nmessage SnapshotFilter {\n enum LoadStatus {\n // Only return compatible snapshots\n CompatibleOnly = 0;\n\n // Return all snapshots.\n All = 1;\n }\n\n // Filter snapshots by load status.\n LoadStatus statusFilter = 1;\n}\n\n// Provides detailed information regarding the snapshot.\nmessage SnapshotDetails {\n enum LoadStatus {\n // The emulator believes that the snapshot is compatible with the emulator\n // that provided this information. The emulator will attempt to load this\n // snapshot when requested.\n //\n // A snapshot is usually compatible when the following statements are true:\n // - The snapshot was taken by the current emulator version. i.e.\n // emulator_build_id in the details field matches the build_id of the\n // emulator that provided this information.\n //\n // - The snapshot was taken on the current running machine, and no hardware\n // changes have taken place between taking and loading the snapshot.\n //\n // - The avd configuration has not changed between when this snapshot was\n // taken and when the snapshot was loaded.\n //\n // - The system images on which the avd is based have not changed.\n Compatible = 0;\n\n // The emulator will not allow loading of the snapshot, as it deems the\n // snapshot to be incompatible. Loading of snapshots can be forced by\n // launching the emulator with the feature \"AllowSnapshotMigration\" enabled.\n Incompatible = 1;\n\n // This snapshot was successfully loaded in the emulator, and was used at\n // the starting point of the current running emulator. The following holds:\n //\n // A loaded snapshot is a compatible snapshot\n // There is at most one snapshot_id that is in the \"Loaded\" state\n Loaded = 2;\n }\n\n // The id of this snapshot. Use this id to load/delete/pull the\n // snapshot.\n string snapshot_id = 1;\n\n // Detailed information about this snapshot. This contains a detailed\n // hardware description of the snapshot. These details are the same\n // as the \"snapshot.pb\" file found in an exported snapshot.\n // Look at the import file for a detailed description of the available\n // fields.\n emulator_snapshot.Snapshot details = 2;\n\n // Provides information about the ability to restore this snapshot.\n LoadStatus status = 3;\n\n // The size of the folder that stores required information to load a snapshot.\n uint64 size = 4;\n}\n\n// A list of on snapshot details.\nmessage SnapshotList {\n repeated SnapshotDetails snapshots = 1;\n}\n\nmessage IceboxTarget {\n // This is the process id to attach to, if this value is not set (0)\n // The process name will be used instead.\n int64 pid = 1;\n\n // The process name to attach to if any, if this is not set the pid will\n // be used. This is usually the application name of your application under\n // test, that is passed in to the am instrument command. It is likely\n // what you will find in your AndroidManifest.xml\n string package_name = 2;\n\n // The name of the snapshot that icebox will create if a snapshot is\n // generated.\n string snapshot_id = 3;\n\n // [Output Only] True if icebox failed to track the given target.\n bool failed = 4;\n\n // [Output Only] Detailed error message that might provide more information.\n string err = 5;\n\n // Maximum number of snapshots the emulator can take during one Icebox run.\n // Set to -1 for unlimited number of snapshots.\n int32 max_snapshot_number = 6;\n}\n\n// list of deleted methods:\n//\n"
},
{
"path": "android_env/proto/state.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\noption java_multiple_files = true;\n\nmessage SaveStateRequest {\n map args = 1;\n}\n\nmessage LoadStateRequest {\n map args = 1;\n}\n\nmessage SaveStateResponse {\n enum Status {\n // Reserved value for unset statuses.\n UNDEFINED = 0;\n // Returned when everything goes well.\n OK = 1;\n // Returned when something internal did not work as expected.\n ERROR = 2;\n }\n Status status = 1;\n // `error_message` is only populated in case of errors.\n string error_message = 2;\n\n // Any additional info returned during the request; e.g., file paths or sizes.\n map additional_info = 3;\n}\n\nmessage LoadStateResponse {\n enum Status {\n // Reserved value for unset statuses.\n UNDEFINED = 0;\n // Returned when everything goes well.\n OK = 1;\n // Returned when there is no state to load.\n NOT_FOUND = 2;\n // Returned when something internal did not work as expected.\n ERROR = 3;\n }\n Status status = 1;\n // `error_message` is only populated in case of errors.\n string error_message = 2;\n\n // Any additional info returned during the request; e.g., file paths or sizes.\n map additional_info = 3;\n}"
},
{
"path": "android_env/proto/task.proto",
"content": "// Copyright 2026 DeepMind Technologies Limited.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\nsyntax = \"proto3\";\n\npackage android_env;\n\nimport \"android_env/proto/adb.proto\";\n\n// An AppScreen identifies a unique configuration that we can observe on the\n// screen of a device.\nmessage AppScreen {\n // Fully-qualified name of the activity.\n string activity = 1;\n\n // A list of regexes to match at each level of the current view hierarchy.\n // The environment uses this list to determine whether the agent has \"exited\"\n // this current task.\n // Example: [\n // \"^DecorView@.*\\[MainActivity\\]$\",\n // \"^android.widget.LinearLayout\\{.*\\}$\",\n // \"^android.widget.FrameLayout\\{.*android\\:id\\/content\\}\",\n // \"^android.widget.RelativeLayout\\{.*\\}\",\n // \"^android.widget.FrameLayout\\{.*app\\:id\\/fragment_holder\\}\",\n // \"^android.widget.RelativeLayout\\{.*\\}\",\n // \"^com.google.example.games.nostalgicracer.views.RaceView3D\\{.*app\\:id\\/gameplay_screen_3d\\}\",\n // ],\n repeated string view_hierarchy_path = 2;\n}\n\n// Waits for `app_screen` to be the current app screen shown to the user.\nmessage WaitForAppScreen {\n AppScreen app_screen = 1;\n // Maximum time in seconds to wait for the activity to become the current one.\n float timeout_sec = 2;\n}\n\nmessage CheckInstall {\n string package_name = 1;\n // Maximum time in seconds to wait.\n float timeout_sec = 2;\n}\n\nmessage Sleep {\n float time_sec = 1;\n}\n\nmessage SuccessCondition {\n int32 num_retries = 1;\n\n oneof check {\n WaitForAppScreen wait_for_app_screen = 2;\n CheckInstall check_install = 3;\n }\n}\n\nmessage SetupStep {\n SuccessCondition success_condition = 1;\n\n oneof step {\n AdbRequest adb_request = 2;\n Sleep sleep = 3;\n }\n}\n\n// A specification of structured observations\n// Analogous to dm_env.specs.Array()\n\nmessage ArraySpec {\n // An identifier for this ArraySpec.\n string name = 1;\n\n // The shape of the multi-dimensional values associated with this ArraySpec,\n repeated int32 shape = 2;\n\n enum DataType {\n INVALID_DATA_TYPE = 0;\n FLOAT = 1;\n DOUBLE = 2;\n INT8 = 3;\n INT16 = 4;\n INT32 = 5;\n INT64 = 6;\n UINT8 = 7;\n UINT16 = 8;\n UINT32 = 9;\n UINT64 = 10;\n BOOL = 11;\n STRING_U1 = 12;\n STRING_U16 = 13;\n STRING_U25 = 14;\n STRING_U250 = 15;\n STRING = 16; // String without max length\n OBJECT = 17;\n }\n\n // Data type of elements we expect to see in an array of this spec.\n DataType dtype = 3;\n}\n\nmessage LogParsingConfig {\n // `filters` are tags used by the app's logging system so that we can\n // identify them in logcat's output. It's the first argument to logging calls\n // such as Log.e(\"ActivityManager\", \"My message\").\n // Example: \"ActivityManager\"\n repeated string filters = 1;\n\n // Regular expressions that define how we can extract RL information such as\n // score, extras and episode end from raw logcat messages.\n message LogRegexps {\n // Regexp expected to match:\n // ...a floating point value which gets accumulated over time.\n // A delta in 'score' corresponds to the reward.\n string score = 1;\n\n // Regexp expected to match:\n // ...a floating point value directly forwarded by the environment.\n repeated string reward = 2;\n\n // Regexp expected to match:\n // ...a signal marking the end of an episode.\n repeated string episode_end = 3;\n\n // Regexp expected to match:\n // ...a string representing pairs of extra names and values.\n repeated string extra = 4;\n\n // Regexp expected to match:\n // ...a dict of extra names and values in json format.\n repeated string json_extra = 5;\n\n // Attaches rewards to arbitrary log messages, for example:\n // {event: \"coin_collected\" reward: 2.3}\n // {event: \"car_crashed\" reward: -1.4}\n message RewardEvent {\n // If `event` is matched, the environment will give `reward`.\n string event = 1;\n\n // Numerical value to give as reward if `event` is matched.\n float reward = 2;\n }\n\n repeated RewardEvent reward_event = 6;\n }\n\n LogRegexps log_regexps = 2;\n}\n\n// Description of a reinforcement learning task to be solved by an agent.\nmessage Task {\n // A globally unique identifier for this task.\n string id = 1;\n\n // A human readable name for this task.\n string name = 2;\n\n // A description of the task.\n string description = 3;\n\n repeated SetupStep setup_steps = 4;\n repeated SetupStep reset_steps = 5;\n\n AppScreen expected_app_screen = 6;\n\n // AndroidEnv resets the episode after `max_episode_sec` is passed since the\n // last reset(). Recommended for time sensitive tasks (e.g. reactive games).\n // Note that this is real time as measured by AndroidEnv and is independent of\n // the speed of simulation of Android.\n // If <= 0.0, this logic is disabled.\n float max_episode_sec = 7;\n\n // The maximum number of interactions in a single episode between the\n // environment and an agent.\n // This setting is appropriate for tasks that are not time-dependent or when\n // the performance of the simulation varies dramatically between runs.\n // If <= 0, this logic is disabled.\n int32 max_episode_steps = 8;\n\n // Defines parameters for parsing messages from logcat.\n LogParsingConfig log_parsing_config = 9;\n\n // NOTE: This field is deprecated and will be removed from this Task\n // definition soon.\n //\n // (Optional): The task may also define extras to help the RL agent.\n // An Extra in AndroidEnv is any information that apps may send to aid the\n // understanding of the task. The type of information sent through this\n // channel is usually something difficult to obtain from raw pixels and may\n // include things such as:\n //\n // - The current board configuration (e.g. of a chess game or a tetris game)\n // - The position of the avatar in a map\n // - Events (e.g. whether a button was pressed or a checkpoint was achieved)\n //\n // Notice that these are entirely optional and may not be available at all.\n // This specification ensures that only extras specified in the Task\n // definition will be passed to the agent, everything else is excluded.\n // The name of an extra must be unique across all extras.\n repeated ArraySpec extras_spec = 10;\n}\n"
},
{
"path": "android_env/wrappers/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/wrappers/a11y/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "android_env/wrappers/a11y/a11y_events.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tools for accessing accessibility events.\"\"\"\n\nfrom collections.abc import Mapping\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env.proto.a11y import a11y_pb2\nimport numpy as np\n\nfrom google.protobuf import any_pb2\n\n\n_A11Y_EVENT_KEY = 'full_event'\n\n\ndef package_events_to_task_extras(\n events: list[a11y_pb2.EventRequest],\n) -> Mapping[str, np.ndarray]:\n if not events:\n return {}\n events = np.stack(events, axis=0)\n return {_A11Y_EVENT_KEY: events}\n\n\ndef extract_events_from_task_extras(\n task_extras: Mapping[str, Any] | None = None,\n) -> list[Mapping[str, str]]:\n \"\"\"Inspects task_extras and extracts all accessibility events detected.\n\n Args:\n task_extras: Task extras forwarded by AndroidEnv. If 'full_event' is not a\n key in task_extras, then this function returns an empty string. Otherwise,\n full_event is expected to be list to be a numpy array with one dimension,\n and contains a list of dictionary describing accessibility events that are\n present in the given task extras. e.g. 'event_type:\n TYPE_WINDOW_CONTENT_CHANGED // event_package_name:\n com.google.android.deskclock // source_class_name:\n android.widget.ImageView'.\n\n Returns:\n List of all events detected\n \"\"\"\n if task_extras is None or _A11Y_EVENT_KEY not in task_extras:\n return []\n\n if (\n not isinstance(task_extras[_A11Y_EVENT_KEY], np.ndarray)\n or task_extras[_A11Y_EVENT_KEY].ndim != 1\n ):\n raise ValueError(\n f'{_A11Y_EVENT_KEY} task extra should be a numpy array with one'\n ' dimension.'\n )\n\n if task_extras[_A11Y_EVENT_KEY].size == 0:\n return []\n\n events = []\n for e in task_extras[_A11Y_EVENT_KEY]:\n if isinstance(e, a11y_pb2.EventRequest):\n events.append(dict(e.event))\n elif isinstance(e, dict):\n events.append(e)\n logging.warning(\n 'The event should come only from the a11y_grpc_wrapper. '\n 'Please verify that the upacking operation has not been '\n 'called twice. See here for full task_extras: %s',\n task_extras,\n )\n elif isinstance(e, any_pb2.Any):\n ev = a11y_pb2.EventRequest()\n new_any = any_pb2.Any()\n new_any.CopyFrom(e)\n new_any.Unpack(ev)\n events.append(dict(ev.event))\n\n else:\n raise TypeError(\n f'Unexpected event type: {type(e)}. See here for full '\n f'task_extras: {task_extras}.'\n )\n\n return events\n\n\ndef keep_latest_event_only(task_extras: dict[str, Any]):\n \"\"\"Removes all a11y events except the last one observed.\"\"\"\n if task_extras is None or 'full_event' not in task_extras:\n return\n\n if (\n not isinstance(task_extras[_A11Y_EVENT_KEY], np.ndarray)\n or task_extras[_A11Y_EVENT_KEY].ndim != 1\n ):\n raise ValueError(\n f'{_A11Y_EVENT_KEY} task extra should be a numpy array with one'\n ' dimension.'\n )\n\n if task_extras[_A11Y_EVENT_KEY].size == 0:\n return []\n\n task_extras[_A11Y_EVENT_KEY] = task_extras[_A11Y_EVENT_KEY][-1:]\n"
},
{
"path": "android_env/wrappers/a11y/a11y_events_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for a11y_events.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.proto.a11y import a11y_pb2\nfrom android_env.wrappers.a11y import a11y_events\nimport numpy as np\n\nfrom google.protobuf import any_pb2\n\n\ndef _event_request(d: dict[str, str]) -> a11y_pb2.EventRequest:\n event_request = a11y_pb2.EventRequest()\n for k, v in d.items():\n event_request.event[k] = v\n return event_request\n\n\ndef _event_request_as_any(d: dict[str, str]) -> any_pb2.Any:\n event_request = _event_request(d)\n response = any_pb2.Any()\n response.Pack(event_request)\n return response\n\n\nclass A11yEventsTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(task_extras={}),\n dict(\n task_extras={'no_full_event': [{'1': '1'}, {'2': '2'}, {'3': '3'}]},\n ),\n dict(\n task_extras={'full_event': np.array([])},\n ),\n dict(\n task_extras={},\n ),\n )\n def test_no_events_in_task_extras(self, task_extras):\n events = a11y_events.extract_events_from_task_extras(task_extras)\n self.assertEmpty(events)\n\n @parameterized.parameters(\n dict(\n task_extras={'full_event': [{'1': '1'}, {'2': '2'}]},\n expected_events=[{'1': '1'}, {'2': '2'}],\n ),\n dict(\n task_extras={'full_event': [{}]},\n expected_events=[{}],\n ),\n dict(\n task_extras={\n 'full_event_wrong_key': [1, 2, 3],\n 'full_event': [{'1': '1'}, {'2': '2'}, {'3': '3'}],\n },\n expected_events=[{'1': '1'}, {'2': '2'}, {'3': '3'}],\n ),\n )\n def test_task_extras(self, task_extras, expected_events):\n event_requests = [_event_request(e) for e in task_extras['full_event']]\n task_extras['full_event'] = np.stack(event_requests, axis=0)\n events = a11y_events.extract_events_from_task_extras(task_extras)\n self.assertEqual(len(events), len(expected_events))\n for i, event in enumerate(expected_events):\n self.assertEqual(len(event), len(expected_events[i]))\n for k, v in event.items():\n self.assertIn(k, expected_events[i])\n self.assertEqual(v, expected_events[i][k])\n\n def test_events_key_has_dict_event_requrests(self):\n event_requests = [\n _event_request({'1': '1'}),\n {'2': '2'},\n _event_request({'3': '3'}),\n ]\n expected_events = [\n {'1': '1'},\n {'2': '2'},\n {'3': '3'},\n ]\n task_extras = {'full_event': np.stack(event_requests, axis=0)}\n events = a11y_events.extract_events_from_task_extras(task_extras)\n self.assertEqual(len(events), len(expected_events))\n for i, event in enumerate(expected_events):\n self.assertEqual(len(event), len(expected_events[i]))\n for k, v in event.items():\n self.assertIn(k, expected_events[i])\n self.assertEqual(v, expected_events[i][k])\n\n def test_events_key_has__event_requrests_packed_as_any(self):\n event_requests = [\n _event_request_as_any({'1': '1'}),\n {'2': '2'},\n _event_request_as_any({'3': '3'}),\n ]\n expected_events = [\n {'1': '1'},\n {'2': '2'},\n {'3': '3'},\n ]\n task_extras = {'full_event': np.stack(event_requests, axis=0)}\n events = a11y_events.extract_events_from_task_extras(task_extras)\n self.assertEqual(len(events), len(expected_events))\n for i, event in enumerate(expected_events):\n self.assertEqual(len(event), len(expected_events[i]))\n for k, v in event.items():\n self.assertIn(k, expected_events[i])\n self.assertEqual(v, expected_events[i][k])\n\n def test_events_key_has_non_event_requrests(self):\n event_requests = [\n _event_request({'1': '1'}),\n 3, # Not an even and not a dict.\n _event_request({'3': '3'}),\n ]\n task_extras = {'full_event': np.stack(event_requests, axis=0)}\n with self.assertRaises(TypeError):\n _ = a11y_events.extract_events_from_task_extras(task_extras)\n\n @parameterized.parameters(\n dict(task_extras={}, expected_extras={}),\n dict(\n task_extras={\n 'no_full_event': 42,\n },\n expected_extras={\n 'no_full_event': 42,\n },\n ),\n dict(\n task_extras={'full_event': np.array([1, 2]), 'no_full_event': 43},\n expected_extras={'full_event': np.array([2]), 'no_full_event': 43},\n ),\n dict(\n task_extras={'full_event': np.array([1, 2, 3])},\n expected_extras={'full_event': np.array([3])},\n ),\n dict(\n task_extras={'full_event': np.array([]), 'no_full_event': 44},\n expected_extras={'full_event': np.array([]), 'no_full_event': 44},\n ),\n )\n def test_keep_latest_only(self, task_extras, expected_extras):\n a11y_events.keep_latest_event_only(task_extras)\n self.assertEqual(len(task_extras), len(expected_extras))\n for k, v in task_extras.items():\n self.assertIn(k, expected_extras)\n if k == 'full_event':\n np.testing.assert_array_equal(v, expected_extras['full_event'])\n else:\n self.assertEqual(v, expected_extras[k])\n pass\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/a11y/a11y_forests.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tools for accessing accessibility events.\"\"\"\n\nfrom collections.abc import Mapping\nfrom typing import Any\n\nfrom android_env.proto.a11y import android_accessibility_forest_pb2\nimport numpy as np\n\nfrom google.protobuf import any_pb2\n\n\n_A11Y_FORESTS_KEY = 'accessibility_tree'\n\n\ndef package_forests_to_task_extras(\n forests: list[android_accessibility_forest_pb2.AndroidAccessibilityForest],\n) -> Mapping[str, np.ndarray]:\n if not forests:\n return {}\n forests = np.stack(forests, axis=0)\n return {_A11Y_FORESTS_KEY: forests}\n\n\ndef task_extras_has_forests(task_extras: Mapping[str, Any]) -> bool:\n \"\"\"Checks that the task_extras has any a11y forest information.\"\"\"\n if _A11Y_FORESTS_KEY not in task_extras:\n return False\n\n payload = task_extras[_A11Y_FORESTS_KEY]\n if not isinstance(payload, np.ndarray) or payload.ndim != 1:\n raise ValueError(\n f'{_A11Y_FORESTS_KEY} task extra should be a numpy array with one'\n f' dimension. payload: {payload}'\n )\n\n if payload.size == 0:\n return False\n\n if any(isinstance(f, any_pb2.Any) for f in payload):\n # Forests were packed as Any.\n return True\n\n return any(\n isinstance(f, android_accessibility_forest_pb2.AndroidAccessibilityForest)\n for f in payload\n )\n\n\ndef convert_to_forest(\n forest: android_accessibility_forest_pb2.AndroidAccessibilityForest\n | any_pb2.Any\n | None,\n) -> android_accessibility_forest_pb2.AndroidAccessibilityForest | None:\n \"\"\"Takes an object and attempts to convert it to a forest.\"\"\"\n if forest is None:\n return None\n\n if isinstance(forest, any_pb2.Any):\n output = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n new_any = any_pb2.Any()\n new_any.CopyFrom(forest)\n new_any.Unpack(output)\n return output\n elif isinstance(\n forest, android_accessibility_forest_pb2.AndroidAccessibilityForest\n ):\n return forest\n else:\n return None\n\n\ndef extract_forests_from_task_extras(\n task_extras: Mapping[str, Any] | None = None,\n) -> list[android_accessibility_forest_pb2.AndroidAccessibilityForest]:\n \"\"\"Inspects task_extras and extracts all accessibility forests detected.\n\n Args:\n task_extras: Task extras forwarded by AndroidEnv. If 'full_event' is not a\n key in task_extras, then this function returns an empty string. Otherwise,\n full_event is expected to be list to be a numpy array with one dimension,\n and contains a list of dictionary describing accessibility forests that\n are present in the given task extras.\n\n Returns:\n List of all forests detected\n \"\"\"\n if task_extras is None or not task_extras_has_forests(task_extras):\n return []\n\n forests = []\n for f in task_extras[_A11Y_FORESTS_KEY]:\n f = convert_to_forest(f)\n if f is not None:\n forests.append(f)\n return forests\n\n\ndef keep_latest_forest_only(task_extras: dict[str, Any]):\n \"\"\"Removes all a11y forests except the last one observed.\"\"\"\n if _A11Y_FORESTS_KEY not in task_extras.keys():\n return\n\n payload = task_extras[_A11Y_FORESTS_KEY]\n if not isinstance(payload, np.ndarray) or payload.ndim != 1:\n raise ValueError(\n f'{_A11Y_FORESTS_KEY} task extra should be a numpy array with one'\n f' dimension. payload: {payload}'\n )\n\n if payload.size == 0:\n return\n\n task_extras[_A11Y_FORESTS_KEY] = payload[-1:]\n"
},
{
"path": "android_env/wrappers/a11y/a11y_forests_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for a11y_forests.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.proto.a11y import android_accessibility_forest_pb2\nfrom android_env.wrappers.a11y import a11y_forests\nimport numpy as np\n\nfrom google.protobuf import any_pb2\n\n\ndef _pack_any(proto_message) -> any_pb2.Any:\n response = any_pb2.Any()\n response.Pack(proto_message)\n return response\n\n\ndef _empty_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n return android_accessibility_forest_pb2.AndroidAccessibilityForest()\n\n\ndef _one_empty_window_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n forest.windows.add()\n return forest\n\n\ndef _two_window_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n window = forest.windows.add()\n window.tree.nodes.add(\n class_name='foo', is_clickable=True, hint_text='Foo hint'\n )\n forest.windows.add()\n return forest\n\n\nclass A11YForestsTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(task_extras={}, expected_forests=[], convert_to_np=[]),\n dict(\n task_extras={'accessibility_tree': []},\n convert_to_np=['accessibility_tree'],\n expected_forests=[],\n ),\n dict(\n task_extras={\n 'not_accessibility_tree': [\n _empty_forest(),\n _one_empty_window_forest(),\n _two_window_forest(),\n ],\n },\n convert_to_np=['not_accessibility_tree'],\n expected_forests=[],\n ),\n dict(\n task_extras={\n 'accessibility_tree': [\n _empty_forest(),\n {'not_a_forest_key': 'nor_a_forest_value'},\n _two_window_forest(),\n ]\n },\n convert_to_np=['accessibility_tree'],\n expected_forests=[_empty_forest(), _two_window_forest()],\n ),\n dict(\n task_extras={\n 'accessibility_tree': [\n {'not_a_forest_key': 'nor_a_forest_value'},\n 3,\n 4,\n {'not_a_forest_key': _empty_forest()},\n ],\n },\n convert_to_np=['accessibility_tree'],\n expected_forests=[],\n ),\n dict(\n task_extras={'accessibility_tree': []},\n convert_to_np=['accessibility_tree'],\n expected_forests=[],\n ),\n dict(\n task_extras={\n 'accessibility_tree_wrong_key': [1, 2, 3],\n 'accessibility_tree': [\n _empty_forest(),\n None,\n None,\n _one_empty_window_forest(),\n _two_window_forest(),\n ],\n },\n convert_to_np=['accessibility_tree', 'accessibility_tree_wrong_key'],\n expected_forests=[\n _empty_forest(),\n _one_empty_window_forest(),\n _two_window_forest(),\n ],\n ),\n dict(\n task_extras={\n 'accessibility_tree_wrong_key': [1, 2, 3],\n 'accessibility_tree': [\n None,\n _pack_any(_empty_forest()),\n _pack_any(_one_empty_window_forest()),\n _pack_any(_two_window_forest()),\n ],\n },\n convert_to_np=['accessibility_tree', 'accessibility_tree_wrong_key'],\n expected_forests=[\n _empty_forest(),\n _one_empty_window_forest(),\n _two_window_forest(),\n ],\n ),\n dict(\n task_extras={\n 'accessibility_tree': [\n _pack_any(_empty_forest()),\n {'not_a_forest_key': 'nor_a_forest_value'},\n None,\n _two_window_forest(),\n None,\n ]\n },\n convert_to_np=['accessibility_tree'],\n expected_forests=[_empty_forest(), _two_window_forest()],\n ),\n )\n def test_task_extras(self, task_extras, expected_forests, convert_to_np):\n for k in convert_to_np:\n if task_extras[k]:\n task_extras[k] = np.stack(task_extras[k], axis=0)\n else:\n task_extras[k] = np.array([])\n forests = a11y_forests.extract_forests_from_task_extras(task_extras)\n self.assertEqual(len(forests), len(expected_forests))\n for idx, f in enumerate(forests):\n self.assertEqual(f, expected_forests[idx])\n\n @parameterized.parameters(\n dict(task_extras={}, expected_extras={}),\n dict(\n task_extras={\n 'no_accessibility_tree': 42,\n },\n expected_extras={\n 'no_accessibility_tree': 42,\n },\n ),\n dict(\n task_extras={'accessibility_tree': []},\n expected_extras={'accessibility_tree': []},\n ),\n dict(\n task_extras={\n 'accessibility_tree': [\n _empty_forest(),\n _one_empty_window_forest(),\n ],\n 'no_accessibility_tree': 43,\n },\n expected_extras={\n 'accessibility_tree': [_one_empty_window_forest()],\n 'no_accessibility_tree': 43,\n },\n ),\n dict(\n task_extras={\n 'accessibility_tree': [\n _empty_forest(),\n _one_empty_window_forest(),\n _two_window_forest(),\n ]\n },\n expected_extras={'accessibility_tree': [_two_window_forest()]},\n ),\n dict(\n task_extras={\n 'accessibility_tree': [],\n 'no_accessibility_tree': 44,\n },\n expected_extras={\n 'accessibility_tree': [],\n 'no_accessibility_tree': 44,\n },\n ),\n )\n def test_keep_latest_only(self, task_extras, expected_extras):\n if 'accessibility_tree' in task_extras:\n if task_extras['accessibility_tree']:\n task_extras['accessibility_tree'] = np.stack(\n task_extras['accessibility_tree'], axis=0\n )\n else:\n task_extras['accessibility_tree'] = np.array([])\n\n a11y_forests.keep_latest_forest_only(task_extras)\n self.assertSameElements(task_extras.keys(), expected_extras.keys())\n for k in task_extras.keys():\n if k == 'accessibility_tree':\n self.assertEqual(len(task_extras[k]), len(expected_extras[k]))\n for idx, f in enumerate(task_extras[k]):\n self.assertEqual(f, expected_extras[k][idx])\n else:\n self.assertEqual(task_extras[k], expected_extras[k])\n pass\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/a11y/a11y_servicer.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Accessibility Servicer implementation.\"\"\"\n\nimport asyncio\nfrom collections.abc import AsyncIterator, Generator, Iterable\nimport threading\n\nfrom absl import logging\nfrom android_env.proto.a11y import a11y_pb2\nfrom android_env.proto.a11y import a11y_pb2_grpc\nfrom android_env.proto.a11y import android_accessibility_forest_pb2\nimport grpc\n\n\nclass A11yServicer(a11y_pb2_grpc.A11yServiceServicer):\n \"\"\"Services the A11yService requests.\"\"\"\n\n def __init__(self, latest_forest_only: bool = False):\n self._received_forests: list[\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n ] = []\n self._received_events: list[a11y_pb2.EventRequest] = []\n self._lock_forests = threading.Lock()\n self._lock_events = threading.Lock()\n self._latest_forest_only = latest_forest_only\n self._paused = True\n\n # A11y Forest bookkeeping.\n self._get_forest = asyncio.Event() # Whether to request a forest.\n self._forest_ready = asyncio.Event() # Whether the forest is ready.\n self._latest_forest: (\n android_accessibility_forest_pb2.AndroidAccessibilityForest | None\n ) = None\n\n def SendForest(\n self,\n request: android_accessibility_forest_pb2.AndroidAccessibilityForest,\n context: grpc.ServicerContext,\n ) -> a11y_pb2.ForestResponse:\n self._process_forest(request)\n return a11y_pb2.ForestResponse()\n\n def SendEvent(\n self,\n request: a11y_pb2.EventRequest,\n context: grpc.ServicerContext,\n ) -> a11y_pb2.EventResponse:\n self._process_event(request)\n return a11y_pb2.EventResponse()\n\n async def Bidi(\n self,\n request_iterator: AsyncIterator[a11y_pb2.ClientToServer],\n context: grpc.aio.ServicerContext,\n ) -> AsyncIterator[a11y_pb2.ServerToClient]:\n \"\"\"Processes incoming ClientToServer requests.\"\"\"\n\n logging.info('Starting A11yServicer.Bidi()')\n\n # Send a dummy message to unblock clients in their loop.\n yield a11y_pb2.ServerToClient()\n\n # This block defines two coroutines:\n #\n # * `read_client_requests()`\n # * `check_forest()`\n #\n # They cooperate with each other and both populate a queue `q` which is\n # consumed in a loop below, which actually yields requests which are sent to\n # the client. The processing finishes when the clients \"closes\" the\n # connection, which causes `read_client_requests()` to put a special value,\n # `STOP_ITERATION`, in the queue.\n\n # Queue for communicating from coroutines to `Bidi()`.\n q = asyncio.Queue()\n\n should_run = True\n\n async def read_client_requests():\n \"\"\"Coroutine for reading client requests.\"\"\"\n\n nonlocal should_run\n async for request in request_iterator:\n field_name = request.WhichOneof('payload')\n match field_name:\n case 'event':\n self._process_event(request.event)\n case 'forest':\n self._latest_forest = request.forest\n self._forest_ready.set()\n self._get_forest.clear() # Reset the `Event`.\n case _:\n logging.error('Unknown field %r', field_name)\n await q.put(a11y_pb2.ServerToClient())\n\n # Send a special value to stop processing this `Bidi` connection.\n await q.put('STOP_ITERATION')\n should_run = False\n\n async def check_forest():\n \"\"\"Coroutine for sending \"get forest\" requests.\"\"\"\n\n nonlocal should_run\n while should_run:\n await self._get_forest.wait()\n await q.put(a11y_pb2.ServerToClient(get_forest={}))\n\n tasks = asyncio.gather(read_client_requests(), check_forest())\n\n while should_run:\n v = await q.get()\n if v == 'STOP_ITERATION':\n break\n else:\n yield v\n\n await tasks\n\n logging.info('Finishing A11yServicer.Bidi()')\n\n async def get_forest(\n self,\n ) -> android_accessibility_forest_pb2.AndroidAccessibilityForest | None:\n \"\"\"Issues a request to get the a11y forest from the client.\"\"\"\n\n self._get_forest.set() # Unblocks coroutine to send a request.\n await self._forest_ready.wait() # Wait for forest to be ready.\n self._forest_ready.clear() # Reset the `Event`.\n return self._latest_forest\n\n def gather_forests(\n self,\n ) -> list[android_accessibility_forest_pb2.AndroidAccessibilityForest]:\n forests = []\n with self._lock_forests:\n forests = self._received_forests\n self._received_forests = []\n return forests\n\n def gather_events(self) -> list[a11y_pb2.EventRequest]:\n events = []\n with self._lock_events:\n events = self._received_events\n self._received_events = []\n return events\n\n def pause_and_clear(self) -> None:\n \"\"\"Temporarily stop receiving events/forests and clear the queue.\n\n Used when resetting the environment; in this case:\n - all events/forests that have been received since last timestep are things\n that happened in the last episode after its `LAST` timestep (so we should\n ignore them, done by clearing the lists).\n - we're about to receive a bunch of events/forests just as a result of\n resetting the environment. We don't want to count these either; thus we\n temporarily stop receiving new ones.\n \"\"\"\n self._paused = True\n with self._lock_forests:\n self._received_forests = []\n with self._lock_events:\n self._received_events = []\n\n def resume(self) -> None:\n \"\"\"Start receiving events/forests (e.g., after a reset).\"\"\"\n self._paused = False\n\n def _process_event(self, event: a11y_pb2.EventRequest) -> None:\n \"\"\"Adds the given event to the internal buffer of events.\"\"\"\n\n if not self._paused:\n with self._lock_events:\n self._received_events.append(event)\n\n def _process_forest(\n self, forest: android_accessibility_forest_pb2.AndroidAccessibilityForest\n ) -> None:\n \"\"\"Adds the given forest to the internal buffer of forests.\"\"\"\n\n if not self._paused:\n with self._lock_forests:\n if self._latest_forest_only:\n self._received_forests = [forest]\n else:\n self._received_forests.append(forest)\n"
},
{
"path": "android_env/wrappers/a11y/a11y_servicer_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for a11y_servicer.\"\"\"\n\nimport asyncio\nfrom collections.abc import AsyncIterator, Iterable\nfrom typing import TypeVar\nfrom unittest import IsolatedAsyncioTestCase, mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env.proto.a11y import a11y_pb2\nfrom android_env.proto.a11y import android_accessibility_forest_pb2\nfrom android_env.wrappers.a11y import a11y_servicer\nimport grpc\n\n\n_T = TypeVar('_T')\n\n\nasync def _aiter(xs: Iterable[_T]) -> AsyncIterator[_T]:\n \"\"\"Utility to make an AsyncIterator from Iterable.\"\"\"\n\n for x in xs:\n yield x\n\n\ndef one_window_one_node_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n window = forest.windows.add()\n node = window.tree.nodes.add()\n node.class_name = 'foo'\n node.is_clickable = True\n node.hint_text = 'Foo hint'\n return forest\n\n\ndef one_window_two_nodes_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n window = forest.windows.add()\n node = window.tree.nodes.add()\n node.class_name = 'bar'\n node.is_clickable = True\n node.hint_text = 'Bar hint'\n node = window.tree.nodes.add()\n node.class_name = 'bar'\n node.is_clickable = False\n node.hint_text = 'Bar hint 2'\n return forest\n\n\ndef empty_dict() -> dict[str, str]:\n return {}\n\n\ndef single_item_dict_with_special_chars() -> dict[str, str]:\n return {'foo': 'bar\\r\\t\\nbaz'}\n\n\nclass A11yServicerTest(parameterized.TestCase, IsolatedAsyncioTestCase):\n\n def test_servicer_sendforest(self):\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer()\n servicer.resume()\n response = servicer.SendForest(one_window_one_node_forest(), mock_context)\n self.assertEqual(response.error, '')\n response = servicer.SendForest(one_window_two_nodes_forest(), mock_context)\n self.assertEqual(response.error, '')\n forests = servicer.gather_forests()\n self.assertLen(forests, 2)\n self.assertEqual(forests[0], one_window_one_node_forest())\n self.assertEqual(forests[1], one_window_two_nodes_forest())\n\n async def test_servicer_bidi_forests(self):\n \"\"\"Checks that the bidirectional interface accepts forests.\"\"\"\n\n # Arrange.\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer()\n\n # Act.\n servicer.resume()\n responses = [\n x\n async for x in servicer.Bidi(\n _aiter([\n a11y_pb2.ClientToServer(\n event=a11y_pb2.EventRequest(\n event=single_item_dict_with_special_chars()\n )\n ),\n a11y_pb2.ClientToServer(forest=one_window_two_nodes_forest()),\n ]),\n mock_context,\n )\n ]\n forest = await servicer.get_forest()\n\n # Assert.\n self.assertEqual(responses[0], a11y_pb2.ServerToClient())\n self.assertEqual(responses[1], a11y_pb2.ServerToClient())\n self.assertIsNotNone(forest)\n self.assertEqual(forest, one_window_two_nodes_forest())\n\n def test_servicer_sendforest_latest_only(self):\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer(latest_forest_only=True)\n servicer.resume()\n response = servicer.SendForest(one_window_one_node_forest(), mock_context)\n self.assertEqual(response.error, '')\n response = servicer.SendForest(one_window_two_nodes_forest(), mock_context)\n self.assertEqual(response.error, '')\n forests = servicer.gather_forests()\n self.assertLen(forests, 1)\n self.assertEqual(forests[0], one_window_two_nodes_forest())\n\n def test_servicer_sendevent(self):\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer()\n servicer.resume()\n response = servicer.SendEvent(\n a11y_pb2.EventRequest(event=empty_dict()), mock_context\n )\n self.assertEqual(response.error, '')\n response = servicer.SendEvent(\n a11y_pb2.EventRequest(event=single_item_dict_with_special_chars()),\n mock_context,\n )\n self.assertEqual(response.error, '')\n events = servicer.gather_events()\n self.assertLen(events, 2)\n self.assertEqual(events[0].event, empty_dict())\n self.assertEqual(events[1].event, single_item_dict_with_special_chars())\n\n async def test_servicer_bidi_events(self):\n \"\"\"Checks that the bidirectional interface accepts events.\"\"\"\n\n # Arrange.\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer()\n\n # Act.\n servicer.resume()\n responses = [\n x\n async for x in servicer.Bidi(\n _aiter([\n a11y_pb2.ClientToServer(\n event=a11y_pb2.EventRequest(event=empty_dict())\n ),\n a11y_pb2.ClientToServer(\n event=a11y_pb2.EventRequest(\n event=single_item_dict_with_special_chars()\n )\n ),\n ]),\n mock_context,\n )\n ]\n events = servicer.gather_events()\n\n # Assert.\n self.assertEqual(responses[0], a11y_pb2.ServerToClient())\n self.assertEqual(responses[1], a11y_pb2.ServerToClient())\n self.assertLen(events, 2)\n self.assertEqual(events[0].event, empty_dict())\n self.assertEqual(events[1].event, single_item_dict_with_special_chars())\n\n def test_servicer_pause_and_clear_pauses(self):\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer()\n servicer.resume()\n servicer.pause_and_clear()\n response = servicer.SendEvent(\n a11y_pb2.EventRequest(event=empty_dict()), mock_context\n )\n self.assertEqual(response.error, '')\n response = servicer.SendForest(one_window_one_node_forest(), mock_context)\n self.assertEqual(response.error, '')\n events = servicer.gather_events()\n self.assertEmpty(events)\n forests = servicer.gather_forests()\n self.assertEmpty(forests)\n\n def test_servicer_pause_and_clear_clears(self):\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n servicer = a11y_servicer.A11yServicer()\n servicer.resume()\n response = servicer.SendEvent(\n a11y_pb2.EventRequest(event=empty_dict()), mock_context\n )\n self.assertEqual(response.error, '')\n response = servicer.SendForest(one_window_one_node_forest(), mock_context)\n self.assertEqual(\n response.error,\n '',\n )\n servicer.pause_and_clear()\n events = servicer.gather_events()\n self.assertEmpty(events)\n forests = servicer.gather_forests()\n self.assertEmpty(forests)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/a11y_grpc_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wraps AndroidEnv to retrieve accessibility messages from gRPC.\"\"\"\n\nfrom concurrent import futures\nimport time\nfrom typing import Any\n\nimport urllib.request\n\nfrom absl import logging\nfrom android_env import env_interface\nfrom android_env.components import action_type as android_action_type_lib\nfrom android_env.proto import adb_pb2\nfrom android_env.proto.a11y import a11y_pb2_grpc\nfrom android_env.wrappers import base_wrapper\nfrom android_env.wrappers.a11y import a11y_events\nfrom android_env.wrappers.a11y import a11y_forests\nfrom android_env.wrappers.a11y import a11y_servicer\nimport dm_env\nimport grpc\nimport numpy as np\nimport portpicker\n\n\ndef _get_accessibility_forwarder_apk() -> bytes:\n logging.info('Downloading accessibility forwarder apk....')\n with urllib.request.urlopen(\n 'https://storage.googleapis.com/android_env-tasks/2024.05.13-accessibility_forwarder.apk'\n ) as response:\n return response.read()\n\n\nclass EnableNetworkingError(ValueError):\n pass\n\n\nclass A11yGrpcWrapper(base_wrapper.BaseWrapper):\n \"\"\"Wrapper which receives A11y events and forests over gRPC.\n\n A11y forest protobufs and event dicts are sent from the Android emulator via\n gRPC from the `AccessibilityForwarder` (for use in developing reward\n functions, etc). This wrapper constructs a server which receives these\n messages and channels them into `task_extras`.\n\n The downside of forwarding this information through gRPC is that no messages\n will be sent if networking is turned off (e.g., if the AVD is in airplane\n mode). To mitigate this problem, the `AccessibilityForwarder` logs an error\n message if it fails to contact the server. This wrapper monitors the logs for\n such error messages, and attempts (in another thread, to not block environment\n transitions) to reconnect the AVD to the network. If this fails to fix the\n problem, this wrapper ends the episode.\n\n This wrapper is implemented to be robust to multiple upstream callers of\n `task_extras`, and to ensure they each receive the same extras at every\n timestep. Thus, the logic is the following:\n * New a11y events/forests are fetched during `reset` and `step`, *not* during\n `task_extras()` calls.\n * If no one has called `task_extras()` since the last `step` or `reset`, the\n extras are accumulated (so that no extras are missed because someone called\n `step()` twice without calling `task_extras()`).\n * If someone *has* called `task_extras()` since last step, the newly fetched\n extras replace the old extras.\n \"\"\"\n\n def __init__(\n self,\n env: env_interface.AndroidEnvInterface,\n disable_other_network_traffic: bool = False,\n install_a11y_forwarding: bool = False,\n start_a11y_service: bool = True,\n enable_a11y_tree_info: bool = False,\n add_latest_a11y_info_to_obs: bool = False,\n a11y_info_timeout: float | None = None,\n max_enable_networking_attempts: int = 10,\n latest_a11y_info_only: bool = False,\n grpc_server_ip: str = '10.0.2.2',\n ):\n \"\"\"Initializes wrapper.\n\n Args:\n env: Environment to wrap.\n disable_other_network_traffic: When True, all network traffic, other than\n the connection to the servicer, is disabled. NOTE: This requires root\n access on the device (i.e. it uses the `su` command). An\n `AdbControllerError` exception will be raised if the underlying command\n fails.\n install_a11y_forwarding: If True, the wrapper handles the installation of\n all packages required for the servicer to collect a11y information.\n start_a11y_service: If True, starts the a11y forwarding services. NOTE:\n The packages must be installed beforehand, e.g., using the\n install_a11y_forwarding flag.\n enable_a11y_tree_info: When False, this wrapper collects only a11y events\n and not a11y tree.\n add_latest_a11y_info_to_obs: When True, the latest observed a11y forest is\n added to the observation.\n a11y_info_timeout: When larger than zero and add_latest_a11y_info_to_obs\n is set to True, the wrapper will wait the corresponding amount of time,\n measured in seconds, to collect the latest a11y forest.\n max_enable_networking_attempts: When the a11y gRPC service fails to\n provide a11y information, we attempt this many times to re-enable the\n networking. If all these attempts fail, fetching task_extras will raise\n an EnableNetworkingError.\n latest_a11y_info_only: When True, the a11y servicer is setup to save only\n the latest tree it has received from the Android app.\n grpc_server_ip: The IP address of the gRPC server which will be\n broadcasted to the AccessibilityForwarder app where it should log the\n a11y info. By default, this is set to the IP address of the AVD's host\n machine which is 10.0.2.2: See\n https://developer.android.com/studio/run/emulator-networking#networkaddresses.\n \"\"\"\n self._env = env\n self._grpc_server_ip = grpc_server_ip\n if install_a11y_forwarding:\n self._install_a11y_forwarding_apk()\n time.sleep(10.0)\n if start_a11y_service:\n self._start_a11y_services()\n time.sleep(3.0)\n if enable_a11y_tree_info:\n self._enable_a11y_tree_logs()\n self._relaunch_count = 0\n self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))\n self._servicer = a11y_servicer.A11yServicer(\n latest_forest_only=latest_a11y_info_only\n )\n a11y_pb2_grpc.add_A11yServiceServicer_to_server(\n self._servicer, self._server\n )\n server_credentials = grpc.local_server_credentials()\n self._port = portpicker.pick_unused_port()\n logging.info('Using port %s', self._port)\n uri_address = f'[::]:{self._port}'\n self._server.add_secure_port(uri_address, server_credentials)\n logging.info('Starting server')\n self._server.start()\n logging.info('Server now running.')\n\n self._max_enable_networking_attempts = max_enable_networking_attempts\n self._reset_enable_networking_attempts()\n\n self._disable_other_network_traffic = disable_other_network_traffic\n self._should_accumulate = False\n self._accumulated_extras = None\n self._add_latest_a11y_info_to_obs = add_latest_a11y_info_to_obs\n self._a11y_info_timeout = a11y_info_timeout\n self._parent_action_spec = self._env.action_spec()\n if self._a11y_info_timeout is not None and self._a11y_info_timeout > 0.0:\n if 'action_type' not in self._parent_action_spec.keys():\n raise ValueError(\n 'action_type not in the parent action spec: '\n f'{self._parent_action_spec}. This is a strong requirement when '\n f'a11y_info_timeout = {a11y_info_timeout} > 0'\n )\n\n def _start_a11y_services(self) -> None:\n \"\"\"Starts the accessibility forwarder services.\n\n Raises:\n RuntimeError: If accessibility service is not started.\n \"\"\"\n start_service_request = adb_pb2.AdbRequest(\n settings=adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SECURE,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(\n key='enabled_accessibility_services',\n value=(\n 'com.google.androidenv.accessibilityforwarder/com.google.'\n 'androidenv.accessibilityforwarder.AccessibilityForwarder'\n ),\n ),\n )\n )\n start_service_response = self._env.execute_adb_call(start_service_request)\n if start_service_response.status != adb_pb2.AdbResponse.Status.OK:\n raise RuntimeError(\n 'Could not start accessibility forwarder '\n 'service: '\n f'{start_service_response}.'\n )\n\n def _install_a11y_forwarding_apk(self) -> None:\n \"\"\"Enables accessibility information forwarding.\"\"\"\n a11y_fwd_apk = _get_accessibility_forwarder_apk()\n # Install and setup the Accesssibility Forwarder.\n install_request = adb_pb2.AdbRequest(\n install_apk=adb_pb2.AdbRequest.InstallApk(\n blob=adb_pb2.AdbRequest.InstallApk.Blob(contents=a11y_fwd_apk),\n )\n )\n install_response = self._env.execute_adb_call(install_request)\n if install_response.status != adb_pb2.AdbResponse.Status.OK:\n raise ValueError(\n f'Could not install accessibility_forwarder.apk: {install_response}.'\n )\n\n def _enable_a11y_tree_logs(self) -> None:\n enable_tree_logs_request = adb_pb2.AdbRequest(\n send_broadcast=adb_pb2.AdbRequest.SendBroadcast(\n action=(\n 'accessibility_forwarder.intent.action.'\n 'ENABLE_ACCESSIBILITY_TREE_LOGS'\n ),\n component=(\n 'com.google.androidenv.accessibilityforwarder/com.google.androidenv.accessibilityforwarder.FlagsBroadcastReceiver'\n ),\n )\n )\n enable_tree_logs_response = self._env.execute_adb_call(\n enable_tree_logs_request\n )\n if enable_tree_logs_response.status != adb_pb2.AdbResponse.Status.OK:\n raise ValueError(\n 'Could not enable accessibility tree logging: '\n f'{enable_tree_logs_response}.'\n )\n\n def _reset_enable_networking_attempts(self) -> None:\n self._enable_networking_attempts_left = self._max_enable_networking_attempts\n self._enabling_networking_future = None\n self._a11y_exception = None\n\n def get_port(self):\n return self._port\n\n def close(self):\n self._server.stop(None)\n logging.info('gRPC server stopped')\n self._env.close()\n\n def attempt_enable_networking(self) -> None:\n \"\"\"Attempts to turn on networking within the Android device.\n\n Attempt to turn on the networking in the Android device, by:\n - turning off airplane mode;\n - turning on the wifi connection.\n \"\"\"\n self.execute_adb_call(\n adb_pb2.AdbRequest(\n settings=adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(\n key='airplane_mode_on', value='0'\n ),\n )\n )\n )\n time.sleep(1.0)\n self.execute_adb_call(\n adb_pb2.AdbRequest(\n generic=adb_pb2.AdbRequest.GenericRequest(\n args=[\n 'shell',\n 'svc',\n 'wifi',\n 'enable',\n ]\n )\n )\n )\n time.sleep(1.0)\n\n def _configure_grpc(self) -> None:\n \"\"\"Configure networking and set the gRPC ip and port on AVD or device.\"\"\"\n\n if self._disable_other_network_traffic:\n self.execute_adb_call(\n adb_pb2.AdbRequest(\n generic=adb_pb2.AdbRequest.GenericRequest(\n args=[\n 'shell',\n 'su',\n '0',\n 'iptables',\n '-A',\n 'OUTPUT',\n '-p',\n 'tcp',\n '-d',\n self._grpc_server_ip,\n '--dport',\n str(self._port),\n '-j',\n 'ACCEPT',\n ]\n )\n )\n )\n time.sleep(3.0)\n self.execute_adb_call(\n adb_pb2.AdbRequest(\n generic=adb_pb2.AdbRequest.GenericRequest(\n args=[\n 'shell',\n 'su',\n '0',\n 'iptables',\n '-A',\n 'OUTPUT',\n '-j',\n 'DROP',\n ]\n )\n )\n )\n time.sleep(3.0)\n\n self.execute_adb_call(\n adb_pb2.AdbRequest(\n settings=adb_pb2.AdbRequest.SettingsRequest(\n name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,\n put=adb_pb2.AdbRequest.SettingsRequest.Put(\n key='no_proxy', value=f'{self._grpc_server_ip}:{self._port}'\n ),\n )\n )\n )\n self.attempt_enable_networking()\n self.execute_adb_call(\n adb_pb2.AdbRequest(\n send_broadcast=adb_pb2.AdbRequest.SendBroadcast(\n action=(\n 'accessibility_forwarder.intent.action.SET_GRPC --ei'\n f' \"port\" {self._port} --es \"host\" {self._grpc_server_ip}'\n ),\n component=(\n 'com.google.androidenv.accessibilityforwarder/com.google.androidenv.accessibilityforwarder.FlagsBroadcastReceiver'\n ),\n )\n )\n )\n\n def _accumulate_and_return_a11y_info(\n self, timer: float | None = None, get_env_observation: bool = True\n ) -> dict[str, Any]:\n \"\"\"Accumulates and returns the latest a11y tree info and observation.\n\n Args:\n timer: If larger than 0, the system will wait this long for a11y info to\n accumulate before it returns a value.\n get_env_observation: If False, the corresponding observation is not\n introduced here.\n\n Returns:\n a dict with a11y forest under key 'a11y_forest'. All other fields will\n provide the observation, if requested.\n \"\"\"\n timer = timer or 0.0\n if timer > 0.0:\n time.sleep(timer)\n\n if get_env_observation:\n # Fetch observation.\n new_ts = self._env.step({\n 'action_type': np.array(\n android_action_type_lib.ActionType.REPEAT,\n dtype=self._parent_action_spec['action_type'].dtype,\n ),\n })\n observation = new_ts.observation\n else:\n observation = {}\n\n extras = self.accumulate_new_extras()\n forests = a11y_forests.extract_forests_from_task_extras(extras)\n if forests:\n observation['a11y_forest'] = forests[-1]\n else:\n observation['a11y_forest'] = None\n return observation\n\n def _fetch_task_extras_and_update_observation(\n self, observation: dict[str, Any], timeout: float = 0.0\n ) -> dict[str, Any]:\n if timeout > 0.0:\n observation = self._accumulate_and_return_a11y_info(\n timeout, get_env_observation=True\n )\n if not self._add_latest_a11y_info_to_obs:\n observation.pop('a11y_forest')\n else:\n new_obs = self._accumulate_and_return_a11y_info(get_env_observation=False)\n if self._add_latest_a11y_info_to_obs:\n observation.update(new_obs)\n return observation\n\n def reset(self) -> dm_env.TimeStep:\n self._reset_enable_networking_attempts()\n self._servicer.pause_and_clear()\n timestep = self._env.reset()\n self._servicer.resume()\n if self._env.stats()['relaunch_count'] > self._relaunch_count:\n self._configure_grpc()\n self._relaunch_count = self._env.stats()['relaunch_count']\n self._accumulated_extras = {}\n timeout = self._a11y_info_timeout or 0.0\n new_observation = self._fetch_task_extras_and_update_observation(\n timestep.observation, timeout\n )\n timestep = timestep._replace(observation=new_observation)\n return timestep\n\n def step(self, action: Any) -> dm_env.TimeStep:\n timeout = float(action.pop('wait_time', self._a11y_info_timeout or 0.0))\n timestep = self._env.step(action)\n new_observation = self._fetch_task_extras_and_update_observation(\n timestep.observation, timeout=timeout\n )\n timestep = timestep._replace(observation=new_observation)\n return timestep\n\n def accumulate_new_extras(self) -> dict[str, Any]:\n new_extras = self._fetch_task_extras()\n if self._should_accumulate:\n for key in new_extras:\n if key in self._accumulated_extras:\n self._accumulated_extras[key] = np.concatenate(\n (self._accumulated_extras[key], new_extras[key]), axis=0\n )\n else:\n self._accumulated_extras[key] = new_extras[key]\n else:\n self._accumulated_extras = new_extras\n self._should_accumulate = True\n return self._accumulated_extras\n\n def _fetch_task_extras(self) -> dict[str, Any]:\n \"\"\"Fetches task_extras from the services.\n\n NOTE: If you want to access the latest a11y information, please use\n accumulate_and_return_a11y_info instead. This function has the side effect\n of clearing the content from the servicer, hence all the a11y info returned\n here won't be accumulated.\n\n Returns:\n A dict with the corresponding task_extras.\n\n Raises:\n EnableNetworkingError: after a fixed number of attempts to revive the a11y\n services by re-enabling the network connection.\n \"\"\"\n base_extras = self._env.task_extras(latest_only=False).copy()\n # If the previous future is done, reset it to the initial state.\n if (\n self._enabling_networking_future is not None\n and self._enabling_networking_future.done()\n ):\n self._enabling_networking_future = None\n self._enable_networking_attempts_left -= 1\n logging.info('Finished enabling networking.')\n\n if (\n self._enabling_networking_future is None\n and 'exception' in base_extras\n and base_extras['exception'].shape[0]\n ):\n self._a11y_exception = base_extras['exception']\n logging.warning(\n 'AccessibilityForwarder logged exceptions: %s', self._a11y_exception\n )\n if self._enable_networking_attempts_left > 0:\n logging.warning(\n 'Attempting to enable networking. %s attempts left.',\n self._enable_networking_attempts_left - 1,\n )\n executor = futures.ThreadPoolExecutor(max_workers=1)\n self._enabling_networking_future = executor.submit(\n self.attempt_enable_networking\n )\n else:\n raise EnableNetworkingError(\n 'A11y service failed multiple times with'\n f' exception.{self._a11y_exception}.'\n )\n\n forests = self._servicer.gather_forests()\n if forests:\n base_extras.update(a11y_forests.package_forests_to_task_extras(forests))\n self._reset_enable_networking_attempts()\n events = self._servicer.gather_events()\n if events:\n base_extras.update(a11y_events.package_events_to_task_extras(events))\n self._reset_enable_networking_attempts()\n return base_extras\n\n def task_extras(self, latest_only: bool = False) -> dict[str, Any]:\n if self._accumulated_extras is None:\n raise RuntimeError('You must call .reset() before calling .task_extras()')\n self._should_accumulate = False\n extras = self._accumulated_extras.copy()\n if latest_only:\n a11y_events.keep_latest_event_only(extras)\n a11y_forests.keep_latest_forest_only(extras)\n return extras\n"
},
{
"path": "android_env/wrappers/a11y_grpc_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for a11y_grpc_wrapper.\"\"\"\n\nimport time\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env import env_interface\nfrom android_env.proto import adb_pb2\nfrom android_env.proto.a11y import a11y_pb2\nfrom android_env.proto.a11y import a11y_pb2_grpc\nfrom android_env.proto.a11y import android_accessibility_forest_pb2\nfrom android_env.wrappers import a11y_grpc_wrapper\nimport dm_env\nimport grpc\nimport numpy as np\n\n\ndef empty_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n return android_accessibility_forest_pb2.AndroidAccessibilityForest()\n\n\ndef one_empty_window_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n _ = forest.windows.add()\n return forest\n\n\ndef one_window_one_node_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n window = forest.windows.add()\n node = window.tree.nodes.add()\n node.class_name = 'foo'\n node.is_clickable = True\n node.hint_text = 'Foo hint'\n return forest\n\n\ndef one_window_two_nodes_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n window = forest.windows.add()\n node = window.tree.nodes.add()\n node.class_name = 'bar'\n node.is_clickable = True\n node.hint_text = 'Bar hint'\n node = window.tree.nodes.add()\n node.class_name = 'bar'\n node.is_clickable = False\n node.hint_text = 'Bar hint 2'\n return forest\n\n\ndef three_windows_forest() -> (\n android_accessibility_forest_pb2.AndroidAccessibilityForest\n):\n forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()\n _ = forest.windows.add()\n window = forest.windows.add()\n node = window.tree.nodes.add()\n node.class_name = 'foo'\n node.is_clickable = True\n node.hint_text = 'hint'\n window = forest.windows.add()\n node = window.tree.nodes.add()\n node.class_name = 'baz'\n node.is_clickable = True\n node.hint_text = 'hint'\n node = window.tree.nodes.add()\n node.class_name = 'foobar'\n node.is_clickable = False\n node.hint_text = 'hint'\n return forest\n\n\ndef empty_dict() -> dict[str, str]:\n return {}\n\n\ndef single_item_dict() -> dict[str, str]:\n return {'foo': 'bar'}\n\n\ndef several_long_items_dict() -> dict[str, str]:\n return {\n 'first_key': 'Lorem ipsum ' * 100,\n 'second_key': 'the beginning is the end is' * 100,\n }\n\n\ndef single_item_dict_with_special_chars() -> dict[str, str]:\n return {'foo': 'bar\\r\\t\\nbaz'}\n\n\ndef _ok_response():\n return adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)\n\n\nclass A11yGrpcWrapperTest(parameterized.TestCase):\n\n def test_server(self):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.task_extras.return_value = {}\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n wrapped_env.reset()\n channel_creds = grpc.local_channel_credentials()\n with grpc.secure_channel(\n f'[::]:{wrapped_env.get_port()}', channel_creds\n ) as channel:\n grpc.channel_ready_future(channel).result()\n stub = a11y_pb2_grpc.A11yServiceStub(channel)\n stub.SendForest(one_window_one_node_forest())\n stub.SendForest(one_window_two_nodes_forest())\n wrapped_env.step({})\n extras = wrapped_env.task_extras(latest_only=False)\n self.assertIn('accessibility_tree', extras)\n self.assertEqual(extras['accessibility_tree'].shape[0], 2)\n\n # tests of fetch_task_extras:\n # exception occurs (ensure attempt to enable networking) and recovers\n # exception occurs and enable networking doesn't help\n # exception occurs twice but with a forest sent between\n\n @parameterized.named_parameters(\n ('no_events_or_forests', [], []),\n (\n 'no_events',\n [],\n [one_window_one_node_forest(), one_window_two_nodes_forest()],\n ),\n ('no_forests', [empty_dict(), single_item_dict()], []),\n (\n 'events_and_forests',\n [empty_dict(), single_item_dict()],\n [one_window_one_node_forest(), one_window_two_nodes_forest()],\n ),\n )\n @mock.patch.object(time, 'sleep', autospec=True)\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_fetch_task_extras(\n self,\n received_events,\n received_forests,\n mock_server,\n mock_add_servicer,\n mock_sleep,\n ):\n del mock_server, mock_add_servicer, mock_sleep\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar', 'baz'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n }\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n wrapped_env.reset()\n for forest in received_forests:\n wrapped_env._servicer.SendForest(forest, mock_context)\n for event in received_events:\n wrapped_env._servicer.SendEvent(\n a11y_pb2.EventRequest(event=event), mock_context\n )\n with mock.patch.object(\n wrapped_env, 'attempt_enable_networking'\n ) as mock_attempt_enable_networking:\n extras = wrapped_env._fetch_task_extras()\n mock_attempt_enable_networking.assert_not_called()\n self.assertIn('foo', extras)\n np.testing.assert_array_equal(extras['foo'], ['bar', 'baz'])\n self.assertIn('some_key', extras)\n np.testing.assert_array_equal(extras['some_key'], ['some_value'])\n if received_events:\n self.assertIn('full_event', extras)\n self.assertLen(extras['full_event'], len(received_events))\n for i, event in enumerate(received_events):\n event = a11y_pb2.EventRequest(event=event)\n np.testing.assert_array_equal(extras['full_event'][i], event)\n else:\n self.assertNotIn('full_event', extras)\n if received_forests:\n self.assertIn('accessibility_tree', extras)\n self.assertLen(extras['accessibility_tree'], len(received_forests))\n for i, forest in enumerate(received_forests):\n np.testing.assert_array_equal(extras['accessibility_tree'][i], forest)\n else:\n self.assertNotIn('accessibility_tree', extras)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_fetch_task_extras_enable_networking(\n self,\n mock_server,\n mock_add_servicer,\n mock_sleep,\n ):\n del mock_server, mock_add_servicer, mock_sleep\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n 'exception': np.array(['fake exception'], dtype='U'),\n }\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n with mock.patch.object(\n wrapped_env, 'attempt_enable_networking'\n ) as mock_attempt_enable_networking:\n extras = wrapped_env._fetch_task_extras()\n self.assertNotIn('accessibility_tree', extras)\n self.assertNotIn('full_event', extras)\n future = wrapped_env._enabling_networking_future\n if future is not None:\n future.result()\n mock_attempt_enable_networking.assert_called_once()\n\n @mock.patch.object(time, 'sleep', autospec=True)\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_fetch_task_extras_enable_networking_twice(\n self,\n mock_server,\n mock_add_servicer,\n mock_sleep,\n ):\n del mock_server, mock_add_servicer, mock_sleep\n mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n }\n\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n wrapped_env.reset()\n\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n 'exception': np.array(['fake exception'], dtype='U'),\n }\n with mock.patch.object(\n wrapped_env, 'attempt_enable_networking'\n ) as mock_attempt_enable_networking:\n extras = wrapped_env._fetch_task_extras()\n self.assertNotIn('accessibility_tree', extras)\n self.assertNotIn('full_event', extras)\n future = wrapped_env._enabling_networking_future\n if future is not None:\n future.result()\n mock_attempt_enable_networking.assert_called_once()\n # Fixed networking; send a forest so the wrapper knows it worked.\n wrapped_env._servicer.SendForest(one_window_one_node_forest(), mock_context)\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n }\n extras = wrapped_env._fetch_task_extras()\n self.assertIn('accessibility_tree', extras)\n self.assertEqual(extras['accessibility_tree'].shape[0], 1)\n self.assertEqual(\n extras['accessibility_tree'][0], one_window_one_node_forest()\n )\n\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n 'exception': np.array(['fake exception'], dtype='U'),\n }\n with mock.patch.object(\n wrapped_env, 'attempt_enable_networking'\n ) as mock_attempt_enable_networking:\n extras = wrapped_env._fetch_task_extras()\n self.assertNotIn('accessibility_tree', extras)\n self.assertNotIn('full_event', extras)\n future = wrapped_env._enabling_networking_future\n if future is not None:\n future.result()\n mock_attempt_enable_networking.assert_called_once()\n\n @mock.patch.object(time, 'sleep', autospec=True)\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_task_extras_raises_a11y_info_exception(\n self, mock_sleep, mock_add_servicer, mock_server\n ):\n del mock_server, mock_add_servicer, mock_sleep\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n }\n\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.reset.return_value = dm_env.restart(observation={'dummy': 42})\n base_env.step.return_value = dm_env.transition(\n observation={'dummy': 42}, reward=0.0\n )\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env,\n add_latest_a11y_info_to_obs=True,\n max_enable_networking_attempts=1,\n )\n wrapped_env.reset()\n\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n 'exception': np.array(['fake exception'], dtype='U'),\n }\n with mock.patch.object(\n wrapped_env, 'attempt_enable_networking'\n ) as mock_attempt_enable_networking:\n extras = wrapped_env._fetch_task_extras()\n self.assertNotIn('accessibility_tree', extras)\n self.assertNotIn('full_event', extras)\n # Wait for the the attempt to finish.\n future = wrapped_env._enabling_networking_future\n if future is not None:\n future.result()\n mock_attempt_enable_networking.assert_called_once()\n # The _fetch_task_extras() call inside the next step will force a restart\n self.assertRaises(\n a11y_grpc_wrapper.EnableNetworkingError, wrapped_env.step, {}\n )\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_configure_grpc(\n self,\n mock_server,\n mock_add_servicer,\n ):\n del mock_server, mock_add_servicer\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.task_extras.return_value = {\n 'foo': np.array(['bar'], dtype='U'),\n 'some_key': np.array(['some_value'], dtype='U'),\n }\n\n base_env.stats.return_value = {'relaunch_count': 1}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n with mock.patch.object(\n wrapped_env, '_configure_grpc'\n ) as mock_configure_grpc:\n wrapped_env.reset()\n mock_configure_grpc.assert_called_once()\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_task_extras_raises_before_reset(\n self, unused_mock_server, unused_mock_add_servicer\n ):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n with self.assertRaisesRegex(\n RuntimeError,\n r'You must call \\.reset\\(\\) before calling \\.task_extras\\(\\)',\n ):\n wrapped_env.task_extras(latest_only=False)\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_extras_accumulate_between_steps(\n self, mock_server, mock_add_servicer\n ):\n del mock_server, mock_add_servicer\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.reset.return_value = dm_env.restart(observation={'dummy': 42})\n base_env.step.return_value = dm_env.transition(\n observation={'dummy': 42}, reward=0.0\n )\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env, add_latest_a11y_info_to_obs=True\n )\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.return_value = {\n 'full_event': np.array(single_item_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(empty_forest(), ndmin=1, dtype=object),\n }\n timestep = wrapped_env.reset()\n self.assertIn('a11y_forest', timestep.observation)\n self.assertEqual(timestep.observation['a11y_forest'], empty_forest())\n wrapped_env._fetch_task_extras.return_value = {\n 'full_event': np.array(empty_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(\n one_window_two_nodes_forest(), ndmin=1, dtype=object\n ),\n }\n timestep = wrapped_env.step({})\n self.assertIn('a11y_forest', timestep.observation)\n self.assertEqual(\n timestep.observation['a11y_forest'], one_window_two_nodes_forest()\n )\n timestep = wrapped_env.step({})\n self.assertIn('a11y_forest', timestep.observation)\n self.assertEqual(\n timestep.observation['a11y_forest'], one_window_two_nodes_forest()\n )\n wrapped_env._fetch_task_extras.return_value = {\n 'full_event': np.array(single_item_dict(), ndmin=1, dtype=object),\n }\n timestep = wrapped_env.step({})\n self.assertIn('a11y_forest', timestep.observation)\n self.assertEqual(\n timestep.observation['a11y_forest'], one_window_two_nodes_forest()\n )\n expected_task_extras = {\n 'full_event': np.array(\n [\n single_item_dict(),\n empty_dict(),\n empty_dict(),\n single_item_dict(),\n ],\n dtype=object,\n ),\n 'accessibility_tree': np.array(\n [\n empty_forest(),\n one_window_two_nodes_forest(),\n one_window_two_nodes_forest(),\n ],\n dtype=object,\n ),\n }\n expected_task_extras_latest = {\n 'full_event': np.array([single_item_dict()], dtype=object),\n 'accessibility_tree': np.array(\n [one_window_two_nodes_forest()], dtype=object\n ),\n }\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n task_extras = wrapped_env.task_extras(latest_only=True)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras_latest['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras_latest['accessibility_tree'],\n )\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_a11y_info_disabled(\n self,\n unused_mock_server,\n unused_mock_add_servicer,\n ):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.action_spec.return_value = {\n 'action_type': dm_env.specs.Array(shape=(), dtype=np.int32)\n }\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.reset.return_value = dm_env.restart(observation={'dummy': 42})\n base_env.step.return_value = dm_env.transition(\n observation={'dummy': 42}, reward=0.0\n )\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env, add_latest_a11y_info_to_obs=False, a11y_info_timeout=1.0\n )\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.return_value = {\n 'accessibility_tree': np.array(empty_forest(), ndmin=1, dtype=object),\n }\n timestep = wrapped_env.reset()\n self.assertNotIn('a11y_forest', timestep.observation)\n timestep = wrapped_env.step({})\n self.assertNotIn('a11y_forest', timestep.observation)\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_a11y_info_with_timer_info_present(\n self,\n unused_mock_server,\n unused_mock_add_servicer,\n ):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.action_spec.return_value = {\n 'action_type': dm_env.specs.Array(shape=(), dtype=np.int32)\n }\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.reset.return_value = dm_env.restart(observation={'dummy': 42})\n base_env.step.return_value = dm_env.transition(\n observation={'dummy': 42}, reward=0.0\n )\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env, add_latest_a11y_info_to_obs=True, a11y_info_timeout=1.0\n )\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.side_effect = [{\n 'accessibility_tree': np.array(empty_forest(), ndmin=1, dtype=object),\n }]\n timestep = wrapped_env.reset()\n self.assertIn('a11y_forest', timestep.observation)\n self.assertEqual(timestep.observation['a11y_forest'], empty_forest())\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_a11y_info_with_timer_task_extra_returned(\n self, unused_mock_server, unused_mock_add_servicer, unused_mock_sleep\n ):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.action_spec.return_value = {\n 'action_type': dm_env.specs.Array(shape=(), dtype=np.int32)\n }\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.reset.return_value = dm_env.restart(observation={'dummy': 42})\n base_env.step.return_value = dm_env.transition(\n observation={'dummy': 42}, reward=0.0\n )\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env, add_latest_a11y_info_to_obs=True, a11y_info_timeout=1.0\n )\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.side_effect = [\n {\n 'accessibility_tree': np.array(\n empty_forest(), ndmin=1, dtype=object\n ),\n },\n ]\n timestep = wrapped_env.reset()\n self.assertIn('a11y_forest', timestep.observation)\n self.assertEqual(timestep.observation['a11y_forest'], empty_forest())\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_a11y_info_with_timer_from_action(\n self, unused_mock_server, unused_mock_add_servicer, mock_sleep\n ):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.action_spec.return_value = {\n 'action_type': dm_env.specs.Array(shape=(), dtype=np.int32)\n }\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.reset.return_value = dm_env.restart(observation={'dummy': 42})\n base_env.step.return_value = dm_env.transition(\n observation={'dummy': 42}, reward=0.0\n )\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env, add_latest_a11y_info_to_obs=True, a11y_info_timeout=0.0\n )\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.side_effect = [\n {\n 'accessibility_tree': np.array(\n empty_forest(), ndmin=1, dtype=object\n ),\n },\n ]\n timestep = wrapped_env.step(action={'wait_time': 1.0})\n self.assertIn('a11y_forest', timestep.observation)\n mock_sleep.assert_called_once()\n self.assertEqual(timestep.observation['a11y_forest'], empty_forest())\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_task_extras_same_between_calls(self, mock_server, mock_add_servicer):\n del mock_server, mock_add_servicer\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n expected_task_extras = {\n 'full_event': np.array(single_item_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(empty_forest(), ndmin=1, dtype=object),\n }\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.return_value = expected_task_extras\n wrapped_env.reset()\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n expected_task_extras = {\n 'full_event': np.array(empty_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(\n one_window_two_nodes_forest(), ndmin=1, dtype=object\n ),\n }\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n wrapped_env._fetch_task_extras.return_value = expected_task_extras\n wrapped_env.step({})\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n @mock.patch.object(\n a11y_pb2_grpc, 'add_A11yServiceServicer_to_server', autospec=True\n )\n @mock.patch.object(grpc, 'server', autospec=True)\n def test_task_extras_clear_if_called_between_step(\n self, mock_server, mock_add_servicer\n ):\n del mock_server, mock_add_servicer\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.stats.return_value = {'relaunch_count': 0}\n base_env.execute_adb_call.return_value = _ok_response()\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(base_env)\n with mock.patch.object(wrapped_env, '_fetch_task_extras'):\n expected_task_extras = {\n 'full_event': np.array(empty_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(empty_forest(), ndmin=1, dtype=object),\n }\n wrapped_env._fetch_task_extras.return_value = expected_task_extras\n wrapped_env.reset()\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n expected_task_extras = {\n 'full_event': np.array(single_item_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(empty_forest(), ndmin=1, dtype=object),\n }\n wrapped_env._fetch_task_extras.return_value = expected_task_extras\n wrapped_env.step({})\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n expected_task_extras = {\n 'full_event': np.array(empty_dict(), ndmin=1, dtype=object),\n 'accessibility_tree': np.array(\n one_window_two_nodes_forest(), ndmin=1, dtype=object\n ),\n }\n wrapped_env._fetch_task_extras.return_value = expected_task_extras\n wrapped_env.step({})\n task_extras = wrapped_env.task_extras(latest_only=False)\n np.testing.assert_equal(\n task_extras['full_event'], expected_task_extras['full_event']\n )\n np.testing.assert_equal(\n task_extras['accessibility_tree'],\n expected_task_extras['accessibility_tree'],\n )\n\n @parameterized.named_parameters(\n ('none_true', False, False, False, 0),\n ('only_install', True, False, False, 1),\n ('only_start', False, True, False, 1),\n ('only_enable_a11y_tree', False, False, True, 1),\n ('install_and_start_no_a11y_tree', True, True, False, 2),\n ('install_and_a11y_tree', True, False, True, 2),\n ('start_and_a11y_tree', False, True, True, 2),\n ('all_true', True, True, True, 3),\n )\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_apk_install_and_start(\n self,\n install_a11y_forwarding: bool,\n start_a11y_service: bool,\n enable_a11y_tree_logs: bool,\n expected_adb_calls: int,\n unused_mock_sleep,\n ):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n\n side_effects = []\n if install_a11y_forwarding:\n side_effects.append(_ok_response()) # install response\n if start_a11y_service:\n side_effects.append(_ok_response()) # start service response\n if enable_a11y_tree_logs:\n side_effects.append(_ok_response()) # enable_tree_request\n\n base_env.execute_adb_call.side_effect = side_effects\n\n _ = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env,\n install_a11y_forwarding=install_a11y_forwarding,\n start_a11y_service=start_a11y_service,\n enable_a11y_tree_info=enable_a11y_tree_logs,\n )\n self.assertEqual(base_env.execute_adb_call.call_count, expected_adb_calls)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_component_and_start(self, unused_mock_sleep):\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n\n side_effects = []\n side_effects.append(_ok_response()) # install response\n side_effects.append(_ok_response()) # start service response\n side_effects.append(_ok_response()) # enable_tree_request\n\n base_env.execute_adb_call.side_effect = side_effects\n\n _ = a11y_grpc_wrapper.A11yGrpcWrapper(\n base_env,\n install_a11y_forwarding=True,\n start_a11y_service=True,\n enable_a11y_tree_info=True,\n )\n\n # call_args returns a tuple of which the first member is a tuple containing\n # the most recent args the mock was called with, and execute_adb_call only\n # has one arg (so [0][0] to access the AdbRequest).\n self.assertEqual(\n base_env.execute_adb_call.call_args[0][0].send_broadcast.component,\n 'com.google.androidenv.accessibilityforwarder/com.google.androidenv.accessibilityforwarder.FlagsBroadcastReceiver',\n )\n\n def test_broadcast_sent_default_grpc_server_ip(self):\n \"\"\"Tests that the broadcast sets the default grpc server ip.\"\"\"\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.stats.return_value = {'relaunch_count': 1}\n\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n env=base_env,\n disable_other_network_traffic=False,\n install_a11y_forwarding=False,\n start_a11y_service=False,\n enable_a11y_tree_info=False,\n )\n wrapped_env.reset()\n\n self.assertStartsWith(\n base_env.execute_adb_call.call_args[0][0].send_broadcast.action,\n 'accessibility_forwarder.intent.action.SET_GRPC',\n )\n\n self.assertIn(\n '--es \"host\" 10.0.2.2',\n base_env.execute_adb_call.call_args[0][0].send_broadcast.action,\n )\n\n @parameterized.parameters(('127.0.0.1',), ('1.2.3.4',), 'localhost')\n def test_broadcast_sent_custom_grpc_server_ip(self, grpc_server_ip):\n \"\"\"Tests that the broadcast sets the custom grpc server ip.\"\"\"\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.stats.return_value = {'relaunch_count': 1}\n\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n env=base_env,\n disable_other_network_traffic=False,\n install_a11y_forwarding=False,\n start_a11y_service=False,\n enable_a11y_tree_info=False,\n grpc_server_ip=grpc_server_ip,\n )\n wrapped_env.reset()\n\n self.assertStartsWith(\n base_env.execute_adb_call.call_args[0][0].send_broadcast.action,\n 'accessibility_forwarder.intent.action.SET_GRPC',\n )\n self.assertIn(\n f'--es \"host\" {grpc_server_ip}',\n base_env.execute_adb_call.call_args[0][0].send_broadcast.action,\n )\n\n def test_broadcast_sent_port(self):\n \"\"\"Tests that the broadcast sets the correct port.\"\"\"\n base_env = mock.create_autospec(\n env_interface.AndroidEnvInterface, instance=True\n )\n base_env.execute_adb_call.return_value = _ok_response()\n base_env.stats.return_value = {'relaunch_count': 1}\n\n wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper(\n env=base_env,\n disable_other_network_traffic=False,\n install_a11y_forwarding=False,\n start_a11y_service=False,\n enable_a11y_tree_info=False,\n )\n wrapped_env.reset()\n\n self.assertStartsWith(\n base_env.execute_adb_call.call_args[0][0].send_broadcast.action,\n 'accessibility_forwarder.intent.action.SET_GRPC',\n )\n self.assertIn(\n f'--ei \"port\" {wrapped_env.get_port()}',\n base_env.execute_adb_call.call_args[0][0].send_broadcast.action,\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/base_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Base class for AndroidEnv wrappers.\"\"\"\n\nfrom typing import Any\n\nfrom absl import logging\nfrom android_env import env_interface\nfrom android_env.proto import adb_pb2\nfrom android_env.proto import state_pb2\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\nclass BaseWrapper(env_interface.AndroidEnvInterface):\n \"\"\"AndroidEnv wrapper.\"\"\"\n\n def __init__(self, env: env_interface.AndroidEnvInterface) -> None:\n self._env = env\n logging.info('Wrapping with %s', self.__class__.__name__)\n\n def reset(self) -> dm_env.TimeStep:\n self._reset_state()\n timestep = self._process_timestep(self._env.reset())\n return timestep\n\n def step(self, action: Any) -> dm_env.TimeStep:\n action = self._process_action(action)\n return self._process_timestep(self._env.step(action))\n\n def task_extras(self, latest_only: bool = True) -> dict[str, np.ndarray]:\n return self._env.task_extras(latest_only=latest_only)\n\n def _reset_state(self):\n pass\n\n def _process_action(self, action: Any) -> Any:\n return action\n\n def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:\n return timestep\n\n def observation_spec(self) -> dict[str, specs.Array]:\n return self._env.observation_spec()\n\n def action_spec(self) -> dict[str, specs.Array]:\n return self._env.action_spec()\n\n def reward_spec(self) -> specs.Array:\n return self._env.reward_spec()\n\n def discount_spec(self) -> specs.Array:\n return self._env.discount_spec()\n\n def _wrapper_stats(self) -> dict[str, Any]:\n \"\"\"Add wrapper specific logging here.\"\"\"\n return {}\n\n def stats(self) -> dict[str, Any]:\n info = self._env.stats()\n info.update(self._wrapper_stats())\n return info\n\n def load_state(\n self, request: state_pb2.LoadStateRequest\n ) -> state_pb2.LoadStateResponse:\n \"\"\"Loads a state.\"\"\"\n return self._env.load_state(request)\n\n def save_state(\n self, request: state_pb2.SaveStateRequest\n ) -> state_pb2.SaveStateResponse:\n \"\"\"Saves a state.\n\n Args:\n request: A `SaveStateRequest` containing any parameters necessary to\n specify how/what state to save.\n\n Returns:\n A `SaveStateResponse` containing the status, error message (if\n applicable), and any other relevant information.\n \"\"\"\n return self._env.save_state(request)\n\n def execute_adb_call(\n self, adb_call: adb_pb2.AdbRequest\n ) -> adb_pb2.AdbResponse:\n return self._env.execute_adb_call(adb_call)\n\n @property\n def raw_action(self) -> Any:\n return self._env.raw_action\n\n @property\n def raw_observation(self) -> Any:\n return self._env.raw_observation\n\n @property\n def raw_env(self) -> env_interface.AndroidEnvInterface:\n \"\"\"Recursively unwrap until we reach the true 'raw' env.\"\"\"\n wrapped = self._env\n if hasattr(wrapped, 'raw_env'):\n return wrapped.raw_env\n return wrapped\n\n def __getattr__(self, attr) -> Any:\n \"\"\"Delegate attribute access to underlying environment.\"\"\"\n return getattr(self._env, attr)\n\n def close(self) -> None:\n self._env.close()\n"
},
{
"path": "android_env/wrappers/base_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.base_wrapper.\"\"\"\n\nfrom unittest import mock\n\nfrom absl import logging\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env.proto import state_pb2\nfrom android_env.wrappers import base_wrapper\n\n\nclass BaseWrapperTest(absltest.TestCase):\n\n @mock.patch.object(logging, 'info')\n def test_base_function_forwarding(self, mock_info):\n base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n wrapped_env = base_wrapper.BaseWrapper(base_env)\n mock_info.assert_called_with('Wrapping with %s', 'BaseWrapper')\n\n fake_ts = 'fake_ts'\n base_env.reset.return_value = fake_ts\n self.assertEqual(fake_ts, wrapped_env.reset())\n base_env.reset.assert_called_once()\n\n fake_ts = 'fake_ts'\n fake_action = 'fake_action'\n base_env.step.return_value = fake_ts\n self.assertEqual(fake_ts, wrapped_env.step(fake_action))\n base_env.step.assert_called_once_with(fake_action)\n\n fake_extras = 'fake_task_extras'\n base_env.task_extras.return_value = fake_extras\n self.assertEqual(fake_extras, wrapped_env.task_extras(latest_only=True))\n base_env.task_extras.assert_called_once_with(latest_only=True)\n\n fake_obs_spec = 'fake_obs_spec'\n base_env.observation_spec.return_value = fake_obs_spec\n self.assertEqual(fake_obs_spec, wrapped_env.observation_spec())\n base_env.observation_spec.assert_called_once()\n\n fake_action_spec = 'fake_action_spec'\n base_env.action_spec.return_value = fake_action_spec\n self.assertEqual(fake_action_spec, wrapped_env.action_spec())\n base_env.action_spec.assert_called_once()\n\n fake_raw_action = 'fake_raw_action'\n type(base_env).raw_action = mock.PropertyMock(return_value=fake_raw_action)\n self.assertEqual(fake_raw_action, wrapped_env.raw_action)\n\n fake_raw_observation = 'fake_raw_observation'\n type(base_env).raw_observation = mock.PropertyMock(\n return_value=fake_raw_observation)\n self.assertEqual(fake_raw_observation, wrapped_env.raw_observation)\n\n load_request = state_pb2.LoadStateRequest(args={})\n expected_response = state_pb2.LoadStateResponse(\n status=state_pb2.LoadStateResponse.Status.OK\n )\n base_env.load_state.return_value = expected_response\n self.assertEqual(wrapped_env.load_state(load_request), expected_response)\n base_env.load_state.assert_called_once_with(load_request)\n\n save_request = state_pb2.SaveStateRequest(args={})\n expected_response = state_pb2.SaveStateResponse(\n status=state_pb2.SaveStateResponse.Status.OK\n )\n base_env.save_state.return_value = expected_response\n self.assertEqual(wrapped_env.save_state(save_request), expected_response)\n base_env.save_state.assert_called_once_with(save_request)\n\n wrapped_env.close()\n base_env.close.assert_called_once()\n\n fake_return_value = 'fake'\n # AndroidEnv::some_random_function() does not exist and calling it should\n # raise an AttributeError.\n with self.assertRaises(AttributeError):\n base_env.some_random_function.return_value = fake_return_value\n\n def test_multiple_wrappers(self):\n base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n wrapped_env_1 = base_wrapper.BaseWrapper(base_env)\n wrapped_env_2 = base_wrapper.BaseWrapper(wrapped_env_1)\n\n wrapped_env_2.close()\n base_env.close.assert_called_once()\n\n def test_raw_env(self):\n base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n wrapped_env_1 = base_wrapper.BaseWrapper(base_env)\n wrapped_env_2 = base_wrapper.BaseWrapper(wrapped_env_1)\n self.assertEqual(base_env, wrapped_env_2.raw_env)\n\n def test_stats(self):\n base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n wrapped_env = base_wrapper.BaseWrapper(base_env)\n base_stats = {'base': 'stats'}\n base_env.stats.return_value = base_stats\n self.assertEqual(base_stats, wrapped_env.stats())\n\n @mock.patch.object(logging, 'info')\n def test_wrapped_stats(self, mock_info):\n base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n\n class LoggingWrapper1(base_wrapper.BaseWrapper):\n\n def _wrapper_stats(self):\n return {\n 'wrapper1': 'stats',\n 'shared': 1,\n }\n\n class LoggingWrapper2(base_wrapper.BaseWrapper):\n\n def _wrapper_stats(self):\n return {\n 'wrapper2': 'stats',\n 'shared': 2,\n }\n\n wrapped_env = LoggingWrapper2(LoggingWrapper1(base_env))\n mock_info.assert_has_calls([\n mock.call('Wrapping with %s', 'LoggingWrapper1'),\n mock.call('Wrapping with %s', 'LoggingWrapper2'),\n ])\n base_stats = {'base': 'stats'}\n base_env.stats.return_value = base_stats\n expected_stats = {\n 'base': 'stats',\n 'wrapper1': 'stats',\n 'wrapper2': 'stats',\n 'shared': 2,\n }\n\n self.assertEqual(expected_stats, wrapped_env.stats())\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/discrete_action_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wraps the AndroidEnv environment to provide discrete actions.\"\"\"\n\nfrom collections.abc import Sequence\nfrom typing import cast\n\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\n_NOISE_CLIP_VALUE = 0.4999\n\n\nclass DiscreteActionWrapper(base_wrapper.BaseWrapper):\n \"\"\"AndroidEnv with discrete actions.\"\"\"\n\n def __init__(\n self,\n env: env_interface.AndroidEnvInterface,\n action_grid: Sequence[int] = (10, 10),\n redundant_actions: bool = True,\n noise: float = 0.1,\n ) -> None:\n super().__init__(env)\n self._parent_action_spec = self._env.action_spec()\n self._assert_base_env()\n self._action_grid = action_grid # [height, width]\n self._grid_size = np.prod(self._action_grid)\n action_types = cast(\n specs.DiscreteArray, self._parent_action_spec['action_type']\n )\n self._num_action_types = action_types.num_values\n self._redundant_actions = redundant_actions\n self._noise = noise\n\n def _assert_base_env(self) -> None:\n \"\"\"Checks that the wrapped env has the right action spec format.\"\"\"\n\n assert len(self._parent_action_spec) == 2\n assert not self._parent_action_spec['action_type'].shape\n assert self._parent_action_spec['touch_position'].shape == (2,)\n\n @property\n def num_actions(self) -> int:\n \"\"\"Number of discrete actions.\"\"\"\n\n if self._redundant_actions:\n return self._grid_size * self._num_action_types\n else:\n return self._grid_size + self._num_action_types - 1\n\n def step(self, action: dict[str, int]) -> dm_env.TimeStep:\n \"\"\"Take a step in the base environment.\"\"\"\n\n return self._env.step(self._process_action(action))\n\n def _process_action(self, action: dict[str, int]) -> dict[str, np.ndarray]:\n \"\"\"Transforms action so that it agrees with AndroidEnv's action spec.\"\"\"\n\n return {\n 'action_type':\n np.array(self._get_action_type(action['action_id']),\n dtype=self._parent_action_spec['action_type'].dtype),\n 'touch_position':\n np.array(self._get_touch_position(action['action_id']),\n dtype=self._parent_action_spec['touch_position'].dtype)\n }\n\n def _get_action_type(self, action_id: int) -> action_type.ActionType:\n \"\"\"Compute action type corresponding to the given action_id.\n\n When `self._redundant_actions` == True the `grid_size` is \"broadcast\" over\n all the possible actions so you end up with `grid_size` discrete actions\n of type 0, `grid_size` discrete actions of type 1, etc. for all action\n types.\n\n When `self._redundant_actions` == False the first `grid_size` actions are\n reserved for \"touch\" and the rest are just added (NOT multiplied) to the\n total number of discrete actions (exactly one of LIFT and REPEAT).\n\n Args:\n action_id: A discrete action.\n Returns:\n action_type: The action_type of the action.\n \"\"\"\n\n if self._redundant_actions:\n assert action_id < self._num_action_types * self._grid_size\n return action_id // self._grid_size\n\n else:\n assert action_id <= self._grid_size + 1\n if action_id < self._grid_size:\n return action_type.ActionType.TOUCH\n elif action_id == self._grid_size:\n return action_type.ActionType.LIFT\n else:\n return action_type.ActionType.REPEAT\n\n def _get_touch_position(self, action_id: int) -> Sequence[float]:\n \"\"\"Compute the position corresponding to the given action_id.\n\n Note: in the touch_position (x, y) of an action, x corresponds to the\n horizontal axis (width), and y corresponds to the vertical axis (height)\n of the screen. BUT, the screen has dimensions (height, width), i.e. the\n first coordinate corresponds to y, and the second coordinate corresponds\n to x. Pay attention to this mismatch in the calculations below.\n\n Args:\n action_id: A discrete action.\n Returns:\n touch_position: The [0,1]x[0,1] coordinate of the action.\n \"\"\"\n\n position_idx = action_id % self._grid_size\n\n x_pos_grid = position_idx % self._action_grid[1] # WIDTH\n y_pos_grid = position_idx // self._action_grid[1] # HEIGHT\n\n noise_x = np.random.normal(loc=0.0, scale=self._noise)\n noise_y = np.random.normal(loc=0.0, scale=self._noise)\n\n # Noise is clipped so that the action will strictly stay in the cell.\n noise_x = max(min(noise_x, _NOISE_CLIP_VALUE), -_NOISE_CLIP_VALUE)\n noise_y = max(min(noise_y, _NOISE_CLIP_VALUE), -_NOISE_CLIP_VALUE)\n\n x_pos = (x_pos_grid + 0.5 + noise_x) / self._action_grid[1] # WIDTH\n y_pos = (y_pos_grid + 0.5 + noise_y) / self._action_grid[0] # HEIGHT\n\n # Project action space to action_spec ranges. For the default case of\n # minimum = [0, 0] and maximum = [1, 1], this will not do anything.\n x_min, y_min = cast(\n specs.BoundedArray, self._parent_action_spec['touch_position']\n ).minimum\n x_max, y_max = cast(\n specs.BoundedArray, self._parent_action_spec['touch_position']\n ).maximum\n\n x_pos = x_min + x_pos * (x_max - x_min)\n y_pos = y_min + y_pos * (y_max - y_min)\n\n return [x_pos, y_pos]\n\n def action_spec(self) -> dict[str, specs.Array]:\n \"\"\"Action spec of the wrapped environment.\"\"\"\n\n return {\n 'action_id':\n specs.DiscreteArray(\n num_values=self.num_actions,\n name='action_id')\n }\n"
},
{
"path": "android_env/wrappers/discrete_action_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.discrete_action_wrapper.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env.components import action_type as action_type_lib\nfrom android_env.wrappers import discrete_action_wrapper\nfrom dm_env import specs\nimport numpy as np\n\nActionType = action_type_lib.ActionType\n\n\ndef _make_array_spec(shape, dtype, name):\n assert len(shape) == 1\n return specs.BoundedArray(\n name=name,\n shape=shape,\n dtype=dtype,\n minimum=np.zeros(shape),\n maximum=(shape[0] - 1) * np.ones(shape), # maximum is inclusive.\n )\n\n\ndef _valid_shape(action):\n assert len(action) == 2, action\n assert not action['action_type'].shape, (\n 'action: %r, shape: %r' %\n (action['action_type'], action['action_type'].shape))\n assert action['touch_position'].shape == (\n 2,), ('action: %r, shape: %r' %\n (action['touch_position'], action['touch_position'].shape))\n\n\ndef _valid_types(action, types):\n for a, t in zip(action.values(), types):\n assert a.dtype == t, '%r is not of dtype %r' % (a, t)\n\n\nclass DiscreteActionWrapperTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self._num_action_types = 3 # Only TOUCH, LIFT, REPEAT.\n self._base_action_spec = {\n 'action_type': specs.DiscreteArray(\n num_values=self._num_action_types, name='action_type'),\n 'touch_position': _make_array_spec(\n shape=(2,), dtype=np.float32, name='touch_position'),\n }\n self.base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n self.base_env.action_spec.return_value = self._base_action_spec\n\n def test_num_actions(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env, action_grid=(3, 3), redundant_actions=True)\n # 27 = 3 * 3 * 2 (H * W * self._num_action_types).\n self.assertEqual(27, wrapped_env.num_actions)\n\n def test_num_actions_non_redundant(self):\n # Check that with `redundant_actions`==False we get an additive term instead\n # of a multiplier in the number of actions.\n non_redudant_wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env, action_grid=(3, 3), redundant_actions=False)\n # 11 = 3 * 3 + 2 (H * W + (self._num_action_types - 1)).\n self.assertEqual(11, non_redudant_wrapped_env.num_actions)\n\n def test_reset(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env, redundant_actions=True)\n fake_timestep = 'ts'\n self.base_env.reset.return_value = fake_timestep\n ts = wrapped_env.reset()\n self.base_env.reset.assert_called_once()\n self.assertEqual(fake_timestep, ts)\n\n def test_step_no_noise(self):\n height = 4\n width = 3\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env,\n action_grid=(height, width),\n noise=0.0,\n redundant_actions=True)\n self.assertEqual(height * width * self._num_action_types,\n wrapped_env.num_actions)\n\n vertical_half_step = 1. / float(height) / 2.\n horizontal_half_step = 1. / float(width) / 2.\n\n delta = 0.0001\n\n # Testing the four corners with each finger position\n def get_verifier(expected_action_type, lower_x, lower_y):\n\n def verifier(x):\n _valid_shape(x)\n _valid_types(x, [np.int32, np.float32])\n self.assertEqual(\n expected_action_type, x['action_type'])\n if lower_y:\n self.assertAlmostEqual(\n vertical_half_step, x['touch_position'][1], delta=delta)\n else:\n self.assertAlmostEqual(\n 1 - vertical_half_step, x['touch_position'][1], delta=delta)\n if lower_x:\n self.assertAlmostEqual(\n horizontal_half_step, x['touch_position'][0], delta=delta)\n else:\n self.assertAlmostEqual(\n 1 - horizontal_half_step, x['touch_position'][0], delta=delta)\n return True\n\n return verifier\n\n action_tests = {\n 0: get_verifier(0, lower_x=True, lower_y=True),\n 2: get_verifier(0, lower_x=False, lower_y=True),\n 9: get_verifier(0, lower_x=True, lower_y=False),\n 11: get_verifier(0, lower_x=False, lower_y=False),\n\n 12: get_verifier(1, lower_x=True, lower_y=True),\n 14: get_verifier(1, lower_x=False, lower_y=True),\n 21: get_verifier(1, lower_x=True, lower_y=False),\n 23: get_verifier(1, lower_x=False, lower_y=False),\n\n 24: get_verifier(2, lower_x=True, lower_y=True),\n 26: get_verifier(2, lower_x=False, lower_y=True),\n 33: get_verifier(2, lower_x=True, lower_y=False),\n 35: get_verifier(2, lower_x=False, lower_y=False),\n }\n\n fake_timestep = 'ts'\n self.base_env.step.return_value = fake_timestep\n\n for action_id, verifier in action_tests.items():\n ts = wrapped_env.step({'action_id': action_id})\n verifier(self.base_env.step.call_args[0][0])\n self.assertEqual(fake_timestep, ts)\n\n def test_step_redundant_actions_invalid_action_id(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env,\n action_grid=(4, 3),\n noise=0.0,\n redundant_actions=True)\n with self.assertRaises(AssertionError):\n _ = wrapped_env.step({'action_id': 36})\n\n def test_step_no_noise_no_redudant_actions(self):\n height = 4\n width = 3\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env,\n action_grid=(height, width),\n noise=0.0,\n redundant_actions=False)\n self.assertEqual(height * width + (self._num_action_types - 1),\n wrapped_env.num_actions)\n\n vertical_half_step = 1. / float(height) / 2.\n horizontal_half_step = 1. / float(width) / 2.\n\n delta = 0.0001\n\n # Testing the four corners with each finger position\n def get_verifier(expected_action_type, lower_x, lower_y):\n\n def verifier(x):\n _valid_shape(x)\n _valid_types(x, [np.int32, np.float32])\n self.assertEqual(expected_action_type, x['action_type'])\n # If the action type == TOUCH, then check the coordinate values.\n if x['action_type'] == ActionType.TOUCH:\n if lower_y:\n self.assertAlmostEqual(\n vertical_half_step, x['touch_position'][1], delta=delta)\n else:\n self.assertAlmostEqual(\n 1 - vertical_half_step, x['touch_position'][1], delta=delta)\n if lower_x:\n self.assertAlmostEqual(\n horizontal_half_step, x['touch_position'][0], delta=delta)\n else:\n self.assertAlmostEqual(\n 1 - horizontal_half_step, x['touch_position'][0], delta=delta)\n return True\n\n return verifier\n\n action_tests = {\n # Touch type actions\n 0: get_verifier(0, lower_x=True, lower_y=True),\n 2: get_verifier(0, lower_x=False, lower_y=True),\n 9: get_verifier(0, lower_x=True, lower_y=False),\n 11: get_verifier(0, lower_x=False, lower_y=False),\n # Actions > grid_size return non-touch actions with (0,0) coordinates.\n 12: get_verifier(1, lower_x=False, lower_y=False),\n 13: get_verifier(2, lower_x=False, lower_y=False),\n }\n\n fake_timestep = 'ts'\n self.base_env.step.return_value = fake_timestep\n\n for action_id, verifier in action_tests.items():\n ts = wrapped_env.step({'action_id': action_id})\n verifier(self.base_env.step.call_args[0][0])\n self.assertEqual(fake_timestep, ts)\n\n def test_step_no_redundant_actions_invalid_action_id(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env,\n action_grid=(4, 3),\n noise=0.0,\n redundant_actions=False)\n with self.assertRaises(AssertionError):\n _ = wrapped_env.step({'action_id': 14})\n\n def test_step_with_noise(self):\n height = 4\n width = 3\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env, action_grid=(height, width), noise=1.0)\n self.assertEqual(height * width * self._num_action_types,\n wrapped_env.num_actions)\n\n vertical_grid_step = 1. / float(height)\n horizontal_grid_step = 1. / float(width)\n\n # Testing the four corners with each finger position\n def get_verifier(expected_up_down, lower_x, lower_y):\n\n def verifier(x):\n _valid_shape(x)\n _valid_types(x, [np.int32, np.float32])\n self.assertEqual(expected_up_down, x['action_type'])\n if lower_y:\n self.assertGreater(vertical_grid_step, x['touch_position'][1])\n else:\n self.assertLess(1 - vertical_grid_step, x['touch_position'][1])\n if lower_x:\n self.assertGreater(horizontal_grid_step, x['touch_position'][0])\n else:\n self.assertLess(1 - horizontal_grid_step, x['touch_position'][0])\n return True\n\n return verifier\n\n action_tests = {\n 0: get_verifier(0, lower_x=True, lower_y=True),\n 2: get_verifier(0, lower_x=False, lower_y=True),\n 9: get_verifier(0, lower_x=True, lower_y=False),\n 11: get_verifier(0, lower_x=False, lower_y=False),\n\n 12: get_verifier(1, lower_x=True, lower_y=True),\n 14: get_verifier(1, lower_x=False, lower_y=True),\n 21: get_verifier(1, lower_x=True, lower_y=False),\n 23: get_verifier(1, lower_x=False, lower_y=False),\n\n 24: get_verifier(2, lower_x=True, lower_y=True),\n 26: get_verifier(2, lower_x=False, lower_y=True),\n 33: get_verifier(2, lower_x=True, lower_y=False),\n 35: get_verifier(2, lower_x=False, lower_y=False),\n }\n\n fake_timestep = 'ts'\n self.base_env.step.return_value = fake_timestep\n\n for action_id, verifier in action_tests.items():\n ts = wrapped_env.step({'action_id': action_id})\n verifier(self.base_env.step.call_args[0][0])\n self.assertEqual(fake_timestep, ts)\n\n def test_parent_spec_type(self):\n base_action_spec = {\n 'action_type': specs.DiscreteArray(\n num_values=self._num_action_types, name='action_type'),\n 'touch_position': _make_array_spec(\n shape=(2,), dtype=np.float64, name='touch_position'),\n }\n base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n base_env.action_spec.return_value = base_action_spec\n\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n base_env, noise=0.0)\n\n fake_timestep = 'ts'\n base_env.step.return_value = fake_timestep\n\n def verifier(x):\n _valid_types(x, [np.int32, np.float64])\n return True\n\n ts = wrapped_env.step({'action_id': 1})\n verifier(base_env.step.call_args[0][0])\n self.assertEqual(fake_timestep, ts)\n\n def test_observation_spec(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env)\n fake_obs_spec = 'fake_obs_spec'\n self.base_env.observation_spec.return_value = fake_obs_spec\n observation_spec = wrapped_env.observation_spec()\n self.base_env.observation_spec.assert_called_once()\n self.assertEqual(fake_obs_spec, observation_spec)\n\n def test_action_spec(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env, action_grid=(4, 5), redundant_actions=True)\n expected_action_spec = {\n 'action_id':\n specs.DiscreteArray(\n num_values=4 * 5 * self._num_action_types, name='action_type')\n }\n self.assertEqual(expected_action_spec, wrapped_env.action_spec())\n\n def test_action_spec_non_redundant(self):\n wrapped_env = discrete_action_wrapper.DiscreteActionWrapper(\n self.base_env, action_grid=(4, 5), redundant_actions=False)\n num_non_touch_actions = self._num_action_types - 1\n expected_action_spec = {\n 'action_id':\n specs.DiscreteArray(\n num_values=4 * 5 + num_non_touch_actions, name='action_type')\n }\n self.assertEqual(expected_action_spec, wrapped_env.action_spec())\n\n def test_assert_base_env_action_spec_too_short(self):\n self.base_env.action_spec.return_value = {\n 'action_type': specs.DiscreteArray(\n num_values=self._num_action_types, name='action_type'),\n }\n with self.assertRaises(AssertionError):\n _ = discrete_action_wrapper.DiscreteActionWrapper(self.base_env)\n\n def test_assert_base_env_action_spec_too_long(self):\n self.base_env.action_spec.return_value = {\n 'action_type': specs.DiscreteArray(\n num_values=self._num_action_types, name='action_type'),\n 'touch_position': _make_array_spec(\n shape=(2,), dtype=np.float32, name='touch_position'),\n 'too_long': _make_array_spec(\n shape=(1,), dtype=np.float32, name='too_long'),\n }\n with self.assertRaises(AssertionError):\n _ = discrete_action_wrapper.DiscreteActionWrapper(self.base_env)\n\n def test_assert_base_env_action_spec_wrong_shapes(self):\n self.base_env.action_spec.return_value = {\n 'action_type': _make_array_spec(\n shape=(2,), dtype=np.float32, name='action_type'),\n 'touch_position': _make_array_spec(\n shape=(1,), dtype=np.float32, name='touch_position')\n }\n with self.assertRaises(AssertionError):\n _ = discrete_action_wrapper.DiscreteActionWrapper(self.base_env)\n\n def test_assert_base_env_ok(self):\n self.base_env.action_spec.return_value = {\n 'action_type': specs.DiscreteArray(\n num_values=self._num_action_types, name='action_type'),\n 'touch_position': _make_array_spec(\n shape=(2,), dtype=np.float32, name='touch_position'),\n }\n _ = discrete_action_wrapper.DiscreteActionWrapper(self.base_env)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/flat_interface_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wraps the AndroidEnv environment to make its interface flat.\"\"\"\n\nfrom typing import Any, cast\n\nfrom android_env import env_interface\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\nRGB_CHANNELS = (0, 1, 2)\n\n\ndef _extract_screen_pixels(obs: np.ndarray) -> np.ndarray:\n \"\"\"Get only screen pixels by removing previous action layer.\"\"\"\n is_grayscale_image = obs.shape[-1] == 2\n if is_grayscale_image:\n return np.expand_dims(obs[..., 0], -1)\n return obs[..., RGB_CHANNELS]\n\n\ndef _get_no_action_observation_spec(\n obs_spec: specs.BoundedArray,\n) -> specs.BoundedArray:\n \"\"\"Create an observation spec without the action layer.\"\"\"\n shape = np.array(obs_spec.shape)\n shape[2] -= 1\n minimum = obs_spec.minimum\n maximum = obs_spec.maximum\n is_scalar = lambda x: np.isscalar(x) or np.ndim(x) == 0\n if not is_scalar(minimum):\n minimum = _extract_screen_pixels(minimum)\n if not is_scalar(maximum):\n maximum = _extract_screen_pixels(maximum)\n return obs_spec.replace(shape=shape, minimum=minimum, maximum=maximum)\n\n\nclass FlatInterfaceWrapper(base_wrapper.BaseWrapper):\n \"\"\"Simple interface for AndroidEnv.\n\n Removes the structure from observations and actions, keeping only the pixel\n observations. Also exposes action as an int32 scalar, making it easier to use\n with conventional discrete agents. This wrapper expects a discretized action\n space.\n \"\"\"\n\n def __init__(\n self,\n env: env_interface.AndroidEnvInterface,\n flat_actions: bool = True,\n flat_observations: bool = True,\n keep_action_layer: bool = True,\n ) -> None:\n super().__init__(env)\n self._flat_actions = flat_actions\n self._flat_observations = flat_observations\n self._keep_action_layer = keep_action_layer\n self._action_name = list(self._env.action_spec())[0]\n self._assert_base_env()\n\n def _assert_base_env(self) -> None:\n base_action_spec = self._env.action_spec()\n assert len(base_action_spec) == 1, self._env.action_spec()\n assert isinstance(base_action_spec, dict)\n assert isinstance(base_action_spec[self._action_name], specs.BoundedArray)\n\n def _process_action(\n self, action: int | np.ndarray | dict[str, Any]\n ) -> int | np.ndarray | dict[str, Any]:\n if self._flat_actions:\n return {self._action_name: action}\n else:\n return action\n\n def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:\n if self._flat_observations:\n step_type, reward, discount, observation = timestep\n # Keep only the pixels.\n pixels = observation['pixels']\n pixels = (\n pixels if self._keep_action_layer else _extract_screen_pixels(pixels)\n )\n return dm_env.TimeStep(\n step_type=step_type,\n reward=reward,\n discount=discount,\n observation=pixels,\n )\n else:\n return timestep\n\n def reset(self) -> dm_env.TimeStep:\n timestep = self._env.reset()\n return self._process_timestep(timestep)\n\n def step(self, action: int) -> dm_env.TimeStep:\n timestep = self._env.step(self._process_action(action))\n return self._process_timestep(timestep)\n\n def observation_spec(self) -> specs.Array | dict[str, specs.Array]: # pytype: disable=signature-mismatch # overriding-return-type-checks\n if self._flat_observations:\n pixels_spec = cast(\n specs.BoundedArray, self._env.observation_spec()['pixels']\n )\n if not self._keep_action_layer:\n return _get_no_action_observation_spec(pixels_spec)\n return pixels_spec\n else:\n return self._env.observation_spec()\n\n def action_spec(self) -> specs.BoundedArray | dict[str, specs.Array]: # pytype: disable=signature-mismatch # overriding-return-type-checks\n if self._flat_actions:\n return self._env.action_spec()[self._action_name] # pytype: disable=bad-return-type\n else:\n return self._env.action_spec()\n"
},
{
"path": "android_env/wrappers/flat_interface_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.flat_interface_wrapper.\"\"\"\n\nfrom typing import cast\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.wrappers import flat_interface_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\ndef _make_array_spec(shape, dtype=np.float32, name=None, maximum=3, minimum=0):\n return specs.BoundedArray(\n shape=shape,\n dtype=dtype,\n name=name,\n maximum=np.ones(shape) * maximum,\n minimum=np.ones(shape) * minimum)\n\n\ndef _make_timestep(observation):\n return dm_env.TimeStep(\n step_type='fake_step_type',\n reward='fake_reward',\n discount='fake_discount',\n observation=observation,\n )\n\n\nclass FlatInterfaceWrapperTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self.action_shape = (1,)\n self.base_action_spec: dict[str, specs.DiscreteArray] = {\n 'action_id': specs.DiscreteArray(name='action_id', num_values=4)\n }\n self.int_obs_shape = (3, 4, 2)\n self.float_obs_shape = (2,)\n self.base_observation_spec = {\n 'pixels': _make_array_spec(\n shape=self.int_obs_shape, dtype=np.uint8, name='pixels'),\n 'obs1': _make_array_spec(\n shape=self.float_obs_shape, dtype=np.float32, name='obs1'),\n }\n # Expected.\n self.expected_observation_spec = _make_array_spec(\n shape=self.int_obs_shape, dtype=np.uint8, name='pixels')\n self.image_obs = np.ones(self.int_obs_shape, dtype=np.uint8)\n self.expected_timestep = _make_timestep(self.image_obs)\n\n # Expected for no new action layer shape.\n expected_new_shape_no_action_layer = (3, 4, 1)\n self.expected_observation_spec_no_action_layer = _make_array_spec(\n shape=expected_new_shape_no_action_layer, dtype=np.uint8, name='pixels')\n self.expected_timestep_no_action_layer = _make_timestep(\n np.ones(expected_new_shape_no_action_layer, dtype=np.uint8))\n\n # Base environment.\n self.other_obs = np.ones(self.float_obs_shape, dtype=np.float32)\n self.base_timestep = _make_timestep({\n 'pixels': self.image_obs,\n 'obs1': self.other_obs})\n self.base_env = mock.create_autospec(dm_env.Environment)\n self.base_env.action_spec.return_value = self.base_action_spec\n self.base_env.observation_spec.return_value = self.base_observation_spec\n self.base_env.reset.return_value = self.base_timestep\n self.base_env.step.return_value = self.base_timestep\n\n def test_reset(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(self.base_env)\n ts = wrapped_env.reset()\n self.base_env.reset.assert_called_once()\n self.assertEqual(self.expected_timestep, ts)\n\n def test_reset_no_action_layer(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(\n self.base_env, keep_action_layer=False)\n ts = wrapped_env.reset()\n self.base_env.reset.assert_called_once()\n self.assertEqual(\n self.expected_timestep_no_action_layer.observation.tolist(),\n ts.observation.tolist())\n\n def test_step(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(self.base_env)\n action = 2\n ts = wrapped_env.step(action)\n\n def verifier(x):\n self.assertIsInstance(x, dict)\n self.assertIsInstance(x['action_id'], int)\n self.assertEqual(x['action_id'], action)\n return True\n verifier(self.base_env.step.call_args[0][0])\n self.assertEqual(self.expected_timestep, ts)\n\n def test_step_no_action_layer(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(\n self.base_env, keep_action_layer=False)\n action = 2\n ts = wrapped_env.step(action)\n\n def verifier(x):\n self.assertIsInstance(x, dict)\n self.assertIsInstance(x['action_id'], int)\n self.assertEqual(x['action_id'], action)\n return True\n\n verifier(self.base_env.step.call_args[0][0])\n self.assertEqual(\n self.expected_timestep_no_action_layer.observation.tolist(),\n ts.observation.tolist())\n\n def test_observation_spec(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(self.base_env)\n observation_spec = wrapped_env.observation_spec()\n self.base_env.observation_spec.assert_called_once()\n self.assertEqual(self.expected_observation_spec, observation_spec)\n\n def test_observation_spec_no_action_layer(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(\n self.base_env, keep_action_layer=False)\n observation_spec = wrapped_env.observation_spec()\n self.base_env.observation_spec.assert_called_once()\n self.assertEqual(self.expected_observation_spec_no_action_layer,\n observation_spec)\n\n def test_action_spec(self):\n wrapped_env = flat_interface_wrapper.FlatInterfaceWrapper(self.base_env)\n action_spec = cast(specs.BoundedArray, wrapped_env.action_spec())\n parent_action_spec = self.base_action_spec['action_id']\n\n self.assertEqual(parent_action_spec.name, action_spec.name)\n self.assertEqual((), action_spec.shape)\n self.assertEqual(np.int32, action_spec.dtype)\n self.assertEqual(0, action_spec.minimum)\n\n def test_bad_action_spec_structured_action(self):\n bad_base_env = mock.create_autospec(dm_env.Environment)\n bad_base_env.action_spec.return_value = {\n 'action_id': _make_array_spec((1,)),\n 'too_many': _make_array_spec((1,))\n }\n with self.assertRaises(AssertionError):\n _ = flat_interface_wrapper.FlatInterfaceWrapper(bad_base_env)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/float_pixels_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Converts pixel observation to from int to float32 between 0.0 and 1.0.\"\"\"\n\nfrom android_env import env_interface\nfrom android_env.components import pixel_fns\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\nclass FloatPixelsWrapper(base_wrapper.BaseWrapper):\n \"\"\"Wraps AndroidEnv for Panultimate agent.\"\"\"\n\n def __init__(self, env: env_interface.AndroidEnvInterface) -> None:\n super().__init__(env)\n self._input_spec = self._env.observation_spec()['pixels']\n self._should_convert_int_to_float = np.issubdtype(self._input_spec.dtype,\n np.integer)\n\n def _process_observation(\n self, observation: dict[str, np.ndarray]\n ) -> dict[str, np.ndarray]:\n if self._should_convert_int_to_float:\n float_pixels = pixel_fns.convert_int_to_float(\n observation['pixels'], self._input_spec\n )\n observation['pixels'] = float_pixels\n return observation\n\n def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:\n step_type, reward, discount, observation = timestep\n return dm_env.TimeStep(\n step_type=step_type,\n reward=reward,\n discount=discount,\n observation=self._process_observation(observation))\n\n def reset(self) -> dm_env.TimeStep:\n return self._process_timestep(self._env.reset())\n\n def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:\n return self._process_timestep(self._env.step(action))\n\n def observation_spec(self) -> dict[str, specs.Array]:\n if self._should_convert_int_to_float:\n observation_spec = self._env.observation_spec()\n observation_spec['pixels'] = specs.BoundedArray(\n shape=self._env.observation_spec()['pixels'].shape,\n dtype=np.float32,\n minimum=0.0,\n maximum=1.0,\n name=self._env.observation_spec()['pixels'].name)\n return observation_spec\n return self._env.observation_spec()\n"
},
{
"path": "android_env/wrappers/float_pixels_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.float_pixels_wrapper.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env.wrappers import float_pixels_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\ndef _make_array_spec(shape, dtype=np.float32, name=None):\n return specs.Array(\n shape=shape,\n dtype=dtype,\n name=name,\n )\n\n\ndef _make_bounded_array_spec(\n shape, dtype=np.float32, name=None, maximum=1.0, minimum=0.0):\n return specs.BoundedArray(\n shape=shape,\n dtype=dtype,\n name=name,\n maximum=maximum,\n minimum=minimum,\n )\n\n\ndef _simple_timestep(obs_shape, obs_type):\n return dm_env.TimeStep(\n step_type=dm_env.StepType.MID,\n reward=3.14,\n discount=0.9,\n observation=(np.ones(shape=obs_shape, dtype=obs_type),),\n )\n\n\nclass FloatPixelsWrapperTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self.pixels_shape = (3, 4)\n base_pixel_spec = _make_array_spec(\n shape=self.pixels_shape, dtype=np.uint8, name='pixels')\n self.other_obs_spec = _make_array_spec(\n shape=(1,), dtype=np.float32, name='other_obs')\n base_observation_spec = {\n 'pixels': base_pixel_spec,\n 'other_obs': self.other_obs_spec\n }\n self.base_env = mock.create_autospec(dm_env.Environment)\n self.base_env.observation_spec.return_value = base_observation_spec\n\n self.base_timestep = dm_env.TimeStep(\n step_type=dm_env.StepType.MID,\n reward=3.14,\n discount=0.9,\n observation={\n 'pixels': np.ones(shape=self.pixels_shape, dtype=np.uint8),\n 'other_obs': [42.2]})\n self.base_env.step.return_value = self.base_timestep\n self.base_env.reset.return_value = self.base_timestep\n\n def test_float_pixels_wrapper_spec(self):\n expected_pixel_spec = _make_bounded_array_spec(\n shape=self.pixels_shape,\n dtype=np.float32,\n name='pixels',\n minimum=0.0,\n maximum=1.0)\n\n wrapped_env = float_pixels_wrapper.FloatPixelsWrapper(self.base_env)\n\n self.assertLen(wrapped_env.observation_spec(), 2)\n self.assertEqual(expected_pixel_spec,\n wrapped_env.observation_spec()['pixels'])\n self.assertEqual(self.other_obs_spec,\n wrapped_env.observation_spec()['other_obs'])\n\n def test_float_pixels_wrapper_step(self):\n wrapped_env = float_pixels_wrapper.FloatPixelsWrapper(self.base_env)\n ts = wrapped_env.step({'fake_action': np.array([1, 2, 3])})\n\n self.assertEqual(self.base_timestep.step_type, ts.step_type)\n self.assertEqual(self.base_timestep.reward, ts.reward)\n self.assertEqual(self.base_timestep.discount, ts.discount)\n self.assertEqual(self.base_timestep.observation['other_obs'],\n ts.observation['other_obs'])\n expected_pixel_value = 1. / 255. # original values are unit8\n expected_pixels = np.ones(\n self.pixels_shape, dtype=np.float32) * expected_pixel_value\n np.testing.assert_equal(expected_pixels, ts.observation['pixels'])\n\n def test_float_pixels_wrapper_reset(self):\n wrapped_env = float_pixels_wrapper.FloatPixelsWrapper(self.base_env)\n ts = wrapped_env.reset()\n\n self.assertEqual(self.base_timestep.step_type, ts.step_type)\n self.assertEqual(self.base_timestep.reward, ts.reward)\n self.assertEqual(self.base_timestep.discount, ts.discount)\n self.assertEqual(self.base_timestep.observation['other_obs'],\n ts.observation['other_obs'])\n expected_pixel_value = 1. / 255. # original values are unit8\n expected_pixels = np.ones(\n self.pixels_shape, dtype=np.float32) * expected_pixel_value\n np.testing.assert_equal(expected_pixels, ts.observation['pixels'])\n\n def test_float_pixels_wrapper_already_float(self):\n base_pixel_spec = _make_array_spec(\n shape=self.pixels_shape, dtype=np.float64, name='pixels')\n base_observation_spec = {\n 'pixels': base_pixel_spec,\n 'other_obs': self.other_obs_spec\n }\n base_env = mock.create_autospec(dm_env.Environment)\n base_env.observation_spec.return_value = base_observation_spec\n\n wrapped_env = float_pixels_wrapper.FloatPixelsWrapper(base_env)\n\n # If the pixels are already float values, then obs_spec does not change.\n self.assertEqual(base_env.observation_spec(),\n wrapped_env.observation_spec())\n\n # The wrapper should not touch the timestep in this case.\n fake_timestep = ('step_type', 'reward', 'discount', 'obs')\n base_env.step.return_value = fake_timestep\n ts = wrapped_env.step({'fake_action': np.array([1, 2, 3])})\n self.assertEqual(fake_timestep, ts)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/gym_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wraps the AndroidEnv to expose an OpenAI Gym interface.\"\"\"\n\nfrom typing import Any\n\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nfrom dm_env import specs\nimport gym\nfrom gym import spaces\nimport numpy as np\n\n\nclass GymInterfaceWrapper(gym.Env):\n \"\"\"AndroidEnv with OpenAI Gym interface.\"\"\"\n\n def __init__(self, env: dm_env.Environment):\n self._env = env\n self.spec = None\n self.action_space = self._spec_to_space(self._env.action_spec())\n self.observation_space = self._spec_to_space(self._env.observation_spec())\n self.metadata = {'render.modes': ['rgb_array']}\n self._latest_observation = None\n\n def _spec_to_space(self, spec: specs.Array) -> spaces.Space:\n \"\"\"Converts dm_env specs to OpenAI Gym spaces.\"\"\"\n\n if isinstance(spec, list):\n return spaces.Tuple([self._spec_to_space(s) for s in spec])\n\n if isinstance(spec, dict):\n return spaces.Dict(\n {name: self._spec_to_space(s) for name, s in spec.items()}\n )\n\n if isinstance(spec, specs.DiscreteArray):\n return spaces.Box(\n shape=(),\n dtype=spec.dtype,\n low=0,\n high=spec.num_values-1)\n\n if isinstance(spec, specs.BoundedArray):\n return spaces.Box(\n shape=spec.shape,\n dtype=spec.dtype,\n low=spec.minimum,\n high=spec.maximum)\n\n if isinstance(spec, specs.Array):\n if spec.dtype == np.uint8:\n low = 0\n high = 255\n else:\n low = -np.inf\n high = np.inf\n return spaces.Box(shape=spec.shape, dtype=spec.dtype, low=low, high=high)\n\n raise ValueError('Unknown type for specs: {}'.format(spec))\n\n def render(self, mode='rgb_array'):\n \"\"\"Renders the environment.\"\"\"\n if mode == 'rgb_array':\n if self._latest_observation is None:\n return\n\n return self._latest_observation['pixels']\n else:\n raise ValueError('Only supported render mode is rgb_array.')\n\n def reset(self) -> np.ndarray:\n self._latest_observation = None\n timestep = self._env.reset()\n return timestep.observation\n\n def step(self, action: dict[str, int]) -> tuple[Any, ...]:\n \"\"\"Take a step in the base environment.\"\"\"\n timestep = self._env.step(action)\n observation = timestep.observation\n self._latest_observation = observation\n reward = timestep.reward\n done = timestep.step_type == dm_env.StepType.LAST\n info = {'discount': timestep.discount}\n return observation, reward, done, info\n"
},
{
"path": "android_env/wrappers/gym_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.gym_wrapper.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env.wrappers import gym_wrapper\nimport dm_env\nfrom dm_env import specs\nfrom gym import spaces\nimport numpy as np\n\n\nclass GymInterfaceWrapperTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self._base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n self._base_env.action_spec.return_value = {\n 'action_type':\n specs.DiscreteArray(\n num_values=3,\n name='action_type'),\n 'touch_position':\n specs.BoundedArray(\n shape=(2,),\n dtype=np.float32,\n minimum=[0.0, 0.0],\n maximum=[1.0, 1.0],\n name='touch_position'),\n }\n self._base_env.observation_spec.return_value = {\n 'pixels':\n specs.Array(\n shape=(480, 320, 3),\n dtype=np.uint8,\n name='pixels'),\n 'timedelta':\n specs.Array(shape=(), dtype=np.int64, name='timedelta'),\n 'orientation':\n specs.Array(\n shape=np.array([4]),\n dtype=np.uint8,\n name='orientation'),\n }\n self._wrapped_env = gym_wrapper.GymInterfaceWrapper(self._base_env)\n self._fake_ts = dm_env.TimeStep(\n step_type=dm_env.StepType.MID,\n observation={'pixels': np.ones(shape=(2, 3))},\n reward=10.0,\n discount=1.0)\n\n def test_render(self):\n self._base_env.step.return_value = self._fake_ts\n _ = self._wrapped_env.step(action=np.zeros(shape=(1,)))\n image = self._wrapped_env.render(mode='rgb_array')\n self.assertTrue(np.array_equal(image, np.ones(shape=(2, 3))))\n\n def test_render_error(self):\n with self.assertRaises(ValueError):\n _ = self._wrapped_env.render(mode='human')\n\n def test_reset(self):\n self._base_env.reset.return_value = dm_env.TimeStep(\n step_type=dm_env.StepType.FIRST,\n observation={'pixels': np.ones(shape=(2, 3))},\n reward=10.0,\n discount=1.0)\n obs = self._wrapped_env.reset()\n self._base_env.reset.assert_called_once()\n self.assertTrue(np.array_equal(obs['pixels'], np.ones(shape=(2, 3))))\n\n def test_step(self):\n self._base_env.step.return_value = self._fake_ts\n obs, _, _, _ = self._wrapped_env.step(action=np.zeros(shape=(1,)))\n self._base_env.step.assert_called_once()\n self.assertTrue(np.array_equal(obs['pixels'], np.ones(shape=(2, 3))))\n\n def test_spec_to_space(self):\n\n spec = specs.Array(\n shape=(2, 3),\n dtype=np.float32)\n space = self._wrapped_env._spec_to_space(spec)\n self.assertEqual(space, spaces.Box(\n low=-np.inf, high=np.inf, shape=spec.shape, dtype=spec.dtype))\n\n spec = specs.BoundedArray(\n shape=(),\n dtype=np.float32,\n minimum=4,\n maximum=5)\n space = self._wrapped_env._spec_to_space(spec)\n self.assertEqual(space, spaces.Box(\n low=4, high=5, shape=spec.shape, dtype=spec.dtype))\n\n spec = specs.DiscreteArray(num_values=4)\n space = self._wrapped_env._spec_to_space(spec)\n self.assertEqual(space, spaces.Box(\n low=0, high=3, shape=(), dtype=np.int32))\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/image_rescale_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wraps the AndroidEnv environment to rescale the observations.\"\"\"\n\nfrom collections.abc import Sequence\nfrom typing import cast\n\nfrom android_env import env_interface\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\nfrom PIL import Image\n\n\n# Taken from https://pillow.readthedocs.io/en/3.2.x/reference/Image.html#PIL.Image.Image.convert\n#\n# This array maps an RGB image to a grayscale image using the ITU-R 709\n# specification which is good for computer displays and HDTV.\nRGB_TO_GRAYSCALE_COEFFICIENTS = [0.2126, 0.7152, 0.0722]\n\n\nclass ImageRescaleWrapper(base_wrapper.BaseWrapper):\n \"\"\"AndroidEnv with rescaled observations.\"\"\"\n\n def __init__(\n self,\n env: env_interface.AndroidEnvInterface,\n zoom_factors: Sequence[float] | None = (0.5, 0.5),\n grayscale: bool = False,\n ) -> None:\n super().__init__(env)\n assert 'pixels' in self._env.observation_spec()\n assert self._env.observation_spec()['pixels'].shape[-1] in [\n 1,\n 3,\n ], 'Number of pixel channels should be 1 or 3.'\n self._grayscale = grayscale\n if zoom_factors is None:\n zoom_factors = (1.0, 1.0)\n # We only zoom the width and height of each layer, and we explicitly do not\n # want to zoom the number of channels so we just multiply it by 1.0.\n self._zoom_factors = tuple(zoom_factors) + (1.0,)\n\n def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:\n observation = timestep.observation\n processed_observation = observation.copy()\n processed_observation['pixels'] = self._process_pixels(\n observation['pixels']\n )\n return timestep._replace(observation=processed_observation)\n\n def _process_pixels(self, raw_observation: np.ndarray) -> np.ndarray:\n # We expect `raw_observation` to have shape (W, H, 3) - 3 for RGB\n new_shape = np.array(\n self._zoom_factors[0:2] * np.array(raw_observation.shape[0:2]),\n dtype=np.int32,\n )[::-1]\n if self._grayscale:\n # When self._grayscale == True, we squash the RGB into a single layer\n image = np.dot(raw_observation, RGB_TO_GRAYSCALE_COEFFICIENTS)\n else:\n image = raw_observation\n return self._resize_image_array(image, new_shape)\n\n def _resize_image_array(\n self, grayscale_or_rbg_array: np.ndarray, new_shape: np.ndarray\n ) -> np.ndarray:\n \"\"\"Resize color or grayscale/action_layer array to new_shape.\"\"\"\n assert new_shape.ndim == 1\n assert len(new_shape) == 2\n resized_array = np.array(\n Image.fromarray(grayscale_or_rbg_array.astype('uint8')).resize(\n tuple(new_shape)\n )\n )\n if resized_array.ndim == 2:\n return np.expand_dims(resized_array, axis=-1)\n return resized_array\n\n def reset(self) -> dm_env.TimeStep:\n timestep = self._env.reset()\n return self._process_timestep(timestep)\n\n def step(self, action) -> dm_env.TimeStep:\n timestep = self._env.step(action)\n return self._process_timestep(timestep)\n\n def observation_spec(self) -> dict[str, specs.Array]:\n parent_spec = self._env.observation_spec().copy()\n parent_pixels = cast(specs.BoundedArray, parent_spec['pixels'])\n out_shape = np.multiply(parent_pixels.shape, self._zoom_factors).astype(\n np.int32\n )\n if self._grayscale:\n # In grayscale mode we want the output shape to be [W, H, 1]\n out_shape[-1] = 1\n parent_spec['pixels'] = specs.BoundedArray(\n shape=out_shape,\n dtype=parent_pixels.dtype,\n name=parent_pixels.name,\n minimum=parent_pixels.minimum,\n maximum=parent_pixels.maximum,\n )\n return parent_spec\n"
},
{
"path": "android_env/wrappers/image_rescale_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.image_rescale_wrapper.\"\"\"\n\nfrom typing import Any\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env.wrappers import image_rescale_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\ndef _simple_spec():\n return specs.BoundedArray(\n shape=np.array([300, 300, 3]),\n dtype=np.uint8,\n name='pixels',\n minimum=0,\n maximum=255)\n\n\ndef _simple_timestep():\n observation = np.ones(shape=[300, 300, 3])\n return dm_env.TimeStep(\n step_type=dm_env.StepType.MID,\n reward=3.14,\n discount=0.9,\n observation={'pixels': observation})\n\n\nclass ImageRescaleWrapperTest(absltest.TestCase):\n\n def test_100x50_grayscale(self):\n fake_timestep = _simple_timestep()\n fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n fake_env.observation_spec.return_value = {'pixels': _simple_spec()}\n fake_env.reset.return_value = fake_timestep\n fake_env.step.return_value = fake_timestep\n\n wrapper = image_rescale_wrapper.ImageRescaleWrapper(\n fake_env, zoom_factors=(1.0 / 3, 1.0 / 6.0), grayscale=True)\n self.assertIsNotNone(wrapper)\n self.assertEqual(wrapper.observation_spec()['pixels'].shape, (100, 50, 1))\n reset_timestep = wrapper.reset()\n reset_image = reset_timestep.observation['pixels']\n self.assertEqual(reset_image.shape, (100, 50, 1))\n step_timestep = wrapper.step(action='fake_action')\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (100, 50, 1))\n\n def test_150x60_full_channels(self):\n fake_timestep = _simple_timestep()\n fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n fake_env.observation_spec.return_value = {'pixels': _simple_spec()}\n fake_env.reset.return_value = fake_timestep\n fake_env.step.return_value = fake_timestep\n\n wrapper = image_rescale_wrapper.ImageRescaleWrapper(\n fake_env, zoom_factors=(1.0 / 2.0, 1.0 / 5.0))\n self.assertIsNotNone(wrapper)\n self.assertEqual(wrapper.observation_spec()['pixels'].shape, (150, 60, 3))\n reset_timestep = wrapper.reset()\n reset_image = reset_timestep.observation['pixels']\n self.assertEqual(reset_image.shape, (150, 60, 3))\n step_timestep = wrapper.step(action='fake_action')\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (150, 60, 3))\n\n def test_list_zoom_factor(self):\n fake_timestep = _simple_timestep()\n fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n fake_env.observation_spec.return_value = {'pixels': _simple_spec()}\n fake_env.reset.return_value = fake_timestep\n fake_env.step.return_value = fake_timestep\n\n wrapper = image_rescale_wrapper.ImageRescaleWrapper(\n fake_env, zoom_factors=[0.5, 0.2])\n self.assertIsNotNone(wrapper)\n self.assertEqual(wrapper.observation_spec()['pixels'].shape, (150, 60, 3))\n reset_timestep = wrapper.reset()\n reset_image = reset_timestep.observation['pixels']\n self.assertEqual(reset_image.shape, (150, 60, 3))\n step_timestep = wrapper.step(action='fake_action')\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (150, 60, 3))\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/last_action_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Extends Android observation with the latest action taken.\"\"\"\n\nfrom typing import cast\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.components import pixel_fns\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\nclass LastActionWrapper(base_wrapper.BaseWrapper):\n \"\"\"Extends Android observations with information about the last action taken.\n\n The position of the last action is denoted by a single white pixel (with a\n value of 255) in a channel of all black pixels (with a value of 0).\n As this wrapper makes use of temporarily stored information about the\n last action taken, it is important to apply on the environment side rather\n than the agent side. Recommended not to apply before an ImageRescaleWrapper,\n to avoid distortion of the single pixel denoting the action position.\n \"\"\"\n\n def __init__(self,\n env: env_interface.AndroidEnvInterface,\n concat_to_pixels: bool = True) -> None:\n \"\"\"Initializes the internal state of this wrapper.\n\n Args:\n env: the environment to wrap.\n concat_to_pixels: If True, will add a channel to the pixel observation.\n If False, will pass the action as an extra observation.\n \"\"\"\n super().__init__(env)\n self._concat_to_pixels = concat_to_pixels\n self._screen_dimensions = self._env.observation_spec()['pixels'].shape[:2]\n\n def _process_timestep(self, timestep: dm_env.TimeStep) -> dm_env.TimeStep:\n observation = timestep.observation.copy()\n processed_observation = self._process_observation(observation)\n return timestep._replace(observation=processed_observation)\n\n def _process_observation(\n self, observation: dict[str, np.ndarray]\n ) -> dict[str, np.ndarray]:\n \"\"\"Extends observation with last_action data.\"\"\"\n processed_observation = observation.copy()\n last_action_layer = self._get_last_action_layer(observation['pixels'])\n if self._concat_to_pixels:\n pixels = observation['pixels'].copy()\n processed_pixels = np.dstack((pixels, last_action_layer))\n processed_observation['pixels'] = processed_pixels\n else:\n processed_observation['last_action'] = last_action_layer\n return processed_observation\n\n def _get_last_action_layer(self, pixels: np.ndarray) -> np.ndarray:\n \"\"\"Makes sure the rescaling doesn't distort the last_action layer.\"\"\"\n\n last_action = self._env.raw_action\n last_action_layer = np.zeros(self._screen_dimensions, dtype=pixels.dtype)\n\n if ('action_type' in last_action and\n last_action['action_type'] == action_type.ActionType.TOUCH):\n touch_position = last_action['touch_position']\n x, y = pixel_fns.touch_position_to_pixel_position(\n touch_position, width_height=self._screen_dimensions[::-1]\n )\n last_action_layer[y, x] = 255\n\n return last_action_layer\n\n def reset(self) -> dm_env.TimeStep:\n timestep = self._env.reset()\n return self._process_timestep(timestep)\n\n def step(self, action) -> dm_env.TimeStep:\n timestep = self._env.step(action)\n return self._process_timestep(timestep)\n\n def observation_spec(self) -> dict[str, specs.Array]:\n parent_spec = self._env.observation_spec().copy()\n parent_pixels = cast(specs.BoundedArray, parent_spec['pixels'])\n shape = parent_pixels.shape\n if self._concat_to_pixels:\n parent_spec['pixels'] = specs.BoundedArray(\n shape=(shape[0], shape[1], shape[2] + 1),\n dtype=parent_pixels.dtype,\n name=parent_pixels.name,\n minimum=parent_pixels.minimum,\n maximum=parent_pixels.maximum)\n else:\n parent_spec.update({\n 'last_action':\n specs.BoundedArray(\n shape=(shape[0], shape[1]),\n dtype=parent_pixels.dtype,\n name='last_action',\n minimum=parent_pixels.minimum,\n maximum=parent_pixels.maximum)\n })\n return parent_spec\n"
},
{
"path": "android_env/wrappers/last_action_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for android_env.wrappers.last_action_wrapper.\"\"\"\n\nfrom typing import Any\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.wrappers import last_action_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\ndef _simple_spec():\n return specs.BoundedArray(\n shape=np.array([120, 80, 3]),\n dtype=np.uint8,\n name='pixels',\n minimum=0,\n maximum=255)\n\n\ndef _simple_timestep():\n observation = np.ones(shape=[120, 80, 3])\n return dm_env.TimeStep(\n step_type=dm_env.StepType.MID,\n reward=3.14,\n discount=0.9,\n observation={'pixels': observation})\n\n\nclass LastActionWrapperTest(absltest.TestCase):\n\n def test_concat_to_pixels(self):\n fake_timestep = _simple_timestep()\n fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n fake_env.observation_spec.return_value = {'pixels': _simple_spec()}\n fake_env.reset.return_value = fake_timestep\n fake_env.step.return_value = fake_timestep\n\n wrapper = last_action_wrapper.LastActionWrapper(\n fake_env, concat_to_pixels=True)\n self.assertIsNotNone(wrapper)\n self.assertEqual(wrapper.observation_spec()['pixels'].shape, (120, 80, 4))\n\n reset_timestep = wrapper.reset()\n reset_image = reset_timestep.observation['pixels']\n self.assertEqual(reset_image.shape, (120, 80, 4))\n last_action_layer = reset_image[:, :, -1]\n self.assertEqual(np.sum(last_action_layer), 0)\n\n action1 = {\n 'action_type': action_type.ActionType.TOUCH,\n 'touch_position': np.array([0.25, 0.75], dtype=np.float32), # (W x H)\n }\n type(fake_env).raw_action = mock.PropertyMock(return_value=action1)\n step_timestep = wrapper.step(action=action1)\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (120, 80, 4)) # (H x W)\n last_action_layer = step_image[:, :, -1]\n self.assertEqual(np.sum(last_action_layer), 255)\n y, x = np.where(last_action_layer == 255)\n self.assertEqual((y.item(), x.item()), (90, 20))\n\n action2 = {\n 'action_type': action_type.ActionType.LIFT,\n 'touch_position': np.array([0.25, 0.75], dtype=np.float32),\n }\n type(fake_env).raw_action = mock.PropertyMock(return_value=action2)\n step_timestep = wrapper.step(action=action2)\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (120, 80, 4))\n last_action_layer = step_image[:, :, -1]\n self.assertEqual(np.sum(last_action_layer), 0)\n\n action3 = {\n 'action_type': action_type.ActionType.TOUCH,\n 'touch_position': np.array([0.25, 1.0], dtype=np.float32),\n }\n type(fake_env).raw_action = mock.PropertyMock(return_value=action3)\n step_timestep = wrapper.step(action=action3)\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (120, 80, 4))\n last_action_layer = step_image[:, :, -1]\n self.assertEqual(np.sum(last_action_layer), 255)\n y, x = np.where(last_action_layer == 255)\n self.assertEqual((y.item(), x.item()), (119, 20))\n\n def test_no_concat_to_pixels(self):\n fake_timestep = _simple_timestep()\n fake_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n fake_env.observation_spec.return_value = {'pixels': _simple_spec()}\n fake_env.reset.return_value = fake_timestep\n fake_env.step.return_value = fake_timestep\n\n wrapper = last_action_wrapper.LastActionWrapper(\n fake_env, concat_to_pixels=False)\n self.assertIsNotNone(wrapper)\n self.assertEqual(wrapper.observation_spec()['pixels'].shape, (120, 80, 3))\n self.assertEqual(wrapper.observation_spec()['last_action'].shape, (120, 80))\n\n reset_timestep = wrapper.reset()\n reset_image = reset_timestep.observation['pixels']\n self.assertEqual(reset_image.shape, (120, 80, 3))\n last_action_layer = reset_timestep.observation['last_action']\n self.assertEqual(np.sum(last_action_layer), 0)\n\n action1 = {\n 'action_type': action_type.ActionType.TOUCH,\n 'touch_position': np.array([0.25, 0.75], dtype=np.float32),\n }\n type(fake_env).raw_action = mock.PropertyMock(return_value=action1)\n step_timestep = wrapper.step(action=action1)\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (120, 80, 3))\n last_action_layer = step_timestep.observation['last_action']\n self.assertEqual(np.sum(last_action_layer), 255)\n y, x = np.where(last_action_layer == 255)\n self.assertEqual((y.item(), x.item()), (90, 20))\n\n action2 = {\n 'action_type': action_type.ActionType.LIFT,\n 'touch_position': np.array([0.25, 0.75], dtype=np.float32),\n }\n type(fake_env).raw_action = mock.PropertyMock(return_value=action2)\n step_timestep = wrapper.step(action=action2)\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (120, 80, 3))\n last_action_layer = step_timestep.observation['last_action']\n self.assertEqual(np.sum(last_action_layer), 0)\n\n action3 = {\n 'action_type': action_type.ActionType.TOUCH,\n 'touch_position': np.array([1.0, 0.75], dtype=np.float32),\n }\n type(fake_env).raw_action = mock.PropertyMock(return_value=action3)\n step_timestep = wrapper.step(action=action3)\n step_image = step_timestep.observation['pixels']\n self.assertEqual(step_image.shape, (120, 80, 3))\n last_action_layer = step_timestep.observation['last_action']\n self.assertEqual(np.sum(last_action_layer), 255)\n y, x = np.where(last_action_layer == 255)\n self.assertEqual((y.item(), x.item()), (90, 79))\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/rate_limit_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Limits interactions with the environment to a given rate.\"\"\"\n\nimport enum\nimport time\n\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nimport numpy as np\n\n\nclass RateLimitWrapper(base_wrapper.BaseWrapper):\n \"\"\"Limits interactions with the environment to a given rate.\"\"\"\n\n class SleepType(enum.IntEnum):\n \"\"\"Determines how the wrapper interacts with the underlying environment.\"\"\"\n\n # The wrapper sleeps before calling `step()` on the underlying environment.\n BEFORE = 0\n\n # The wrapper sleeps after calling `step()` on the underlying environment.\n AFTER = 1\n\n # The wrapper first calls `step()`, obtaining a TimeStep which is ignored,\n # then it sleeps, and then it calls `step(REPEAT)` to obtain a TimeStep\n # that's as fresh as possible.\n #\n # Note that for both BEFORE and AFTER_WITH_REPEAT, the _total_ amount of\n # time inside this wrapper may go beyond the rate specified in `rate`\n # because the sleep does not account for the time taken by step().\n AFTER_WITH_REPEAT = 2\n\n def __init__(self,\n env: env_interface.AndroidEnvInterface,\n rate: float,\n sleep_type: SleepType = SleepType.AFTER_WITH_REPEAT):\n \"\"\"Initializes this wrapper.\n\n Args:\n env: The underlying environment to which this wrapper is applied.\n rate: The desired rate in Hz to interact with the environment. If <=0.0,\n this wrapper will be disabled.\n sleep_type: This determines how the wrapper will interact with the\n underlying AndroidEnv environment.\n \"\"\"\n super().__init__(env)\n self._assert_base_env()\n self._last_step_time = None\n self._max_wait = 1.0 / rate if rate > 0.0 else 0.0\n self._sleep_type = sleep_type\n\n def _assert_base_env(self):\n \"\"\"Checks that the wrapped env has the right action spec format.\"\"\"\n parent_action_spec = self._env.action_spec()\n assert len(parent_action_spec) == 2\n assert not parent_action_spec['action_type'].shape\n assert parent_action_spec['touch_position'].shape == (2,)\n\n def reset(self):\n timestep = self._env.reset()\n self._last_step_time = time.time()\n return timestep\n\n def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:\n \"\"\"Takes a step while maintaining a steady interaction rate.\"\"\"\n\n # If max_wait is non-positive, the wrapper has no effect.\n if self._max_wait <= 0.0:\n return self._env.step(action)\n\n if self._sleep_type == RateLimitWrapper.SleepType.BEFORE:\n self._wait()\n\n timestep = self._env.step(action)\n if timestep.last():\n return timestep\n\n if self._sleep_type == RateLimitWrapper.SleepType.AFTER_WITH_REPEAT:\n for k in action.keys():\n if k.startswith('action_type'):\n action[k] = np.array(action_type.ActionType.REPEAT, dtype=np.uint8)\n self._wait()\n first_reward = timestep.reward or 0.0\n timestep = self._env.step(action)\n second_reward = timestep.reward or 0.0\n # Accumulate rewards over the two steps taken.\n timestep = timestep._replace(reward=first_reward + second_reward)\n\n elif self._sleep_type == RateLimitWrapper.SleepType.AFTER:\n self._wait()\n\n self._last_step_time = time.time()\n\n return timestep\n\n def _wait(self) -> None:\n if self._max_wait > 0.0 and self._last_step_time is not None:\n time_since_step = time.time() - self._last_step_time\n sec_to_wait = self._max_wait - time_since_step\n if sec_to_wait > 0.0:\n time.sleep(sec_to_wait)\n"
},
{
"path": "android_env/wrappers/rate_limit_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for rate_limit_wrapper.\"\"\"\n\nimport time\nfrom typing import Any, Protocol\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.wrappers import rate_limit_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\ndef _get_base_env():\n env = mock.create_autospec(env_interface.AndroidEnvInterface)\n env.action_spec.return_value = {\n 'action_type':\n specs.DiscreteArray(\n num_values=len(action_type.ActionType),\n name='action_type'),\n 'touch_position':\n specs.BoundedArray(\n shape=(2,),\n dtype=np.float32,\n minimum=[0.0, 0.0],\n maximum=[1.0, 1.0],\n name='touch_position'),\n }\n return env\n\n\nclass _FnWithTimestamps(Protocol):\n \"\"\"A function with `timestamp` and `timestamps` attributes.\"\"\"\n\n timestamp: float\n timestamps: list[float]\n\n\ndef _with_timestamp(fn: Any) -> _FnWithTimestamps:\n return fn\n\n\nclass RateLimitWrapperTest(parameterized.TestCase):\n\n @parameterized.named_parameters(\n ('zero_rate', 0),\n ('negative_rate', -50),\n )\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_disabled(self, rate, mock_sleep):\n \"\"\"With a non-positive rate, this wrapper should do nothing.\"\"\"\n env = _get_base_env()\n wrapper = rate_limit_wrapper.RateLimitWrapper(env, rate=rate)\n _ = wrapper.reset()\n mock_sleep.assert_not_called()\n _ = wrapper.step({\n 'action_type': np.array(action_type.ActionType.LIFT, dtype=np.uint8),\n 'touch_position': np.array([0.123, 0.456])\n })\n mock_sleep.assert_not_called()\n # When the wrapper is disabled, base step should only be called once.\n env.step.assert_called_once()\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_enabled(self, mock_sleep):\n \"\"\"When enabled, the wrapper should sleep for a period in [0, 1/rate].\"\"\"\n\n env = _get_base_env()\n env.step.return_value = dm_env.transition(reward=None, observation=None)\n wrapper = rate_limit_wrapper.RateLimitWrapper(env, rate=1/33.33)\n\n _ = wrapper.reset()\n mock_sleep.assert_not_called() # It should never sleep during reset().\n\n # Step for 100 steps.\n for _ in range(100):\n _ = wrapper.step({\n 'action_type':\n np.array(action_type.ActionType.LIFT, dtype=np.uint8),\n 'touch_position':\n np.array([0.123, 0.456])\n })\n\n # Check that there are 100 calls and that they're all within [0, 1/rate].\n self.assertLen(mock_sleep.call_args_list, 100)\n for call in mock_sleep.call_args_list:\n args, unused_kwargs = call\n sleep_time = args[0]\n self.assertBetween(sleep_time, 0.0, 33.33)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_enabled_sleep_type_before(self, mock_sleep):\n \"\"\"When sleep_type==BEFORE, sleep should come before step().\"\"\"\n\n env = _get_base_env()\n wrapper = rate_limit_wrapper.RateLimitWrapper(\n env,\n rate=1/33.33,\n sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.BEFORE)\n\n _ = wrapper.reset()\n mock_sleep.assert_not_called() # It should never sleep during reset().\n\n @_with_timestamp\n def _sleep_fn(sleep_time):\n _sleep_fn.timestamp = time.time()\n self.assertBetween(sleep_time, 0.0, 33.33)\n\n mock_sleep.side_effect = _sleep_fn\n\n def _step_fn(action):\n self.assertEqual(\n action['action_type'],\n np.array(action_type.ActionType.LIFT, dtype=np.uint8))\n _step_fn.timestamps.append(time.time())\n return dm_env.transition(reward=None, observation=None)\n\n _step_fn.timestamps = []\n\n env.step.side_effect = _step_fn\n\n _ = wrapper.step({\n 'action_type': np.array(action_type.ActionType.LIFT, dtype=np.uint8),\n 'touch_position': np.array([0.123, 0.456])\n })\n\n self.assertLen(_step_fn.timestamps, 1)\n # We expect sleep to have been executed BEFORE a single `step()`.\n self.assertGreaterEqual(_step_fn.timestamps[0], _sleep_fn.timestamp)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_enabled_sleep_type_after(self, mock_sleep):\n \"\"\"When sleep_type==AFTER, sleep should come after step().\"\"\"\n\n env = _get_base_env()\n wrapper = rate_limit_wrapper.RateLimitWrapper(\n env,\n rate=1/33.33,\n sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.AFTER)\n _ = wrapper.reset()\n mock_sleep.assert_not_called() # It should never sleep during reset().\n\n @_with_timestamp\n def _sleep_fn(sleep_time):\n _sleep_fn.timestamp = time.time()\n self.assertBetween(sleep_time, 0.0, 33.33)\n\n mock_sleep.side_effect = _sleep_fn\n\n def _step_fn(action):\n self.assertEqual(\n action['action_type'],\n np.array(action_type.ActionType.LIFT, dtype=np.uint8))\n _step_fn.timestamps.append(time.time())\n return dm_env.transition(reward=None, observation=None)\n\n _step_fn.timestamps = []\n\n env.step.side_effect = _step_fn\n\n _ = wrapper.step({\n 'action_type': np.array(action_type.ActionType.LIFT, dtype=np.uint8),\n 'touch_position': np.array([0.123, 0.456])\n })\n\n # We expect sleep to have been executed AFTER a single `step()`.\n self.assertLen(_step_fn.timestamps, 1)\n self.assertLessEqual(_step_fn.timestamps[0], _sleep_fn.timestamp)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_enabled_sleep_type_after_with_repeat(self, mock_sleep):\n \"\"\"When sleep_type==AFTER_WITH_REPEAT, sleep should be between 2 steps().\"\"\"\n\n env = _get_base_env()\n wrapper = rate_limit_wrapper.RateLimitWrapper(\n env,\n rate=1/33.33,\n sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType\n .AFTER_WITH_REPEAT)\n\n _ = wrapper.reset()\n mock_sleep.assert_not_called() # It should never sleep during reset().\n\n @_with_timestamp\n def _sleep_fn(sleep_time):\n _sleep_fn.timestamp = time.time()\n self.assertBetween(sleep_time, 0.0, 33.33)\n\n mock_sleep.side_effect = _sleep_fn\n\n @_with_timestamp\n def _step_fn(action):\n # On even calls the action should be the actual agent action, but on odd\n # calls they should be REPEATs.\n if len(_step_fn.timestamps) % 2 == 0:\n self.assertEqual(\n action['action_type'],\n np.array(action_type.ActionType.LIFT, dtype=np.uint8))\n else:\n self.assertEqual(\n action['action_type'],\n np.array(action_type.ActionType.REPEAT, dtype=np.uint8))\n _step_fn.timestamps.append(time.time())\n return dm_env.transition(reward=1.0, observation=None)\n\n _step_fn.timestamps = []\n\n env.step.side_effect = _step_fn\n\n timestep = wrapper.step({\n 'action_type': np.array(action_type.ActionType.LIFT, dtype=np.uint8),\n 'touch_position': np.array([0.123, 0.456])\n })\n\n # When the wrapper is enabled, base step should be called twice.\n self.assertEqual(env.step.call_count, 2)\n\n # `step()` should be called twice: before `sleep()` and after it.\n self.assertLen(_step_fn.timestamps, 2)\n self.assertGreaterEqual(_sleep_fn.timestamp, _step_fn.timestamps[0])\n self.assertLessEqual(_sleep_fn.timestamp, _step_fn.timestamps[1])\n # Rewards should accumulate over the two step() calls\n self.assertEqual(timestep.reward, 2.0)\n\n @mock.patch.object(time, 'sleep', autospec=True)\n def test_enabled_sleep_type_after_with_repeat_last(self, mock_sleep):\n \"\"\"If the first step is a LAST, second step should not be taken.\"\"\"\n\n env = _get_base_env()\n wrapper = rate_limit_wrapper.RateLimitWrapper(\n env,\n rate=1/33.33,\n sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType\n .AFTER_WITH_REPEAT)\n\n _ = wrapper.reset()\n mock_sleep.assert_not_called() # It should never sleep during reset().\n\n env.step.return_value = dm_env.termination(reward=None, observation=None)\n\n _ = wrapper.step({\n 'action_type': np.array(action_type.ActionType.LIFT, dtype=np.uint8),\n 'touch_position': np.array([0.123, 0.456])\n })\n\n # Second step call should be skipped.\n env.step.assert_called_once()\n mock_sleep.assert_not_called()\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "android_env/wrappers/tap_action_wrapper.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wraps the AndroidEnv environment to provide tap actions of a given duration.\"\"\"\n\nfrom collections.abc import Sequence\nfrom typing import Any\n\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.wrappers import base_wrapper\nimport dm_env\nimport numpy as np\n\n\nclass TapActionWrapper(base_wrapper.BaseWrapper):\n \"\"\"AndroidEnv with tap actions.\"\"\"\n\n def __init__(self,\n env: env_interface.AndroidEnvInterface,\n num_frames: int = 5,\n touch_only: bool = False) -> None:\n super().__init__(env)\n assert 'action_type' in env.action_spec()\n self._touch_only = touch_only\n self._num_frames = num_frames\n self._env_steps = 0\n\n def stats(self) -> dict[str, Any]:\n \"\"\"Returns a dictionary of metrics logged by the environment.\"\"\"\n logs = self._env.stats()\n logs.update({'env_steps': self._env_steps})\n return logs\n\n def _process_action(\n self, action: dict[str, np.ndarray]\n ) -> Sequence[dict[str, np.ndarray]]:\n if self._touch_only:\n assert action['action_type'] == 0\n touch_action = action.copy()\n touch_action['action_type'] = np.array(\n action_type.ActionType.TOUCH\n ).astype(self.action_spec()['action_type'].dtype)\n actions = [touch_action] * self._num_frames\n lift_action = action.copy()\n lift_action['action_type'] = np.array(action_type.ActionType.LIFT).astype(\n self.action_spec()['action_type'].dtype\n )\n actions.append(lift_action)\n\n else:\n if action['action_type'] == action_type.ActionType.TOUCH:\n actions = [action] * self._num_frames\n lift_action = action.copy()\n lift_action['action_type'] = np.array(\n action_type.ActionType.LIFT\n ).astype(self.action_spec()['action_type'].dtype)\n actions.append(lift_action)\n else:\n actions = [action] * (self._num_frames + 1)\n\n return actions\n\n def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep:\n \"\"\"Takes a step in the environment.\"\"\"\n self._env_steps += self._num_frames + 1\n actions = self._process_action(action)\n total_reward = 0.0\n for idx in range(len(actions)):\n step_type, reward, discount, observation = self._env.step(actions[idx])\n if reward:\n total_reward += reward\n if step_type == dm_env.StepType.LAST:\n return dm_env.TimeStep(\n step_type=step_type,\n reward=total_reward,\n discount=discount,\n observation=observation)\n return dm_env.TimeStep(\n step_type=step_type,\n reward=total_reward,\n discount=discount,\n observation=observation)\n\n def action_spec(self) -> dict[str, dm_env.specs.Array]:\n if self._touch_only:\n return {\n 'action_type':\n dm_env.specs.DiscreteArray(num_values=1, name='action_type'),\n 'touch_position':\n self._env.action_spec()['touch_position'],\n }\n else:\n return self._env.action_spec()\n"
},
{
"path": "android_env/wrappers/tap_action_wrapper_test.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for tap_action_wrapper.\"\"\"\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\nfrom android_env import env_interface\nfrom android_env.components import action_type\nfrom android_env.wrappers import tap_action_wrapper\nimport dm_env\nfrom dm_env import specs\nimport numpy as np\n\n\ndef _make_array_spec(shape, dtype, name):\n return specs.BoundedArray(\n name=name,\n shape=shape,\n dtype=dtype,\n minimum=np.zeros(shape),\n maximum=np.ones(shape), # maximum is inclusive.\n )\n\n\nclass TapActionWrapperTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self._base_action_spec = {\n 'action_type': specs.DiscreteArray(\n num_values=3, name='action_type'),\n 'touch_position': _make_array_spec(\n shape=(2,), dtype=np.float32, name='touch_position'),\n }\n self.base_env = mock.create_autospec(env_interface.AndroidEnvInterface)\n self.base_env.action_spec.return_value = self._base_action_spec\n\n def test_process_action_repeat(self):\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=3)\n action = {\n 'action_type': np.array(action_type.ActionType.REPEAT, dtype=np.int32),\n 'touch_position': np.array([0.5, 0.5], dtype=np.float32),\n }\n actions = wrapped_env._process_action(action)\n self.assertLen(actions, wrapped_env._num_frames + 1)\n self.assertEqual(action, actions[-1])\n\n def test_process_action_lift(self):\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=3)\n action = {\n 'action_type': np.array(action_type.ActionType.LIFT, dtype=np.int32),\n 'touch_position': np.array([0.5, 0.5], dtype=np.float32),\n }\n actions = wrapped_env._process_action(action)\n self.assertLen(actions, wrapped_env._num_frames + 1)\n self.assertEqual(action, actions[-1])\n\n def test_process_action_touch(self):\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=3)\n action = {\n 'action_type': np.array(action_type.ActionType.TOUCH, dtype=np.int32),\n 'touch_position': np.array([0.5, 0.5], dtype=np.float32),\n }\n actions = wrapped_env._process_action(action)\n self.assertLen(actions, wrapped_env._num_frames + 1)\n self.assertEqual(\n actions[-1]['action_type'], np.array(action_type.ActionType.LIFT)\n )\n\n def test_reset(self):\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=5)\n fake_timestep = 'ts'\n self.base_env.reset.return_value = fake_timestep\n ts = wrapped_env.reset()\n self.base_env.reset.assert_called_once()\n self.assertEqual(fake_timestep, ts)\n\n def test_step(self):\n # Arrange.\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=5)\n fake_timestep = dm_env.TimeStep(\n step_type='fake_type',\n reward=0.0,\n discount=1.0,\n observation='fake_obs')\n self.base_env.step.return_value = fake_timestep\n self.base_env.stats.return_value = {}\n\n # Act.\n ts = wrapped_env.step({\n 'action_type': np.array(action_type.ActionType.REPEAT, dtype=np.int32),\n 'touch_position': np.array([0.5, 0.5], dtype=np.float32),\n })\n stats = wrapped_env.stats()\n\n # Assert.\n self.assertEqual(wrapped_env._num_frames+1, self.base_env.step.call_count)\n self.assertIsInstance(ts, dm_env.TimeStep)\n self.assertIsInstance(stats, dict)\n self.assertIn('env_steps', stats)\n self.assertEqual(stats['env_steps'], 6)\n\n def test_observation_spec(self):\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=5)\n fake_obs_spec = 'fake_obs_spec'\n self.base_env.observation_spec.return_value = fake_obs_spec\n observation_spec = wrapped_env.observation_spec()\n self.base_env.observation_spec.assert_called_once()\n self.assertEqual(fake_obs_spec, observation_spec)\n\n def test_action_spec(self):\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=5)\n self.base_env.action_spec.return_value = self._base_action_spec\n action_spec = wrapped_env.action_spec()\n self.base_env.action_spec.assert_called()\n self.assertEqual(self.base_env.action_spec(),\n action_spec)\n\n def test_stats(self):\n \"\"\"Checks that returned stats have expected properties.\"\"\"\n\n # Arrange.\n self.base_env.stats.return_value = {\n 'some_key': 12345,\n 'another_key': 5.4321,\n }\n wrapped_env = tap_action_wrapper.TapActionWrapper(\n self.base_env, num_frames=5\n )\n\n # Act.\n stats = wrapped_env.stats()\n\n # Assert.\n self.assertIsInstance(stats, dict)\n # Original entries should still be present.\n self.assertIn('some_key', stats)\n self.assertEqual(stats['some_key'], 12345)\n self.assertIn('another_key', stats)\n self.assertEqual(stats['another_key'], 5.4321)\n # TapActionWrapper inserts its own `env_steps`.\n self.assertIn('env_steps', stats)\n self.assertEqual(stats['env_steps'], 0)\n\n\nif __name__ == '__main__':\n absltest.main()\n"
},
{
"path": "docs/emulator_guide.md",
"content": "# AndroidEnv - Emulator Setup Guide\n\nIn this document we provide a step-by-step guide for creating a virtual Android\ndevice with Android Studio. After creating an AVD\n([Android Virtual Device](https://developer.android.com/studio/run/managing-avds))\nyou will be able to connect it to an AndroidEnv instance and you're ready to go.\n\nTo get started, you will need to download\n[Android Studio](https://developer.android.com/studio) - an IDE widely used by\nAndroid developers.\n\n## Install an SDK Platform Package\n\nAndroid Studio comes with the Android Software Development Toolkit (SDK) which,\namong others, allows you to install different versions of Android. Click on\n**Tools** > **SDK Manager** and select the SDK version that you would like to\nuse.\n\n\n\nWe recommend that you set the `Android SDK Location` to be in your home\ndirectory (for example, on Linux the default one is `~/Android/Sdk`, while on\nmacOS - `~/Library/Android/sdk`). You can always find the SDK location in\nAndroid Studio under **Preferences** > **Appearance & Behavior** > **System\nSettings** > **Android SDKs** > _Android SDK Location_.\n\nIf you set the the custom `Android SDK Location`, make note of it - you will\nneed it for connecting AndroidEnv to your AVD.\n\n\n\n## Create an AVD\n\nNow it is time to create a virtual device (AVD). Go to **Tools** > **AVD\nManager**.\n\n\n\nIn the pop-up window you will find an option to **Create Virtual Device**.\n\n\n\nConfigure the virtual device. You can select the model or choose from more\nadvanced settings (refer to the\n[Android docs](https://developer.android.com/studio/run/managing-avds) for\nstep-by-step instructions).\n\n\n\nName your AVD and take note of this value. It will be neccessary for connecting\nAndroidEnv to this virtual device.\n\n\n\nOnce you are done, you will see the new AVD show up in the **AVD Manager**.\nClick on **View details** to inspect some of its properties.\n\n\n\nTake note of the `AVD Path`. This value will be neccessary for connecting\nAndroidEnv to this device. We recommend that you set this to be your home\ndirectory (for instance, on Linux or macOS it may be `~/.android/avd`).\n\n\n\n## Ready to use\n\nWith SDK and AVD both set up, you are now ready to use this emulated device with\nAndroidEnv. Don't forget to take note of the following three values: your AVD\nname, the AVD path, and the SDK path. For example, on Linux they may be:\n\n```\n--avd_name=my_avd\n--avd_package_path=~/.android/avd\n--android_sdk_root=~/Android/Sdk\n```\n\nNext, once you have set up the AVD, follow the\n[Task steps](instructions.md#the-task) in the\n[Running the environment guide](instructions.md) to finish setting up\nAndroidEnv.\n\nHowever, if you want to just interact with the newly created device, click on\nthe run button next to your AVD in the **AVD Manager** in Android Studio (this\nstep is optional).\n\n\n\nYou will see an emulator window pop up. You can interact with it by clicking on\nthe screen.\n\n\n\nThere are many other features in Android Studio that let you customize your\ndevice. For example, you can create custom images with pre-installed\napplications or configured settings.\n"
},
{
"path": "docs/environment.md",
"content": "# AndroidEnv - Environment features\n\nAndroidEnv is a complex environment that, while offering an almost endless range\nof possibilites for RL research and investigation, poses multiple kinds of\nchallenges simultaneously. In this document we outline AndroidEnv's main\nfeatures that render it such a unique learning environment.\n\n## Real-time environment\n\nAndroidEnv is built on top of an emulated Android device, allowing the agent to\ncommunicate with the emulator through touch actions. Android emulators are\ncreated independently from our environment implementation and simulate real\nAndroid devices in the most realistic manner. This simulation runs real-time,\nindependently of the agent, meaning that the simulaton will not wait for agent\ninput between frames. This aspect of the environment renders the RL setup\nsimilar to a robotics problem, where the challenges of real-time interaction and\nconsequent noise in observations have to be overcome. Please note there is\ncurrently no straightforward way to slow down the simulation either.\n\n## Action space\n\n![](\"images/classic_2048.gif\")
\n\nPerhaps one of the most interesting features of AndroidEnv is its large and\ncomplex action interface. The raw action space of the environment consists of a\ntuple `(x,y) in [0,1]x[0,1]` determining the location of the action on the\nAndroid screen, and a discrete value `ActionType in {LIFT, TOUCH, REPEAT}`\nindicating whether the agent wants to touch the screen at this chosen location\nor not. This action space is uniform across all tasks/apps.\n\n**Gestures.** The complexity of the interface arises from the fact that\nindividual raw actions on their own do not neccessarily trigger a meaningful\nchange in the environment. Most Android applications are designed such that they\ncan be controlled/navigated through common touchscreen gestures such as pressing\nbuttons, swiping, scrolling, pinching, drag and drop etc. Each of these can be\nthought of as particular sequences of raw actions: for example, *touching* the\nscreen at a particular location, then immediately *lifting* the imaginary finger\nmight be interpreted as a *press of a button*; while a sequence of *touches*\naligned in a vertical line might be interpreted as *scrolling*. We note that\nAndroidEnv does not support multitouch actions at the moment, but it is a\npossible feature to add.\n\nIt is important to point out that it is out of the environment's control to\ndetermine how particular sequences of raw actions get interpreted by the Android\nsimulator - much like when humans interact with physical devices, a certain\ngesture on the screen might be interpreted differently if it is performed at a\nslightly different angle or speed.\n\nTap | Double Tap | Touch & Hold | Flick Left | Flick Right | Scroll (H) | Scroll (V) | Drag & Drop\n-------------------------------------------------------- | ---------------------------------------------------------------------- | ------------------------------------------------------------------------ | ---------------------------------------------------------------------- | ------------------------------------------------------------------------ | ------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------- | -----------\n |  |  |  |  |  |  | \n\n**Wrappers.** It is possible to alter the raw action space of the environment by\napplying [wrappers](#wrappers). For example one might discretize the action\nspace by splitting the screen up into a grid of a desired size; restrict the\nActionType to *touch* only; or fix certain gesture skills. We note here that\nthese wrappers, again, will not alter how the particular sequence of performed\nraw actions gets interpreted by the Android simulator.\n\n## Observation space\n\nThe observation space of AndroidEnv consists of three main components:\n(`pixels`, `timedelta`, `orientation`), the most notable of these being\n`pixels`. The original screen size will depend on the type of emulator used, but\ngiven that it will correspond to real device screen sizes, this will usually be\nquite large (of course, this can be scaled down, e.g. with wrappers). The\n`timedelta` component captures the amount of time passed since the last\nobservation was fetched. The `orientation`, even though it does not affect the\nlayout of the RGB image in the observation, might carry relevant information for\nthe agent. For example, if there is text on the screen, it is important to know\nhow it is oriented. Again, a benefit of this observation space is that it is\nuniform across all tasks. As mentioned above, observations often carry spatial\ncues and are suggestive of the kind of actions/gestures that are meaningful to\nperform in a given state.\n\n## Task extras\n\nOn top of the default observations (`pixels`, `timedelta`, `orientation`), some\ntasks might expose additional structured observations after each step. An\n*extra* in AndroidEnv is any information that an app may send to aid the\nunderstanding of the task. The type of information sent through this channel is\nusually something difficult to obtain from raw pixels and may include meaningful\ninformation such as:\n\n* The current board configuration (e.g. of a chess game or of a tetris game)\n in matrix or string form.\n* The position of the avatar in a map.\n* Events such as whether a button was pressed or whether a checkpoint was\n achieved.\n\nNote that these are entirely optional and may not be available at all.\n\nTo request extras from the environment, you can call `env.task_extras()` after\neach `env.step()`, which will return a dictionary of all the extra observations\nobserved during the previous step (or an empty dict is there's none available).\nFor example:\n\n```python\nfor _ in range(FLAGS.n_steps):\n\n action = agent.select_action(timestep.observation)\n timestep = env.step(action)\n logging.info('observation: %s', timestep.observation)\n logging.info('extra observations: %s', env.task_extras())\n```\n\nPlease note however that the env might not return extras at every timestep, only\nwhen something meaningful happened (e.g. only when a button was pressed, or when\nthe state of the board has changed).\n\nWhen integrating your own APK as a new task for the environment, you can define\nyour own extras by following the instructions\n[here](tasks_guide.md#log-messages-and-custom-apks).\n\n## Wrappers\n\nAndroidEnv's action- and observation spaces can be altered by applying suitable\nwrappers. While custom wrappers can be built easily, we have provided a number\nof useful wrappers that demonstrate their usage:\n\n* `discrete_action_wrapper`: Discretizes the action space into an `nxk` grid.\n* `flat_interface_wrapper`: Removes the dictionary structure from the\n observation and action specs.\n* `float_pixels_wrapper`: Projects the pixel RGB values from the integer range\n `[0, 255]` to the float range `[0, 1]`.\n* `image_rescale_wrapper`: Resizes the pixel observations by the selected\n ratio.\n* `gym_wrapper`: Changes the environment interface from\n [dm_env](https://github.com/deepmind/dm_env) to\n [OpenAI](https://gym.openai.com/) gym interface.\n* `last_action_wrapper`: Extends the observation with a one-hot encoded\n location of the previously taken action, in order to aid agents without\n built-in memory.\n\n## Internal structure of AndroidEnv\n\nThe chart below gives an overview of the internal workings of the system,\nillustrating how different classes interact with each other and what their\nindividual roles are. See the source code for more details.\n\n\n"
},
{
"path": "docs/example_tasks.md",
"content": "# AndroidEnv - Available tasks\n\nThis page gives a detailed overview of the example tasks provided with\nAndroidEnv. The purpose is to give researchers an idea of the different kinds of\nchallenges that AndroidEnv poses.\n\nTo use any of these tasks in your own experiments, click on **Download** to\ndownload a ZIP file containing textprotos and the corresponding APKs. After\ndownloading, move the `.apk` and `.textproto` files to a directory of your\nchoice and take note of their path. This information is needed for\n[running](instructions.md#create-the-env) an AndroidEnv instance with the given\ntask.\n\n\n\n| App / Game | Interface | Time reactive | Multi-level | Rewards | Extras | Download |\n| --------------------------------------------------- | -------------------- | -------------- | ----------------------------- | ---------- | ------------ | -------- |\n| [Vokram (MDP)](#vokram) | Tapping (buttons) | No | Yes (4 variants) | Dense | Yes | [Download](https://storage.googleapis.com/android_env-tasks/mdp.tar.gz) |\n| [Apple Flinger](#apple-flinger) | Drag & drop | No | Yes (6 variants) | Dense | No | [Download](https://storage.googleapis.com/android_env-tasks/apple_flinger.tar.gz) |\n| [Blockinger](#blockinger) | Tapping (buttons) | Yes | No | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/blockinger.tar.gz) |\n| [Catch](#catch) | Touch | Yes | No | Dense | Yes | [Download](https://storage.googleapis.com/android_env-tasks/catch_the_ball.tar.gz) |\n| [Classic 2048](#classic-2048) | Swiping | No | No | Dense | Yes | [Download](https://storage.googleapis.com/android_env-tasks/classic_2048.tar.gz) |\n| [Dodge](#dodge) | Tapping | Yes | No | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/dodge.tar.gz) |\n| [DroidFish (Chess)](#droidfish) | Tapping | No | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/droidfish.tar.gz) |\n| [FlappyDroid](#flappydroid) | Tapping | Yes | Yes (2 levels) | Dense | No | [Download](https://storage.googleapis.com/android_env-tasks/systemui_egg_land.tar.gz) |\n| [FloodIt](#floodit) | Tapping (buttons) | No | Yes (4 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/floodit.tar.gz) |\n| [Frozen Bubble](#frozen-bubble) | Dragging, tapping | No | No | Sparse | No | [Download](https://storage.googleapis.com/android_env-tasks/frozen_bubble.tar.gz) |\n| [Memory Game](#memory-game) | Tapping | No | Yes (6 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/memory_game.tar.gz) |\n| [Minesweeper](#minesweeper) | Tapping | No | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/minesweeper.tar.gz) |\n| [Nostalgic Racer](#nostalgic-racer) | Touch | Yes | Yes (2 variants) | Dense | Yes | [Download](https://storage.googleapis.com/android_env-tasks/nostalgic_racer.tar.gz) |\n| [Open Sudoku](#open-sudoku) | Tapping (buttons) | No | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/open_sudoku.tar.gz) |\n| [Perfection](#perfection) | Drag & drop | No | Yes (3 game types) | Dense | Yes | [Download](https://storage.googleapis.com/android_env-tasks/perfection.tar.gz) |\n| [Rocket Sleigh](#rocket-sleigh) | Tapping | Yes | No | Dense | No | [Download](https://storage.googleapis.com/android_env-tasks/rocket_sleigh.tar.gz) |\n| [Pong](#pong) | Drag | Yes | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/pong.tar.gz) |\n| [SGT Puzzles - Blackbox](#sgt-puzzles-blackbox) | Tapping | No | Yes (4 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Bridge](#sgt-puzzles-bridge) | Drag & drop | No | Yes (5 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Cube](#sgt-puzzles-cube) | Tapping | No | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Dominosa](#sgt-puzzles-dominosa) | Tapping | No | Yes (5 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Fifteen](#sgt-puzzles-fifteen) | Tapping | No | Yes (4 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Flip](#sgt-puzzles-flip) | Tapping | No | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/apple_flinger.tar.gz) |\n| [SGT Puzzles - Flood](#sgt-puzzles-flood) | Tapping | No | Yes (3 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Galaxies](#sgt-puzzles-galaxies) | Tapping | No | Yes (6 sizes) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Guess](#sgt-puzzles-guess) | Tapping | No | Yes (4 levels) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Inertia](#sgt-puzzles-inertia) | Tapping | No | Yes (2 sizes) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Light Up](#sgt-puzzles-light-up) | Tapping | No | Yes (5 sizes) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Loopy](#sgt-puzzles-loopy) | Tapping | No | Yes (3 sizes) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [SGT Puzzles - Net](#sgt-puzzles-net) | Tapping | No | Yes (5 sizes) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/sgtpuzzles.tar.gz) |\n| [Shattered Pixel Dungeon](#shattered-pixel-dungeon) | Tapping | Yes | Yes (4 variants) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/shattered_pixel_dungeon.tar.gz) |\n| [Simple Solitaire](#simple-solitaire) | Drag & drop | No | Yes (19 tasks) | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/simple_solitaire.tar.gz) |\n| [Snake](#snake) | Tapping (buttons) | Yes | No | Sparse | Yes | [Download](https://storage.googleapis.com/android_env-tasks/aosp_samples_snake.tar.gz) |\n| [Vector Pinball](#vector-pinball) | Tapping | Yes | Yes (5 variants) | Sparse | No | [Download](https://storage.googleapis.com/android_env-tasks/vector_pinball.tar.gz) |\n\n\n\n## Vokram\n\nVokram is our in-house implementation of an Android app that displays a\nMarkov-Decision-Process (MDP) graph as buttons on the screen which the agent\nmust use to select its actions. The observation is simply the color of the\nbackground, and the actions are the buttons themselves which are presented in\ndifferent colors.\n\n* **mdp_0000**: This is a task that presents the agent with two colored, but\n unlabeled buttons on the screen. Pressing one of the buttons gives the agent\n a reward of `-1` and redraws the buttons on the screen. The other button\n gives a reward of `+1` and terminates the episode. The color of the buttons\n is the same throughout the episode. The sizes of the buttons are randomized\n at each screen draw. Pressing anywhere else on the screen gives a reward of\n zero. The task lasts up to 60 seconds, at which point the episode is\n restarted. The underlying dynamics governing the buttons is a simple 2-state\n 2-action Markov Decision Process (MDP). The MDP is an intentionally simple\n environment that can be used to debug agents.\n\n* **mdp_0001**: This is similar to `mdp_0000` but it's even simpler. It\n presents the agent with a single button which gives a reward of `+1` and\n terminates the episode when pressed. This task can be used for example to\n train agents to click buttons.\n\n* **mdp_0002**: In this task there are two buttons, pressing either of which\n will terminate the episode with a return of `+1`.\n\n* **mdp_0003**: An equivalent of `mdp_0000` with rewards reversed: the episode\n ends when the wrong button is clicked, and carries on with a new set of\n buttons when the correct one is clicked.\n\n\n Extras returned
\n\n* `actions`:\n - Set of all buttons present, e.g. `['A', 'B']`.\n - Returned when any button is pressed.\n - Has `shape=[2], dtype=STRING_U1`.\n* `clicks`:\n - Character representing the button pressed.\n - Returned when any button is pressed.\n - Has `shape=[1], dtype=STRING_U1`.\n* `buttons`:\n - Coordinates of the top left and bottom right corners of each button, e.g\n `[[x_a_0, y_a_0, x_a_1, y_a_1], [x_b_0, y_a_0, x_b_1, y_b_1]]`.\n - Returned when any button is pressed.\n - Has `shape=[2, 4], dtype=INT32`.\n\n \n\n**mdp_0000** | **mdp_0001** | **mdp_0002** | **mdp_0003**\n------------------------------------------------ | ------------------------------------------------ | ------------------------------------------------ | ------------\n |  |  | \n\n## Apple Flinger\n\nA clone of Angry Birds. Even though the game offers many levels, we currently\nexpose six levels. See the original github repo for more info:\nhttps://gitlab.com/ar-/apple-flinger.\n\n\n Extras returned
\n Returns no extras.\n \n\n**apple_flinger_M_1_1** | **apple_flinger_M_1_2** | **apple_flinger_M_1_18**\n---------------------------------------------------------------------- | ---------------------------------------------------------------------- | ------------------------\n |  | \n\n**apple_flinger_M_2_1** | **apple_flinger_M_2_2** | **apple_flinger_M_2_18**\n----------------------------------------------------------------------- | ---------------------------------------------------------------------- | ------------------------\n) |  | \n\n## Blockinger\n\nThis is a Tetris clone implemented with on-screen controls. See the original\ngithub repo for more info: https://github.com/tasioleiva/Blockinger.git.\n\n\n Extras returned
\n\n* `down_pressed`, `left_pressed`, `right_pressed`, `rotate_right_pressed`,\n `drop_pressed`:\n - Indicates that said button has been pressed.\n - Returned when said button has been pressed.\n - Has `shape=[1], dtype=INT32`.\n* `current_board`:\n - One-hot encoded state of the board.\n - Has `shape=[18, 10], dtype=INT32`.\n* `current_line`, `cleared_lines`:\n - Index of the relevant line.\n - Has `shape=[1], dtype=INT32`.\n* `current_piece`, `next_piece`:\n - Index representing the type of piece.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n\n\n## Catch\n\nClassic Catch game.\n\n\n Extras returned
\n\n* `ball`:\n - `x, y` coordinates of the ball.\n - Returned every timestep.\n - Has `shape=[2], dtype=INT32`.\n* `paddle`:\n - `x, y` coordinates of the paddle.\n - Returned every timestep.\n - Has `shape=[2], dtype=INT32`.\n* `paddle_width`:\n - Width of the paddle.\n - Returned every timestep.\n - Has `shape=[1], dtype=INT32`.\n* `lives`:\n - Number of lives left.\n - Returned every timestep.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n\n\n## Classic 2048\n\nThis is an Android implementation of a popular game in the 2010s. See the\noriginal github repo for more info: https://github.com/tpcstld/2048.git.\n\n\n Extras returned
\n\n* `grid`:\n - State of the board.\n - Returned when the board changes.\n - Has `shape=[4, 4], dtype=INT32`.\n* `direction`:\n - Index representing the direction of the last swipe (between 0-3).\n - Returned when the swipe prompted a board change.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n\n\n## Dodge\n\nGuide the ball from the red line to the green line without getting hit by the\nfloating dots.\n\n\n Extras returned
\n\n* `lives`:\n - Number of lives left.\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n* `level`:\n - Current level.\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n\n\n## DroidFish\n\nStandard chess game. You can choose whether to play as a specific player\n(black/white), or have the player colour randomly assigned at the beginning of\neach episode. The numbers 1, 10 and 100 indicate the level of difficulty. Take a\nlook at a few sample moves below to get an idea of roughly how well the bot\nplays for each level of difficulty. You can see that the 1% and 10% bots often\nmake very obvious mistakes. See the original github repo for more info:\nhttps://github.com/peterosterlund2/droidfish.\n\n\n Extras returned
\n\n* `board`:\n - State of the board, representing pieces by indices. No piece - 0\n - White pieces - 1: king, 2: queen, 3: rook, 4: bishop, 5: knight, 6: pawn\n - Black pieces - 7: king, 8: queen, 9: rook, 10: bishop, 11: knight, 12:\n pawn\n - Returned when the board changes.\n - Has `shape=[8, 8], dtype=INT32`.\n* `selection`:\n - Coordinate of selected piece (between 0-64, -1 if selection is removed)\n - Returned when a piece is selected (or unselected).\n - Has `shape=[1], dtype=INT32`.\n* `moved`:\n - Coordinates \"from\" and \"to\" cells when a piece is moved (between 0-64)\n - Returned when a piece is moved.\n - Has `shape=[2], dtype=INT32`.\n* `invalid`:\n - Coordinates \"from\" and \"to\" cells of an invalid move attempt (between\n 0-64)\n - Returned upon invalid move request\n - Has `shape=[2], dtype=INT32`.\n\n \n\n**droidfish_black_1** | **droidfish_black_10** | **droidfish_black_100**\n------------------------------------------------------------------ | -------------------------------------------------------------------- | -----------------------\n |  | \n\n**droidfish_white_1** | **droidfish_white_10** | **droidfish_white_100**\n------------------------------------------------------------------ | -------------------------------------------------------------------- | -----------------------\n |  | \n\n**droidfish_random_1** | **droidfish_random_10** | **droidfish_random_100**\n-------------------------------------------------------------------- | ---------------------------------------------------------------------- | ------------------------\n |  | \n\n## FlappyDroid\n\nA clone of the well-known game Flappy Birds.\n\n\n Extras returned
\n Returns no extras.\n \n\n**systemui_egg_land_default** | **systemui_egg_land_half_speed**\n---------------------------------------------------------------------------------- | --------------------------------\n | \n\n## FloodIt\n\nFloodIt is a game where the player needs to fill the board with a single color.\nThe dynamics of the game are driven by a few colorful buttons at the bottom of\nthe screen, which when pressed cause the currently active region to change its\ncolor to the color of the pressed button. When this active region changes color\nit absorbs neighboring squares that have the same color, thus expanding the\nactive region. The active region starts as a single square at the top-left\ncorner of the board. The game gives a single reward at the end of the game if\nthe player manages to fill the entire board with the same color within the\nmaximum number steps, otherwise the reward is just zero.\n\nThis is a very hard-exploration game because the game does not give intermediate\nrewards until the very end of the game, and the number of possible moves is\nincredibly big.\n\nYou can find another implementation of this game in the task\n[SGT Puzzles - Flood](#sgt-puzzles-flood).\n\n\n Extras returned
\n\n* `board`:\n - State of the board, representing colours by their indices.\n - 0: purple, 1: blue, 2: green, 3: yellow, 4: red, 5: pink\n - Returned when the board changes.\n - Has `shape=[board_size, board_size], dtype=INT32`.\n* `clicked`:\n - Index of the colour clicked (between 0-5)\n - Returned when said colour is clicked.\n - Has `shape=[1], dtype=INT32`.\n* `flipped`:\n - The number of new cells that just got merged into the big blob (0 or\n more)\n - Returned when the board state changes\n - Has `shape=[1], dtype=INT32`.\n\n \n\n**floodit_easy** | **floodit_medium** | **floodit_hard**\n-------------------------------------------------------- | ------------------------------------------------------------ | ----------------\n |  | \n\n### Task `mdp_flood_it`\n\nCustom task created for pretraining agents to locate and press FloodIt buttons\non the screen.\n\n\n\n## Frozen Bubble\n\nShoot the coloured bubbles in a direction of your choice. Groups of bubbles with\nthe same colour will drop. Remove all bubbles from the board before the time\nruns out. See the original github repo for more info:\nhttps://github.com/robinst/frozen-bubble-android.git.\n\n\n Extras returned
\n Returns no extras.\n \n\n\n\n## Memory Game\n\nClassic memory game. Find the pairs of images. See the original github repo for\nmore info: https://github.com/sromku/memory-game/.\n\n\n Extras returned
\n\n* `flip`:\n - Index of the card flipped\n - Returned when a card is clicked.\n - Has `shape=[1], dtype=INT32`.\n* `cards`:\n - Number of cards still on the board.\n - Returned upon finding a pair.\n - Has `shape=[1], dtype=INT32`.\n* `remained`:\n - Number of cards remaining at the end of the episode.\n - Returned when an episode is over.\n - Has `shape=[1], dtype=INT32`.\n* `stars`:\n - Number of stars achieved at the end of the game.\n - Returned upon finishing the game.\n - Has `shape=[1], dtype=INT32`.\n* `achieved`:\n - Score obtained by the end of the episode.\n - Returned when an episode is over.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n**memory_game_animals_beginner** | **memory_game_animals_easy** | **memory_game_monsters_medium**\n---------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- | -------------------------------\n |  | \n\n**memory_game_monsters_hard** | **memory_game_emojis_hardest** | **memory_game_emojis_master**\n---------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ | -----------------------------\n |  | \n\n## Minesweeper\n\nThis is an Android implementation of a popular game on Desktop in the 1990s. See\nthe original github repo for more info: https://gitlab.com/ar-/apple-flinger.\n\n\n Extras returned
\n\n* `hidden`:\n - Number of hidden cells.\n - Returned whenever the board changes.\n - Has `shape=[1], dtype=INT32`.\n* `revealed`:\n - Number of revealed cells.\n - Returned whenever the board changes.\n - Has `shape=[1], dtype=INT32`.\n* `bombs`:\n - Number of bombs in the game.\n - Returned whenever the board changes.\n - Has `shape=[1], dtype=INT32`.\n* `click`:\n - Coordinates of the cell clicked (row, column).\n - Returned whenever the board changes.\n - Has `shape=[2], dtype=INT32`.\n* `grid`:\n - State of the board.\n - -1 = hidden, -2 = marked, 9 = bomb, 0-8 = number of nearby bombs\n - Returned whenever the board changes.\n - Has `shape=[grid_height, grid_width], dtype=INT32`.\n\n \n\n**minesweeper_easy** | **minesweeper_medium** | **minesweeper_hard**\n---------------------------------------------------------------- | -------------------------------------------------------------------- | --------------------\n |  | \n\n## Nostalgic Racer\n\nNostalgicRacer is a racing game that offers Atari-like graphics and controls.\nThe objective is to maximize the score which increases as the car moves forward\nand by collecting coins and speed-ups.\n\n\n Extras returned
\n Returns no extras.\n \n\n### Task `nostalgic_racer`\n\nThe player can only control whether the car should move left, move right or stay\nput by touching on the screen. If the touch is on the right pane the car moves\nright, if the touch is on the left pane the car moves left and no touches leaves\nthe car in the same position. Pressing for too little time moves the car by\nminiscule amounts, with effects similar to staying put.\n\n### Task `nostalgic_racer_2d`\n\nThis is the same underlying game as NostalgicRacer with the same objective.\nHowever, the interface is very different. The observation is given as a 2D view\nfrom the top with no perspective and the touchscreen determines the position\nthat the car should move to (sideways).\n\n**nostaligc_racer** | **nostalgic_racer_2d**\n-------------------------------------------------------------- | ----------------------\n | \n\n## Open Sudoku\n\nClassic Sudoku game with different levels of difficulty. The board is randomised\nover a set of 30 boards for each level. See the original github repo for more\ninfo: https://github.com/ogarcia/opensudoku.\n\n\n Extras returned
\n\n* `value`:\n - Number pressed (between 1-9, 0 if the \"delete\" button is pressed).\n - Returned upon clicking said button.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n**open_sudoku_easy** | **open_sudoku_medium** | **open_sudoku_hard**\n---------------------------------------------------------------- | -------------------------------------------------------------------- | --------------------\n |  | \n\n## Perfection\n\nDrag the items corresponding to the targets with the same shape.\n\n\n Extras returned
\n\n* `moving`:\n - The ID of the piece being dragged on the screen or 0.\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n* `todo`:\n - Number of pieces yet to be moved to a hole.\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n* `done`:\n - Number of pieces correctly moved to a hole.\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n**perfection_1_circle_static** | **perfection_1_cube_static** | **perfection_1_plus_static** | **perfection_1_triangle_static**\n------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------- | --------------------------------\n |  |  | \n\n**perfection_default** | **perfection_4_colors_square_static** | **perfection_4_pieces_static**\n-------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | ------------------------------\n |  | \n\n## Rocket Sleigh\n\nA Flappy Bird-like game where you have to collect christmas presents while\navoiding trees. The sleigh is powered by a rocket that needs to recharge over\ntime after you use up its fuel.\n\n\n Extras returned
\n Returns no extras.\n \n\n\n\n## Pong\n\nClassic Pong game.\n\n\n Extras returned
\n\n* `ball`:\n - The ball coordinates: [left, top, right, bottom].\n - Returned when its value changes.\n - Has `shape=[4], dtype=INT32`.\n* `computer`:\n - The computer paddle coordinates: [left, top, right, bottom].\n - Returned when its value changes.\n - Has `shape=[4], dtype=INT32`.\n* `human`:\n - The human paddle coordinates: [left, top, right, bottom].\n - Returned when its value changes.\n - Has `shape=[4], dtype=INT32`.\n* `collision`:\n - Indicates collision of paddle and ball: (0=no collision, 1=collision).\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n* `state`:\n - The current state of the game: (0=pause, 1=ready, 2=running, 3=lose,\n 4=win).\n - Returned when its value changes.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n**pong_easy** | **pong_default** | **pong_hard**\n-------------------------------------------------- | -------------------------------------------------------- | -------------\n |  | \n\n## SGT Puzzles - Blackbox\n\nThere's an invisible laser beam originating from each of the cells at the edge\nof the grid. There are also a given number of balls inside the grid, hidden from\nthe player whose aim is to guess where those balls are. The player can figure\nout where those balls might be by looking at how they *deflect* the laser beams.\nClicking on an edge cell the player can reveal information about how that\nparticular laser beam travels. Click on a cell; if the cell reveals an `H`, it\nmeans the straight laser beam leaving this cell hits a ball frontally. If the\ncell reveals a *number* along with another cell with the same number, that means\nthe laser beam originating in the first cell ends up getting absorbed in the\ncorresponding pair cell. If the cell reveals an `R`, that means the laser beam\nwas *reflected*: either its origin and the cell it gets absorbed in is the\n*same*, or the beam gets bent before entering the grid. See the description\nbelow.\n\nThe balls affect the travel of the laser beam in the following way:\n\n* If a laser beam hits it straight, it gets absorbed. This is denoted by the\n letter `H`.\n* If a laser beam hits *its corner*, the beam gets deflected by 90 degrees.\n* If a laser beam hits a balls corner right at the edge of the grid, i.e.\n before it enters the grid, it is considered *reflected*.\n* If a laser beam enters the same cell that it originally left, it is\n considered *reflected* too.\n\nOnce the player has placed the given number of balls on the screen, a green dot\nappears that allows the player to check if their solution was correct. See the\noriginal github repo for more info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `balls`:\n - The number of balls in the game arena.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `guesses`:\n - The number of guessed balls made by the agent.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `wrong`:\n - 1 if the guesses are wrong, 0 otherwise.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `lasers`:\n - The number of lasers in the grid\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - Representation of the grid cells:\n - In the arena: `G`=guessed ball `' '`=empty\n - In the range: `[0-9]`=the number of lasers, `H`=beam hit, `R`=beam\n reflected,\n - `?` unknown, `' '` for the corners\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=STRING_U1`.\n\n \n\n**blackbox_3x3_1_ball** | **blackbox_5x5_3_balls** | **blackbox_8x8_5_balls** | **blackbox_10x10_5_balls**\n--------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- | --------------------------\n |  |  | \n\n## SGT Puzzles - Bridge\n\nConnect nodes on the board so that each number denotes the degree of the given\nvertex. Edges are not allowed to cross each other and the graph has to be\nconnected. Edges have to be horizontal or vertical, and there can be at most two\nparallel bridges between any pair of nodes. See the original github repo for\nmore info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `islands`:\n - Number of nodes on the board.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - Representation of the current state of the board.\n - `[0-9]=island, ' '=empty`\n - `'|'=vertical line, '\"'=double vertical line, '!'=wrong vertical line`\n - `'-'=horizontal line, '='=double horizontal line, '~'=wrong horizontal\n line`\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=STRING_U1`.\n\n \n\n**bridge_7x7_easy** | **bridge_7x7_medium** | **bridge_7x7_hard** | **bridge_10x10_medium** | **bridge_15x15_medium**\n------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------- | -----------------------\n |  |  |  | \n\n## SGT Puzzles - Cube\n\nThere are six coloured squares that you have to collect with a moving cube. If\nthe cube rolls on top of a coloured cell, the colour gets attached to that side\nof the cube; and if a coloured side of the cube rolls on an empty cell, the\ncolour is removed from the cube. The goal is to have all six sides of the cube\ncoloured. See the original github repo for more info:\nhttps://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `current`:\n - Index of the current grid cell the cube is on.\n - Returned whenever the cube moves.\n - Has `shape=[1], dtype=INT32`.\n* `previous`:\n - Index of the previous grid cell the cube was on.\n - Returned whenever the cube moves.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - The grid state (0 = dark, 1 = blue)\n - Returned whenever the cube moves.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n* `face_count`:\n - The number of dark faces on the cube.\n - Returned whenever the cube moves.\n - Has `shape=[1], dtype=INT32`.\n* `face_colour_count`:\n - The number of blue faces on the cube.\n - Returned whenever the cube moves.\n - Has `shape=[1], dtype=INT32`.\n* `faces`:\n - The cube faces (0 = dark, 1 = blue)\n - Returned whenever the cube moves.\n - Has `shape=[6], dtype=INT32`.\n\n \n\n**cube_c3x3** | **cube_c4x4** | **cube_c8x8**\n------------------------------------------------------------------------ | ------------------------------------------------------------------------ | -------------\n |  | \n\n## SGT Puzzles - Dominosa\n\nPlace 2x1 size dominoes on the board such that the full board is covered, making\nsure that no two dominoes have the same pair of numbers on them. There needs to\nbe exactly one of (0, 0), (0, 1) (0, 2), ... (1, 1), (1, 2) etc. See the\noriginal github repo for more info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `numbers`:\n - Numbers as they appear in the grid.\n - Returned whenever the grid changes.\n - Has `shape=[height, width], dtype=INT32`.\n* `grid`:\n - Letters representing the dominoes currently placed on the board.\n - 'R=right, L=left, T=top, B=bottom'\n - Returned whenever the grid changes.\n - Has `shape=[height, with], dtype=INT32`.\n* `clash`:\n - Represents clashes on the board (i.e. if two dominoes have the same\n pair)\n - '1=clash, 0=no clash'\n - Returned whenever the grid changes.\n - Has `shape=[height, width], dtype=INT32`.\n\n \n\n**dominosa_1** | **dominosa_3** | **dominosa_3a** | **dominosa_6** | **dominosa_9**\n-------------------------------------------------------------------------- | -------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | -------------------------------------------------------------------------- | --------------\n |  |  |  | \n\n## SGT Puzzles - Fifteen\n\nOrder the tiles in increasing order, starting from the top left corner. See the\noriginal github repo for more info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `grid`:\n - Current state of the grid.\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n* `empty`:\n - Index of the single empty cell in the grid.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `movecount`:\n - Number of moves made so far.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n\n \n\n**fifteen_2x2** | **fifteen_3x3** | **fifteen_4x4** | **fifteen_6x6**\n---------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | ---------------\n |  |  | \n\n## SGT Puzzles - Flip\n\nClicking on a cell will flip the colour of some of its neighbours, which are\ndetermined by the symbol in the cell. The goal is to make all the cells have the\nsame colour. See the original github repo for more info:\nhttps://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `light`:\n - The number of light cells.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `dark`:\n - The number of dark cells\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `moves`:\n - The number of moves made by the player.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - State of the board (0 = dark, 1 = light).\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n* `gridMatrix`:\n - The grid matrix of square neighbours (-1 = outside, 1 = neighbour, 0 =\n not neighbour)\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size, 3, 3s], dtype=INT32`.\n\n \n\n**flip_3x3c** | **flip_4x4c** | **flip_5x5r**\n------------------------------------------------------------------------ | ------------------------------------------------------------------------ | -------------\n |  | \n\n## SGT Puzzles - Flood\n\nFloodIt is a game where the player needs to fill the board with a single color.\nThe dynamics of the game are driven by colored areas of the board, which when\npressed cause the currently active region to change its color to the color of\nthe pressed button. When this active region changes color it absorbs neighboring\nsquares that have the same color, thus expanding the active region. The active\nregion starts as a single square at the top-left corner of the board. The game\ngives a single reward at the end of the game if the player manages to fill the\nentire board with the same color within the maximum number steps, otherwise the\nreward is just zero. See the original github repo for more info:\nhttps://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `board`:\n - State of the board, representing colours by their indices.\n - 0: red, 1: yellow, 2: green, 3: blue, 4: orange, 5: purple,\n - 6: brown, 7: light blue, 8: light green, 9: pink\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n\n \n\n**sgtpuzzles_flood_3x3_easy** | **sgtpuzzles_flood_12x12_medium** | **sgtpuzzles_flood_16x16_hard**\n---------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | -------------------------------\n |  | \n\n## SGT Puzzles - Galaxies\n\nSplit the grid up into centrally symmetric areas. The centre of symmetry for\neach area is denoted by a dot on the grid. See the original github repo for more\ninfo: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `dot`:\n - Number of dots on the board.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - String representation of the board:\n - `o`=dot, `' '`=empty, `+`, `-`, `|` = cell corners.\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=STRING_U1`.\n\n \n\n**galaxies_3x3_normal** | **galaxies_5x5_normal** | **galaxies_7x7_normal**\n--------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- | -----------------------\n |  | \n\n**galaxies_7x7_unreasonable** | **galaxies_10x10_normal** | **galaxies_15x15_normal**\n--------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- | -------------------------\n |  | \n\n## SGT Puzzles - Guess\n\nThe computer has thought of a sequence of colours that you have to guess. Fill\nthe top row with colours of your choice, and wait for the computer to give you\nfeedback about your sequence. It will show a black dot for each colour that is\nplaced in the correct position, and a white dot for each that is present in the\nhidden sequence, but not at the position your guess. Try to figure out the\nhidden sequence before you run out of guesses! See the original github repo for\nmore info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `peg`:\n - Indices representing the colours selected in the latest row.\n - Returned after the row is completed and evaluated.\n - Has `shape=[row_length], dtype=INT32`.\n* `feedback`:\n - Evaluation of the latest guess (0: incorrect, 1: correct place, 2:\n correct colour)\n - Returned after the row is completed and evaluated.\n - Has `shape=[row_length], dtype=INT32`.\n\n \n\n**guess_basic** | **guess_quick** | **guess_standard** | **guess_super**\n---------------------------------------------------------------------------- | ----------------------------------------------------------------- | ----------------------------------------------------------------------- | ---------------\n |  |  | \n\n## SGT Puzzles - Inertia\n\nCollect all the blue diamonds on the board without colliding into a bomb. You\ncan move the ball in the 8 main directions (including the diagonals). The ball\nwill keep on moving in that direction until it hits a wall, a bomb, a diamond or\na circle. Circles and diamonds have grip, i.e. it will stop the ball from\ncontinuing to move in the direction it was going towards. See the original\ngithub repo for more info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `gems`:\n - Current number of gems still on the board.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `distancemoved`:\n - Number of cells just moved.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - Symbols of grid cells (b=blank, g=gem, m=mine, s=stop, w=wall)\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n\n \n\n**inertia_5x5** | **inertia_10x10**\n----------------------------------------------------------------- | -----------------\n | \n\n## SGT Puzzles - Light Up\n\nYou have a grid of squares. Some are empty (black) and some are *walls* (grey);\nsome of the walls are numbered. Your aim is to *light up* all the empty squares\nby placing light bulbs in some of them. The numbers denote how many bulbs' light\nhits these directly in a straight sight. Meanwhile, no two bulbs should light up\neach other (i.e. only one bulb allowed in straight sight). See the original\ngithub repo for more info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `grid`:\n - String representation of the board:\n - '#' = blocked, ' ' = empty/black, 'L' = light, 'l' = illuminated,\n - 'X' = impossible, 'number' = number of bulbs hitting this cell.\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=STRING_U1`.\n\n \n\n**light_up_3x3_easy** | **light_up_5x5_easy** | **light_up_7x7_easy** | **light_up_10x10_tricky** | **light_up_14x14_easy**\n---------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------- | -----------------------\n |  |  |  | \n\n## SGT Puzzles - Loopy\n\nDraw a closed loop along the edges of the grid. A number in a cell denotes the\nnumber of edges adjacent to that cell. The loop cannot intersect itself. See the\noriginal github repo for more info: https://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `grid`:\n - String representation of the board:\n - The grid lines and cells:\n - `.` = dots (cell corners)\n - `0-9` = number on cell face or ` ` for empty face\n - `?` = unknown (default),`x` = no line,`-` = `|` = line,`~` = `/` = error\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=STRING_U1`.\n\n \n\n**loopy_3x3_easy** | **loopy_5x5_easy** | **loopy_7x7_easy** | **loopy_7x7_normal** | **loopy_7x7_hard**\n---------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- | ------------------\n |  |  |  | \n\n## SGT Puzzles - Net\n\nThere are a number of light bulbs, wires and a single light source in the\nmiddle. Connect the wires such that all bulbs are lit up, without loose ends or\nloops in the wiring. You can rotate the tiles by clicking on them. If the task\nname has a *w* suffix in it then it is allowed to connect wires on opposing\nedges of the grid. See the original github repo for more info:\nhttps://github.com/chrisboyle/sgtpuzzles.\n\n\n Extras returned
\n\n* `active`:\n - The number of active/completed cells.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `total`:\n - The total number of cells\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - The grid cells represented by numbers.\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n* `gridCompleted`:\n - The grid cells' active/completed status (0 = false, 1 = true).\n - Returned whenever the grid changes.\n - Has `shape=[grid_size, grid_size], dtype=INT32`.\n\n \n\n**net_3x3** | **net_5x5** | **net_7x7w** | **net_9x9** | **net_11x11w**\n-------------------------------------------------------------------- | -------------------------------------------------------------------- | ---------------------------------------------------------------------- | -------------------------------------------------------------------- | --------------\n |  |  |  | \n\n## Shattered Pixel Dungeon\n\nShattered Pixel Dungeon is a Roguelike RPG, with pixel art graphics and lots of\nvariety and replayability. Every game is unique, with four different playable\ncharacters, randomized levels and enemies, and over 150 items to collect and\nuse. The game is simple to get into, but has lots of depth. Strategy is required\nif you want to win! See the original github repo for more info:\nhttps://github.com/00-Evan/shattered-pixel-dungeon.git.\n\n\n Extras returned
\n\n* `action`:\n - The action just completed by the hero.\n - Returned whenever an action is taken.\n - Has `shape=[1], dtype=STRING_U25`.\n* `dst`:\n - TThe destination of the action.\n - Returned whenever an action is taken.\n - Has `shape=[1], dtype=INT32`.\n* `level`:\n - The level reached.\n - Returned whenever a new level is reached.\n - Has `shape=[1], dtype=INT32`.\n* `depth`:\n - The depth of the level reached.\n - Returned whenever a new floor is reached.\n - Has `shape=[1], dtype=INT32`.\n* `deepestFloor`:\n - The deepest reached floor statistic.\n - Returned whenever a new floor is reached.\n - Has `shape=[1], dtype=INT32`.\n* `gold`:\n - The gold level reached.\n - Returned whenever gold is acquired.\n - Has `shape=[1], dtype=INT32`.\n* `totalGold`:\n - The gold collected statistic.\n - Returned whenever gold is acquired.\n - Has `shape=[1], dtype=INT32`.\n* `addedGold`:\n - The gold acquired at the most recent step.\n - Returned whenever gold is acquired.\n - Has `shape=[1], dtype=INT32`.\n* `newlyVisited`:\n - Number of new squares uncovered.\n - Returned whenever new squares are uncovered.\n - Has `shape=[1], dtype=INT32`.\n* `damageDealt`:\n - Damage dealt by hero.\n - Returned whenever the hero deals damage.\n - Has `shape=[1], dtype=INT32`.\n* `damageTaken`:\n - Damage taken by hero.\n - Returned whenever the hero takes damage.\n - Has `shape=[1], dtype=INT32`.\n* `HP`:\n - Hit Points (Health) of hero.\n - Returned whenever its value changes.\n - Has `shape=[1], dtype=INT32`.\n* `exp`:\n - Experience gained by hero.\n - Returned whenever the hero gains experience points.\n - Has `shape=[1], dtype=INT32`.\n* `dew`:\n - How many dew (an immediately consumed item) the user picks up.\n - Returned whenever a dew is picked up.\n - Has `shape=[1], dtype=INT32`.\n* `heal`:\n - Amount the hero is healed.\n - Returned whenever the hero is healed.\n - Has `shape=[1], dtype=INT32`.\n* `itemPickup`:\n - The name of an item that is picked up.\n - Returned whenever an item is picked up.\n - Has `shape=[1], dtype=STRING_U25`.\n* `itemDrop`:\n - The name of an item that is dropped.\n - Returned whenever an item is dropped.\n - Has `shape=[1], dtype=STRING_U25`.\n* `destroy`:\n - The name of the enemy that is killed.\n - Returned whenever an enemy is killed.\n - Has `shape=[1], dtype=STRING_U25`.\n* `search`:\n - Whether the user clicked the \"Search\" button.\n - Returned whenever the button is clicked.\n - Has `shape=[1], dtype=INT32`.\n* `wait`:\n - Whether the user clicked the \"Wait\" button.\n - Returned whenever the button is clicked.\n - Has `shape=[1], dtype=INT32`.\n* `openedInventory`:\n - Whether the user clicked the \"Inventory\" button.\n - Returned whenever the button is clicked.\n - Has `shape=[1], dtype=INT32`.\n\n \n\nhuntress | mage | rogue | warrior\n------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | -------\n |  |  | \n\n## Simple Solitaire\n\nThis is an Android implementation of\n[Solitaire](https://en.wikipedia.org/wiki/Solitaire) card games. We currently\nsupport 19 variants listed here in alphabetical order. Note that the full\ntask_IDs take the form `simple_solitaire_aces_up`,\n`simple_solitaire_calculation` etc. See the original github repo for more info:\nhttps://github.com/TobiasBielefeld/Simple-Solitaire.git.\n\n\n Extras returned
\n\n* `card`:\n - The new visible card `[kind, suit]`:\n - kind: `a=ace, k=king, q=queen, j=jack, x=10, 2-9=digit`.\n - suit: `c=clubs, d=diamonds, h=hearts, s=spades`.\n - Returned when a card is moved.\n - Has `shape=[2], dtype=STRING_U1`.\n* `stack_i`:\n - A non-empty stack of visible cards `[kind, suit]`.\n - `i` different extras (`stack_0`, `stack_1`, `...`), one corresponding to\n each stack.\n - Returned when a card is moved.\n - Has `shape=[52, 2], dtype=STRING_U1`.\n\n \n\n**aces_up** | **calculation** | **canfield** | **forty_eight** | **freecell**\n-------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- | ------------\n |  |  |  | \n\n**golf** | **grandfathers_clock** | **gypsy** | **klondike** | **maze**\n-------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- | --------\n |  |  |  | \n\n**mod3** | **napoleons_tomb** | **pyramid** | **simple_simon** | **spider**\n-------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | ----------\n |  |  |  | \n\n**spiderette** | **tri_peaks** | **vegas** | **yukon**\n-------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------- | ---------\n |  |  | \n\n## Snake\n\nClassic Snake game.\n\n\n Extras returned
\n\n* `move`:\n - The desired direction of movement: `0`=left, `1`=up, `2`=down,\n `3`=right.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `direction`:\n - The direction of the snake: `1`=north, `2`=south, `3`=east, `4`=west.\n - Returned whenever the grid changes.\n - Has `shape=[1], dtype=INT32`.\n* `grid`:\n - The grid cells: `x`=border, `' '`=empty, `s`=snake, `a`=apple.\n - Returned whenever the grid changes.\n - Has `shape=[13, 19], dtype=STRING_U1`.\n\n \n\n\n\n## Vector Pinball\n\nA simple vector-based Pinball game with realistic physics. See the original\ngithub repo for more info: https://github.com/dozingcat/Vector-Pinball.\n\n\n Extras returned
\n Returns no extras.\n \n\n**vector_pinball_table_1** | **vector_pinball_table_2** | **vector_pinball_table_3**\n---------------------------------------------------------------------------- | ---------------------------------------------------------------------------- | --------------------------\n |  | \n\n**vector_pinball_table_4** | **vector_pinball_table_5**\n---------------------------------------------------------------------------- | --------------------------\n | \n"
},
{
"path": "docs/instructions.md",
"content": "# AndroidEnv - Running the environment\n\nIn order to create an AndroidEnv instance you will need to provide two main\ncomponents: a [simulator](#the-simulator) and a [task](#the-task). In the\nfollowing sections you will learn how you can create them.\n\n### The simulator\n\nFirst, you will need to provide your Android virtual device (AVD) that the\nenvironment (and through it, the agent) can communicate with. While this can\nalso be a physical device, in most cases you will need a virtual emulated\ndevice. There are many ways to emulate an AVD - in our examples, we will use\n[Android Studio](https://developer.android.com/studio) to create one.\n\n1. In Android Studio, create a virtual device by following this step-by-step\n [guide](emulator_guide.md).\n2. Follow the steps below to attach the AVD to your environment.\n\n### The task and examples with games and other apps\n\nA `task` is a particular definition of an RL problem that the agent will be\ninteracting with. A `task` may include critical RL information, such as what the\nrewards are, when the episodes are supposed to terminate, and what the reset\nprocedures are that the environment should perform upon episode termination\n(e.g. start or relaunch an app, clear cache, etc.). This information is packaged\ninto a `Task()` proto message, which gets passed passed to AndroidEnv.\n\n* For ready-made example tasks provided with AndroidEnv, check out the\n [Available tasks](example_tasks.md), featuring Vokram (with Markov Decision\n Process (MDP)), Pong, DroidFish (a chess clone), Blockinger (a tetris\n clone), and more.\n\n* See the [Tasks guide](tasks_guide.md) for details on features and\n capabilities of tasks, as well as how to create custom ones.\n\n### Create the environment\n\nAfter setting up the simulator and creating a task, you may find the\n[`loader.load()`](https://github.com/deepmind/android_env/blob/main/android_env/loader.py)\nfunction handy for creating an environment instance by providing relevant\narguments, such as:\n\n* `task_path`: the path pointing to the `.textproto` file describing the\n desired task.\n* `avd_name`: the name of the AVD specified when your created it in Android\n Studio.\n* `android_avd_home` (Optional): the path to where the AVD is installed.\n (default value: `~/.android/avd`).\n* `android_sdk_root` (Optional): the root directory of the Android SDK.\n (default value: `~/Android/Sdk`).\n* `emulator_path` (Optional): the path to the emulator binary. (default:\n `~/Android/Sdk/emulator/emulator`).\n* `adb_path` (Optional): the path to the ADB\n ([Android Debug Bridge](https://developer.android.com/studio/command-line/adb)).\n (default value: `~/Android/Sdk/platform-tools/adb`).\n\nFor the AVD name and path, in Android Studio, go to **Tools** > **AVD Manager**,\nright click on your virtual device, and select **View Details**, where you will\nfind the `avd_name` and its path.\n\nFor the Android SDK location, in Android Studio, go to **Preferences** >\n**Appearance & Behavior** > **System Settings** > **Android SDKs** and note the\n_Android SDK Location_. In the SDK folder, you will find `/emulator/emulator` as\nwell as the ADB path (`/platform-tools/adb`).\n\nYour example configuration may look like this, depending on how you set up your\nemulator:\n\n```python\nfrom android_env import loader\n\nenv = loader.load(\n avd_name='my_avd',\n android_avd_home='/Users/username/.android/avd',\n android_sdk_root='/Users/username/Library/Android/sdk',\n emulator_path='/Users/username/Library/Android/sdk/emulator/emulator',\n adb_path='/Users/username/Library/Android/sdk/platform-tools/adb',\n task_path='/Users/username/android_env/my_tasks/my_task.textproto',\n)\n```\n\n## Example RL agent scripts\n\nThe `examples` directory contains a few simple example agent setups, such as:\n\n* [`run_random_agent.py`](https://github.com/deepmind/android_env/blob/main/examples/run_random_agent.py):\n Runs a simple loop performing randomly selected actions in the environment.\n* [`run_acme_agent.py`](https://github.com/deepmind/android_env/blob/main/examples/run_acme_agent.py):\n Runs a training loop with an\n [Acme](https://deepmind.com/research/publications/Acme) DQN agent,\n implemented in the popular DeepMind RL framework. This will require to\n install the [`acme`](https://github.com/deepmind/acme) dependency.\n* [`run_human_agent.py`](https://github.com/deepmind/android_env/blob/main/examples/run_human_agent.py):\n Creates a [`pygame`](https://www.pygame.org) instance that lets a human user\n interact with the environment and observe environment mechanics, such as\n rewards or task extras. You will need to install the [PyGame] dependency.\n\nFor instance, here is how you can run\n[`run_random_agent.py`](https://github.com/8bitmp3/android_env/blob/main/examples/run_random_agent.py)\nin a folder where you have your APK file, such as\n[Apple Flinger](https://github.com/deepmind/android_env/blob/main/docs/example_tasks.md#apple-flinger)\nfrom\n[Example tasks](https://github.com/deepmind/android_env/blob/main/docs/example_tasks.md).\n(The downloaded TAR file contains the APK file and `.textproto` definitions.)\n\n```shell\npython3 run_random_agent.py \\\n--avd_name='my_avd' \\\n--android_avd_home=/Users/username/.android/avd \\\n--android_sdk_root=/Users/username/Library/Android/sdk \\\n--emulator_path=/Users/username/Library/Android/sdk/emulator/emulator \\\n--adb_path=/Users/username/Library/Android/sdk/platform-tools/adb \\\n--num_steps=1000 \\\n--task_path=/Users/username//apple_flinger_M_1_1.textproto\n```\n"
},
{
"path": "docs/tasks_guide.md",
"content": "# AndroidEnv - Tasks\n\nWith AndroidEnv we provide a mechanism for easily defining RL tasks for the\nagent to learn. This includes various types of information such as what app/game\nit should train on, what rewards the environment returns, or the start state\ndistribution and the episode end criteria.\n\n## Task structure\n\nA *task* definition is captured in the form of a `Task()` proto message. These\nare most easily created by writing a `.textproto` file, then parsing it into a\nproto message. In this section you can find a detailed description about the\ntypes of information that make up a task, and an example demonstrating exactly\nhow to put these into code.\n\n\n Expand this tab to view the main types of information captured in these messages:
\n\n* `id`: An ID used to identify the task.\n\n* `setup_steps`: These are steps the environment will perform right after\n launching the simulator. Possible steps include:\n\n * `install_apk`: Installs an application from a specified path to the APK\n file.\n * `start_activity`: Launches the requested app/activity.\n * `rotate`: Sets the orientation of the device (landscape/portrait).\n\n* `reset_steps`: These are steps the environment will perform right at the\n beginning of a new RL episode. Possible steps include:\n\n * `force_stop`: Stops a given app.\n * `start_activity`: Launches the requested app/activity.\n * `start_screen_pinning`: Restricts the agent's interaction to a\n particular activity through\n [screen pinning](https://support.google.com/android/answer/9455138?hl=en),\n meaning the agent will not be able to quit the given app.\n * `clear_cache`: Clears the cache of a given app.\n\n* `success_conditions`: For each success condition defined, the environment\n will make sure that these conditions were met after finishing `setup_steps`\n and `reset_steps`. They might include conditions such as:\n\n * `check_install`: Makes sure that the request app was successfully\n installed.\n * `wait_for_app_screen`: Waits until the request app was successfully\n launched.\n\n* `expected_app_screen`: If this value is set to a particular activity, the\n environment will periodically check if the agent is still interacting with\n said activity, making sure it has not accidentally quit the application we\n want it to be training on.\n\n* `max_episode_sec`: Puts a time limit on the episodes, triggering an episode\n reset if the current episode has lasted too long.\n\n* `max_duration_steps`: Puts a step limit on the episodes, triggering an\n episode reset once the agent has reached the specified limit.\n\n* `log_parsing_config`: If the environment is parsing logcat messages, this\n field will determine what information it should listen for using regular\n expressions.\n\n * `filters`: The environment filters log messages for these labels which\n signify that such messages were meant to be parsed by AndroidEnv.\n * `log_regexps`: Once a log message was identified as relevant using the\n filters, the environment parses its contents using these regular\n expressions. For example, an application might be sending log messages\n of the form `reward: 1.0`, then the task will capture this info using\n the regexp `^[Rr]eward: ([-+]?[0-9]*\\\\.?[0-9]*)$`.\n\n \n\n\n Expand this tab to see what an example `.textproto` file might look like in practice:
\n\n```python\nid: \"classic_2048\"\nname: \"Classic 2048 - Default\"\ndescription: \"Slide numbered tiles on a grid to combine them to create a tile with the number 2048\"\npackage_name: \"com.tpcstld.twozerogame\"\nfull_activity_name: \"com.tpcstld.twozerogame/com.tpcstld.twozerogame.MainActivity\"\n\n# Perform these upon launching the environment\nsetup_steps: [\n {\n # Install the 2048 app\n adb_call: {\n install_apk: {\n filesystem: {\n path: path/to/classic_2048.apk\n }\n }\n }\n # Check if it was installed correctly\n success_condition: {\n check_install: {\n package_name: \"com.tpcstld.twozerogame\"\n timeout_sec: 10.0\n }\n }\n },\n # Orient the screen in portait mode\n { adb_call: { rotate: { orientation: PORTRAIT_0 } } }\n]\n\n# Perform these upon episode resets\nreset_steps: [\n\n # Stop the 2048 app\n { adb_call: { force_stop: { package_name: \"com.tpcstld.twozerogame\" } } },\n { adb_call: { clear_cache: { package_name: \"com.tpcstld.twozerogame\" } } },\n\n # Start the 2048 app\n {\n adb_call: {\n start_activity: {\n full_activity: \"com.tpcstld.twozerogame/com.tpcstld.twozerogame.MainActivity\"\n extra_args: [\n \"--ez\", '\"RL_TASK_ENABLED\"', '\"true\"',\n \"--es\", '\"RL_TASK_GAME_CONFIG\"', '\"{}\"'\n ]\n }\n }\n\n # Wait until the app has launched successfully\n success_condition: {\n wait_for_app_screen: {\n app_screen: {\n activity: \"com.tpcstld.twozerogame/com.tpcstld.twozerogame.MainActivity\"\n view_hierarchy_path: [\n ]\n }\n timeout_sec: 10.0\n }\n num_retries: 10\n }\n },\n\n # Make sure the agent cannot quit the 2048 app\n {\n adb_call: {\n start_screen_pinning: {\n full_activity: \"com.tpcstld.twozerogame/com.tpcstld.twozerogame.MainActivity\"\n }\n }\n }\n]\n\n# Periodically check if the agent has accidentally quit the app\nexpected_app_screen: {\n activity: \"com.tpcstld.twozerogame/com.tpcstld.twozerogame.MainActivity\"\n view_hierarchy_path: []\n}\n\nmax_episode_steps: 500\n\n# Capture expected format of log messages\nlog_parsing_config: {\n filters: [\"AndroidRLTask:V\"]\n log_regexps: {\n score: \"^[Ss]core: ([-+]?[0-9]*\\\\.?[0-9]*)$\"\n reward: \"^[Rr]eward: ([-+]?[0-9]*\\\\.?[0-9]*)$\"\n episode_end: \"^episode[ _]end$\"\n extra: \"^extra: (?P[^ ]*)[ ]?(?P.*)$\"\n json_extra: \"^json_extra: (?P.*)$\"\n }\n}\n```\n\n \n\n## Log messages and custom APKs\n\nYou might have noticed that tasks often rely on log messages exposed by the\nAndroid system, which AndroidEnv can intercept and translate into items such as\nrewards, episode end signals or task extras.\n\nOne way to define rewards is by using\n`log_parsing_config.LogRegexps.RewardEvent` messages in the task proto. These\nconsist of a regular expression and a numeric value indicating the intended\nreward. If the regexp is matched in any of the lines of the logcat stream, the\nagent will receive the given reward. It is also possible to have multiple of\nthese RewardEvents, allowing us to give rewards for different log messages. The\nsame applies for episode end signals: logcat messages that match the regexps\ndefined in `log_parsing_config.LogRegexps.episode_end` will trigger an episode\nreset.\n\nOf course, applications might not send suitable messages by default, so in order\nto have access to such messages, we often add them to the apps' source code to\nmatch our expectations. For example, in the case of the 2048 app, we find in the\ngame's source code the exact lines where the score is computed, and add a line\nto log this value in the format that is expected by the textproto (or\nconversely, make sure the textproto matches the format you specified here). For\nexample:\n\n```java\n// Make sure thet LOG_FILTER matches 'filters' in the textproto\npublic static final String LOG_FILTER = \"AndroidRLTask\";\n\n// Make sure that the corresponding part of 'log_regexps' will match this string\nLog.i(LOG_FILTER, String.format(Locale.ENGLISH, \"reward: %r\", reward_value))\n```\n\nYou can take a look at example APKs extended with log messages in the example\ntasks (see the section below).\n\n## Example tasks\n\nAlong with the environment implementation we provide a set of example task\ndefinitions. These were chosen so that they would demonstrate the large variety\nof different challenges (e.g. app navigtion, puzzle games, time-reactive games,\nadventure games, card games...) and corresponding interfaces (e.g. button\npressing, swiping, drag-and-drop...) available in AndroidEnv. You can find a\nlist and detailed description of each of these tasks in\n[example_tasks.md](example_tasks.md).\n"
},
{
"path": "examples/__init__.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
},
{
"path": "examples/run_acme_agent.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Acme DQN agent interacting with AndroidEnv.\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nimport acme\nfrom acme import specs\nfrom acme import wrappers as acme_wrappers\nfrom acme.agents.tf import dqn\nfrom acme.tf import networks\nfrom android_env import loader\nfrom android_env.components import config_classes\nfrom android_env.wrappers import discrete_action_wrapper\nfrom android_env.wrappers import float_pixels_wrapper\nfrom android_env.wrappers import image_rescale_wrapper\n\n# Simulator args\nflags.DEFINE_string('avd_name', None, 'Name of AVD to use.')\nflags.DEFINE_string('android_avd_home', '~/.android/avd', 'Path to AVD.')\nflags.DEFINE_string('android_sdk_root', '~/Android/Sdk', 'Path to SDK.')\nflags.DEFINE_string('emulator_path',\n '~/Android/Sdk/emulator/emulator', 'Path to emulator.')\nflags.DEFINE_string('adb_path',\n '~/Android/Sdk/platform-tools/adb', 'Path to ADB.')\n\n# Environment args\nflags.DEFINE_string('task_path', None, 'Path to task textproto file.')\n\n# Experiment args\nflags.DEFINE_integer('num_episodes', 100, 'Number of episodes.')\n\nFLAGS = flags.FLAGS\n\n\ndef apply_wrappers(env):\n \"\"\"Applies a series of wrappers to the environment.\"\"\"\n env = discrete_action_wrapper.DiscreteActionWrapper(env, action_grid=(10, 10))\n env = image_rescale_wrapper.ImageRescaleWrapper(\n env, zoom_factors=(0.25, 0.25))\n env = float_pixels_wrapper.FloatPixelsWrapper(env)\n env = acme_wrappers.SinglePrecisionWrapper(env)\n return env\n\n\ndef main(_):\n\n config = config_classes.AndroidEnvConfig(\n task=config_classes.FilesystemTaskConfig(path=FLAGS.task_path),\n simulator=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n emulator_path=FLAGS.emulator_path,\n android_sdk_root=FLAGS.android_sdk_root,\n android_avd_home=FLAGS.android_avd_home,\n avd_name=FLAGS.avd_name,\n run_headless=FLAGS.run_headless,\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path=FLAGS.adb_path\n ),\n ),\n )\n with loader.load(config) as env:\n\n env = apply_wrappers(env)\n env_spec = specs.make_environment_spec(env)\n\n agent = dqn.DQN(\n environment_spec=env_spec,\n network=networks.DQNAtariNetwork(\n num_actions=env_spec.actions.num_values),\n batch_size=10,\n samples_per_insert=2,\n min_replay_size=10)\n\n loop = acme.EnvironmentLoop(env, agent)\n loop.run(num_episodes=FLAGS.num_episodes)\n\n\nif __name__ == '__main__':\n logging.set_verbosity('info')\n logging.set_stderrthreshold('info')\n flags.mark_flags_as_required(['task_path', 'avd_name'])\n app.run(main)\n"
},
{
"path": "examples/run_human_agent.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Loads an interactive session where a human acts on behalf of an agent.\"\"\"\n\nimport time\nfrom typing import Any\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom android_env import loader\nfrom android_env.components import action_type\nfrom android_env.components import config_classes\nfrom android_env.components import pixel_fns\nimport dm_env\nimport numpy as np\nimport pygame\n\n# Simulator args.\nflags.DEFINE_string('avd_name', None, 'Name of AVD to use.')\nflags.DEFINE_string('android_avd_home', '~/.android/avd', 'Path to AVD.')\nflags.DEFINE_string('android_sdk_root', '~/Android/Sdk', 'Path to SDK.')\nflags.DEFINE_string('emulator_path',\n '~/Android/Sdk/emulator/emulator', 'Path to emulator.')\nflags.DEFINE_string('adb_path',\n '~/Android/Sdk/platform-tools/adb', 'Path to ADB.')\nflags.DEFINE_boolean('run_headless', True, 'Optionally turn off display.')\n\n# Environment args.\nflags.DEFINE_string('task_path', None, 'Path to task textproto file.')\n\n# Pygame args.\nflags.DEFINE_list('screen_size', '480,720', 'Screen width, height in pixels.')\nflags.DEFINE_float('frame_rate', 1.0/30.0, 'Frame rate in seconds.')\n\nFLAGS = flags.FLAGS\n\n\ndef _get_action_from_event(\n event: pygame.event.Event, screen: pygame.Surface, orientation: int\n) -> dict[str, Any]:\n \"\"\"Returns the current action by reading data from a pygame Event object.\"\"\"\n\n act_type = action_type.ActionType.LIFT\n if event.type == pygame.MOUSEBUTTONDOWN:\n act_type = action_type.ActionType.TOUCH\n\n return {\n 'action_type':\n np.array(act_type, dtype=np.int32),\n 'touch_position':\n _scale_position(event.pos, screen, orientation),\n }\n\n\ndef _get_action_from_mouse(\n screen: pygame.Surface, orientation: int\n) -> dict[str, Any]:\n \"\"\"Returns the current action by reading data from the mouse.\"\"\"\n\n act_type = action_type.ActionType.LIFT\n if pygame.mouse.get_pressed()[0]:\n act_type = action_type.ActionType.TOUCH\n\n return {\n 'action_type':\n np.array(act_type, dtype=np.int32),\n 'touch_position':\n _scale_position(pygame.mouse.get_pos(), screen, orientation),\n }\n\n\ndef _scale_position(position: np.ndarray, screen: pygame.Surface,\n orientation: int) -> np.ndarray:\n \"\"\"AndroidEnv accepts mouse inputs as floats so we need to scale it.\"\"\"\n\n scaled_pos = np.divide(position, screen.get_size(), dtype=np.float32)\n if orientation == 1: # LANDSCAPE_90\n scaled_pos = scaled_pos[::-1]\n scaled_pos[0] = 1 - scaled_pos[0]\n return scaled_pos\n\n\ndef _accumulate_reward(\n timestep: dm_env.TimeStep,\n episode_return: float) -> float:\n \"\"\"Accumulates rewards collected over the course of an episode.\"\"\"\n\n if timestep.reward and timestep.reward != 0:\n logging.info('Reward: %s', timestep.reward)\n episode_return += timestep.reward\n\n if timestep.first():\n episode_return = 0\n elif timestep.last():\n logging.info('Episode return: %s', episode_return)\n\n return episode_return\n\n\ndef _render_pygame_frame(surface: pygame.Surface, screen: pygame.Surface,\n orientation: int, timestep: dm_env.TimeStep) -> None:\n \"\"\"Displays latest observation on pygame surface.\"\"\"\n\n frame = timestep.observation['pixels'][:, :, :3] # (H x W x C) (RGB)\n frame = pixel_fns.transpose_pixels(frame) # (W x H x C)\n frame = pixel_fns.orient_pixels(frame, orientation)\n\n pygame.surfarray.blit_array(surface, frame)\n pygame.transform.smoothscale(surface, screen.get_size(), screen)\n\n pygame.display.flip()\n\n\ndef main(_):\n\n pygame.init()\n pygame.display.set_caption('android_human_agent')\n\n config = config_classes.AndroidEnvConfig(\n task=config_classes.FilesystemTaskConfig(path=FLAGS.task_path),\n simulator=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n emulator_path=FLAGS.emulator_path,\n android_sdk_root=FLAGS.android_sdk_root,\n android_avd_home=FLAGS.android_avd_home,\n avd_name=FLAGS.avd_name,\n run_headless=FLAGS.run_headless,\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path=FLAGS.adb_path\n ),\n ),\n )\n with loader.load(config) as env:\n\n # Reset environment.\n first_timestep = env.reset()\n orientation = np.argmax(first_timestep.observation['orientation'])\n\n # Create pygame canvas.\n screen_size = list(map(int, FLAGS.screen_size)) # (W x H)\n obs_shape = env.observation_spec()['pixels'].shape[:2] # (H x W)\n\n if (orientation == 1 or orientation == 3): # LANDSCAPE_90 | LANDSCAPE_270\n screen_size = screen_size[::-1]\n obs_shape = obs_shape[::-1]\n\n screen = pygame.display.set_mode(screen_size) # takes (W x H)\n surface = pygame.Surface(obs_shape[::-1]) # takes (W x H)\n\n # Start game loop.\n prev_frame = time.time()\n episode_return = 0\n\n while True:\n if pygame.key.get_pressed()[pygame.K_ESCAPE]:\n return\n\n all_events = pygame.event.get()\n for event in all_events:\n if event.type == pygame.QUIT:\n return\n\n # Filter event queue for mouse click events.\n mouse_click_events = [\n event for event in all_events\n if event.type in [pygame.MOUSEBUTTONDOWN, pygame.MOUSEBUTTONUP]\n ]\n\n # Process all mouse click events.\n for event in mouse_click_events:\n action = _get_action_from_event(event, screen, orientation)\n timestep = env.step(action)\n episode_return = _accumulate_reward(timestep, episode_return)\n _render_pygame_frame(surface, screen, orientation, timestep)\n\n # Sample the current position of the mouse either way.\n action = _get_action_from_mouse(screen, orientation)\n timestep = env.step(action)\n episode_return = _accumulate_reward(timestep, episode_return)\n _render_pygame_frame(surface, screen, orientation, timestep)\n\n # Limit framerate.\n now = time.time()\n frame_time = now - prev_frame\n if frame_time < FLAGS.frame_rate:\n time.sleep(FLAGS.frame_rate - frame_time)\n prev_frame = now\n\n\nif __name__ == '__main__':\n logging.set_verbosity('info')\n logging.set_stderrthreshold('info')\n flags.mark_flags_as_required(['avd_name', 'task_path'])\n app.run(main)\n"
},
{
"path": "examples/run_random_agent.py",
"content": "# coding=utf-8\n# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Example script demonstrating usage of AndroidEnv.\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom android_env import loader\nfrom android_env.components import config_classes\nfrom dm_env import specs\nimport numpy as np\n\nFLAGS = flags.FLAGS\n\n# Simulator args.\nflags.DEFINE_string('avd_name', None, 'Name of AVD to use.')\nflags.DEFINE_string('android_avd_home', '~/.android/avd', 'Path to AVD.')\nflags.DEFINE_string('android_sdk_root', '~/Android/Sdk', 'Path to SDK.')\nflags.DEFINE_string('emulator_path',\n '~/Android/Sdk/emulator/emulator', 'Path to emulator.')\nflags.DEFINE_string('adb_path',\n '~/Android/Sdk/platform-tools/adb', 'Path to ADB.')\nflags.DEFINE_bool('run_headless', False,\n 'Whether to display the emulator window.')\n\n# Environment args.\nflags.DEFINE_string('task_path', None, 'Path to task textproto file.')\n\n# Experiment args.\nflags.DEFINE_integer('num_steps', 1000, 'Number of steps to take.')\n\n\ndef main(_):\n\n config = config_classes.AndroidEnvConfig(\n task=config_classes.FilesystemTaskConfig(path=FLAGS.task_path),\n simulator=config_classes.EmulatorConfig(\n emulator_launcher=config_classes.EmulatorLauncherConfig(\n emulator_path=FLAGS.emulator_path,\n android_sdk_root=FLAGS.android_sdk_root,\n android_avd_home=FLAGS.android_avd_home,\n avd_name=FLAGS.avd_name,\n run_headless=FLAGS.run_headless,\n ),\n adb_controller=config_classes.AdbControllerConfig(\n adb_path=FLAGS.adb_path\n ),\n ),\n )\n with loader.load(config) as env:\n\n action_spec = env.action_spec()\n\n def get_random_action() -> dict[str, np.ndarray]:\n \"\"\"Returns a random AndroidEnv action.\"\"\"\n action = {}\n for k, v in action_spec.items():\n if isinstance(v, specs.DiscreteArray):\n action[k] = np.random.randint(low=0, high=v.num_values, dtype=v.dtype)\n else:\n action[k] = np.random.random(size=v.shape).astype(v.dtype)\n return action\n\n _ = env.reset()\n\n for step in range(FLAGS.num_steps):\n action = get_random_action()\n timestep = env.step(action=action)\n reward = timestep.reward\n logging.info('Step %r, action: %r, reward: %r', step, action, reward)\n\n\nif __name__ == '__main__':\n logging.set_verbosity('info')\n logging.set_stderrthreshold('info')\n flags.mark_flags_as_required(['avd_name', 'task_path'])\n app.run(main)\n"
},
{
"path": "pyproject.toml",
"content": "[build-system]\nrequires = [\n \"setuptools\",\n \"wheel\"\n]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"android-env\"\nversion = \"1.2.2\"\ndescription = \"AndroidEnv environment and library for training agents.\"\nauthors = [{name = \"DeepMind\"}]\nlicense = {file = \"LICENSE\"}\nreadme = {text = \"Read the README at https://github.com/deepmind/android_env for more information.\", content-type = \"text/plain\"}\nkeywords = [\"Android\", \"OS\", \"reinforcement-learning\"]\nrequires-python = \">=3.10\"\ndependencies = [\n \"absl-py>=0.1.0\",\n \"dm_env\",\n \"grpcio\",\n \"numpy>=1.21\",\n \"portpicker>=1.2.0\",\n \"protobuf>=2.6\",\n \"pygame\",\n]\n\n[project.optional-dependencies]\ntesting = [\n \"gym\",\n \"pillow\",\n \"pytype\",\n]\nacme = [\"dm-acme\"]\ngym = [\"gym\"]\n\n[project.urls]\nrepository = \"https://github.com/deepmind/android_env\"\ndeepmind = \"https://www.deepmind.com/publications/androidenv-the-android-learning-environment\"\narxiv = \"https://arxiv.org/abs/2105.13231\"\n"
},
{
"path": "setup.py",
"content": "# Copyright 2026 DeepMind Technologies Limited.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Simple package definition for using with `pip`.\"\"\"\n\nimport importlib\nimport os\n\nimport setuptools\nfrom setuptools import find_packages\nfrom setuptools import setup\nfrom setuptools.command.build_ext import build_ext\nfrom setuptools.command.build_py import build_py\n\n_ROOT_DIR = os.path.dirname(os.path.abspath(__file__))\n\n# Tuple of proto message definitions to build Python bindings for. Paths must\n# be relative to root directory.\n_ANDROID_ENV_PROTOS = (\n 'android_env/proto/adb.proto',\n 'android_env/proto/emulator_controller.proto',\n 'android_env/proto/snapshot.proto',\n 'android_env/proto/snapshot_service.proto',\n 'android_env/proto/state.proto',\n 'android_env/proto/task.proto',\n 'android_env/proto/a11y/a11y.proto',\n 'android_env/proto/a11y/android_accessibility_action.proto',\n 'android_env/proto/a11y/android_accessibility_forest.proto',\n 'android_env/proto/a11y/android_accessibility_node_info.proto',\n 'android_env/proto/a11y/android_accessibility_node_info_clickable_span.proto',\n 'android_env/proto/a11y/android_accessibility_tree.proto',\n 'android_env/proto/a11y/android_accessibility_window_info.proto',\n 'android_env/proto/a11y/rect.proto',\n)\n\n\nclass _GenerateProtoFiles(setuptools.Command):\n \"\"\"Command to generate protobuf bindings for AndroidEnv protos.\"\"\"\n\n descriptions = 'Generates Python protobuf bindings for AndroidEnv protos.'\n user_options = []\n\n def initialize_options(self):\n pass\n\n def finalize_options(self):\n pass\n\n def run(self):\n # Import grpc_tools here, after setuptools has installed setup_requires\n # dependencies.\n from grpc_tools import protoc # pylint: disable=g-import-not-at-top\n\n with importlib.resources.as_file(\n importlib.resources.files('grpc_tools').joinpath('_proto')\n ) as path:\n grpc_protos_include = str(path)\n\n for proto_path in _ANDROID_ENV_PROTOS:\n proto_args = [\n 'grpc_tools.protoc',\n '--proto_path={}'.format(grpc_protos_include),\n '--proto_path={}'.format(_ROOT_DIR),\n '--python_out={}'.format(_ROOT_DIR),\n '--pyi_out={}'.format(_ROOT_DIR),\n '--grpc_python_out={}'.format(_ROOT_DIR),\n os.path.join(_ROOT_DIR, proto_path),\n ]\n if protoc.main(proto_args) != 0:\n raise RuntimeError('ERROR: {}'.format(proto_args))\n\n\nclass _BuildExt(build_ext):\n \"\"\"Generate protobuf bindings in build_ext stage.\"\"\"\n\n def run(self):\n self.run_command('generate_protos')\n build_ext.run(self)\n\n\nclass _BuildPy(build_py):\n \"\"\"Generate protobuf bindings in build_py stage.\"\"\"\n\n def run(self):\n self.run_command('generate_protos')\n build_py.run(self)\n\nsetup(\n packages=find_packages(exclude=['examples']),\n package_data={'': ['proto/*.proto']}, # Copy protobuf files.\n include_package_data=True,\n setup_requires=['grpcio-tools'],\n cmdclass={\n 'build_ext': _BuildExt,\n 'build_py': _BuildPy,\n 'generate_protos': _GenerateProtoFiles,\n },\n)\n"
}
]