DeFlow: Decoupling Behavior-Prior Modeling and Value Maximization for Offline Policy Extraction

In plain words

A robot or game agent learns from a fixed log of past play, and it has to pick actions the log actually supports. Modern generative models are good at copying a varied set of past moves, but turning such a model into a policy that also chases a reward is slow and awkward. The usual route runs the reward signal back through every step of the generator, which is costly and unstable.

DeFlow splits the job in two. One part is the generator itself, which is trained only to copy the logged behaviour and is then left alone. The other part is a small add-on network that looks at a move the generator proposed and nudges it a little toward moves that score better. A penalty keeps those nudges small, which is the authors' stand-in for staying near moves the log showed; it is a soft rule, not a hard guarantee. Because the generator never has to be trained through its own step-by-step process, the setup stays simpler. The authors report that it is competitive with recent methods, strongest on some multi-step manipulation tasks and when the agent keeps learning online after the offline phase.

TL;DR

• Offline RL with a multi-step flow-matching policy: the flow is trained only on logged actions and kept as a behaviour prior §3.2.
• A small residual MLP f_phi(s, a_base) adds a correction to each sampled action and is trained by DDPG-style Q-ascent with a stopped-gradient base action §3.2 alg. 1.
• A log-space dual update on a penalty weight keeps the mean squared correction near a per-task budget delta; no gradient flows through the ODE solver §3.2 app. B.1.
• Author-reported: highest of the listed methods on cube-double (40, vs 29 for FQL and 25 for SORL) and puzzle-3x3 (43, vs 30 for FQL); SORL is stronger on antmaze-large, humanoidmaze-medium, antsoccer and scene tab. 1 §4.2.
• The official code differs from the paper in defaults and in one puzzle budget; see Paper vs code before reproducing.

Problem

Offline RL has to improve a policy without trusting actions the dataset barely covers §1 §2.1. On multimodal benchmarks such as OGBench, unimodal actors cannot represent the behaviour distribution, while flow-matching and diffusion policies can §1.

The paper calls the difficulty an expressivity-optimizability dilemma. Differentiating a Q-objective through an iterative generator needs backpropagation through the solver §1 §3.1. Shortcut methods such as FQL and SORL avoid that cost by distilling or parameterising a one-step or few-step policy, but the authors argue that this can compress a multimodal flow and makes a scalar BC weight a sensitive trade-off §1 §3.1 §5.

The question the paper asks is whether the multi-step prior can stay intact while a small local module does the value improvement §1.

Method

Policy structure

The action is the sum of a fixed base proposal and a learned residual. The base is a flow sample, with the gradient stopped, and the residual is an MLP that sees the state and that base action §3.2:

a = \text{sg}[\mu_\psi(s, z)] + f_\phi\big(s,\ \text{sg}[\mu_\psi(s, z)]\big), \quad z \sim \mathcal{N}(0, I) §3.2

Here sg is stop-gradient, mu_psi is the Euler-integrated flow, and the displayed action is taken before clipping to the action bounds §3.2. The flow is trained only by flow matching, so Q-learning gradients never change the prior §3.2. The residual is conditioned on the sampled action itself, not only on the state, so two proposals from the same state can receive different corrections §3.2 §5.

Flow prior

The flow uses the conditional flow-matching regression with a straight path between a noise sample and a logged action, and it is integrated with an Euler scheme §2.2 app. B alg. 1. The residual has the same architecture as the flow, a four-layer MLP with width 512 app. B.

Trust-region objective

The refinement is posed as a maximisation of Q over the residual, with a budget on the mean squared residual per action coordinate, where d is the action dimension §3.2:

\max_{f_\phi}\ \mathbb{E}_{s,z}\big[Q(s,\ a_{\text{flow}} + \Delta a)\big] \quad \text{s.t.}\quad \mathbb{E}_{s,z}\big[\lVert \Delta a \rVert_2^2 / d\big] \le \delta §3.2

The paper stresses that delta is a squared-displacement budget on average, not a per-action radius §3.2.

Lagrangian and dual update

To solve the constrained problem, the Q term is divided by the stopped batch mean of |Q| (Q-normalisation) and the penalty weight alpha is written as exp(eta). The refiner minimises the normalised negative Q plus sg[alpha] * ||Delta a||^2 / d, and eta is updated by a surrogate dual step that increases alpha when the measured residual exceeds delta §3.2 app. B.1.

The paper describes this as a surrogate dual update, not an exact solution of the non-convex problem, and it notes that minibatch regularisation gives no hard pointwise bound §3.2. Its appendix gives a Wasserstein-2 bound between the refined and base distributions for a fixed state, W_2^2 <= E||f||^2, and the average over states is only bounded when the population constraint holds app. A.

Algorithm

The paper's Algorithm 1 is a loop over minibatches. The steps are the Bellman update for the critic, the flow-matching update, then the residual and dual updates, as written below alg. 1 §3.2.

# DeFlow update  [alg. 1]
# Input:  minibatch (s, a, r, s', mask) from D; flow v_psi; residual f_phi
# Output: updated Q_xi, v_psi, f_phi and penalty exponent eta

def deflow_update(batch):
    # 1. Critic target uses the refined policy at s'  [alg. 1] [§2.1]
    next_actions = act(batch.next_states)               # (B, d), clipped
    next_values = target_q(batch.next_states, next_actions)  # ensemble
    target = batch.rewards + discount * batch.mask * next_values
    update(critic, squared_error(critic(batch.states, batch.actions), target))

    # 2. Flow prior learns only from logged actions  [alg. 1] [§2.2]
    noise = sample_normal(batch.size, action_dim)        # (B, d)
    time = sample_uniform(batch.size, 1)                 # (B, 1)
    point = (1 - time) * noise + time * batch.actions    # straight path
    update(flow, squared_error(flow(batch.states, point, time),
                               batch.actions - noise))

    # 3. Base proposal from the flow, with no gradient  [alg. 1] [§3.2]
    base = stop_gradient(sample_flow(batch.states))      # (B, d)

    # 4. Residual and its squared size  [alg. 1] [§3.2]
    shift = residual(batch.states, base)                 # (B, d)
    penalty = mean(square(shift))                        # scalar

    # 5. Q term, normalised by the stopped batch mean of |Q|  [§3.2] [app. B]
    action = clip(base + shift, -1, 1)                   # (B, d)
    q = critic_mean(batch.states, action)                # (B,)
    q_loss = -mean(q) / stop_gradient(mean(abs(q)))

    # 6. Residual update with a penalty weight alpha = exp(eta)  [§3.2]
    alpha = stop_gradient(exp(eta))
    update(residual, q_loss + alpha * penalty)

    # 7. Dual step on eta: grow alpha while the shift exceeds the budget  [§3.2]
    dual_loss = eta * (budget - stop_gradient(penalty))
    update(eta, dual_loss)

The sample_flow step is Euler integration with M steps, with the final clip the paper applies to the base proposal alg. 1 app. B.

# Base proposal from the flow  [alg. 1] [§3.2]
# Input:  states (B, obs); Output: actions (B, d)

def sample_flow(states):
    # Start from Gaussian noise  [alg. 1]
    actions = sample_normal(states.size, action_dim)     # (B, d)
    # Euler steps from t=0 to t=1  [alg. 1] [§2.2]
    for step in range(flow_steps):                       # M steps
        time = step / flow_steps
        actions = actions + flow(states, actions, time) / flow_steps
    return clip(actions, -1, 1)

At evaluation the agent calls act, which is sample_flow followed by the residual and a clip, so the executed action is clip(base + residual(base)) alg. 1 deflow.py.

Hyperparameters

Task-specific values come from the paper's appendix table. The table lists FQL's BC coefficient alpha next to DeFlow's budget delta for each task, which is the one place where the paper gives per-task settings app. B.1. Values below are grouped by task family, one row per family, with all five or six tasks sharing a value unless noted app. B.1.

The paper states that IAV (intrinsic action variance) is only a rough guide to delta, not a validated predictor, and that the per-task table overrides the generic heuristic app. B.1.

General settings, with the source of each:

Results

Author-reported. Setup: OGBench primary benchmark (50 state-based single-task tasks and 5 visual tasks) and 18 D4RL tasks, 73 offline tasks in total, averaged over 8 seeds for state tasks and 4 for pixel tasks, following the OGBench protocol with a fixed number of gradient steps §4.1. Baselines are BC, IQL, ReBRAC, FAWAC, FBRAC, IFQL, FQL and SORL (where aligned OGBench numbers exist). The paper flags that SORL is not compute-matched §4.1 §4.3. Bold means within 95% of the row's best §4.1 tab. 1.

Table 1: offline results (score, mean ± std) tab. 1

Source: tab. 1 §4.2. The "–" entries mean the SORL number is not reported for that row, and they are not used when picking the row best tab. 1. The paper's appendix gives the task-level version of the non-SORL baselines app. C.

Table 2: offline-to-online results (offline → online) tab. 2

Baselines are IQL, ReBRAC, Cal-QL, RLPD, IFQL and FQL; DeFlow reuses the offline configuration and appends online transitions to its replay §4.1 tab. 2.

Source: tab. 2 §4.2. Figure 3 plots the aggregated O2O results over 8 seeds fig. 3; its numbers are not given in text beyond the table.

The paper reports that DeFlow is competitive with FQL on many O2O tasks and is strongest on cube-double and puzzle-4x4, and weaker on some cloned Adroit tasks §4.2.

Table 3: compute trade-offs tab. 3

Here M is the number of flow-solver steps, and C_v, C_f, C_Q are the costs of one vector-field call, the residual module and one critic call §4.4 tab. 3.

Table 4: timing in ms per step on an NVIDIA RTX 3090 tab. 4

Source: tab. 4 §4.4 app. D. The paper says these timings describe this implementation only; they are not a speed claim against SORL or a scaling law §4.4.

Diagnostics (not headline results). These come from the appendix and are single-setting probes, not benchmark claims §4.2 app. C:

• Residual attribution, same checkpoints, 50 episodes each: cube-double task1, base flow only 10 ± 2 vs 62 ± 12 with the residual; puzzle-3x3 task2, 0.5 ± 0.9 vs 35 ± 11 app. C.3.
• Mixed human/expert dataset: FQL 21 ± 3, DeFlow 37 ± 5. Removing action conditioning from the residual drops the score from 40 ± 5 to 5 ± 3 app. C.1.
• Residual budget sweep on a manipulation diagnostic: delta 0.0005 gives 55 ± 8; 0.001 gives 50 ± 7; 0.005 gives 27 ± 4; 0.01 gives 8 ± 5; 0.02 gives 4 ± 2 app. C.1.
• Combining DeFlow with floq's value function: cube-double-play 58 ± 5 (DeFlow + floq) vs 40 ± 5 (DeFlow); puzzle-3x3-play 50 ± 3 vs 43 ± 4 app. C.1.
• Residual depth, four seeds: cube-double task1 gives 81 ± 9 (one layer), 74 ± 23 (two layers), 69 ± 10 (four); puzzle-3x3 task2 gives 6 ± 3, 55 ± 7, 38 ± 11; scene task2 gives 52 ± 12, 70 ± 13, 59 ± 17. The paper does not claim a monotone depth effect app. C.4.

Paper vs code

The official repository is github.com/muzhancun/DeFlow at commit 7c4b9f1ff41b91a3e166c0d98a8338587cc45dec. The paper's project page links it, and its owner is the first author. The repository has no LICENSE file, so cite it by path and line only and do not reproduce code.

The code is built on FQL and floq README.md app. B. The main agent is agents/deflow.py; the floq combination is agents/deflowvf.py, which this card did not review in detail deflow.py.

Discrepancies, each with both sides:

1. Q-normalisation is off by default in the code. The paper applies Q-normalisation to the refinement loss §3.2. The code's default is normalize_q_loss=False deflow.py, but every README command for the main results passes --agent.normalize_q_loss=True README.md README.md. A run with default flags is not the paper's setup.
2. Puzzle budget differs. The paper's appendix table gives delta = 0.0005 for puzzle-3x3-play and puzzle-4x4-play app. B.1 and the capacity table uses 0.0005 for puzzle-3x3 app. C.4. The README commands use --agent.target_divergence=0.005 for both README.md README.md. The discrepancy is about a factor of ten. This card follows the paper; check with the authors before reproducing the puzzle numbers.
3. Discount is not in the paper's main text. The code default is 0.99 deflow.py. The README sets 0.995 for antmaze-giant, humanoidmaze and antsoccer README.md README.md README.md. The paper's appendix capacity table also lists 0.995 for humanoidmaze-medium and antsoccer-arena app. C.4.
4. Target aggregation. The code default is mean deflow.py. The README uses min for antmaze-large, antmaze-giant and the pen, door, hammer and relocate D4RL tasks README.md README.md. The paper's capacity table lists min for antmaze-large and pen-cloned app. C.4.
5. Offline step counts. The code default is 1,000,000 main.py deflow.py. The README uses 500,000 for D4RL and visual tasks README.md. The paper states 1,000,000 for OGBench capacity runs and 500,000 for pen-cloned app. D. It does not state the step count for the main table.
6. Initial penalty. The code comment next to the Lagrange setup says the initial value 2.3 gives alpha ≈ 10 deflow.py. The code then sets init_value=1.0 deflow.py. The paper says eta starts at 1 app. B.1. The code matches the paper; the comment is stale.
7. Flow steps. The code default is 10 Euler steps deflow.py. The paper's main text does not state the step count; its capacity runs use ten app. C.4.
8. Items in code, not in paper. The critic has two ensemble members with LayerNorm deflow.py deflow.py. The target update is tau = 0.005 deflow.py. The paper mentions a target network but gives no coefficient §2.1.
9. Test-time scaling variants exist but are not in the main results. The code has sample_actions_max_q and sample_actions_iterative_refinement deflow.py. The paper says the main results use no candidate ranking §4.4 app. B.1. The evaluation path for the main results calls sample_actions evaluation.py. Note that sample_actions_max_q picks argmin of the mean critic value deflow.py. For a best-of-N selection that looks like the lowest-value candidate, so do not reuse that function without checking.
10. Clipping of alpha is disabled. The code's clip on alpha is commented out deflow.py. The paper does not describe a clip.
11. Fixed-penalty branch doubles the coefficient. The paper notes that the fixed-coefficient code path multiplies the configured coefficient twice, giving c² app. C.2. The code confirms this: distill_loss already includes alpha deflow.py, and the actor loss multiplies by alpha again deflow.py. The main results use the adaptive branch, which is the default deflow.py.

Pitfalls

• delta is a mean squared budget per coordinate. It is not a radius, and its scale depends on the action units and dimension. The code's penalty is the mean over batch and action dimensions, so it is already divided by d deflow.py §3.2. Do not change the normalisation without rescaling delta.
• Q-normalisation changes what alpha means. The Q term is divided by the stopped batch mean of |Q| §3.2 deflow.py. Without it, the same delta gives a different balance between value and penalty. All reported runs use it §3.2 app. B.1.
• The dual step is an Adam update, not a plain gradient step. eta uses the shared Adam learning rate of 3e-4 app. B.1 deflow.py. Its initial value is 1, so alpha starts near 2.72 app. B.1. alpha has no upper or lower clip in the code deflow.py.
• The budget is enforced on minibatches only. The paper states that minibatch regularisation gives no hard pointwise bound and no guarantee of support membership §3.2. The W2 bound in the appendix holds for a fixed state; the bound averaged over states, E_s W_2^2 <= d * delta, needs the population constraint to hold and is not a per-state guarantee app. A.
• The penalty is measured before clipping, the action is clipped after. The paper's appendix argues that projecting the refined action onto the action box cannot increase its distance to the base proposal, so the bound survives clipping app. A alg. 1.
• Stochastic base action. Each call to act samples new flow noise, so the critic target and the actor gradient each see a different base action deflow.py deflow.py.
• Hidden costs. Sampling still costs M vector-field evaluations per action, so the flow is not one-step. The residual adds C_f per evaluation §4.4 tab. 3.
• The SORL comparison is not compute-matched. SORL can use shortcut step counts and critic ranking at test time, and DeFlow here does not §4.1 §4.3. The paper's wins on cube-double and puzzle are not a general claim.
• Several headline wins depend on the penalty. The paper's own ablation shows that the fixed-coefficient sweep fails at some settings in both the offline and online phases app. C.2. The adaptive penalty removes some tuning but not the choice of delta §3.2 §6.
• Attribution is conditional on the trained prior. The residual-off comparison uses the same trained checkpoints and does not vary prior quality app. C.3.
• The floq combination is a diagnostic. The paper calls the DeFlow + floq numbers a diagnostic, and the implementation is in agents/deflowvf.py, which is not reviewed here app. C.1 deflowvf.py.
• Budget selection history is not recorded. The paper says the records do not give a complete count of candidate trials, so the reported configuration is reproducible but not an audit of selection cost app. B.1.

Lineage

• Builds on: flow matching for the prior §2.2 and DDPG-style deterministic policy gradients for the refiner §1. The code is built on FQL and floq README.md app. B. The paper's bibliography gives floq as arXiv:2509.06863.
• Compared against: BC, IQL, ReBRAC, FAWAC, FBRAC, IFQL, FQL, SORL, Cal-QL and RLPD §4.1 tab. 1 tab. 2.
• Related designs: residual policy learning §5, decoupled large policies with small RL heads §1, Q-score matching §5, and value-gradient methods that backpropagate through sampling §5.
• Superseded by: not known from the content of this wiki.