🤖 machine learning

Offline Reinforcement Learning with Generative Trajectory Policies

This paper introduces Generative Trajectory Policies (GTPs), a unified offline reinforcement learning framework that bridges the gap between slow diffusion models and fast consistency models by learning continuous-time ODE-based trajectory solution maps, achieving state-of-the-art performance on D4RL benchmarks including perfect scores on challenging AntMaze tasks.

Original authors: Xinsong Feng, Leshu Tang, Chenan Wang, Haipeng Chen

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xinsong Feng, Leshu Tang, Chenan Wang, Haipeng Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to walk through a complex maze. You have a giant video library of other robots attempting this task, but you can't let your new robot practice in the real world; it must learn entirely by watching these old videos. This is the challenge of Offline Reinforcement Learning.

The problem is that the videos show many different ways to move. Sometimes the robot turns left, sometimes right, sometimes it stumbles, and sometimes it glides perfectly. A simple robot brain might get confused and try to "average" these moves, resulting in a robot that just spins in circles. It needs a brain that can understand all these different possibilities at once.

The Old Dilemma: Slow vs. Stupid

In the past, researchers tried to build these "smart brains" using two main types of tools, but they faced a frustrating trade-off:

The Slow Artist (Diffusion Models): Imagine an artist who creates a picture by starting with a blank canvas covered in static noise and slowly, step-by-step, refining it until a perfect image appears. This method is incredibly good at capturing complex details and different styles (multi-modal behaviors). However, it takes a long time to paint one picture. If your robot needs to make a decision every second, this slow process is too sluggish to be useful.
The Fast Sketcher (Consistency Models): Now imagine a sketcher who tries to draw the final picture in one or two quick strokes. It's incredibly fast, but because it skips the refinement steps, the drawing often looks blurry or wrong. The robot moves quickly but crashes into walls because its understanding of the world is too simple.

The New Solution: The "Generative Trajectory Policy" (GTP)

The authors of this paper, Xinsong Feng and colleagues, say: "Why choose between slow and fast? Let's build a robot brain that is both."

They introduce a new method called Generative Trajectory Policies (GTP). To understand how it works, imagine the robot's movement not as a single jump from point A to point B, but as a continuous movie of its journey.

The Unified View: The authors realized that both the "Slow Artist" and the "Fast Sketcher" are actually just looking at the same movie from different angles. They are both trying to solve a mathematical equation (an Ordinary Differential Equation, or ODE) that describes how the robot moves from chaos (noise) to order (a perfect action).
The Magic Trick: Instead of watching the movie frame-by-frame (slow) or guessing the ending instantly (fast), GTP learns the entire map of the movie at once. It learns the "solution map," which tells it exactly how to get from any point in the noise to the perfect action, no matter how many steps it takes.

How They Made It Work (The Two Hacks)

Learning this "entire movie map" from scratch is hard. If the robot guesses wrong early on, it gets confused and keeps making worse guesses. The authors added two clever "hacks" to make this practical:

The "Cheat Sheet" (Score Approximation):
Normally, to learn the movie map, the robot would have to simulate the whole movie over and over again to check its work. This is computationally expensive.
- The Analogy: Imagine trying to learn a song by playing the whole symphony to check if you hit the right note. The authors say, "No, just look at the sheet music." They found a way to use a simple, direct formula (a "cheat sheet") to check the robot's progress without simulating the whole movie. This makes training fast and stable, preventing the robot from spiraling into confusion.
The "Coach's Whistle" (Value-Driven Guidance):
Just copying the videos isn't enough; the robot needs to learn better moves than the ones in the videos.
- The Analogy: Imagine a coach watching the video library. When the robot in the video makes a move that leads to a high score, the coach blows a whistle and says, "Do that more!" When the robot makes a bad move, the coach says, "Don't do that."
- The authors built a system where a "critic" (the coach) watches the robot's training and weights the lessons. It tells the robot to pay extra attention to the high-scoring moves in the videos, effectively teaching it to improve beyond just copying.

The Results

The team tested this new robot brain on a famous set of challenges called D4RL, which includes tasks like walking on two legs and navigating complex mazes (AntMaze).

The Outcome: GTP beat almost every other method. It was fast enough to be practical but smart enough to handle the most difficult, multi-step puzzles.
The Highlight: On the notoriously difficult "AntMaze" tasks (where a robot must navigate a huge, winding maze), GTP achieved perfect scores on several levels, something previous methods struggled to do.

Summary

In short, the authors bridged the gap between "slow but smart" and "fast but dumb." By viewing robot movement as a continuous journey and using a "cheat sheet" to learn it efficiently, they created a policy that is both highly expressive (can handle complex, multi-option behaviors) and computationally efficient. They proved that you don't have to sacrifice speed for intelligence in offline learning.

Technical Summary: Offline Reinforcement Learning with Generative Trajectory Policies

Problem Statement

Offline Reinforcement Learning (RL) faces a fundamental challenge: learning a policy from a static, pre-collected dataset without further environment interaction. This setting requires the policy to generalize from limited, often narrow experience to an unpredictable world. A critical bottleneck in this domain is the trade-off between policy expressiveness and computational efficiency.

Traditional offline RL methods often rely on simple function approximators (e.g., Gaussian policies) that struggle to capture the complex, multi-modal behaviors present in real-world datasets. Conversely, recent generative approaches, such as Diffusion Policies, offer high expressiveness by modeling complex action distributions but suffer from slow, iterative sampling processes that impose a significant computational burden. To mitigate this, Consistency Policies were introduced to accelerate inference (often to one or two steps), but this simplification frequently leads to degraded policy performance. The central research question addressed by this work is: Is it possible to design a policy class that achieves both high policy expressiveness and computational efficiency?

Methodology

1. Unified ODE Framework

The authors propose a unifying perspective that views modern generative models—including Diffusion Models, Flow Matching, Consistency Models (CMs), and Consistency Trajectory Models (CTMs)—as specific instances of learning a continuous-time generative trajectory governed by an Ordinary Differential Equation (ODE).

The reverse process is described by the ODE:
$\frac{dx_t}{dt} = f(x_t, t)$
where $f(x, t)$ defines a deterministic trajectory from a noise distribution $x_T$ to a data sample $x_0$ . The authors introduce a Flow Map $\Phi(x_t, t, s)$ , which maps a state at time $t$ to a state at time $s$ . They argue that existing methods are merely approximations of specific aspects of this flow map.

To learn this map, they define two complementary training objectives:

Instantaneous Flow Loss (Local Anchor): Enforces correctness for infinitesimal steps by matching the learned map to the vector field (or denoiser) at the limit $s \to t$ . This unifies diffusion-style denoising and flow-matching velocity estimation.
Trajectory Consistency Loss (Global Regulator): Enforces self-consistency across long, multi-step jumps, ensuring that the composition of intermediate steps equals the direct jump.

2. Generative Trajectory Policies (GTP)

Building on this framework, the paper introduces Generative Trajectory Policies (GTPs), a new policy paradigm that learns the entire solution map of the underlying ODE. Unlike methods confined to slow, high-fidelity sampling or fast, low-fidelity shortcuts, GTPs enable flexible, multi-step, deterministic generation.

3. Practical Adaptations for Offline RL

Translating the GTP paradigm into a functional offline RL algorithm requires addressing three specific challenges: prohibitive computational cost, training instability, and misaligned objectives. The authors introduce two key theoretically principled adaptations:

Stable Score Approximation:
- Challenge: Learning the ODE trajectory typically requires numerically solving the ODE backward using an inner-loop solver for every sample, which is computationally intractable and unstable due to error propagation from early-stage inaccurate estimates.
- Solution: The authors replace the learned vector field with a closed-form surrogate anchored to the offline data sample: $\tilde{f}(x_t, t) = (x_t - x)/t$ .
- Theoretical Guarantee: Theorem 4.1 proves that using this surrogate changes the training objective only by $O(h^p)$ (where $h$ is the step size), making the practical objective asymptotically equivalent to the ideal one. This eliminates the need for costly multi-step ODE integration during training while ensuring stability.
Value-Driven Guidance:
- Challenge: Standard generative objectives match the data distribution (Behavior Cloning), which does not inherently lead to policy improvement.
- Solution: The authors formalize a value-weighted training objective. By deriving the optimal solution to a KL-regularized policy optimization problem, they show that training a generative policy to match the optimal policy is equivalent to solving a weighted generative objective where the weights are the exponentiated advantage function: $w(s, a) \propto \exp(\eta A(s, a))$ .
- Implementation: To ensure stability, advantage weights are normalized and negative values are truncated. This allows the policy to preferentially imitate high-advantage actions while remaining anchored to the data support.

Key Contributions

Generative Trajectory Policies (GTP): A new, highly expressive policy paradigm derived from a unifying framework connecting modern generative models to continuous-time ODE trajectories.
Theoretical Adaptations: Two key adaptations for practical offline RL:
- A score approximation technique that provides theoretical guarantees for efficiency and stability, removing the need for iterative solvers during training.
- A variational framework for value-driven policy improvement that bridges the gap between imitation learning and policy optimization.
Empirical Validation: Extensive experiments demonstrating that GTP achieves state-of-the-art performance on D4RL benchmarks, outperforming prior generative and offline RL methods.

Experimental Results

The authors evaluated GTP on the D4RL benchmark, covering Gym (locomotion) and AntMaze (sparse-reward, long-horizon) domains.

Behavior Cloning (BC) Performance: As a pure imitation learner (without value guidance), GTP-BC achieved state-of-the-art results in 11 out of 15 tasks. Notably, it significantly outperformed Diffusion-BC and Consistency-BC, particularly in the challenging AntMaze suite (average score 66.3 vs. 44.1 for the next best).
Policy Improvement: The full GTP actor-critic algorithm set new state-of-the-art records for generative policies.
- On Gym tasks, GTP achieved an average return of 89.0, surpassing the previous best (Diffusion-QL at 87.9).
- On AntMaze tasks, GTP achieved an average of 80.6, significantly outperforming Diffusion-QL (69.6) and QGPO (78.3).
- Notable Achievement: GTP achieved a perfect score of 100.0 on the notoriously difficult antmaze-umaze task.
Efficiency: Ablation studies confirmed that GTP does not rely on long ODE trajectories for performance; reducing the sampling horizon from 5 to 2 steps maintained high performance while substantially reducing inference latency.

Significance and Claims

The paper claims that GTP successfully bridges the gap between expressiveness and efficiency in offline RL. By viewing generative policies through the lens of a unified ODE framework and learning the full solution map, GTP avoids the computational bottlenecks of iterative diffusion models and the performance degradation of simplified consistency models.

The authors emphasize that their approach provides a principled foundation for designing generative policies. The introduction of score approximation and value-driven guidance makes this theoretically rich framework practically viable for offline RL. The results suggest that learning the full continuous-time trajectory provides a powerful inductive bias for capturing complex, temporally extended behaviors, particularly in sparse-reward, long-horizon environments where other methods struggle.

The paper concludes that while inference is fast, reducing the substantial training time of this model class remains an important avenue for future research. The work opens a promising direction for harnessing continuous-time dynamics in RL, offering a balance between robustness and performance that was previously unattainable.