SE-GA: Memory-Augmented Self-Evolution for GUI Agents 

This is the implementation code for the paper "SE-GA: Memory-Augmented Self-Evolution for GUI Agents". This work has been accepted by ICML 2026.

---

Environment Setup

Requirements

• Python >= 3.11
• PyTorch >= 2.8.0 (with CUDA)
• NVIDIA GPUs

Install Dependencies

pip install -r requirements.txt

or

pip install -e .

Or install individually:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers>=4.46.0 peft accelerate deepspeed
pip install vllm
pip install swanlab  # optional, for logging

Environment Variables

For TTME memory (if using OpenAI-compatible embeddings):

export EMBEDDING_API_KEY=your-api-key
export EMBEDDING_BASE_URL=your-base-url
export EMBEDDING_MODEL=text-embedding-3-small

If you experience slow downloads from HuggingFace, set the mirror endpoint:

export HF_ENDPOINT=https://hf-mirror.com

You can also add this to your config YAML:

env:
  HF_ENDPOINT: "https://hf-mirror.com"

---

MASE-Training Pipeline

Pipeline Overview

The training pipeline consists of three sequential stages:

Stage	Name	Method	Goal	Model Input
Grounding SFT	Grounding SFT	Supervised Fine-Tuning	Learn basic GUI action format	Screenshot + Task
Planning SFT	Planning SFT	Supervised Fine-Tuning	Learn reasoning with thinking	Screenshot + Task + History
RL	RL	Group Relative Policy Optimization	Optimize action quality via RL	Screenshot + Task + History

The model progressively learns:

1. Grounding SFT: Simple action generation ({"action_type": "CLICK", "action_info": [x, y]})
2. Planning SFT: Structured thinking before action (<thinking>...<analysis>...<reasoning>...<instruction>...</thinking><answer>...</answer>)
3. RL: Improved action quality through reinforcement learning with custom reward functions

Configuration

All training parameters are managed in a single YAML file: configs/default.yaml.

The configuration is organized into five sections:

model:           # Shared model settings (dtype, attention, image resolution)
env:             # Logging project name
grounding_sft:   # Grounding SFT: model_id, hyperparameters, data paths
planning_sft:    # Planning SFT: model_id, hyperparameters, data paths
grpo:            # GRPO RL: model_id, hyperparameters, LoRA config, data paths

Quick Start

One-click full pipeline:

bash run.sh

This runs all stages sequentially: Grounding SFT → Planning SFT → RL.

Using Python directly:

python -m gui_sega.run --config configs/default.yaml

Running Individual Stages

# Run only SFT (Grounding + Planning sequentially)
bash run_sft.sh

# Run only Grounding SFT
bash run_sft.sh --stage grounding_sft

# Run only Planning SFT
bash run_sft.sh --stage planning_sft

# Run only GRPO RL
bash run_rl.sh

# Specify GPUs
bash run.sh --gpus 0,1 --num_gpus 2

# Use custom config
bash run.sh --config configs/my_experiment.yaml

---

TTME-Inference:

The Test-Time Memory Extension (TTME) module provides a hierarchical memory system for GUI agents during inference, addressing failures in long-horizon tasks caused by constrained context windows.

Three-Layer Architecture

┌─────────────────────────────────────────────────────┐
│                    TTMEMemory                        │
│                                                      │
│  ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│  │   Episodic   │ │   Semantic   │ │ Experiential │ │
│  │   Memory     │ │   Memory     │ │   Memory     │ │
│  │   (M_EPI)    │ │   (M_SEM)    │ │   (M_EXP)    │ │
│  ├──────────────┤ ├──────────────┤ ├──────────────┤ │
│  │ Short-term   │ │ Long-term    │ │ Historical   │ │
│  │ working      │ │ knowledge    │ │ task         │ │
│  │ memory       │ │ repository   │ │ strategies   │ │
│  ├──────────────┤ ├──────────────┤ ├──────────────┤ │
│  │ <o_k, a_k,   │ │ <k_sem, d>   │ │ <τ, g(τ),   │ │
│  │  o_{k+1}>    │ │              │ │  k_intent,   │ │
│  │              │ │              │ │  k_task>     │ │
│  ├──────────────┤ ├──────────────┤ ├──────────────┤ │
│  │ Sliding      │ │ Embedding    │ │ Hybrid       │ │
│  │ Window       │ │ Similarity   │ │ Retrieval    │ │
│  │ (horizon H)  │ │ (Top-K)      │ │ (λ·text +    │ │
│  │              │ │              │ │  (1-λ)·visual│ │
│  └──────────────┘ └──────────────┘ └──────────────┘ │
│                                                      │
│  Output: M_retrieved = C_epi + C_sem + C_exp        │
└─────────────────────────────────────────────────────┘

Layer	Function	Data Structure	Retrieval
Episodic	Track recent actions within current task	<observation, action, next_observation>	Sliding window (horizon H)
Semantic	Store universal interaction rules	<description, embedding>	Cosine similarity Top-K
Experiential	Recall strategies from similar past tasks	<trajectory, summary, intent_emb, task_emb>	Hybrid: λ·text_sim + (1-λ)·visual_sim

Run the full example:

python -m gui_sega.ttme.memory_example

---

Citation

If you find this work helpful, please consider citing our paper:

@misc{jin2026segamemoryaugmentedselfevolutiongui,
      title={SE-GA: Memory-Augmented Self-Evolution for GUI Agents}, 
      author={Shilong Jin and Lanjun Wang and Zhuosheng Zhang},
      year={2026},
      eprint={2605.16883},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.16883}, 
}

Acknowledgements

This project builds upon TRL (HuggingFace) and Qwen2.5-VL (Alibaba Cloud).