Latest / Embodied AI 101 / Q-Guided Flow: Test-Time Gradient Guidance of Flow Policies

Transcript

Q-Guided Flow: Test-Time Gradient Guidance of Flow Policies · Embodied AI 101

• Q-Guided Flow: Test-Time Gradient Guidance of Flow Policies. Modern robotic control algorithms often rely on expressive models – for example, diffusion or flow-based generative policies – to represent complex, multimodal action distributions. Training such policies via pure imitation can scale stably to high-dimensional, real-world tasks. However, using them in reinforcement learning (RL) for policy improvement has proven much trickier: naively adding actor-critic optimization to diffusers or flow models tends to require special objectives or unstable backpropagation through the denoising process. In “Q-Guided Flow,” Zhou et al. (2026) take a different approach. Instead of modifying the model during training, Q-Guided Flow (QGF) leaves the flow-based policy fixed after standard behavioral cloning and uses a separately trained value function to steer the policy at inference time. In other words, all the policy improvement happens at test time, by guiding the sampling process toward higher-Q actions. This clever scheme – essentially performing lightweight “planning” by gradient ascent in action space guided by a critic – yields surprisingly strong results. In simulated benchmarks of long-horizon, goal-conditioned tasks, QGF significantly outperforms existing test-time methods and even matches or beats many fully-trained RL baselines, without any costly actor updates. Crucially, it achieves this with higher sample and compute efficiency: the flows are pretrained by imitation, and the only online costs are computing a gradient at each step of the action denoising. In practice this makes the method more stable, easier to scale, and more data-efficient than conventional RL, as the authors demonstrate. In this article we unpack Q-Guided Flow in detail. We first explain the motivation and the high-level idea of test-time policy improvement for diffusion/flow policies. Then we describe exactly how QGF works – how the reference flow is trained, how a value function is learned, and how at inference time the critic’s gradient is incorporated into the flow’s denoising steps. Next, we highlight the surprising technical contributions: two key approximations that make the guidance both efficient and effective. We summarize the extensive experimental evaluation on offline RL benchmarks (especially the OGBench suite of long-horizon tasks), showing that QGF markedly improves performance in challenging control problems. Finally, we discuss the implications for robotics: by combining rich generative policies with test-time optimization, Q-Guided Flow opens a new way to get high performance from limited data, particularly in long-horizon, multi-step tasks where conventional RL struggles. Why Test-Time Policy Improvement?. To appreciate what Q-Guided Flow does, it helps to step back and consider the tension between imitation learning and reinforcement learning in modern continuous control. On one hand, diffusion and flow-based policies (inspired by generative models) have proven remarkably effective for imitation and planning. These models can learn very flexible, multimodal distributions over actions conditioned on the current state, trained by simple supervised objectives on logged data. On the other hand, when we want those policies to improve beyond imitation – to optimize cumulative reward – standard RL methods become cumbersome. Doing gradient-based policy optimization through a diffusion process is notoriously unstable: you either need tricky training objectives or backpropagate gradients through the many denoising steps, which breaks gradient flow and scalability. This has led to a tantalizing question: Can we keep the stability of imitation training, and push the policy toward higher reward only at test time? In other words, train a rich policy by supervised learning (behavioral cloning) and train a separate value function (critic), then at inference time use the critic to adjust the policy outputs without any further weight updates. QGF takes exactly this approach. Zhou et al. note that we already know how to do stable imitation with flows or diffusers; the issue is how to inject value information into the policy. They propose to do so via a novel gradient guidance during sampling, analogous to classifier-guided diffusion in vision, but using a Q-function as the guide. To make this concrete, imagine you have a pre-trained flow policy $\hat\pi(a|s)$ (from behavioral cloning on a dataset of expert or logged trajectories) and a critic $Q(s,a)$ trained by temporal-difference (TD) learning on the same data (or a combination of data and rollouts). At inference time, given a state $s$, you normally sample an initial noise vector $a_0 \sim \mathcal{N}(0,I)$ and run it through the flow’s denoising process to get an action $a_T$ ~ the model’s suggested action. QGF tweaks this process: at each denoising step, it computes the gradient of the critic with respect to the current (partially denoised) action, and nudges the action in the direction of higher value. Intuitively, the critic pushes the flow to choose actions that the Q-function likes, while the flow ensures those actions stay near the expert distribution. Zhou et al. argue that this “critic-guided sampling” can approximate the solution to a KL-regularized RL problem, whose optimal policy satisfies $$\pi(a|s) ;\propto; \hat\pi(a|s)\exp\bigl(\tfrac{1}{\beta}Q(s,a)\bigr).$$ Taking gradients of that “improved” policy shows that its log-probability gradient is $$\nabla_a\log\pi(a|s) = \nabla_a\log\hat\pi(a|s) + \tfrac{1}{\beta}\nabla_a Q(s,a).$$ So one natural plan is: during the flow’s denoising from noise to an action, add the term $\tfrac{1}{\beta}\nabla_a Q(s,a)$ to the score/velocity field at each step. This is indeed what Q-Guided Flow attempts, with some important twists to make it work in practice. In effect, QGF is a form of test-time policy optimization: policy training is frozen (standard BC for the flow, standard TD for the Q), and then actions are refined by ascending the value gradient during generation. As the authors put it, this approach avoids the instability of actor-critic training while still yielding strong performance. Naive Gradient Guidance: Pitfalls. Before describing QGF in detail, it’s worth understanding why the simplest idea – just plugging in the critic gradient – doesn’t work out-of-the-box. There are two obvious approaches, and both have serious drawbacks:. OOD gradient (Out-of-Distribution query). One could simply take the noisy action $a_t$ at some intermediate time $t$ in the denoising process, compute $\nabla_a Q(s,a_t)$, and add that to the flow’s velocity. On paper this matches the earlier gradient form. But in practice, this asks the critic to output a meaningful gradient at a point $a_t$ that is very different from any action seen in training: the critic was only trained on clean, final actions from the dataset. Noisy $a_t$ are far outside its support, so its gradient there is essentially garbage. Zhou et al. demonstrate an illustrative 1D example where following this naive “OOD gradient” consistently pushes actions toward the wrong mode or sub-optimal action, regardless of how strong the guidance is. In short, $\nabla_a Q(s,a_t)$ at a noisy $a_t$ is not reliable, so this simple trick “misdirects the flow” and fails to improve the policy. BPTT gradient (Backpropagate through denoising). A more principled variant is this: given the noisy action $a_t$, one could compute its fully denoised version $a_1 = \mathrm{ODE}(a_t)$ by running the entire remainder of the flow. Then note that $Q(s,a_1)$ is well-defined and differentiable, so we define $Q(s,a_t) := Q(s,a_1)$. The gradient $\nabla_{a_t}Q(s,a_1)$ can then be computed by backpropagating through the whole denoising process (using the chain rule). This avoids querying the critic on weird inputs, since we evaluate Q only at the final, legitimate action $a_1$. It is theoretically the “correct” gradient of $Q$ with respect to the noisy action. However, it has fatal drawbacks: backpropagating through an entire diffusion/flow pipeline is computationally intensive and extremely unstable. Small changes in $a_t$ can cause wildly different $a_1$, leading to highly noisy gradients. In practice, Zhou et al. find that using this full “BPTT gradient” makes the sampling process erratic and diverges for even modest guidance weights. In short, neither the “noisy Q gradient” nor the “full backprop” approach worked well. QGF's ingenuity lies in a middle path: a lightweight, first-order approximation of the denoising that suffices to keep the actions in-distribution, combined with a trick to limit variance. This achieves nearly the same effect as the naive ideal, but much more stably. The Q-Guided Flow Algorithm. Pretraining: Flow Policy and Critic. Q-Guided Flow operates in an offline RL setting. The assumption is you have a dataset of transitions $(s,a,s',r)$ (for example, robot demonstrations or logged experience) and you want to derive a high-quality policy. The approach is:. Train a reference flow policy by behavioral cloning. The flow policy is a model $v_\theta(s,a,t)$ – a “velocity field” conditioned on state $s$, action value $a$, and time index $t\in[0,1]$ – that defines a continuous-time denoising process (think of it like the “score network” of a diffusion model, or a continuous normalizing flow). The flow is trained via standard supervised learning: we match its generated actions to the dataset actions. Concretely, at training time one samples $a_0 \sim \mathcal{N}(0,I)$ and iteratively applies the velocity $v_\theta$ to gradually transform $a_0$ into an action $a_1$ that should match the expert action. The training minimizes a regression loss so that the denoised $a_1$ reproduces the data. In effect, $\hat\pi(a|s)$ is a generative model that can sample expert-like actions by running noise through this learned flow. Train a value function critic. Simultaneously, we train a Q-function $Q_\phi(s,a)$ by standard offline RL or TD learning on the same data (or, optionally, on data plus additional fine-tuning). For example, one might use Implicit Q-Learning (IQL) or another offline RL algorithm to learn a Q that approximates cumulative reward. Importantly, this Q is trained on the clean (final) actions from the data, not on intermediate noisy actions. Once the critic is trained, we freeze its weights. After these two steps, we have a fixed flow model representing the policy and a fixed critic. Crucially, we have not yet tried to make the policy high-return beyond what the BC data provided; that will happen only at inference. Inference: Gradient Guidance. At test time, given some state $s$, we want to sample an improved action that (hopefully) has higher Q-value than a pure BC sample. The QGF inference procedure is as follows. We discretize time $t$ in $T$ increments from $0$ (pure noise) to $1$ (clean action). Initialize $a_0 \sim \mathcal{N}(0,I)$. Then repeat for $t = 0, 1/T, 2/T, \dots, 1-1/T$:. Approximate a clean action: First take one Euler step of the flow’s vector field to get an estimate of the nearly-denoised action: $$ a' ;=; a_t + (1 - t),v_\theta(s,,a_t,,t). $$ This $a'$ is intended as a “short guess” for what the final action would be if we ran the whole denoising to completion from $a_t$. (It’s just one step, not the full ODE integration.). Get the critic gradient: Compute the gradient of the Q-function wrt this approximate action: $g = \nabla_{a'} Q_\phi(s, a')$. In practice the authors set aside the Jacobian for now and treat $g$ as if it were aligned with the true gradient at $a_t$. Update the action: Finally, we update the noisy action $a_t$ by a step that mixes the flow’s natural denoising velocity with the critic gradient: $$ a_{t+1} ;=; a_t ;+; \tfrac{1}{T}\Bigl(v_\theta(s, a_t, t) ;+; \tfrac{1}{\beta},g\Bigr). $$ Here $\frac1{T}$ is the step size (since we have $T$ total steps), and $1/\beta$ is a tunable “guidance strength” scalar that weights how strongly to follow the critic. Conceptually, the term $v_\theta$ moves us one step down the flow’s learned denoising path, while the term $(1/\beta)g$ nudges the action in the direction of increasing Q-value. This loop is repeated until $t = 1$. At the end, we output the final action $a_T$. In pseudo-code borrowed from the authors:. This process is illustrated in the QGF project overview: at each “denoising time”, the policy’s velocity is supplemented by the critic gradient pointing to higher-value regions. Because we never change $\theta$ or $\phi$ during this loop, no further learning occurs; all the policy improvement is implicit in the sampling. Two key design choices make this work in practice. First, QGF never actually backpropagates the gradient of the entire denoising chain (it does not propagate the gradient through $v_{\theta}$). It only computes $g = \nabla_{a'} Q(s,a')$ at the one-step approximation $a'$. Second, when relating this gradient in $a'$ back to changes in $a_t$, QGF ignores the Jacobian $\partial a'/\partial a_t$. In other words, it treats the critic gradient at $a'$ as if it were aligned with $a_t$ without rescaling. Concretely, this means the algorithm simply uses $g$ directly. The authors find that multiplying by the Jacobian (which involves differentiating through $v_\theta$) actually makes things worse – it amplifies noise and drastically increases gradient variance. Ironically, the “crude” simplifications turn out to be virtues: by not forcing the action exactly onto the data manifold and by dropping the ill-conditioned Jacobian, QGF gets a lower-variance, more reliable gradient estimate. Why the Approximations Help. It may seem surprising that these shortcuts – one Euler step and identity Jacobian – could outperform the more exact calculations. Yet the empirical analysis shows they do. Intuitively, the one-step approximation allows the policy to deviate slightly from the strict imitation distribution. If the data contained multiple modes (e.g. a human could take several distinct valid actions), a full deterministic denoising would tie you exactly to those modes. The first-order approximate allows the policy to “choose” among modes when guided by the critic, rather than faithfully reproducing everything in the dataset. Likewise, using the identity Jacobian avoids magnifying gradient noise; early denoising steps are very rough, so the true Jacobian can be very large or small, introducing instability. Simply ignoring it leads to a smoother, lower-variance gradient. Zhou et al. quantify this with two experiments. They measure gradient stability by computing the cosine similarity between $\nabla Q(s,a_t)$ and $\nabla Q(s,a_t+\epsilon)$ under a small perturbation $\epsilon$. QGF’s estimator (with the two approximations) has the highest cosine similarity (close to 1), far above the naive and backprop gradients. They also simulate an “optimization” by running denoising with each gradient method and checking the final Q-value. QGF’s output actions have the highest Q-values, almost matching an ideal “best-of-N” oracle. In contrast, the naive OOD gradient often fails to raise $Q$ at all. These results show that the two approximations aren’t just computational shortcuts – they are key to making the guidance effective. In short, Q-Guided Flow’s recipe is: take a quick peek at what the next step would be (one-step Euler), ask the critic for a direction there, then move in that direction, and repeat. Although approximate, this works much better than any other simple alternative. Experimental Evaluation. The Q-Guided Flow paper backs up its method with extensive experiments, focusing on high-dimensional, long-horizon tasks from the recently introduced OGBench suite of offline control tasks. OGBench includes a variety of environments (maze navigation, complex manipulators, multi-object puzzles, etc.) with different goals; it was specifically designed to benchmark offline and goal-conditioned RLアルorithms. Zhou et al. use both single-task and goal-conditioned offline settings. In all cases they compare to two groups of baselines:. Test-time RL baselines: These are methods that, like QGF, do not train a new policy but try to optimize action values at inference. Examples include simple best-of-$N$ sampling (draw many candidate actions and pick the one with highest Q), or previous value-guided sampling schemes. All such methods keep a behavioral clone policy and a critic, but differ in how they combine them. (The paper compares to existing methods from prior work, though details are not all public; the emphasis is that QGF substantially beats each of these prior test-time strategies.). Training-time (offline RL) baselines: These are state-of-the-art offline RL algorithms that train a policy via standard RL objectives (often actor-critic) on the offline data. Think IQL, MARWIL, AWAC, CRR, and others. These represent the “gold standard” in the offline setting, but their training can be unstable or require costly hyperparameter tuning for diffusion policies. The main result is compelling: QGF substantially outperforms all prior test-time methods by a large margin, and even matches the best training-time algorithms. In a suite of 20 standard OGBench tasks, each run for 500k critic training steps, QGF achieves the highest average return among test-time schemes. Notably, the best fully-trained baseline (“EDP” in their plots) also happens to use a first-order Euler trick – but even that strong baseline is only roughly on par with QGF. QGF’s pure advantage is that it uses no extra RL training on the policy; after 500k steps it can generate much better actions simply by smarter sampling. The authors also emphasize compute efficiency. One might think: if you have a critic, why not do a brute-force best-of-$N$ (BFN) search – sample $N$ actions from the flow and take the highest-Q one? Indeed, BFN can improve performance, but it’s very expensive. Zhou et al. show that QGF with a small guidance weight outperforms best-of-4 sampling, at a tiny fraction of the compute. Combining QGF with BFN is also effective: QGF+BFN ($N=16$) achieves similar returns to vanilla BFN-16, but QGF’s gradients make it more efficient, so it costs much less in raw FLOPs. In other words, by embedding gradient guidance into the sampling, QGF beats naive sampling even with fewer attempts. Crucially for robotics, QGF shines as tasks get harder. In longer-horizon goal-conditioned tasks (like multi-step puzzles or large mazes), QGF is consistently the best-performing method as difficulty increases. The variant with exact Jacobians (QGF-Jacobian) lags behind, confirming that the approximations were indeed important. Perhaps most strikingly, when model capacity is increased, QGF’s performance continues to climb: moving from a small model to a 3.2M-parameter flow yields about a 4× increase in returns. Meanwhile, the training-time methods plateau or even collapse with larger models. This suggests QGF scales robustly with model size, whereas actor-critic approaches often become unstable when too large. Even the choice of critic benefits from QGF’s flexibility. The authors report that using a stronger TD-trained critic (they call it QAM, a particular offline algorithm) in place of the default IQL critic further boosts results. Remarkably, “QGF+QAM” outperforms the QAM baseline itself. In other words, even with an already good policy from QAM, applying QGF-style guidance with that critic improves performance even more, showing that QGF is largely agnostic to how the critic was trained and can amplify any value information it contains. Across the board, these experiments illustrate that Q-Guided Flow is a compelling new alternative. It consistently beats prior test-time policies (often by large margins) and closes the gap to the best training-time ones, while being dramatically cheaper and more stable. The consistency across many seeds and tasks indicates the method’s reliability. The authors share detailed plots: for example, a chart of results at 1 million training steps shows QGF at the top of the leaderboard on 25 goal-conditioned tasks, and compute plots show its big win in efficiency. Implications for Robotic Control. What does Q-Guided Flow mean for robotics and embodied AI? Several implications stand out:. Sample and data efficiency: QGF makes very effective use of existing data. Unlike typical RL, it doesn’t require additional environment interactions once the dataset is collected. The required “learning” is just offline behavioral cloning and critic fitting, both of which can be done with standard supervised or offline updates. Then all improvement happens at inference. This suggests much higher “sample efficiency” in practice: if we gather a modest demonstration dataset, QGF will wring far more performance out of it than imitation alone, without needing millions of new trials. As the authors note, the approach is “much cheaper to run” than full RL training, and the offline experiments indeed use only 500k–1M steps of critic learning to match or exceed baselines that require lengthy training. Long-horizon, multi-modal tasks: The most striking gains came in tasks like maze navigation with many steps, multi-object manipulation puzzles, and other goal-conditioned benchmarks. In these problems, precision and planning matter, and pure behavior cloning often fails to reach the goal reliably. By contrast, QGF’s critic-in-the-loop lets the agent steer step-by-step towards the goal. In practice, this means a robot using QGF can adjust each action online to maximize eventual success, guided by the learned value function. For example, in an assembly task requiring 20 sequential motions, QGF can subtly alter early actions to avoid getting trapped, something BC alone wouldn’t do. The authors explicitly show that as the horizon lengthens and goals become farther, QGF stays ahead of everything else. This bodes well for complex robotics tasks like multi-stage planning or search. Stability of training: Because QGF avoids backpropagating through the flow policy, it sidesteps many instabilities that plague deep RL. Training a diffusion-based policy via RL would require either an off-policy actor-critic or in-distribution constraints; both are known to be finicky. Here, the policy itself never “sees” rewards during training, so it never diverges from the demonstration behavior (except by intentional test-time guidance). This means we can use very large models and flexible architectures for the policy without worrying about RL collapse. Indeed, Zhou et al. report that QGF scales up smoothly with model capacity. For roboticists, that means we can exploit big neural networks to model subtle jittery commands or vision-based conditioning, and still safely refine the behavior with a value function. Test-time planning and adaptation: The QGF process is akin to a brief “planning” phase at inference: sampling an action is no longer a single forward pass, but a small optimization loop. This does add latency, but the authors argue it’s quite manageable and highly parallelizable. In robotics, a planner query might afford some milliseconds for gradient updates, especially if actions are continuous and high-precision. Moreover, because the critic gradient is taken with respect to $a'$, even early steps are likely to push toward high-value modes, so one might get decent performance even with modest $T$ steps. Importantly, one could imagine extending QGF online: after taking action $a_T$, the resulting state $s'$ and any reward could be used to refine the critic (off-line updates), and the next action would again be guided. This could allow fast adaptation to changes in the environment or task by reusing the same policy. In essence, QGF makes the policy “algorithmic” – it has a built-in optimization routine guided by $Q$. Combining generative policies with value-based control: Conceptually, QGF is part of a broader trend of mixing generative models (for plausible actions or behaviors) with classical value or model-based planning. It’s analogous to classifier-guided diffusion in vision, but here the “classifier” is a reward critic. For robotics, this suggests a unified view: use powerful generative tools to model what kinds of actions are valid or expert-like, and use a value function to discriminate which of those actions maximize reward. QGF is a clean realization of that idea. The success of QGF invites further exploration: maybe we can extend similar test-time guidance to other modalities (e.g., diffusion policies over trajectories) or integrate with learned world models. Practical considerations: Of course, QGF’s reliance on the critic means the value function quality is crucial. If $Q$ is poorly estimated, the guidance may push actions astray (though dropping the Jacobian and using approximations seems to mitigate some risk). The experiments use offline RL critics which are as good as can be hoped from the available data. In a real robot scenario, one might need to collect some rollouts or use cautious update rules to ensure $Q$ is reliable. That said, the flexibility to plug in any critic is a strength: even basic TD fits can give improvements, and better critics give even more. The analysis shows that QGF + a stronger TD critic (like QAM) easily beats using a weaker IQL critic. So improving the critic with domain knowledge will only help QGF. In summary, Q-Guided Flow is a promising framework for robotic control when data is limited and tasks are complex. By deferring policy improvement to inference time, it sidesteps many RL pitfalls while still capturing reward information. As long as a good critic is available, one can treat the flow policy as a rich suggestion mechanism and “correct” it on-the-fly. The result is a method that is both practical and powerful. Conclusion. Zhou et al.’s Q-Guided Flow represents a clever new way to do reinforcement learning with generative policies. The core idea – improve a pre-trained diffusion/flow policy at test time via gradient guidance from a critic – turns out to be extremely effective. Technically, the paper’s main contributions are the specific approximations that make this work: a one-step Euler approximation of the denoising and the decision to set the Jacobian to identity. These choices produce a low-variance gradient that reliably points toward higher-value actions. The methodology neatly avoids the usual instabilities and high costs of actor-critic training, instead leveraging stable imitation learning and standard value fitting. Experimentally, QGF delivers on its promise. On a suite of challenging offline RL benchmarks (including goal-conditioned tasks with long horizons), it outperforms all prior test-time approaches by a wide margin and rivals state-of-the-art training-time algorithms. It also scales favorably with model capacity and additional computational budget. Furthermore, it is framework-agnostic: one can plug in any flow or diffusion policy and any critic, and often see immediate gains. The projects’ code release and detailed ablations show that QGF is not a fluke but a robust technique. For the roboticist or ML researcher, Q-Guided Flow suggests a new paradigm: think of policies as samplers guided by higher-level knowledge, rather than monotonically optimized functions. This opens up possibilities to combine learned dynamics, multimodal action representations, and planning in new ways. In practice, the method can give robots better long-horizon behavior from the same data, and may shorten the RL training loop to just a critic update. As one of the authors’ key observations says, we can get the benefit of policy improvement without the instability of pushing the policy network itself. Looking ahead, it will be interesting to see how QGF ideas generalize. Could one apply Q-guidance to diffusion models in other domains (e.g. text or images) for conditional generation? In robotics, could this be extended to simultaneous multi-step trajectory generation, or combined with model predictive control? The notion of test-time RL brings to mind analogies with model-predictive control or on-the-fly re-planning, but here it leverages a learned policy and critic instead of a known model. Ultimately, Q-Guided Flow enriches the set of tools for embodied AI: when stable training is hard, sometimes it pays to simply train less and plan more – even during execution. Sources: Zhiyuan Zhou et al. introduce and analyze QGF in their 2026 paper; their project page and paper details the algorithm and results. In particular, the method’s description and pseudocode come from the authors’ overview, and the experimental claims are drawn from their reported results. The high-level motivation and comparisons to prior methods are summarized in the abstract and introduction.