# SwitchMem

**Learning How to Remember: Memory Modality Selection for GUI Agents**

SwitchMem selects **Text or Image memory at each interaction step** for a
vision-language GUI agent. It learns when each representation is useful and
passes the selected memory to a frozen actor, helping the agent reuse past
demonstrations across changing GUI states.

[Running experiments](#running-experiments) · [Detailed usage](docs/usage.md)

![SwitchMem overview: offline preference-target construction and a pairwise router that selects one memory modality for the frozen actor.](assets/figures/switchmem_overview.png)

_Figure 5 from the paper. Text and Image are two representations of the same
retrieved demonstration. Offline training learns from their effect on the
actor's hidden states; inference provides the actor with one selected memory._

## Why choose memory modality at each step?

The same demonstration can help differently depending on whether its past
GUI state is described in text or shown as a screenshot. Text is useful for
some interaction steps, while images preserve visual details needed at others.
Across AITW and LearnGUI, neither representation wins consistently.

<p align="center">
  <img src="assets/figures/modality_preferences.png" width="560" alt="Among non-tied steps, Text and Image win 48.0% and 52.0% on AITW, and 47.7% and 52.3% on LearnGUI.">
</p>

_Figure 2. Both modalities are preferred on substantial subsets of steps._

![Left: a step-wise oracle outperforms the best fixed modality on AITW and LearnGUI. Right: the D-CoE modality-preference signal becomes stronger in later actor layers.](assets/figures/oracle_and_hidden_states.png)

_Figures 3–4. The oracle gap motivates adaptive selection. Layer-wise probing
shows that modality preference is reflected in the actor's hidden states; the
probe uses examples with at least five-of-seven preference agreement._

## How SwitchMem works

1. **Retrieve once.** Represent the same demonstration as Text and Image,
   keeping its instruction and demonstrated action fixed.
2. **Construct preference targets.** During offline training, compare the
   actor's hidden states under the two memory conditions. The signed
   Image-minus-Text shift, called Directional Chain-of-Embedding (D-CoE), feeds
   a standardized PCA/logistic preference probe.
3. **Train a pairwise router.** A shared MLP scores the current state with
   each memory candidate and learns their relative preference using MSE.
4. **Select and act.** Apply the frozen routing threshold and generate the
   next action with the selected memory. The actor weights stay frozen.

The implementation includes preference-target construction, five-fold router
training, offline AITW/LearnGUI evaluation, and online MobileMiniWoB++ transfer
calibration and rollouts.

## Results

The paper reports the following results with the same frozen **Qwen3.5-9B**
actor for the memory-augmented methods:

| Memory strategy |   AITW (%) | LearnGUI (%) | MobileMiniWoB++ (%) |
| --------------- | ---------: | -----------: | ------------------: |
| No memory       |     41.920 |       23.485 |              58.696 |
| Text            |     48.760 |       35.638 |              53.261 |
| Image           |     49.300 |       37.441 |              57.609 |
| Both            |     47.080 |       26.919 |              60.870 |
| **SwitchMem**   | **51.420** |   **38.131** |          **63.043** |

_Table 1, selected rows. AITW and LearnGUI measure step-level action scores;
MobileMiniWoB++ measures terminal task success. These are results reported in
the paper, not outputs of the local smoke test._

![LearnGUI analyses: routing-baseline comparison, SwitchMem versus ratio-matched random routing, and latency versus action score.](assets/figures/routing_analysis.png)

_Figure 7. Adaptive selection improves over simpler routing baselines and
ratio-matched random routing on LearnGUI. The final panel compares latency
and performance; the paper measures latency on an RTX PRO 6000 GPU._

## Quick start

Use **Python 3.10+** and run commands from the repository root:

```sh
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python smoke_test.py
```

This CPU check runs a small training job and exercises all three inference
paths with synthetic data and mocked services. It needs no model server,
benchmark downloads, or emulator. A successful run reports `OK`.

## Running experiments

Real experiments use the training and inference scripts below. Prepare your
benchmark records, screenshots, cached features, and sampled utilities, then
connect the required model services. The paper uses Qwen3.5-9B as the actor
and Qwen3-VL-Embedding-2B for retrieval and memory embedding.

Dataset preparation, feature/utility collection, and service startup are
external preparation steps. The [usage guide](docs/usage.md) specifies the
input schemas, service contracts, options, and output files.

### Train a router

```sh
python train.py \
  --data data/train.npz \
  --output artifacts/router.pt \
  --device cpu
```

Use `--device cuda:0` for an available CUDA GPU. Training consumes frozen
actor representations and Text/Image memory embeddings, and saves a bundled
five-fold router for offline inference.
See [training inputs and options](docs/usage.md#2-train-the-offline-router).

### Evaluate AITW and LearnGUI

Start an actor service that exposes `/v1/chat/completions`, then run:

```sh
python infer_aitw.py \
  --samples data/aitw/samples.jsonl \
  --bank data/aitw/memory.jsonl \
  --features data/aitw/features.npz \
  --router artifacts/router.pt \
  --actor-url http://127.0.0.1:8000 \
  --output outputs/aitw/predictions.jsonl

python infer_learngui.py \
  --samples data/learngui/samples.jsonl \
  --features data/learngui/features.npz \
  --router artifacts/router.pt \
  --actor-url http://127.0.0.1:8000 \
  --output outputs/learngui/predictions.jsonl
```

Each runner saves predicted actions, selected modalities, router scores, and
benchmark action scores. See [offline input formats and actor requirements](docs/usage.md#3-aitw-and-learngui-offline-inference).

### Evaluate MobileMiniWoB++

The online runner retrieves a memory for each screen, probes the actor's
current hidden state, selects a modality, and executes the generated action
in Android. Run **`calibrate`**, then **`evaluate`** with different suite seeds.

This path requires an actor with hidden-state output, an embedding service,
an AndroidLab environment, Docker/ADB, a successful-demonstration bank, and
compatible pretrained transfer weights. Its `--source-router` directory uses
five weight files plus `report.json`; it has a different architecture and
checkpoint format from the offline `router.pt`.

See the [complete MobileMiniWoB++ setup and commands](docs/usage.md#4-mobileminiwob-transfer-inference).

## Repository layout

| Path                                             | Purpose                                                                                                  |
| ------------------------------------------------ | -------------------------------------------------------------------------------------------------------- |
| [train.py](train.py)                             | Preference-target construction, grouped five-fold router training, and threshold selection.              |
| [model.py](model.py)                             | Shared pairwise MLP, modality selection, and offline ensemble loading.                                   |
| [common.py](common.py)                           | Offline input loading, ID joins, image encoding, actor requests, and prediction output.                  |
| [infer_aitw.py](infer_aitw.py)                   | AITW prompts, action parsing, coordinate conversion, and action matching.                                |
| [infer_learngui.py](infer_learngui.py)           | LearnGUI prompts, action parsing, and continuous action scoring.                                         |
| [infer_mobileminiwob.py](infer_mobileminiwob.py) | Transfer-router loading, memory retrieval, calibration, Android actions, and terminal reward evaluation. |
| [mobile_common.py](mobile_common.py)             | Mobile prompts, interaction history, GUI XML descriptions, and action parsing.                           |
| [smoke_test.py](smoke_test.py)                   | Synthetic CPU checks for training, inference, parsers, and CLI help.                                     |
| [requirements.txt](requirements.txt)             | Python dependencies.                                                                                     |
| [docs/usage.md](docs/usage.md)                   | Full installation, input schemas, execution commands, and troubleshooting.                               |
| [assets/figures/](assets/figures/)               | Figures from the paper used in this README.                                                |

For all CLI options, run the relevant script with `--help`. The
[troubleshooting section](docs/usage.md#cli-help-and-common-setup-errors)
covers missing inputs, feature IDs, fold construction, checkpoints, and
service connections.