WM-R1: Training GUI Agents to Reason and leverage World Models with Reinforcement Learning
WM-R1 replaces real environments with world models, significantly enhancing GUI agents' reasoning capabilities.
Key Findings
Methodology
WM-R1 framework replaces real environments with world models for training, using GRPO for policy optimization. World models handle state transitions, allowing agents to reason about candidate actions' consequences before execution. This method supports massively parallelized and granular trajectory generation, introducing a multi-dimensional rule-based reward to optimize task success, trajectory efficiency, and world model utilization.
Key Results
- WM-R1 achieved a 39.8% success rate on AndroidWorld, outperforming UI-R1 by 9.0 points, demonstrating its significant advantage in long-horizon tasks.
- On GUI-Odyssey, WM-R1 achieved a success rate of 31.6%, surpassing UI-R1's 21.5%, showing superiority in cross-app navigation tasks.
- WM-R1 improved icon grounding on ScreenSpot-Pro by 15.1 points, indicating its improvement in GUI grounding tasks.
Significance
By eliminating reliance on real environments, WM-R1 reduces training costs and instability, providing an efficient training method for GUI agents on mobile platforms. Its superior performance in long-horizon tasks and complex environments indicates its potential in automating complex digital workflows.
Technical Contribution
WM-R1 introduces world models as a core component of the training environment, offering a novel training perspective compared to traditional GRPO methods. By simulating state transitions, WM-R1 achieves more efficient policy optimization and stronger reasoning capabilities.
Novelty
WM-R1 is the first framework to train GUI agents entirely within a simulated environment, embedding world models into the reasoning process to evaluate candidate actions' consequences before execution, showing significant innovation compared to existing methods.
Limitations
- WM-R1 may perform poorly in extremely complex or dynamically changing environments due to limited predictive capabilities of world models.
- Over-reliance on world models in some cases may lead to wasted computational resources.
Future Work
Future research could explore more complex world models to further improve WM-R1's performance in dynamic environments. Additionally, combining other reinforcement learning strategies may bring further performance enhancements.
AI Executive Summary
WM-R1 is a novel reinforcement learning framework designed to train GUI agents on mobile platforms by replacing real environments with world models. Traditional reinforcement learning methods typically require extensive real-environment interactions, leading to high resource costs and instability, especially in GUI scenarios. WM-R1 eliminates the need for real environments by using world models as the source of state transitions, supporting massively parallelized and granular trajectory generation.
In experiments, WM-R1 demonstrated outstanding performance on Android mobile benchmarks, significantly outperforming GRPO-only baselines. This indicates WM-R1's superiority in long-horizon tasks and complex environments, particularly in tasks requiring multi-step reasoning. WM-R1's success lies in its ability to evaluate the consequences of candidate actions before execution, optimizing task success, trajectory efficiency, and world model utilization.
Despite WM-R1's significant performance improvements, challenges remain in handling extremely complex or dynamically changing environments. Future research could explore more complex world models and other reinforcement learning strategies to further enhance performance. WM-R1's emergence offers new possibilities for automating complex digital workflows on mobile platforms.
Deep Analysis
Background
In recent years, the demand for automated GUI agents has increased with the proliferation of mobile devices. Traditional reinforcement learning methods typically rely on extensive real-environment interactions for training these agents, leading to high resource costs and instability. While some research has attempted to enhance reasoning capabilities through world models, most methods use world models only at inference time, without integrating them into the training process.
Core Problem
In GUI scenarios, reinforcement learning faces three major bottlenecks: high computational cost, strong irreversibility, and significant environmental noise. These issues limit agents' performance in complex environments, especially in tasks requiring multi-step reasoning.
Innovation
WM-R1 eliminates reliance on real environments by embedding world models into the training process. Its core innovations include: 1) using world models for state transitions, 2) evaluating candidate actions' consequences before execution, and 3) supporting massively parallelized and granular trajectory generation.
Methodology
- �� Use Qwen2.5-VL-3/7B as the base model, trained on a dataset combining AndroidCode, GUI-Odyssey, and GUI-R1.
- �� Employ GRPO for policy optimization, driven by a multi-dimensional rule-based reward.
- �� World models handle state transitions, allowing agents to reason about candidate actions' consequences before execution.
Experiments
Experiments were conducted on benchmarks such as AndroidWorld, GUI-Odyssey, and AndroidControl to evaluate WM-R1's performance in long-horizon tasks and GUI grounding tasks. The Code2World-8B model was used to generate renderable HTML, rendered to screenshots using a headless browser.
Results
WM-R1 achieved a 39.8% success rate on AndroidWorld, outperforming UI-R1 by 9.0 points. On GUI-Odyssey, WM-R1 achieved a success rate of 31.6%, surpassing UI-R1's 21.5%. WM-R1 improved icon grounding on ScreenSpot-Pro by 15.1 points.
Applications
WM-R1 can be used to automate complex digital workflows, especially on mobile platforms. Its superior performance in long-horizon tasks and complex environments makes it valuable in mobile app development and testing.
Limitations & Outlook
Despite WM-R1's significant performance improvements, challenges remain in handling extremely complex or dynamically changing environments. Additionally, over-reliance on world models may lead to wasted computational resources. Future research could explore more complex world models and other reinforcement learning strategies to further enhance performance.
Plain Language Accessible to non-experts
Imagine a factory where robots need to complete a series of tasks. Traditional methods require robots to constantly try in the real environment, which is time-consuming and expensive. WM-R1 is like equipping robots with a virtual simulator, allowing them to practice in a virtual environment until they find the best solution. This way, when robots perform tasks in the real environment, they can be more efficient and accurate.
ELI14 Explained like you're 14
Hey there, imagine you're playing a super complex mobile game. This game has many levels, and each level requires you to make the right choices to pass. Now, imagine having a super smart assistant that can try out your choices in a simulated world before you make them, and then tell you which choice is the best! That's what WM-R1 does! It helps computer programs practice in a virtual world so they can perform better in the real world! Isn't that cool?
Glossary
Reinforcement Learning
A machine learning method that learns optimal strategies through interaction with the environment.
Used in WM-R1 to train GUI agents.
World Model
A model that simulates environment state transitions to predict future states.
WM-R1 uses world models to replace real environments for training.
GRPO (Group Relative Policy Optimization)
A policy optimization algorithm that updates policies by computing advantages relative to a group.
WM-R1 uses GRPO for policy optimization.
GUI Agent
An automated tool capable of performing tasks in a graphical user interface.
The subject of training in WM-R1.
Parallelization
Executing multiple computational tasks simultaneously to improve efficiency.
WM-R1 supports massively parallelized training.
Open Questions Unanswered questions from this research
- 1 How to improve world model prediction accuracy in dynamically changing environments remains an open question.
- 2 Existing methods perform poorly in handling extremely complex tasks, requiring further research.
Applications
Immediate Applications
Mobile App Testing
WM-R1 can be used to automate the testing process of mobile apps, improving testing efficiency and accuracy.
Long-term Vision
Complex Workflow Automation
WM-R1 has potential in automating complex digital workflows, potentially transforming future work methods.
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 .