DISTRL: AN ASYNCHRONOUS DISTRIBUTED REINFORCEMENT LEARNING FRAMEWORK FOR ON-DEVICE CONTROL AGENTS

DISTRL:
AN ASYNCHRONOUS DISTRIBUTED
REINFORCEMENT LEARNING FRAMEWORK
FOR ON-DEVICE CONTROL AGENTS

Taiyi Wang^{1 2 * †}, Zhihao Wu^{3 *}, Jianheng Liu⁴, Jianye Hao³, Jun Wang⁴, Kun Shao^{3 †}

¹ University of Cambridge
² Powersense Technology Limited
³ Huawei Noah's Ark Lab
⁴ University College London
^*Equal Contribution
^†Corresponding authors: Taiyi.Wang@cl.cam.ac.uk, shaokun2@huawei.com

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 3X improvement in training efficiency and enables training data collection 2.4X 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.

System Design

DistRL is an asynchronous distributed reinforcement learning framework that decouples trajectory collection from policy learning. The framework consists of a central Host Learner equipped with GPUs for policy training, and multiple heterogeneous Workers running Android emulators or connected to mobile devices. This separation optimizes efficiency by aligning tasks with appropriate hardware - CPU-intensive environment interactions on worker machines and GPU-accelerated policy training on the host server. The host maintains a Circular Replay Buffer for storing trajectories and uses a FIFO Queue to manage incoming experiences from workers efficiently.

The workers operate in parallel through multi-threading, with each thread executing the latest policy received from the Host Learner. Workers collect trajectories during environment interactions and asynchronously send them back to the host for training. The framework demonstrates excellent scalability - with two 96 vCPU machines supporting up to 32 concurrent emulators and nearly linear scaling with worker performance.

A-RIDE: The Backbone of DistRL

Our method employs advantage-based estimations to refine policy gradient updates, as an extension of Generalized Advantage Estimation (GAE) by Schulman et al. (2015), to better suit asynchronous, distributed environments common in device control tasks.

Advantage Computation

With the trajectory-level rewards and state value estimates obtained, we compute the advantage function A(s_t, a_t) using one-step estimation:

A(s_t, a_t) = Q(s_t,a_t) - V(s_t) = r(s_t, a_t) + γV(s_t+1) - V(s_t)

which correctly represents the advantage function as per the policy gradient theorem. Here, r(s_t, a_t) includes signals of immediate rewards. The advantage and value functions are further modified by off-policy corrections.

Policy Optimization with Robust Regularization

Finally, the policy network optimizes the following loss:

ℒ = −𝔼_μ[ρ_t A(s_t, a_t) log π(a_t|s_t)] − β 𝔼_μ[ℍ(π(a_t|s_t))] + λ 𝔼_μ[𝒫_invalid(a_t)],

where ρ_t = π(a_t|s_t)/μ(a_t|s_t) is the importance sampling ratio between the target policy π and the behavior policy μ, ℍ is the entropy term for prevention of overfitting, 𝒫_invalid(a_t) imposes a penalty on actions deemed invalid based on task-specific criteria, β controls the strength of entropy regularization, and λ modulates the penalty's influence. The penalty is assigned using validation through pre-trained LLMs like Gemini, ensuring that inappropriate actions are penalized. This formulation encourages the agent to explore a diverse set of actions while constraining it to generate valid and meaningful commands, thereby enhancing both exploration and policy robustness when dealing with online non-stationarities.

Off-policy correction: Retrace

To enhance the estimation of the state-value function V(s_t) in off-policy and asynchronous settings, we apply Retrace(λ) corrections directly to V(s_t). The Retrace algorithm extends the TD(λ) method to off-policy learning by incorporating importance sampling ratios and a trace decay parameter λ. Specifically, we update V(s_t) using the correction term δ_t, computed as:

V(s_t) ← V(s_t) + δ_t

δ_t = Σ_k=t^H γ^k-t (∏_i=t+1^k c_i) [r_k + γV(s_k+1) - V(s_k)]

where c_i = λ min(1, ρ_i) with λ ∈ [0,1] being the trace decay parameter, and ρ_i is the importance sampling ratio as mentioned before. This correction term effectively adjusts the value estimates using future rewards and importance sampling, enabling off-policy learning while mitigating variance due to importance weights. By applying Retrace(λ), we improve the estimation of V(s_t) in off-policy settings, enhancing the stability and convergence of the value network.

Distributed Prioritized Experience Replay (DPER)

To improve sample efficiency, we employ Distributed Prioritized Experience Replay (DPER). For each trajectory τ = {(s_t, a_t, r_t, s_t+1)}_t=0^H, we compute the priority p(τ) as:

p(τ) = w₁|δ̄| + w₂ρ̄ + w₃ℍ̄

where |δ̄| is the average absolute temporal-difference (TD) error over the trajectory, calculated as δ_t = r_t + γV(s_t+1) - V(s_t); ρ̄ is the average importance sampling ratio ρ_t; and ℍ̄ is the average policy entropy, ℍ_t = -log π(a_t|s_t), encouraging exploration by encouraging policy uncertainty, thus avoiding early convergence to suboptimal policies during training in dynamic environments. The weights w₁, w₂, and w₃ balance the contributions of each component, which is selected by grid-search (See Appendix: Hyperparameters). Trajectories with higher priorities are replayed more frequently, focusing learning on the most informative experiences. Priorities are periodically updated based on the latest policy, recalculating them to focus learning on the most informative experiences, ensuring continual adaptation to evolving behavior policies.

Results

Training performance (32 emulators) between the current SoTA (DigiRL) and DistRL, highlighting the enhanced efficiency of DistRL’s distributed framework during online training. (a) Wall-clock time comparison; (b) Training efficiency comparison; (c) Accumulated trajectories collection ability comparison; (d) Scalability of different training frameworks

Framework Type	Framework Name	General		Web Shopping
Framework Type	Framework Name	Training	Test	Training	Test
Prompting	AppAgent + GPT-4v	41.4	43.0	31.2	35.2
Prompting	AppAgent + Gemini	39.1	45.3	30.5	32.0
Learning	AutoUI	38.3	40.6	42.2	44.5
	DigiRL (single, online)	64.6 ± 1.5	59.9 ± 2.1	63.3 ± 1.5	59.6 ± 3.1
	DigiRL (multi)	67.7 ± 1.3	61.2 ± 2.4	64.5 ± 1.1	59.9 ± 2.8
	DistRL (Ours)	75.5 ± 0.2	73.2 ± 1.1	69.8 ± 0.5	68.5 ± 1.7

Main comparisons regarding the success rate of different agents across various settings. Each experiment is repeated three times and the mean and standard deviation are reported. Results are evaluated with our autonomous evaluator with the 128 user instructions in the train and test set.

For full results and more details, please refer to our paper.

BibTeX

@article{wang2024distrl, title={DistRL: An Asynchronous Distributed Reinforcement Learning Framework for On-Device Control Agents}, author={Wang, Taiyi and Wu, Zhihao and Liu, Jianheng and Hao, Jianye and Wang, Jun and Shao, Kun}, journal={arXiv preprint arXiv:2410.14803}, year={2024} }