결과가 갈리는 지점 학습: 확률적 청크 마스킹을 통한 효율적인 VLA 강화학습 (RL)

강화학습 (Reinforcement Learning, RL)은 작업 성공을 직접적으로 최적화함으로써 시각-언어-행동 (Vision-Language-Action, VLA) 정책이 학습 분포를 넘어 일반화할 수 있도록 하지만, 사후 학습 (post-training)은 계산 비용이 많이 듭니다. 이에 대한 자연스러운 대응은 더 빠른 시뮬레이터와 월드 모델 (world models)을 통해 롤아웃 수집 (rollout collection) 속도를 높이는 것이었습니다. GRPO 기반 VLA RL 연구에서 우리는 주요 비용이 다른 곳에 있음을 발견했습니다. 실험 결과, 롤아웃 수집은 단계당 실제 시간 (wall-clock time)의 약 21%만을 차지하는 반면, 그래디언트 계산 (gradient computation)이 약 78%를 차지합니다. 그래디언트 비용이 지배적인 이유는 이 계산의 상당 부분이 학습에 거의 기여하지 않는 단계에 소비되기 때문입니다. GRPO의 학습 신호는 어드밴티지 분산 (advantage variance)에 의해 구동됩니다. 즉, 성공한 롤아웃과 실패한 롤아웃이 갈라지는 단계에서만 학습 신호가 생성됩니다. 그러나 GRPO는 롤아웃 내의 모든 청크 (chunk)에 동일한 어드밴티지를 할당합니다. 그 결과, 액터 업데이트 (actor-update) 계산이 사전 학습 (pre-training) 및 지도 미세 조정 (supervised fine-tuning) 이후 정책이 이미 처리할 수 있는 단계를 포함하여 궤적 (trajectory) 전체에 걸쳐 균등하게 소비됩니다.

본 논문은 궤적당 확률적으로 선택된 소수의 청크 서브셋에 그래디언트 계산을 할당하는 GRPO의 드롭인 (drop-in) 수정 방식인 확률적 청크 마스킹 (Probabilistic Chunk Masking, PCM)을 제시합니다. PCM은 단계별 그래디언트 분산 (per-phase gradient variance)의 롤아웃 유도 프록시 (proxy)인 성공-실패 액션 분산 (success-failure action variance)을 사용하여 의미론적 단계 (semantic phases)를 점수화하고, 온라인으로 업데이트되는 단계별 유지 확률 (phase-level keep probabilities)을 통해 고정된 청크 예산을 샘플링합니다. 우리는 그래디언트 계산이 유용한 지점을 결정하는 양으로서 단계별 그래디언트 분산을 공식화하고, 성공-실패 액션 분산이 이를 측정 가능한 프록시로 제공함을 보여줍니다. PCM은 보상 모델 (reward model)이나 학습된 크리틱 (learned critic)을 필요로 하지 않습니다. 세 가지 LIBERO 벤치마크에서 PCM은 표준 GRPO의 최종 성공률과 일치하면서도, 궤적 청크의 20% 미만을 역전파 (backpropagating)하면서 2.38배의 실제 시간 단축, 4.8배 빠른 그래디언트 업데이트, 그리고 60% 낮은 피크 활성화 메모리 (peak activation memory)를 달성했습니다.