DistRL: An Asynchronous Distributed Reinforcement Learning Framework for On-Device Control Agent

Authors: Taiyi Wang, Zhihao Wu, Jianheng Liu, Jianye HAO, Jun Wang, Kun Shao
Primary area: applications to robotics, autonomy, planning
Keywords: Mobile Agent, LLM, Reinforcement Learning, Fine Tuning, Distributed Training

Abstract:
On-device control agents, especially on mobile devices, are responsible for operating mobile devices to fulfill users' requests, enabling seamless and intuitive interactions. Integrating Multimodal Large Language Models (MLLMs) into these agents enhances their ability to understand and execute complex commands, thereby improving user experience. However, fine-tuning MLLMs for on-device control presents significant challenges due to limited data availability and inefficient online training processes. This paper introduces DistRL, a novel framework designed to enhance the efficiency of online RL fine-tuning for mobile device control agents. DistRL employs centralized training and decentralized data acquisition to ensure efficient fine-tuning in the context of dynamic online interactions. Additionally, the framework is backed by our tailor-made RL algorithm, which effectively balances exploration with the prioritized utilization of collected data to ensure stable and robust training. Our experiments show that, on average, DistRL delivers a 3$\times$ improvement in training efficiency and enables training data collection 2.4$\times$ faster than the leading synchronous multi-machine methods. Notably, after training, DistRL achieves a 20\% relative improvement in success rate compared to state-of-the-art methods on general Android tasks from an open benchmark, significantly outperforming existing approaches while maintaining the same training time. These results validate DistRL as a scalable and efficient solution, offering substantial improvements in both training efficiency and agent performance for real-world, in-the-wild device control tasks.

Decision: accept (poster)
Reviewer scores: 6, 5, 8, 8

Reviews and discussion:
<thread>
creation date:
1730244889188

thread_type:
official_review

<thread_object>
creator: reviewer_oskg

<content>
summary:
This paper focuses on the domain of on-device control agents and presents DistRL, an asynchronous distributed reinforcement learning framework for device control agents. This paper designs a centralized training and decentralized data acquisition system design to ensure the efficiency of online RL fine-tuning given potentially heterogenous hardware specifications. It also proposes a tailor-made RL algorithm in the context of distributed device control training. Both the system design and new RL algorithm shows improvements over the prior state-of-the-art DigiRL in a popular device-control benchmark AitW.

strengths:
The paper addresses a realistic concern of asynchronomous data collection in applying RL fine-tuning for device-control agents, so it is highly relevant. 

The proposed system-level optimizations can be valuable to the open-source community. I hope the authors can clean up and release the infrastructure for their experiments.

The proposed RL algorithm and system design seems to make a difference on popular AitW benchmark.

weaknesses:
The writing of the paper, in particular the introduction and methodology, can be improved. In fact, the current state of writing has caused me trouble for understanding the contributions of the paper. A significant amount of re-writing is required before the paper can be published in a top-tier venue like ICLR. Please see my questions for more details.

The organization of the section 5 methodology makes it hard to understand what parts are novel in the paper and what parts are adapted from existing literature. From the current writing, it seems to be a combination of existing components from the prior literature. Additional motivations are needed to explain what is novel in the tailor-made RL algorithm.

questions:
Line 53-70 seems repetitive.

Line 71 fine-tuning(type of light training) -> fine-tuning

Paragrapho one and two of intro seem to be a bit verbose in introducing the background.

Line 78-85, it seems unnecessary to introduce the notion of RLHF here cuz this paper does not have human preferences.

Line 86-90 is confusing, it seems to describe challenges of exisiting offline datasets such as AitW and Android-Control, but why is it relevant for online RL fine-tuning?

Line 98-104 can be best put in the related works.

There are too many bolded text in the introduction that can tire the readers. I would suggest removing the bold font for related methods. Same applies to the experiment section.

Line 121-122, te claim of "DistRL is the first work to scale autonomous, online RL fine-tuning for mobile device control in a distributed environment." may be incorrect, DigiRL also supports distributed data collection. I understand that DistRL supports more distributed RL optimizations, but it might be good to update this sentence.

Line 247-257, would be great to defer algorithm-related designs to the next section as this section mainly talks about system design.

There can be more details for the method. In particular, do you maintain a separate model for $Q$ and $\rho$? If so, how are those models trained? How is $\Pcal_{invalid}$ calculated? How is $A(s_t, a_t)$ calculated? More motivations can also be included. For example, why does DPER improve sample efficiency? 5.3 includes a training loss for the trajectory value function $V_{traj}$, where is it used?

Figure 5 (b) needs more explanations, what does "Proportion Above 60%" mean?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730400466820

thread_type:
official_review

<thread_object>
creator: reviewer_w8my

<content>
summary:
This paper provides a way to train device control agents using reinforcement learning. Specifically, it proposes a setup of asynchronous data collection using parallelized data collection on emulators (or real devices), combined with a centralized learning algorithm. The algorithm for training is V-trace, combined with prioritized replay buffer and an additional component encouraging exploration. The system runs on multiple workers and uses a single host learner.

The paper evaluates results on AitW using model-based evaluation. Compared to existing literature/SoTA the DistRL approach achieves a high success rate in less wall time by applying a distributed, asynchronous training approach, which results in faster emulation speed (traj/min). In the General test set, DistRL achieves a success rate a relative improvement of over 19% vs previous work (DigiRL) in the Web Shopping test set, DistRL attains a success rate  outperforming DigiRL by over 14%. The work conducts ablation analysis.

strengths:
The work achieves impressive results on a challenging task. It provides a path to scale up data collection and enable future work to tackle even more complex problems. Implementing this distributed setup and making it work is a valuable contribution. This work is particularly relevant to future researchers where researchers can easily spin up many workers on cloud platforms. 

The paper is well-written, with technical depth and there are enough details describing the methodology and approach. The paper does a good job of comparing their results to existing works and agents.

It is good the authors run multiple experiments per and average the results to provide more stable results. It is also good that the work includes an ablation study.

weaknesses:
In order for this work to be a major contribution, I feel it should be open-sourced since a core contribution is the engineering work around centralized training and decentralized data acquisition, which requires complex communication between different systems. Open-sourcing will greatly help the community to take advantage of the contributions here and with reproducibility.

The core RL techniques used are relatively standard. There's no fundamental advance in machine learning theory or methodology. However combining IMPALA with prioritized experience replay seems to be novel, and it’s implementation is highly non-trivial.

If both DigiRL (single) and DigiRL (multi) have the same amount of emulators, why are the performance numbers between them different? Ultimately, both consist of 32 emulators, if I understand correctly, operating with synchronous training. Please clarify how they are different and why it's relevant to include both. I also think the number of emulators should be displayed more prominently as a key experimental parameter (rather than only in Figure 5)

Spelling/Writing: "ensuring that only contextually appropriate actions are penalized". I believe this should be “inappropriate actions are penalized”?

questions:
“The evaluator receives the current observation, composed of the task description and a screenshot of the device, and outputs a reward signal.” – Have the authors validated that this does a good enough job of success validation, versus, say using the entire trajectory?

To me, “On-device End-to-End Agent Performance Evaluation” implies that the inference is done on the device, but this appears not to be the case. I understand it could be done in theory because the model is small enough, but it’s not actually on device as far as I can tell. Is there another way to describe it?

## Comments
(You don’t have to address these; they are intended to be helpful only.)

For future work, you may consider SeedRL to further scale the approach. https://arxiv.org/pdf/1910.06591#page=5.27

Future work could evaluate on online benchmarks such as AndroidWorld to ensure that results are easily comparable for future researchers, since model-based evaluation is harder to reproduce due to costs and because of differences in model outputs, even for the same base model.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730620251173

thread_type:
official_review

<thread_object>
creator: reviewer_fxpf

<content>
summary:
This paper introduces DistRL, an asynchronous distributed reinforcement learning (RL) framework designed to fine-tune on-device control agents, particularly for mobile environments. The key contributions include a scalable architecture that decouples data acquisition from policy learning, improving training efficiency and robustness in dynamic environments. The paper claims significant improvements in training speed and success rates compared to state-of-the-art methods, validated through experiments on Android tasks.

strengths:
1. The decoupling of data acquisition from policy learning allows the framework to handle heterogeneous devices efficiently, which is crucial for real-world applications.
2. The asynchronous nature of DistRL enables faster data collection and policy updates compared to synchronous methods, which is demonstrated by a claimed 3x improvement in training speed.
3. The focus on mobile device control agents addresses a practical problem with significant industry relevance, particularly in automating complex user interactions on smartphones.

weaknesses:
1.  There is a significant discrepancy between the model descriptions in different sections (T5-based multimodal model mentioned in Sec 1 vs GPT-2/ROBERTA in A.5.1). This inconsistency undermines the paper’s credibility and makes it difficult to understand what model was actually used. If different models were used for different tasks or components (e.g., policy learning vs critic evaluation), the authors should explain why and provide a rationale for these choices. If there were intentional changes in model usage (e.g., using GPT-2 for policy learning and ROBERTA for critic evaluation), the authors should explain why these variations were necessary and how they impacted performance.

2. It is unclear how the proposed method differs from IMPALA from the methodology perspective. I understand that applying them to a new on-device setting is trivial, but it is confusing to me what are the major differences from the algorithm side. It seems that A-RIDE builds upon existing methods like Retrace and GAE but does not introduce substantial innovations beyond adapting these techniques to a distributed setting. Additionally, the author mentioned in Sec. 2.3 that:
> Furthermore, we develop a novel reinforcement learning algorithm specifically designed to handle the stochasticity and scalability inherent in mobile device environments.
However, I did not see any specific techniques and algorithm improvements that address the stochasticity and scalability issues.
If there are indeed major algorithm changes, it would be beneficial to add a baseline of IMPALA and IMPACT.

3.  Additionally, the training and testing set used are not clearly described. I am a bit confused about the difference between the training and testing sets -- are there overlaps of tasks between them? How is the expert-curated and AndroidWorld training set selected? This questions the generalization issues of the experiments.

4. The evaluation metrics and result plots are a bit confusing. It is unclear whether the plots in Fig. 6 are for the training task success rate or the testing task's. Additionally, adding the testing task success rate plots instead of just training one in Fig. 5 would be beneficial as well.

5. The x-axis in Fig. 5 (a) is the time. From Fig. 5 (c) we can see there are significantly more trajectories collected than baselines. I acknowledge that the implementations of the proposed method is much more efficient than baselines, but I am also curious about what the plot looks like if the x-axis is the trajectory-num or the step-num, which is commonly reported in the RL literature. It would be super helpful to understand the upper bounds of the baseline algorithms given the same number of collected data. It is also interesting to know the synchronized algorithm performance trend.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730687231166

thread_type:
official_review

<thread_object>
creator: reviewer_vpis

<content>
summary:
The paper presents DistRL, an asynchronous distributed reinforcement learning framework designed explicitly for on-device control agents in mobile environments. DistRL aims to optimize training efficiency and agent performance for complex tasks, leveraging centralized training with decentralized data acquisition. By introducing a unique off-policy algorithm, A-RIDE, DistRL effectively manages asynchronous data collection and prioritizes valuable interactions, addressing challenges like delayed updates and non-stationary data. Experiments demonstrate its significant performance and efficiency improvements over existing frameworks, especially for real-world mobile tasks, positioning DistRL as a scalable solution in distributed reinforcement learning for mobile devices.

strengths:
- DistRL employs a decoupled architecture that enhances scalability and robustness by distributing data acquisition across heterogeneous mobile devices, making it efficient for large-scale deployment. Moreover, they propose an A-RIDE algorithm adapted for asynchronous environments, balancing exploration and stability, which improves sample efficiency and task success rates.
- The paper provides comprehensive experiment results that significantly outperform baselines, showcasing superior success rates in complex tasks. Ablation studies highlight the importance of its components, like DPER and retrace correction.
- The paper is well-written and easy to follow.

weaknesses:
- Lots of hyperparameters are required. (e.g., Equation 1, 2, DPER). Do you have any tips or evidence for determining them, and are the values sensitive to different types of environments?
- Although the experiments cover general tasks, expanding the testing to a broader range of mobile applications would offer more substantial evidence of DistRL’s adaptability. (from additional tasks from AITW and AndroidWorld)
- Is VLM Evaluator (Gemini)'s performance in providing reward signals credible? Did you encounter any failure cases where VLM misleadingly provides success/failure signals? AndroidWorld provides manually defined dense rewards. Have you tried training DistRL with such a reward?
- While the asynchronous design improves adaptability, more detail on DistRL's response to drastic changes in mobile environments or new, unseen tasks would strengthen the discussion of its robustness in dynamic settings.

questions:
- Any results from INSTALL and GOOGLEAPPS of AITW? / Why not compare results in AndroidWorld tasks?
- In real-world cases, many users (workers) have different types of devices and different OS versions. How does DistRL handle scenarios where the mobile device encounters significant UI changes or unexpected app updates during training?
- What are the trade-offs between computational overhead and empirical performance (success rate) when using more environments (workers) and Distributed Prioritized Experience Replay (DPER)?
- Could the framework be adapted for other operating systems, such as iOS, and if so, what modifications would be necessary?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731952450731

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overview of Our Response & Open Source Concerns

comment:
We sincerely appreciate all the valuable suggestions and acknowledge the reviewers' agreement that our contributions are significant within the application domain. We aim to present our work as a cohesive and integrated system for everyone's use. **Our work is driven by real-world practice and substantial needs for developing efficient, asynchronous RL fine-tuning systems for LLM agents**.

## **Commitment to Open Source and Transparency**

**In adherence to anonymity policies, we are committed to releasing the complete codebase, model weights, necessary prompts, datasets, and product demos for DistRL post reviews shortly.** 

Interested parties and reviewers can refer to the temporary code repository here: `https://anonymous.4open.science/r/DistRL-C41F/`.

## **Revisions Based on Valuable Feedback**

We have thoroughly revised our paper to incorporate all the insightful comments received (**Main revisions were marked with titles highlighted in blue**). The main revisions include, but are not limited to:

1. **Additional Experiments:**
	We have included the requested extra experiments in both the appendix and the main body of the paper, particularly focusing on the validation results of our evaluators (See Section 6.3).

2. **Broader Discussions on Generalization:**
	Expanded discussions on the broader scopes of generalization abilities across various applications have been added as per the reviewers' suggestions (See Section A.6.2, A.7).

3. **Enhanced Methodology Section**
The methodology section (Section 5) has been revised to emphasize our novel algorithm in comparison to existing works. Additional details have been provided in Section A.4 to clarify our approach further.

4. **Improved Presentation and Corrections**
We have meticulously corrected typos, replaced citations with concise forms, removed redundant bold texts for better readability, and enhanced the overall presentation quality of the paper. (Mainly from Section 1 to Section 2)

Some references I mentioned in the responses below:

 - [3] Hao Bai, Yifei Zhou, Mert Cemri, Jiayi Pan, Alane Suhr, Sergey Levine, and Aviral Kumar. Digirl: Training in-the-wild device-control agents with autonomous reinforcement learning. arXiv preprint arXiv:2406.11896, 2024.

 - [6] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al. Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. In International conference on machine learning, pp. 1407–1416. PMLR, 2018.

 - [42] Zhuosheng Zhan and Aston Zhang. You only look at screens: Multimodal chain-of-action agents. arXiv preprint arXiv:2309.11436, 2023.

We believe these revisions address the reviewers' concerns comprehensively and significantly improve the quality and clarity of our work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313226556

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response - Discussion Summary

comment:
Dear PC, SAC, AC, and Reviewers,

We sincerely appreciate your valuable suggestions and are grateful for the consensus among the reviewers regarding the significance of our contributions within the application domain. Our goal is to present our work as a cohesive and integrated system for widespread use, driven by real-world practices and the substantial need for developing efficient, distributed, and asynchronous reinforcement learning (RL) fine-tuning frameworks for large language model (LLM) agents on mobile devices.

After extensive discussions, we have addressed most of the concerns raised by the reviewers. Reviewers **vpis, oSkG, and w8MY** have expressed a very positive attitude toward our work, acknowledging its strengths, including its novelty, strong performance, and comprehensive experiments and ablation studies, as recognized by all reviewers.

Possibly due to the busy schedule, reviewer **fxPf** did not provide feedback on our responses, in which we have provided detailed experiments and responses. However, we believe they can resolve the issues raised by Reviewer **fxPf**.


**Commitment to Open Source and Transparency**

In adherence to anonymity policies, we are committed to releasing the complete codebase, model weights, necessary prompts, datasets, and product demonstrations for DistRL shortly after the review process concludes.

We extend our heartfelt thanks to the reviewers for their thorough evaluations and invaluable feedback throughout the discussion period. We greatly appreciate the recognition and thoughtful assessment of our paper.


Best regards,
Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
**summary**

The paper introduces DistRL, an asynchronous distributed RL framework designed for training on-device control agents in mobile environments. Combining centralized training with decentralized data collection, DistRL addresses challenges like delayed updates and non-stationary data through its novel off-policy algorithm, A-RIDE, which incorporates asynchronous data collection, prioritized replay, and exploration-driven components. The scalable architecture decouples policy learning from data collection, enhancing training efficiency and robustness across heterogeneous hardware. Evaluations on the AitW benchmark demonstrate significant improvements over the state-of-the-art DigiRL, achieving faster training, higher emulation speeds, and success rate improvements of over 19% in the General test set and 14% in the Web Shopping test set.

---

**strengths**

* Novel idea of decoupling architecture and A-RIDE algorithm for asynchronous environments.
* Comprehensive results: Experimental results demonstrate significant performance improvements over baselines, showcasing superior success rates in complex tasks. Also, the asynchronous nature of the framework accelerates data collection and policy updates, achieving a 3x improvement in training speed compared to synchronous methods.
* Practical topic: The framework addresses a real-world problem of mobile device control, focusing on automating complex user interactions with high industry relevance.
* The paper is well-written, providing detailed descriptions of the methodology and clear comparisons with existing works.

---

**weaknesses**

* Evaluation ambiguities: The evaluation metrics and result plots are somewhat unclear. The reliance on VLM Evaluator (Gemini) for reward signals raises questions about its accuracy. The paper does not discuss potential failure cases where VLM may provide misleading success/failure signals.
* Given that a significant contribution of the work lies in its engineering, including centralized training and decentralized data acquisition, open-sourcing the implementation would enhance reproducibility and provide value to the community.
* Incremental innovation: The RL techniques employed, such as combining IMPALA with prioritized experience replay, are relatively standard and do not represent a fundamental advance in machine learning theory. However, their implementation in a distributed setting is novel and non-trivial.

---

**decision**

I think the paper makes a nice contribution that the community will find valuable. However, I encourage the authors to think carefully about how to reflect on the comments from reviewers (+open-sourcing the framework) in the camera-ready version.