Offline Reinforcement Learning with Generative Trajectory Policies

AI summary

Overview

Research area: Offline reinforcement learning (RL) and generative modeling, specifically the use of diffusion, flow matching, and consistency models as policy classes.

Technical level: Advanced. The paper assumes familiarity with actor-critic offline RL, stochastic differential equations / probability flow ODEs, score matching, and behavior cloning.

Scope: The paper proposes a unifying Ordinary Differential Equation (ODE) view of modern generative models, derives a policy paradigm called Generative Trajectory Policies (GTPs), and evaluates it on the D4RL benchmark.

What This Paper Is About

Offline RL agents must learn policies from a fixed, pre-collected dataset with no further environment interaction, which makes policy expressiveness critical for capturing the multi-modal behaviors in real data. Existing generative policies face a trade-off: diffusion policies are expressive but require slow, iterative sampling, while consistency policies are fast but often degrade in quality. The paper's goal is to design a policy class that achieves both expressiveness and computational efficiency by learning the entire solution map of a continuous-time generative ODE.

Key Contributions

1. A unifying ODE framework. The authors argue that diffusion models, Consistency Models, Consistency Trajectory Models (CTMs), various forms of Flow Matching, Shortcut Models, and Mean Flows can all be understood as instances of learning a continuous-time generative trajectory governed by an ODE. This reframes model design around two axes: defining the vector field, and solving the ODE efficiently by learning its solution (flow) map.

2. Generative Trajectory Policies (GTPs). Building on that foundation, the paper defines GTP as a policy class that generates actions by learning the entire solution map of the underlying ODE, enabling flexible multi-step deterministic generation rather than only slow high-fidelity sampling or fast low-fidelity shortcuts.

3. Two theoretically grounded adaptations for offline RL. The first is a score approximation (Theorem 4.1) that replaces the self-referential Inst Map with a closed-form surrogate anchored to the offline sample, with the practical and ideal objectives shown to differ by O(h^p). The second is a variational, advantage-weighted objective (Theorem 4.4) that turns generative imitation into value-driven policy improvement.

4. Empirical validation on D4RL. GTP is reported to achieve state-of-the-art performance among generative policies, with perfect scores on several AntMaze tasks. Code is available at https://github.com/wmd3i/gtp.

Main Findings

• Generative models are unified by one ODE. The true flow map Φ(x_t, t, s) = x_t + ∫_t^s f(x_τ, τ) dτ provides a single representation under which diffusion denoisers estimate its infinitesimal form and consistency models enforce its compositional structure.

• Two complementary training objectives emerge. The Instantaneous Flow Loss (a local anchor) enforces correctness at the limit s → t and recovers the diffusion denoiser and the flow-matching velocity target; the Trajectory Consistency Loss (a global regulator) enforces Φ(x_t, t, s) ≈ Φ(Φ(x_t, t, u), u, s) for t > u > s.

• Score approximation is nearly lossless. Theorem 4.1 shows that replacing the true vector field f* with the surrogate f̃(x_t, t) = (x_t − x)/t changes the training objective by |L_prac(θ) − L_ideal(θ)| = O(h^p), so the two objectives coincide in expectation as h → 0. This removes the need for multi-step ODE integration, with intermediate points obtained by a one-step perturbation x_u = x + u·z.

• Advantage weighting is the theoretically correct value signal. Theorem 4.4 derives π*(a|s) ∝ π_BC(a|s) exp(η A(s,a)) with A(s,a) = Q(s,a) − V(s), giving the weighted generative objective max_θ E[exp(η A(s,a)) ℓ_gen(π_θ; a|s)].

• Strong behavior cloning results. In the pure BC setting (η = 0), GTP-BC achieves state-of-the-art performance in 11 out of 15 tasks. Average Gym score is 82.3 versus D-BC 76.3 and C-BC 69.7; average AntMaze score is 66.3 versus C-BC 44.1 and D-BC 41.2.

• State-of-the-art offline RL results. With full actor-critic training, GTP reaches the highest Gym average (89.0), ahead of D-QL (87.9), BDM (87.3), QGPO (86.6), and C-AC (85.1). On AntMaze, GTP averages 80.6 versus IDQL-A 79.1 and QGPO 78.3.

• Perfect AntMaze scores. GTP achieves a perfect 100 ± 0 on antmaze-umaze, and reports 94.2 ± 2.0 on antmaze-md and 71.0 ± 4.9 on antmaze-ld, tasks where several baselines score near zero.

• Both components are validated by ablation. On hopper-medium-expert-v2, GTP scores 112.2 ± 0.6 in 4.26 hours; removing score approximation scores 99.7 ± 1.7 in 5.23 hours; a linear Q-learning actor loss with λ = 0.01 scores 111.4 ± 0.9 in 5.08 hours, while λ = 0.1 and λ = 1.0 diverge.

• Efficiency is not purchased with fidelity. Diffusion policies and GTP used K = 5 sampling steps and consistency policies used K = 2 in the evaluation, with each policy evaluated over 10 episodes for Gym tasks and 100 episodes for other tasks.

Methodology in Plain English

The starting point is the observation that many modern generative models can be written as an ODE that carries a sample from noise at time T to data at time 0. Solving that ODE numerically is what makes diffusion slow, so the authors instead learn the map that jumps directly from one point on the trajectory to another. They parameterize this map using a surrogate function that estimates the endpoint by extrapolating the average velocity over an interval, and then recover the true flow map by linear interpolation.

Two training signals keep this map honest. A local signal forces the infinitesimal step to behave like a correct denoiser, and a global signal forces composed jumps to agree with one another, so that a long jump equals the sum of its shorter pieces. Running this naively in offline RL is expensive and unstable, because the supervision targets themselves come from the model's own early, inaccurate predictions. The fix is to replace those self-generated targets with a closed-form expression built directly from the offline data sample, which the authors prove costs only O(h^p) in the objective and can be computed without any ODE solver. Finally, to move beyond plain imitation, the actor loss is reweighted by exponentiated advantages from a double Q-network critic. In practice those weights are normalized by the advantage standard deviation and any negative advantages are clipped to zero, which keeps optimization stable. Actions are then produced by starting from Gaussian noise and applying the learned map over K timesteps.

Why This Matters

Impact on research. The paper reframes a crowded landscape of generative policies as variations on one ODE solution map, offering a design space rather than a single architecture. It also separates two failure modes that had been entangled in the literature: the expressiveness-versus-efficiency trade-off, and the choice between generative imitation and value-based improvement.

Real-world applications. The paper does not report deployed systems or non-benchmark applications. The task domains it evaluates are simulated navigation and locomotion (AntMaze and Gym). Settings where offline policy learning from fixed logged data is relevant, and to which the approach could plausibly extend, include:

• Robotics, where collecting new interaction data is costly or unsafe.
• Healthcare treatment recommendation from historical patient records.
• Autonomous driving from logged fleet data.
• Industrial control and recommendation systems trained on archived logs.

Industry relevance. By combining fast inference with high reported returns, GTP targets the practical objection that diffusion policies are too slow to deploy under latency constraints — a key barrier to using offline RL in production. The reported training times (4.26 hours versus 5.23 hours without the approximation on one task) and the use of only K = 5 sampling steps speak directly to compute budgets.

Future Directions

• Beyond D4RL. The paper states that additional results on OGBench and visual-observation variants are deferred to Appendix D, leaving the question of how GTP scales to higher-dimensional observations open in the main text.

• Tightening the theory. The O(h^p) bound in Theorem 4.1 relies on Lipschitz and bounded-moment assumptions; whether guarantees can be extended to end-to-end offline RL performance, rather than only to the discrepancy between two training objectives, is unresolved.

• Hyperparameter sensitivity. The advantage-weighting scheme has coefficients η, λ_Flow, and ε that require tuning, and the ablation shows that an alternative value-loss formulation diverges at λ = 0.1 and 1.0. How these interact across task suites is not fully characterized.

• Which generative family is best under this umbrella. Since diffusion, flow matching, consistency, shortcut, and mean-flow models are all instances of the same flow map, comparing their practical trade-offs within the GTP framework is a natural next experiment.

Target Audience

Researchers and graduate students working on offline RL, generative policies, or diffusion and flow-based decision making; practitioners who need rapid policy inference under fixed logged data; and readers interested in unifying theory for continuous-time generative models. The paper's combination of ODE-level formalism and D4RL benchmarking makes it most useful to those already comfortable with actor-critic algorithms and score-based generative modeling.

Authors’ abstract

Generative models have emerged as a powerful class of policies for offline reinforcement learning (RL) due to their ability to capture complex, multi-modal behaviors. However, existing methods face a stark trade-off: slow, iterative models like diffusion policies are computationally expensive, while fast, single-step models like consistency policies often suffer from degraded performance. In this paper, we demonstrate that it is possible to bridge this gap. The key to moving beyond the limitations of individual methods, we argue, lies in a unifying perspective that views modern generative models, including diffusion, flow matching, and consistency models, as specific instances of learning a continuous-time generative trajectory governed by an Ordinary Differential Equation (ODE). This principled foundation provides a clearer design space for generative policies in RL and allows us to propose Generative Trajectory Policies (GTPs), a new and more general policy paradigm that learns the entire solution map of the underlying ODE. To make this paradigm practical for offline RL, we further introduce two key theoretically principled adaptations. Empirical results demonstrate that GTP achieves state-of-the-art performance on D4RL benchmarks - it significantly outperforms prior generative policies, achieving perfect scores on several notoriously hard AntMaze tasks.