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 · Detailed usage

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.

Figure 2. Both modalities are preferred on substantial subsets of steps.

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.

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:

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 specifies the input schemas, service contracts, options, and output files.

Train a router

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.

Evaluate AITW and LearnGUI

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

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.

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.

Repository layout

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

For all CLI options, run the relevant script with --help. The troubleshooting section covers missing inputs, feature IDs, fold construction, checkpoints, and service connections.