Source-linked AI summary
WM-R1: Training GUI Agents to Reason and leverage World Models with Reinforcement Learning
Yu Han, Tianwen Qian
TL;DR
Mobile GUI reinforcement learning is costly and unstable because it relies on extensive real-environment interaction. WM-R1 trains entirely from world-model state transitions and uses those predictions during reasoning, achieving stronger benchmark performance than GRPO-only and inference-time simulation baselines.
Problem
Existing mobile GUI RL methods require extensive real-environment interaction, creating high computational cost, irreversibility, and environment noise.
Method
WM-R1 replaces real Android transitions with world-model simulation during GRPO training and embeds predicted consequences into chain-of-thought action reasoning.
Results
WM-R1-trained agents outperform GRPO and inference-time simulation baselines across reported mobile GUI benchmarks and model scales.
Takeaways & Limitations
World-model-simulated training provides a practical path toward generalizable mobile GUI reasoning agents without real-environment interaction.
Takeaways & Limitations
WM-R1 depends on world-model fidelity and may struggle with highly dynamic or interactive elements, where inaccurate predictions can produce suboptimal policies.
Abstract
from arXiv · showhide
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 .
Introduction
WM-R1 addresses the cost, irreversibility, and noise of real-environment RL for mobile GUI agents by replacing Android interactions with world-model simulation. It integrates simulated outcomes into reasoning and combines task success, trajectory length, and world-model use in training rewards.
- Existing GUI RL methods require extensive real-environment interaction, which is slower than inference and introduces irreversibility and transition noise.
- WM-R1 replaces the real Android environment with world models that generate simulated GUI state transitions during training.
- World-model simulation enables massively parallelized, step-level trajectory generation without real-environment interaction.
- WM-R1 embeds world-model predictions into chain-of-thought through <call_wm>, allowing candidate actions to be simulated and refined before commitment.
- The training loop uses GRPO with a composite reward for task success, trajectory length, and world-model call count.
4. Experiments
Experiments evaluate WM-R1-trained agents on AndroidWorld, GUI-Odyssey, and AndroidControl at 3B and 7B model scales, comparing them with GRPO baselines.
- WM-R1-trained agents outperform GRPO baselines on AndroidWorld, GUI-Odyssey, and AndroidControl at both 3B and 7B scales.
Related Work
Related work spans visual GUI agents, reinforcement-learning methods, and world models for GUI interaction. Existing world-model approaches are categorized by whether they predict text or DOM updates, pixels, or executable code.
- GUI agents: Visual GUI-agent research includes spatial-aware understanding, explicit grounding supervision, and unified vision-language-action prediction.
- Reinforcement learning for GUI agents: PPO uses a value network, whereas GRPO computes advantages relative to sampled trajectory groups and has been extended to vision-language models and GUI agents.
- World models for GUI: World-model GUI methods include text-based language or DOM prediction, pixel-based screenshot synthesis, and code-based renderable HTML.
Preliminaries
Standard GUI RL observes screenshots and interaction histories, samples actions, receives environment transitions and rewards, and uses GRPO for policy optimization. WM-R1 replaces those transitions with world-model predictions while retaining group-relative optimization and simulated trajectory generation.
- Standard RL: A GUI agent maps a screenshot and interaction history to an action such as clicking, typing, or scrolling, after which the environment produces the next state and sparse reward.
- GRPO: GRPO computes each trajectory’s advantage from its reward relative to the group mean and standard deviation, avoiding a separate value network.
- GRPO: The policy update uses an importance-sampling ratio, while β controls the KL-divergence penalty relative to a reference policy.
- WM-R1: WM-R1 generates multiple candidate trajectories in a world-model-simulated environment, scores them with a rule-based reward model, then computes GRPO advantages and updates the policy.
- WM-R1: During reasoning, the agent can invoke the world model to predict a next state and refine a candidate action before producing the final action.
Method: WM-R1
WM-R1 trains mobile GUI agents entirely through a learned world model, integrating simulated action consequences into chain-of-thought reasoning. GRPO optimizes a composite reward for task success, trajectory efficiency, and world-model usage.
- World-model training: WM-R1 replaces the real Android environment with a learned world model for fully simulated, massively parallel agent training.Trajectory generation no longer requires real-environment interaction.
- World-model reasoning: During chain-of-thought, the agent can invoke <call_wm> to simulate a proposed action before producing the final action.The simulated screenshot is returned as a system observation within the same thinking block.
- Trajectory generation: Each rollout iterates between observing a simulated screenshot, reasoning with optional world-model calls, simulating consequences, and repeating until termination or Tmax.If no world-model call is made, the reasoning block ends and the agent outputs an action.
- Composite reward: The composite reward combines task success, a length penalty, and a world-model usage term as R = α · Rsuccess + β · RL + γ · RWM.This reward design targets correctness, efficiency, and simulation usage jointly.
- Success reward: Success is judged by an LLM against a reference trajectory, with binary reward for semantic equivalence of the action sequence and effects.Exact match is defined by semantic equivalence rather than literal action-string identity.
- Efficiency and utilization: The length reward penalizes unnecessarily long trajectories using a budget mixing reference length and maximum allowed length, while RWM regulates simulation-call usage relative to average calls.GRPO advantages favor trajectories that use simulation effectively without unnecessary or poorly timed calls.
Experiments
WM-R1 is evaluated on curated Android GUI tasks, long-horizon interaction, grounding, generalization, training dynamics, and component ablations. Across these experiments, world-model-simulated training improves performance while agents increasingly use simulation selectively.
- Evaluation Setup: 2000 examples are sampled from three Android GUI sources after removing trivially easy and impossibly hard tasks.The retained tasks target intermediate complexity, where world-model feedback is considered most informative.
- Long-Horizon Tasks: WM-R1-7B reaches 39.8 on AndroidWorld and 31.6 on GUI-Odyssey, surpassing UI-R1-7B by +9.0 and +10.1 points, respectively.At 3B scale, WM-R1-3B scores 29.5 on AndroidWorld versus 18.6 for UI-R1-3B.
- Long-Horizon Tasks: +11.1 on AndroidWorld at 7B is WM-R1’s gain over inference-time Code2World augmentation, with Hard-subset gains of +9.6 on AndroidWorld and +8.8 on GUI-Odyssey.The corresponding AndroidWorld gain at 3B is +8.9.
- Generalization: +16.0 is WM-R1’s average OOD improvement relative to zero-shot, compared with +8.7 for UI-R1.The paper attributes this pattern to world-model-based reasoning and a more diverse simulated training distribution.
- Training Dynamics: 39.8% success by episode 15 gives WM-R1 a 9.0-point advantage over UI-R1, while its policy entropy remains higher at convergence, approximately 3.2 versus 2.6.WM-R1 also reaches 31.5% by episode 5 and shows lower variance across seeds.
- Action Prediction and Grounding: WM-R1 improves mobile grounding across benchmarks, including 90.4 icon accuracy on ScreenSpot-V2 and 87.8 AndroidControl grounding at 7B.At 3B, ScreenSpot-Pro icon and text grounding improve over zero-shot by +15.1 and +14.2 points.
- Ablation Studies: Removing chain-of-thought reasoning causes the largest ablation degradation, −5.2 on AndroidWorld and −4.9 on its Hard subset.Length reward contributes a 3.3-point gain, while the WM call reward has the smallest individual impact but stabilizes simulator use.
- World Model Utilization Dynamics: WM calls decline from approximately 4 per trajectory during exploration to approximately 1.8 during selective utilization as average reward rises to 0.52.The intermediate calibration phase reduces calls to approximately 2.5 per trajectory.
Conclusion
WM-R1 trains GUI agents entirely in a world-model-simulated environment, using simulated transitions during GRPO training rather than only at inference. The resulting agents improve over zero-shot and GRPO baselines on mobile benchmarks.
- WM-R1 uses world models as the source of all state transitions during GRPO training, rather than as post-hoc inference-time proposal evaluators.
- World-model-simulated training enables massively parallel rollouts and integrates simulated visual feedback into the agent’s chain-of-thought reasoning.
- WM-R1-trained agents improve over zero-shot and GRPO baselines on mobile benchmarks.
Additional Training Details
The additional training details specify the datasets, frozen world model, distributed infrastructure, composite reward, and evaluation procedures used for WM-R1.
- Distributed training setup: Distributed training uses Ray actors, FSDP for memory-efficient policy training, and vLLM for rollout sequence generation.
- World model configuration: Training uses a frozen Code2World-8B world model that renders next GUI states as 1080 × 2400 screenshots from screenshots and actions.
- Data curation: The training data combines AndroidCode, GUI-Odyssey, and GUI-R1, filters by task difficulty, and randomly samples 2000 examples.
- Reward function: The composite reward combines task success, trajectory length, and world-model call usage, with weights α = 1.0, β = 0.5, and γ = 0.1.
- Evaluation: The evaluation protocols are applied uniformly across methods, and LLM judgments agree with programmatic matching on 97.0% of a 100-task deterministic subset.
WM-R1 Training Algorithm
The WM-R1 training loop generates multiple rollouts per task, lets the agent query the world model while reasoning, and updates the policy using composite rewards and GRPO advantages.
- Rollout generation: For each batch of task instructions, WM-R1 generates multiple rollouts over steps up to a maximum horizon.
- World-model interaction: When a thought contains <call_wm>, the loop extracts an action, queries the world model, and appends the simulated observation to the trajectory.
- Action selection: The agent’s thought must end with <action> for the loop to extract the final action; otherwise, the trajectory is marked complete.
- Policy update: After rollouts finish, WM-R1 computes composite rewards, derives group-relative GRPO advantages, and updates the policy πθ.
Parameter Sensitivity Analysis
WM-R1 is moderately sensitive to hyperparameters, with default settings forming a stable configuration. Additional experiments show favorable cost-performance tradeoffs, cross-architecture gains, and important scope boundaries.
- Sensitivity overview: WM-R1 shows no catastrophic degradation under reasonable hyperparameter variation, and the default configuration is a stable optimum across dimensions.
- Computational cost: 11.2 wall-clock hours are required for WM-R1-7B to complete 15 training episodes on 2000 tasks using one NVIDIA H200.
- Computational cost: +9.0 points on AndroidWorld are obtained by WM-R1-7B over UI-R1-7B, despite WM-R1 requiring 1.44× the training time.
- Inference cost: WM-R1 adds approximately 2.3 seconds of inference latency per 7B trajectory through an average of 1.8 world-model calls.
- Qwen3-VL transfer: +9.3 and +10.0 points are achieved by WM-R1-8B over UI-R1-8B on AndroidWorld and GUI-Odyssey, respectively.
- Limitations: WM-R1 depends on world-model fidelity and currently assumes discrete actions, limiting direct coverage of dynamic interfaces and continuous controls.