zhouzypaul/qgf

This repository is the research code for Q-Guided Flow (QGF), an algorithm in the field of reinforcement learning (RL), a branch of AI where a system learns to make decisions by receiving feedback on whether its actions achieve a goal. The problem QGF addresses is a common one in offline RL: you have a large dataset of recorded behavior, and you want to learn a policy (a decision-making system) that does better than what was recorded, without collecting additional experience by interacting with the environment. Most approaches to this combine training the policy and a critic (a system that estimates how good an action is) together in a loop, which can be unstable or slow. QGF separates these two concerns. During training, it independently trains a flow-matching policy using behavioral cloning (learning to imitate the recorded actions) and trains a critic using TD-learning (a standard technique for estimating action value). Flow-matching is a type of generative model that produces outputs by progressively refining a noisy starting point through a learned denoising process. At inference time, rather than retraining the policy, QGF uses gradients from the critic to steer the denoising process toward actions the critic rates more highly, achieving policy improvement at test time without any joint actor-critic training. The codebase is written in JAX, a Python library for fast numerical computation on GPUs and TPUs. Experiments use the OGBench benchmark, a suite of offline RL environments based on simulated robot manipulation tasks (cube stacking, scene arrangement, puzzle solving). The repository also includes implementations of several baseline RL algorithms for comparison, and scripts that generate SLURM batch jobs for running the full set of paper experiments on a compute cluster. The paper associated with this code is available on arXiv (2606.11087). The project does not have a license file mentioned in the README.