Guía por Gradientes en Tiempo de Prueba de Políticas de Flujo en Aprendizaje por Refuerzo

Resumen

Las políticas de control continuo expresivas, como los modelos de difusión y flujo, constituyen la base de los avances recientes en el escalado del aprendizaje por imitación para el control de robots simulados y reales. Si bien se sabe que escalan de manera estable en el entorno supervisado de aprendizaje por imitación, su integración en tuberías de aprendizaje por refuerzo (RL) para la mejora de políticas ha resultado más difícil. A menudo se requieren objetivos de entrenamiento especializados o la retropropagación a través de procesos de eliminación de ruido, lo que provoca problemas conocidos de estabilidad y afecta la escalabilidad. En este artículo estudiamos si esquemas simples de mejora de políticas únicamente en tiempo de prueba, manteniendo intacto el entrenamiento supervisado estable de políticas, pueden ser una alternativa competitiva que evite estos problemas. Con este fin, proponemos QGF (Flujo Guiado por Q), un algoritmo de RL que realiza la optimización de políticas completamente en tiempo de prueba. QGF funciona entrenando previamente tanto una política de flujo de referencia (mediante un objetivo estándar de clonación conductual) como un crítico de función de valor y, en tiempo de prueba, utiliza el gradiente del valor para guiar la política de referencia hacia la generación de acciones de mayor valor, sin necesidad de aprendizaje adicional de políticas. Empíricamente, QGF supera a métodos previos de RL en tiempo de prueba en benchmarks de RL fuera de línea con una sola tarea y condicionados por objetivos, con espacios de acción de alta dimensionalidad, y es competitivo con algoritmos de vanguardia en tiempo de entrenamiento, siendo mucho más económico de ejecutar. Además, muestra un escalado favorable con el tamaño del modelo al evitar la inestabilidad del entrenamiento actor-crítico, ofreciendo una alternativa práctica y efectiva de algoritmo de RL con políticas expresivas.

English

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.