论文精选

PCM:概率性分块掩码加速VLA强化学习2.38倍

Learn Where Outcomes Diverge: Efficient VLA RL via Probabilistic Chunk Masking

精选理由

做VLA机器人强化学习的团队终于有了省算力的方案——PCM直接替换GRPO就能省60%内存、快2倍多,效果还不打折,建议做后训练优化的点开看看。

AI 摘要

该论文发现GRPO算法在VLA策略强化学习中,梯度计算占78%时间,而大部分计算浪费在策略已掌握的阶段。为此提出概率性分块掩码(PCM),通过成功-失败动作方差识别关键阶段,仅对少量分块进行梯度更新。PCM无需额外奖励模型,在LIBERO基准上保持相同成功率的同时,实现2.38倍加速、4.8倍梯度更新加速和60%峰值内存降低。

原文 · arXiv cs.LG

Learn Where Outcomes Diverge: Efficient VLA RL via Probabilistic Chunk Masking

Reinforcement learning (RL) allows vision-language-action (VLA) policies to generalize beyond their training distribution by optimizing directly for task success, but post-training is computationally expensive. A natural response has been to speed rollout collection through faster simulators and world models. In GRPO-based VLA RL, we find that the dominant cost lies elsewhere: gradient computation accounts for approximately 78% of wall-clock time per step in our runs, while rollout collection accounts for only 21%. Gradient cost dominates because much of this computation is spent on phases that contribute little to learning. GRPO's learning signal is driven by advantage variance: only phases where successful and failed rollouts diverge produce learning signal. However, GRPO assigns the same advantage to every chunk in a rollout. As a result, actor-update compute is spent uniformly across the trajectory, including phases the policy already handles after pre-training and supervised fine-tuning. This paper presents Probabilistic Chunk Masking (PCM), a drop-in modification to GRPO that allocates gradient computation to a small, probabilistically selected subset of chunks per trajectory. PCM scores semantic phases using success-failure action variance, a rollout-derived proxy for per-phase gradient variance, and samples a fixed chunk budget with online-updated phase-level keep probabilities. We formalize per-phase gradient variance as the quantity determines where gradient computation is useful and show that success-failure action variance provides a measurable proxy for it. PCM requires no reward model or learned critic. On three LIBERO benchmarks, PCM matches the final success rate of standard GRPO while achieving 2.38 times wall-clock speedup, 4.8 times faster gradient updates, and 60% lower peak activation memory, while backpropagating through fewer than 20% of trajectory chunks.

PCM:概率性分块掩码加速VLA强化学习2.38倍 · AI 热点