Yonghoon Dong, Kyungmin Lee, Changyeon Kim, Jaehyuk Kim, Jinwoo Shin
5 min
Abstract
Off-policy reinforcement learning of pretrained flow policies remains challenging due to the instability of optimization arising from the multi-step sampling process. Recently, Q-learning with Adjoint Matching (QAM) addressed this issue by reformulating into a memoryless stochastic optimal control (SOC) problem with a learned critic. However, QAM inherits a fundamental fragility of critic-guided improvement: small critic errors are amplified when critics are ill-conditioned, often leading to model collapse. This paper introduces Trust Region Q-Adjoint Matching (TRQAM), a stable off-policy fine-tuning algorithm that adaptively controls the path-space KL with pretrained flow policies through projected dual descent. Specifically, we optimize the trust-region parameter $λ$ in SOC dynamics, and theoretically show that the path-space KL can be represented by a closed-form function of $λ$. As a result, our method can precisely control the exact deviation from pretrained flow policies, achieving stable off-policy RL. Through experiments on 50 OGBench tasks, TRQAM consistently outperforms prior arts in both offline RL and offline-to-online RL. In particular, TRQAM achieves an overall success rate of 68% in offline RL, substantially improves the strongest baseline at 46%.
Sam: Exactly. The thermostat is wired into the engine, not just stuck on the dashboard. The technical way this works involves something called the "diffusion coefficient"—which is essentially a measure of how much randomness is in the AI's path. By mathematically linking that randomness to the trust-region value, the researchers guarantee that the AI's behavior can only deviate from its original training by a controlled, predictable amount. There's a mathematical theorem—Girsanov's theorem—that proves this relationship holds exactly, not just approximately.
Alex: So it's a structural guarantee, not just a suggestion.
Sam: Yes. And there's a second layer of protection on top of that. To make sure the AI isn't just changing in a controlled way but also changing in a *useful* way, the researchers track something called "path-space KL divergence." Think of it as a budget for the total difference between the AI's new behavior and its old behavior, measured across the entire learning journey—not just at one moment in time. The AI has to spend that budget wisely. It can't just thrash around randomly within the allowed range; it has to use its freedom to actually improve.
Alex: So it's not just about capping how much it changes—it's about making sure the changes are genuinely productive?
Sam: Right. And to keep all of this calibrated in real time, the system runs an adaptive loop. It constantly monitors whether the model is drifting toward the edge of its allowed range. If it is, the loop automatically tightens the trust-region—it turns down the thermostat. This means the fine-tuning process is self-correcting. It doesn't require a human to step in and adjust things manually.
Alex: What does that look like in practice? Does the paper test this on actual tasks?
Sam: It does. The paper tests TRQAM on tasks like navigation—where an AI has to find its way through an environment—and manipulation, where a robotic system has to interact with physical objects. The evidence suggests that by internalizing these constraints rather than bolting them on afterward, the AI maintains a meaningfully higher success rate on these tasks compared to methods that use after-the-fact penalties. The paper is careful not to overstate this—it frames these as promising results rather than definitive proof—but the pattern across the tests is consistent.
Alex: It's a notable shift in philosophy, really. Instead of training the AI and then trying to rein it in, you're building the reins into the training itself.
Sam: That's a fair summary. The underlying insight is that stability and improvement don't have to be in tension. If you design the fine-tuning process carefully enough, you can have both—a model that gets better at new tasks without losing the foundation it was built on. Whether this approach holds up across a wider range of models and tasks is still an open question, but the framework itself is a meaningful step toward more reliable AI fine-tuning.
Alex: Thanks for walking us through that, Sam. And thanks to everyone listening to ResearchPod.