π€ AI Summary
This work addresses the challenge that off-policy samples generated by heterogeneous models can compromise the convergence of Generalized Reward-Weighted Policy Optimization (GRPO) in decentralized collaborative reinforcement learning. To overcome this, the paper proposes the first framework enabling efficient co-training of GRPO with heterogeneous agents, introducing a Filtered Truncated Importance Sampling (F-TIS) mechanism. F-TIS effectively corrects off-policy bias while preserving convergence guarantees and enhancing generalization. By leveraging lightweight communication to aggregate diverse local data, the method achieves convergence performance comparable to purely on-policy training across various heterogeneity settings, and demonstrates up to a 12% performance improvement on out-of-distribution tasks.
π Abstract
Reinforcement learning methods such as GRPO have seen great popularity in LLM post-training. In GRPO, models produce completions to a set of prompts, which are rewarded, and the policy is updated towards the relatively high reward completions. Due to the auto-regressive nature of models, the generation phase of such style of training can be extremely time consuming. As a solution, prior work has sought to distribute the inference step across many nodes, working parallel. These works assume primarily homogeneous models in the training in order to keep samples as close to on-policy as possible. This assumption may be impractical in decentralized systems, where parties with various computes and preferences may wish to collaborate on the same task. Thus, decentralized training requires an approach that can handle heterogeneous models - different models collaborating on the same tasks. However, this leads to highly off-policy samples presented during training, which prior work has identified that off-policy samples can hurt GRPO convergence. To enable heterogeneity, we propose Filtered Truncated Importance Sampling (F-TIS) - a GRPO-style training paradigm that can use off-policy samples to improve local model's learning. Our framework allows various models to collaborate in the same RL training run while being communication efficient. We extensively evaluate F-TIS in various heterogeneous setups and we show that it exhibits identical final model convergence to purely on-sample training. Furthermore, we observe in some setups better generalization on out-of-distribution tasks than on-policy training, increasing model's performance by up to 12\%.