arXiv:2605.16154cs.LGcs.RO2026-05

通过概率掩码只计算关键片段梯度,让视觉语言动作强化学习更快更省显存

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

论文配图:Learn Where Outcomes Diverge: Efficient VLA RL via Probabilistic Chunk Masking
图 1 · 摘自论文原文
  • 只对成功与失败轨迹差异大的片段计算梯度,其余忽略
  • 在三个LIBERO任务上达成相同成功率,速度提升2.38倍
  • 无需奖励模型或价值函数,适合资源受限的RL部署场景

强化学习使视觉-语言-动作(VLA)策略能泛化到训练分布之外,但后训练阶段计算成本高昂。在基于GRPO的VLA强化学习中,我们发现主要开销来自梯度计算:每步约78%的运行时间用于梯度计算,而采样仅占21%。这是因为大量计算浪费在政策已掌握的阶段。GRPO的优化信号依赖优势方差:只有成功与失败轨迹出现分歧的阶段才提供有效学习信号,但其对整个轨迹的所有片段赋予相同优势。本文提出概率片段掩码(PCM),作为GRPO的即插即用改进,仅对每条轨迹中概率性选择的小部分片段进行梯度计算。PCM利用成功-失败动作方差(一种从轨迹中推导出的代理指标)评估语义阶段,并以在线更新的保留概率采样固定数量的片段。我们形式化了每段梯度方差的概念,并证明成功-失败动作方差是其可测量的代理。PCM无需奖励模型或学习的价值函数。在三个LIBERO基准上,PCM达到与标准GRPO相当的最终成功率,实现2.38倍的墙钟速度提升、4.8倍更快的梯度更新和60%更低的峰值激活内存,同时反向传播的片段少于20%。

原文摘要 · Abstract (English)

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.

强化学习效率优化视觉语言动作梯度压缩

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。