Reversal Q-Learning (RQL) from Scratch: Making Denoising an RL Action Without the Horizon Tax

A long read that strings together flow policies, the extended MDP, flow reversal, the reparameterization policy gradient, DDPG/SAC, and RQL’s value function — into one line. Every key idea is paired with a live, interactive visualization (there are eleven). Drag and play as you read.

Note: this is a teaching-style reconstruction based on the RQL project page (Oberai, Park, Levine, 2026). The code is illustrative PyTorch-flavored pseudocode meant to explain the idea — it is not the official implementation; some engineering details (discount bookkeeping, network shapes) favor clarity and may differ from the original.

TL;DR

RQL’s problem: do reinforcement learning on a flow (diffusion-style) policy using offline data.

• A flow policy doesn’t emit an action in one shot — it starts from noise and refines it over $F$ “denoise” steps. Expressive, but awkward to put under RL.
• A natural idea is to treat each denoise step as an RL action. That lets you train with stable, local, one-step gradients — but it multiplies the value-learning horizon by $F$ (off-policy RL’s “curse of horizon”), and the offline data has no intermediate steps in it.
• RQL’s key move is flow reversal: integrate the ODE backwards from the dataset’s final action to fill in the missing steps, building “virtual trajectories.” Because those intermediate transitions are deterministic, multi-step returns across them are zero-variance and unbiased — so the stretched horizon is free.
• In practice: learn a value $V(s, x^f, f)$ defined on partial actions, then nudge each denoise step with a reparameterization (pathwise) gradient, plus a BC term to anchor the policy to the data.

Let’s take it apart, one piece at a time.

1. The setting: a flow policy + offline RL

We’re in a standard MDP — state $s$, action $a$, reward $r$, next state $s’$, discount $\gamma$ — with a batch of offline data: transitions $(s, a, r, s’)$. We want the best policy we can get, and the policy is a flow model.

Why a flow policy? Because the distribution of good actions is often multi-modal and complex. A plain Gaussian policy can only represent one mode; a flow/diffusion model can fit arbitrarily complex distributions. The cost is that generating an action takes many steps — which is the headache the rest of the post resolves.

2. What a flow policy actually is

A flow policy turns noise into an action through a chain of deterministic refinement steps. Start from noise $x^0 \sim \mathcal{N}(0, I)$ and integrate an ODE given by a network $v_\theta$:

The final $x^F$ is the action $a$. Discretized into $F$ Euler steps, each step is

Here is the fact that runs through the whole post: given the velocity field, the path from any $x^f$ to $x^F$ is fully determined. The only randomness is the initial noise $x^0$. Drag the demo: the noise cloud flows to the modes; re-roll the noise and different $x^0$ flow to different modes — but each path is deterministic.

# flow policy forward (generate one action), conceptual
def sample_action(state, v_net, F):
    x = torch.randn(action_dim)        # x^0 ~ N(0, I)
    for f in range(F):                 # F denoise steps
        x = x + v_net(state, x, f)     # x^{f+1} = x^f + v(s, x^f, f)
    return x                           # x^F, executed in the environment

Note the execution flow: sample noise → run $F$ denoise steps → get $x^F$ → execute once. In the environment this is one action; the environment-level horizon doesn’t change at all. We’ll use this repeatedly.

3. Two routes: treat the flow as a black box, or treat each step as an action

To do RL on this policy, the core difficulty is: how do you take the gradient w.r.t. $\theta$ that pushes the action toward higher value? There are two fundamentally different routes.

Route A: treat the whole flow as a black box

Learn one value $Q(s, x^F)$ on the final action only. The MDP is unchanged, the horizon is unchanged. Sounds great — until policy improvement, where you need

and $x^F$ is the output of $F$ integration steps. Differentiating through it means backpropagating through all $F$ steps (BPTT). Each step multiplies the gradient by its Jacobian, so the gradient at the input is roughly $\lambda^F$ — exploding if $\lambda>1$, vanishing if $\lambda<1$. Drag $F$ and $\lambda$:

To dodge BPTT, the black-box route spawns workarounds, each with a wound: BPTT (unstable), one-step distillation (squashes the flow to one step, killing expressiveness), weighted regression (no gradient at all, just value-weighted regression — empirically weak).

Route B: treat each denoise step as an RL action

“Unroll” the generation into an extended MDP: one environment step becomes $F$ decision steps. The single action node $a_t$ splits into $x^0 \to x^1 \to \cdots \to x^F (= a_t) \to s’$. Now you can learn a value $V(s, x^f, f)$ at every intermediate step and update each step with a local one-step gradient — no BPTT, full flow expressiveness retained. Hit “unroll”:

4. The core tension: horizon inflation (training ≠ execution) + no data for intermediate steps

A common confusion (mine, at first): “but at execution time you denoise once and execute $x^F$ — how can the horizon change?” The answer: the horizon problem isn’t in execution, it’s in training-time value learning.

Once you adopt route B and bootstrap value through the intermediate steps, the bootstrap chain length gets multiplied by $F$:

Off-policy RL is very sensitive to long horizons — value error compounds along the bootstrap chain (the curse of horizon). Multiplying it by $F$ is bad. Note this is a virtual horizon inside value learning, not the environment horizon. Slide $F$ and $\gamma$:

A second problem piles on: the offline data only has $(s, x^F, r, s’)$ — it never recorded the intermediate denoise steps $x^0, \dots, x^{F-1}$. You don’t even have samples to train the extended MDP. The tension:

RQL keeps Route B’s policy-improvement upside while killing its horizon and data problems.

5. Solution part 1: flow reversal — manufacture virtual trajectories from the data

Recall the key fact: denoise transitions are deterministic. So even though the data has no intermediate steps, we can take the dataset’s final action $x^F$ and integrate the velocity field backwards to recover any intermediate step:

This builds a “virtual trajectory” $x^0 \to \cdots \to x^F$, perfectly fitting the extended framework. These trajectories are deterministic (ODE-defined, no randomness) and on-policy (the current flow field, run in reverse). Drag the dataset action $x^F$; the reversal follows:

# flow reversal: recover intermediate steps from a dataset action x_F
def flow_reversal(state, x_F, v_net, F):
    xs = [None] * (F + 1); xs[F] = x_F; x = x_F
    for f in reversed(range(F)):       # F-1, ..., 0
        x = x - v_net(state, x, f)     # x^f = x^{f+1} - v(s, x^f, f)
        xs[f] = x
    return xs                          # [x^0, x^1, ..., x^F]

Because the intermediate transitions are deterministic, multi-step returns across them are unbiased and zero-variance — you pay no off-policy cost for that extra length. That’s what makes stable off-policy RL possible in the extended framework.

6. The insight: within one action, every step has the same value

This is RQL’s most elegant — and most overlooked — point, and it answers “where do the intermediate-step values even come from.”

The value $V$ is one shared network that takes the step index $f$ as input: $V(s, x^f, f)$. Its TD target is:

The target contains no $f$. Whether you ask about $V(s, x^0, 0)$, $V(s, x^5, 5)$, or $V(s, x^F, F)$, the regression target is the same number. Why? Because the denoise steps between $x^f$ and $x^F$ are (1) deterministic, (2) reward-free, and crucially (3) un-discounted — $\gamma$ fires only on the real environment transition, never inside denoising. So the value is constant along the deterministic chain:

In words: a partial action’s value is just the $Q$-value of the complete action it will become. Flip the toggle to see the wrong version (discounting every denoise step decays the value $F\times$ faster and inflates the horizon):

So if the values are equal, why does $V$ still take $f$?

Because the policy update needs not the scalar but its gradient $\nabla_{x^f} V$. The objects $x^f$ at different $f$ are very different — near $f=0$ almost pure noise, near $f=F$ nearly a finished action. The same coordinate means something different at each noise level, so even though the on-trajectory scalar is equal, the shape/slope of $V(s, \cdot, f)$ differs across $f$. The network needs $f$ to point the gradient correctly at each level. Slide $f$: the dot’s height (the value) barely moves; the slope changes a lot.

# value training: inputs from reversal, target independent of f
def value_loss(batch, V_net, V_target, v_net, F, gamma):
    s, a, r, s_next = batch.s, batch.a, batch.r, batch.s_next
    xs = flow_reversal(s, a, v_net, F)              # [x^0, ..., x^F]
    x_next_0 = torch.randn_like(a)                  # x'^0 ~ N(0, I)
    with torch.no_grad():
        target = r + gamma * V_target(s_next, x_next_0, 0)   # no f
    loss = 0.0
    for f in range(F + 1):
        loss = loss + huber(V_net(s, xs[f], f) - target)     # regress every step to the SAME target
    return loss / (F + 1)

7. The tool: the reparameterization (pathwise) policy gradient

RQL’s policy update is essentially the “stepwise” version of this classic trick, so let’s nail it down. Policy improvement maximizes $J(\theta) = \mathbb{E}{a \sim \pi\theta}[Q(s, a)]$. The snag: $a$ is sampled, and sampling isn’t differentiable. Two routes:

• Score-function (REINFORCE): $\nabla_\theta J = \mathbb{E}[Q(s,a)\nabla_\theta \log \pi_\theta(a \mid s)]$. Uses only the value of $Q$ (black box), but is high variance and ignores $Q$’s slope.
• Reparameterization (pathwise): write $a = g_\theta(s, \epsilon)$ with $\epsilon$ fixed noise. Now $g_\theta$ is deterministic and differentiable, and the gradient flows straight through: $\nabla_\theta J = \mathbb{E}\epsilon[\nabla_a Q \cdot \nabla\theta g_\theta]$. Low variance, uses the first-order slope.

The mental model: freeze the dice ($\epsilon$), then push a single ball along a known smooth rail. Score-function keeps re-rolling the dice and only checks where points land (jittery); reparameterization freezes one die, reads $\nabla_a Q$ at that one point, and asks how to move $\mu$ to slide it uphill. Switch modes and watch the gradient arrow:

8. What DDPG and SAC do

Both are off-policy actor-critics whose actor is updated by $\nabla_a Q$ — the pathwise gradient. DDPG’s actor is deterministic $\mu_\theta(s)$; its update is $\nabla_\theta J = \mathbb{E}[\nabla_a Q|{a=\mu\theta(s)}\cdot \nabla_\theta \mu_\theta]$. SAC makes the policy stochastic with an entropy bonus and writes the action as a reparameterized squashed Gaussian $a = \tanh(\mu_\theta + \sigma_\theta \odot \epsilon)$, then uses the same $\nabla_a Q$ through it. The crucial shared trait: the actor is one forward pass, so the gradient reaches the parameters in a single cheap hop. That single hop is the picture below — drag the action; hit ascend and it walks uphill along $\nabla_a Q$:

9. RQL’s actor update: make the SAC/DDPG trick “stepwise”

Now assemble everything. RQL’s policy loss is:

Read it: at each partial action $x^f$, take one step $v$ to get $x^{f+1}$, and raise the value after that step, $V(\cdot, f+1)$. The gradient passes through that one call to $v$ — a single-step pathwise gradient. The BC term anchors the policy near the data (essential offline). The punchline is the gradient flow: DDPG/SAC actors are one forward pass (1 hop); RQL is $F$ steps but one hop per step, never BPTT through the chain. Hit backprop:

And here is that update actually running — the stepwise pathwise gradient bending a flow trajectory toward a high-value mode, one cheap hop per step:

# RQL policy update: stepwise pathwise gradient + BC
def policy_loss(batch, V_net, v_net, F, bc_coef):
    s, a = batch.s, batch.a
    xs = flow_reversal(s, a, v_net, F)        # partial actions via reversal
    value_term = 0.0
    for f in range(F):
        x_f = xs[f].detach()                  # gradient flows only through this one v
        x_next = x_f + v_net(s, x_f, f)       # x^{f+1}
        value_term = value_term - V_net(s, x_next, f + 1).mean()
    bc_term = 0.0
    for f in range(F):
        v_target = (xs[f + 1] - xs[f]).detach()             # the data's own step
        bc_term = bc_term + ((v_net(s, xs[f].detach(), f) - v_target) ** 2).mean()
    return value_term / F + bc_coef * bc_term / F

10. The full training loop

def rql_train_step(batch, V_net, V_target, v_net, opt_V, opt_v, F, gamma, bc_coef, tau):
    loss_V = value_loss(batch, V_net, V_target, v_net, F, gamma)
    opt_V.zero_grad(); loss_V.backward(); opt_V.step()
    loss_v = policy_loss(batch, V_net, v_net, F, bc_coef)
    opt_v.zero_grad(); loss_v.backward(); opt_v.step()
    with torch.no_grad():                      # Polyak target update
        for p, p_t in zip(V_net.parameters(), V_target.parameters()):
            p_t.mul_(1 - tau).add_(tau * p)

for step in range(num_steps):
    batch = replay.sample(batch_size)          # offline, off-policy
    rql_train_step(batch, V_net, V_target, v_net, opt_V, opt_v, F, gamma, bc_coef, tau)

Everything is fully offline and off-policy; the intermediate steps are conjured by reversal on the fly — the dataset never stores denoise sequences.

11. Why it beats BPTT / distillation / weighted regression

RQL gets both sides at once: like DDPG/SAC it pushes the policy with stable first-order $\nabla V$ gradients, stepwise (no BPTT, no distillation); and via reversal + “equal values within an action,” the stretched horizon is free (zero-variance, unbiased, no effective-horizon inflation). The only cost is needing a value defined on partial actions — which is exactly what can be trained unbiasedly. The authors report the best offline-RL performance against 19 SOTA flow-RL algorithms across 50 tasks.

12. One line to remember

Treating denoise steps as RL actions stretches the value-learning horizon; but because the flow is deterministic, you can reverse-solve the intermediate steps from offline data, building zero-variance multi-step returns that cancel the horizon penalty; then a value function on partial actions lets you apply a cheap, stable, stepwise reparameterization gradient — that is RQL.

References

• Aditya Oberai, Seohong Park, Sergey Levine. Reversal Q-Learning. 2026. Project page: aober.ai/rql
• Background: DDPG (Lillicrap et al., 2016, arXiv:1509.02971); TD3 (Fujimoto et al., 2018, arXiv:1802.09477); SAC (Haarnoja et al., 2018, arXiv:1801.01290); Flow Matching (Lipman et al., 2023, arXiv:2210.02747).

Reminder: the code here is a teaching reconstruction for explaining the mechanism, not the official implementation; for the real discount bookkeeping, architecture, and regularizers, see the original paper and code.