WM-R1: Training GUI Agents to Reason and leverage World Models with Reinforcement Learning

Yu Han; Tianwen Qian

cs.AIAug 27, 2026

WM-R1: Training GUI Agents to Reason and leverage World Models with Reinforcement Learning

Authors: Yu Han, Tianwen Qian

Organizations: School of Computer Science and Technology, East China Normal University

Abstract

GUI agents trained with reinforcement learning (RL) have showcased strong environment learning capabilities on mobile platforms. However, RL typically demands extensive real-environment interactions, leading to high resource costs and instability, especially in GUI scenarios. To address these, we propose WM-R1, the first reinforcement learning framework that trains mobile GUI agents with world models instead of real environments. Specifically, world models serve as the source of state transitions during all rollouts, replacing the real Android environment within the training loop. WM-R1 also embeds world models directly into the thinking process, enabling agents to reason about the consequences of candidate actions before committing to the final action. Crucially, WM-R1 eliminates the need for real-environment interaction, supports massively parallelized and step-level granularized trajectory generation grounded in world models, and introduces a multi-dimensional rule-based reward that jointly optimizes task success, trajectory efficiency, and world model utilization. For efficient training, we curate a high-quality dataset of 2000 challenging tasks. Experiments on Android mobile benchmarks demonstrate that WM-R1-trained agents significantly outperform GRPO-only baselines and inference-time simulation methods. Code is available at https://github.com/genalyu/WM-R1 .

Figures & tables

Figure 1: Comparison of training paradigms. Left: Standard GRPO generates trajectories through real-environment interaction, incurring high computational cost, environment latency, and irreversibility. Right: WM-R1 replaces the environment with world models, enabling fully simulated trajectory generation.

Figure 2: WM-R1 training framework. For each task, the policy model generates multiple candidate trajectories with chain-of-thought reasoning in a world-model-simulated environment. At each reasoning step, the agent can invoke the world model to predict the consequence of a candidate action. The simulated next state becomes part of the trajectory. After all trajectories complete, a rule-based Reward Model computes scores, followed by GRPO advantage computation and policy update.

	AndroidWorld		GUI-Odyssey		ScreenSpot-Pro		ScreenSpot-V2		AndroidControl
Method	Full	Hard	Full	Hard	Icon	Text	Icon	Text	Type	Gnd
3B scale
Qwen2.5-VL-3B (ZS)	15.2	11.8	8.5	5.2	28.4	38.6	64.5	88.2	76.2	69.5
Qwen2.5-VL-3B + CW	20.6	16.3	14.2	9.8	31.2	41.5	67.8	89.1	78.5	72.3
GUI-R1-3B	17.4	13.5	10.2	6.8	35.6	44.3	76.8	90.5	82.3	75.6
UI-TARS-3B	19.8	15.6	13.4	8.9	39.2	48.5	80.4	91.8	84.7	77.8

Table 1: Main results across all benchmarks. Bold: best per scale. AndroidWorld and GUI-Odyssey report success rate (%); Hard subsets contain tasks requiring

\geq

5 ground-truth steps. ScreenSpot-Pro/V2 report mobile grounding accuracy (%). AndroidControl reports action type prediction (Type) and coordinate grounding (Gnd) accuracy (%). CW: Code2World inference-time augmentation.

Figure 3: Training curves. WM-R1 reaches 39.8% success rate by episode 15, outperforming UI-R1 (30.8%) by 9.0 points with lower variance.

Figure 4: Policy entropy. Both methods show a decrease-then-recover pattern, but WM-R1 maintains higher entropy throughout, indicating sustained exploration from the noise-free world model.

Config	AW	AW-H	$\Delta$
WM-R1 (full)	39.8	32.7	—
$-$ CoT	34.6	27.8	$-$ 5.2 / $-$ 4.9
$-$ $R_{L}$	36.5	30.1	$-$ 3.3 / $-$ 2.6
$-$ $R_{\text{WM}}$	37.2	30.8	$-$ 2.6 / $-$ 1.9

Table 2: Ablation study on AndroidWorld (7B). Chain-of-thought reasoning is the most critical component; removing it causes the largest degradation (

-

5.2 on AW,

-

4.9 on AW-H).

R_{L}

and

R_{\text{WM}}

contribute smaller but complementary gains.

Figure 5: WM utilization dynamics. WM calls per trajectory (blue, left axis) decrease from

\sim

4.2 to

\sim

1.8 while average episode reward (red, right axis) increases from 0.15 to 0.52. Three phases emerge: exploration (ep. 1–4), calibration (ep. 5–10), and selective utilization (ep. 11–15).

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Base model	Qwen2.5-VL-3B / 7B
World model	Code2World-8B (frozen)
GPUs per node	1 $\times$ NVIDIA H200 (141GB)
Number of nodes	1
World model instances	1
Rollout batch size	2

Appendix

Table 3: WM-R1 training hyperparameters.

Parameter	Value	AndroidWorld
Reward $\alpha$ (success)	0.5	37.2
	1.0 (default)	39.8
	2.0	38.5
Reward $\beta$ (length)	0.1	36.8
	0.3	38.1
	0.5 (default)	39.8

Appendix

Table 4: Parameter sensitivity analysis on AndroidWorld (Qwen2.5-VL-7B). The default configuration is highlighted in each row.

Method	GPU	Train hours	Env required	Inference +overhead
UI-R1-7B (no env)	1 $\times$ H200	7.8	No	—
MobileGUI-RL-7B	1 $\times$ H200	36.5	Yes	—
WM-R1-3B	1 $\times$ H200	8.4	No	+1.6s/traj
WM-R1-7B	1 $\times$ H200	11.2	No	+2.3s/traj

Appendix

Table 5: Computational cost comparison. “GPU” denotes hardware used; “Train hours” is wall-clock training time for the full 15-episode curriculum on the 2000-task dataset; “Env required” indicates whether a live Android environment is needed during training; “Inference +overhead” reports the additional latency per trajectory from world model calls at inference time relative to a standard generate-and-act baseline. UI-R1 (no env) refers to the offline GRPO variant that trains on pre-collected trajectories without real-environment interaction during training.

Method	AW	AW-H	GO	GO-H
4B scale
Qwen3-VL-4B (ZS)	24.8	19.5	14.2	9.1
Qwen3-VL-4B + CW	31.2	25.4	21.5	15.8
UI-R1-4B	26.5	21.2	17.8	12.4
WM-R1-4B	38.2	31.5	30.2	23.4
8B scale

Appendix

Table 6: WM-R1 results with Qwen3-VL base models on long-horizon benchmarks. AW: AndroidWorld, GO: GUI-Odyssey, H: Hard subset. Bold: best per scale.

RL for GUI Agents

Reinforcement learning for GUI agents trains screen-based AI to complete multi-step tasks across mobile, web, and desktop interfaces.

See all 48 papers

Generalization in Online Reinforcement Learning for Mobile Agents
Mar 8, 2026Li Gu, Zihuan Jiang, Zhixiang Chi +5RL for GUI Agents Mobile GUI Agents
Mila – Québec AI Institute · Concordia University · University of Toronto +3
GSAR: Goal-State-Anchor Rewards for Mobile GUI Agents with Self-Evolving Data Synthesis
Aug 24, 2026Long Zhang, Yuhan Chen, Chaoran Zhang +7RL for GUI Agents Reward Design for RL
Wuhan University · Xiaomi Inc.
How Mobile World Model Guides GUI Agents?
May 11, 2026Weikai Xu, Kun Huang, Yunren Feng +10World Model Learning GUI Agents
1Nanyang Technological University · 2MiLM Plus, Xiaomi Inc. · 3Independent Researchers +3

Mar 8, 2026J/K move · Enter open · S save

Generalization in Online Reinforcement Learning for Mobile Agents

Li Gu, Zihuan Jiang, Zhixiang Chi, Huan Liu+4

Mila – Québec AI Institute · Concordia University · University of Toronto+3

RL for GUI Agents Mobile GUI Agents Generalization in RL

Open paper AskPDF arXiv

Graphical user interface (GUI)-based mobile agents automate digital tasks on mobile devices by interpreting natural-language instructions and interacting with the screen. While recent methods apply reinforcement learning (RL) to train vision-language-model(VLM) agents in interactive environments with a primary focus on performance, generalization remains underexplored due to the lack of standardized benchmarks and open-source RL systems. In this work, we formalize the problem as a Contextual Markov Decision Process (CMDP) and introduce AndroidWorld-Generalization, a benchmark with three increasingly challenging regimes for evaluating zero-shot generalization to unseen task instances, templates, and applications. We further propose an RL training system that integrates Group Relative Policy Optimization (GRPO) with a scalable rollout collection system, consisting of containerized infrastructure and asynchronous execution % , and error recovery to support reliable and efficient training. Experiments on AndroidWorld-Generalization show that RL enables a 7B-parameter VLM agent to surpass supervised fine-tuning baselines, yielding a 26.1% improvement on unseen instances but only limited gains on unseen templates (15.7%) and apps (8.3%), underscoring the challenges of generalization. As a preliminary step, we demonstrate that few-shot adaptation at test-time improves performance on unseen apps, motivating future research in this direction. To support reproducibility and fair comparison, we open-source the full RL training system, including the environment, task suite, models, prompt configurations, and the underlying infrastructure \footnote{https://github.com/zihuanjiang/AndroidWorld-Generalization}.

Abstract

RL for GUI Agents

Explore similar work

Generalization in Online Reinforcement Learning for Mobile Agents

GSAR: Goal-State-Anchor Rewards for Mobile GUI Agents with Self-Evolving Data Synthesis

How Mobile World Model Guides GUI Agents?