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

• URL: http://arxiv.org/abs/2605.16154v1
• Date: Fri, 15 May 2026 16:33:59 GMT
• Title: Learn Where Outcomes Diverge: Efficient VLA RL via Probabilistic Chunk Masking
• Abstract summary: This paper presents Probabilistic Chunk Masking (PCM), a drop-in modification to GRPO that allocates gradient to a small, probabilistically selected subset of chunks per trajectory.<n>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.
• Score: 5.238545250784642
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL [18.67514175891835]
  Diffusion reinforcement learning (RL) has recently achieved significant success in post-training image and video generative models.<n>Most RL methods, including DanceGRPO and FlowGRPO, recompute selected timesteps with gradient tracking after rollout.<n>We present LeanGRPO by restructuring the data-parallel layout and introducing two recompute-free training schedules for trajectory-logprob computation RL.
  arXiv  Detail & Related papers  (2026-09-03T08:24:56Z)
• When Do Larger Batches Help Scale LLM Reinforcement Learning? [13.656143976571007]
  We study tradeoffs in reinforcement learning for large language models.<n>We find that larger batches improve generation throughput by up to 2.29x on fixed hardware.
  arXiv  Detail & Related papers  (2026-08-29T14:32:09Z)
• Max-Q Selective Imitation for Human-in-the-Loop Online Robot Learning [1.96243636752331]
  Human-in-the-loop (HIL) online reinforcement learning for real robots must absorb human interventions quickly while continuing to improve beyond the human prior.<n>We present a training method for this setting based on two components.<n>emphMC Q-chunk critic regresses chunk-level action values onto Monte Carlo returns from a replay buffer.<n>emphmax-Q selective imitation updates the actor by imitating, at each state, the higher-$Q$ action between the current policy action and a buffer sample.
  arXiv  Detail & Related papers  (2026-08-15T07:13:59Z)
• Learning as Reasoning Unfolds: Progressive Rollout Allocation for Efficient Reinforcement Learning [23.987185167160117]
  VarIance Guided Online Rollout allocation (VIGOR)<n>VarIance Guided Online Rollout allocation (VIGOR)<n>VarIance Guided Online Rollout allocation (VIGOR)
  arXiv  Detail & Related papers  (2026-07-24T06:04:04Z)
• Learning from Own Solutions: Self-Conditioned Credit Assignment for Reinforcement Learning with Verifiable Rewards [49.1203423784326]
  We propose SC- GRPO (Self-Conditioned GRPO), which uses KL divergence mentioned before as a multiplicative weight on GRPO gradients.<n>Across five benchmarks spanning math, code, and agentic tasks, SC- GRPO consistently outperforms 8.1% over GRPO and 5.9% over DAPO with stronger OOD performance.
  arXiv  Detail & Related papers  (2026-06-17T08:26:02Z)
• DISA: Offline Importance Sampling for Distribution-Matching LLM-RL [56.9445657766829]
  We introduce DISA, which moves this calibration problem outside the RL loop.<n> DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate.<n>On two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline Flow.
  arXiv  Detail & Related papers  (2026-05-17T07:14:44Z)
• Learning-Zone Energy: Online Data Selection for Efficient RL Post-Training [28.00739954235118]
  We propose a theoretically grounded, fully online data selection framework that concentrates on the model's active learning frontier.<n>A forward pruner with replay further reduces wall-clock time cost by skipping rollout generation for persistently solved prompts.<n>Our method retains only 40% of the training data per step yet matches or surpasses full-data baselines.
  arXiv  Detail & Related papers  (2026-05-16T14:01:12Z)
• Off-Policy Value-Based Reinforcement Learning for Large Language Models [25.962820072445222]
  ReVal is a Bellman-update-based method that combines stepwise signals capturing internal consistency with trajectory-level signals derived from outcome verification.<n>On DeepSeek-R1-Distill-1.5B, ReVal improves training efficiency and achieves improvement of 2.7% in AIME24 and 4.5% in out-of-domain benchmark GPQA over GRPO.
  arXiv  Detail & Related papers  (2026-03-24T15:55:02Z)
• Recurrent-Depth VLA: Implicit Test-Time Compute Scaling of Vision-Language-Action Models via Latent Iterative Reasoning [38.78160379823724]
  Recurrent-Depth VLA (RD-VLA) is an architecture that achieves computational adaptivity via latent iterative refinement rather than explicit token generation.<n>RD-VLA employs a recurrent, weight-tied action head that supports arbitrary inference depth with a constant memory footprint.
  arXiv  Detail & Related papers  (2026-02-08T07:21:01Z)
• Just-In-Time Reinforcement Learning: Continual Learning in LLM Agents Without Gradient Updates [53.3717573880076]
  We introduce Just-In-Time Reinforcement Learning (JitRL), a training-free framework that enables test-time policy optimization without any gradient updates.<n>JitRL maintains a dynamic, non-parametric memory of experiences and retrieves relevant trajectories to estimate action advantages on-the-fly.<n>Experiments on WebArena and Jericho demonstrate that JitRL establishes a new state-of-the-art among training-free methods.
  arXiv  Detail & Related papers  (2026-01-26T14:16:51Z)
• Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling [75.36692892951018]
  Increasing the batch size during training is a promising strategy to accelerate large language model pretraining.<n>This work develops a principled framework for batch-size scheduling.<n>It introduces Seesaw: whenever a standard scheduler would halve the learning rate, Seesaw instead multiplies it by $1/sqrt2$ and doubles the batch size.
  arXiv  Detail & Related papers  (2025-10-16T14:17:38Z)
• Agentic Reinforcement Learning with Implicit Step Rewards [92.26560379363492]
  Large language models (LLMs) are increasingly developed as autonomous agents using reinforcement learning (agentic RL)<n>We introduce implicit step rewards for agentic RL (iStar), a general credit-assignment strategy that integrates seamlessly with standard RL algorithms.<n>We evaluate our method on three challenging agent benchmarks, including WebShop and VisualSokoban, as well as open-ended social interactions with unverifiable rewards in SOTOPIA.
  arXiv  Detail & Related papers  (2025-09-23T16:15:42Z)
• VinePPO: Refining Credit Assignment in RL Training of LLMs [66.80143024475635]
  We propose VinePPO, a straightforward approach that leverages the flexibility of language environments to compute unbiased Monte Carlo-based estimates.<n>Our method consistently outperforms PPO and other baselines across MATH and GSM8K datasets in less wall-clock time.
  arXiv  Detail & Related papers  (2024-10-02T15:49:30Z)
• Learnable Distribution Calibration for Few-Shot Class-Incremental Learning [122.2241120474278]
  Few-shot class-incremental learning (FSCIL) faces challenges of memorizing old class distributions and estimating new class distributions given few training samples. We propose a learnable distribution calibration (LDC) approach, with the aim to systematically solve these two challenges using a unified framework.
  arXiv  Detail & Related papers  (2022-10-01T09:40:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.