강화 학습에서 플로우 정책의 테스트 시간 기울기 안내

Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning

디퓨전 모델 및 플로우 모델과 같은 표현력이 뛰어난 연속 제어 정책은 시뮬레이션 및 실제 로봇 제어를 위한 모방 학습의 확장 분야에서 최근 상당한 발전을 이루는 데 중요한 역할을 합니다. 이러한 모델들은 지도 기반 모방 학습 환경에서는 안정적으로 확장되는 것으로 알려져 있지만, 정책 개선을 위한 강화 학습(RL) 파이프라인에 통합하는 것은 더욱 어려운 과제로 밝혀졌습니다. 이는 종종 특수한 훈련 목표가 필요하거나 디노이징 과정을 통해 역전파를 수행해야 하며, 이로 인해 안정성 문제가 발생하고 확장성에 영향을 미칩니다. 본 논문에서는 테스트 시간에서만 간단한 정책 개선 방식을 사용하여 안정적인 지도 기반 정책 훈련을 유지하면서 이러한 문제를 회피하는 경쟁력 있는 대안이 될 수 있는지에 대한 질문을 연구합니다. 이를 위해 우리는 QGF(Q-Guided Flow)라는 RL 알고리즘을 제안합니다. QGF는 참조 플로우 정책(표준 행동 복제 목표를 통해 사전 훈련됨)과 가치 함수 평가기를 함께 학습하고, 테스트 시간에 가치 기울기를 사용하여 추가적인 정책 학습 없이 참조 정책이 더 높은 가치를 갖는 동작을 생성하도록 안내합니다. 실험적으로, QGF는 고차원 액션 공간을 가진 단일 작업 및 목표 지향적 오프라인 RL 벤치마크에서 기존의 테스트 시간 RL 방법보다 뛰어난 성능을 보이며, 최첨단 훈련 시간 알고리즘과 경쟁하면서 훨씬 저렴하게 실행됩니다. 또한, QGF는 액터-크리틱 훈련의 불안정성을 피함으로써 모델 크기에 대한 유리한 확장성을 보여주며, 표현력이 뛰어난 정책을 갖춘 실용적이고 효과적인 대안 RL 알고리즘입니다.

Expressive continuous control policies, such as diffusion and flow models, form the backbone of recent advances in scaling imitation learning for simulated and real robot control. While they are known to scale stably in the supervised imitation learning setting, incorporating them into reinforcement learning (RL) pipelines for policy improvement has proven more difficult. It often requires specialized training objectives or backpropagating through denoising processes, which cause well-known issues with stability and affect scalability. In this paper we study the question of whether simple policy improvement schemes at test time alone, leaving stable supervised policy training intact, can be a competitive alternative which sidesteps these issues. To this end, we propose QGF (Q-Guided Flow), an RL algorithm that performs policy optimization entirely at test time. QGF works by pre-training both a reference flow policy (via a standard behavioral cloning objective) and a value function critic and, at test time, using the value gradient to guide the reference policy to generate higher-value actions without any additional policy learning. Empirically, QGF outperforms prior test-time RL methods on single-task and goal-conditioned offline RL benchmarks with high-dimensional action spaces, and is competitive with state-of-the-art training-time algorithms while being much cheaper to run. Moreover, it exhibits favorable scaling with model size by avoiding the instability of actor-critic training, offering a practical and effective alternative RL algorithm with expressive policies.