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

TL;DR

This paper introduces Probabilistic Chunk Masking (PCM), a method that reduces gradient computation in vision-language-action reinforcement learning by selectively focusing on informative trajectory segments, leading to significant speedups.

Contribution

PCM is a novel modification to GRPO that allocates gradient computation to a subset of trajectory chunks based on success-failure variance, improving efficiency without sacrificing success rates.

Findings

PCM achieves 2.38x wall-clock speedup over standard GRPO.

Fewer than 20% of trajectory chunks are backpropagated through with PCM.

PCM reduces peak activation memory by 60% while maintaining success rates.

Abstract

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.