TLDR: Generative Trajectory Policies (GTPs) are a new method for offline reinforcement learning that unifies various generative models (like diffusion and consistency models) under a continuous-time ODE framework. This allows GTPs to learn complex, multi-modal behaviors efficiently, overcoming the trade-off between policy expressiveness and computational speed. Key innovations include a stable score approximation for efficient training and a value-driven guidance mechanism for policy improvement. GTPs achieve state-of-the-art performance on D4RL benchmarks, demonstrating superior imitation and policy improvement capabilities.
Offline Reinforcement Learning (RL) is a field where an artificial agent learns how to make decisions from a pre-collected dataset, without interacting further with its environment. This setup is crucial for many real-world applications where direct interaction might be costly or dangerous. However, it presents a significant challenge: how can an agent learn a robust policy from limited, often narrow, experiences and generalize it to an unpredictable world?
A core issue in offline RL is the need for policies that are expressive enough to capture the rich and often varied patterns of behavior found in real-world datasets. Traditional methods often struggle with this, leading to policies that might choose actions not present in the training data, resulting in inaccurate value estimates and poor performance. This has led to a growing interest in generative models, like diffusion models, which are excellent at representing complex, multi-modal distributions.
However, existing generative policies face a dilemma: powerful models like diffusion policies are computationally expensive and slow due to their iterative sampling process. On the other hand, faster, single-step models, such as consistency policies, often compromise on performance. This creates a fundamental trade-off between how expressive a policy can be and how computationally efficient it is.
Introducing Generative Trajectory Policies (GTPs)
A new research paper, “Offline Reinforcement Learning with Generative Trajectory Policies,” introduces a novel approach that aims to bridge this gap. The authors, Xinsong Feng, Leshu Tang, Chenan Wang, and Haipeng Chen, propose Generative Trajectory Policies (GTPs), a new policy paradigm that unifies modern generative models—including diffusion, flow matching, and consistency models—under a single, continuous-time framework. This framework views these models as learning a generative trajectory governed by an Ordinary Differential Equation (ODE).
By learning the entire solution map of this underlying ODE, GTPs offer a more flexible and powerful way to generate actions. Unlike previous methods confined to either slow, high-fidelity sampling or fast, low-fidelity shortcuts, GTPs enable flexible, multi-step, deterministic generation that can achieve high performance even with a few sampling steps.
Key Innovations for Practical Offline RL
To make GTPs practical for offline RL, the researchers introduced two theoretically principled adaptations:
1. Efficient and Stable Training via Score Approximation: Learning an ODE trajectory typically requires on-trajectory supervision, which involves numerically solving the ODE multiple times. This is computationally intensive and prone to instability, especially early in training when the model’s estimates are inaccurate. GTPs address this by replacing the complex, self-referential supervision with a closed-form surrogate anchored to the offline data sample. This “score approximation” makes training both efficient and robust, avoiding the vicious cycle of bad targets leading to bad updates.
2. Value-Driven Guidance for Policy Improvement: Generative models usually aim to match the data distribution, which in offline RL is akin to behavior cloning (imitating observed actions). While useful, this doesn’t inherently lead to policy improvement. GTPs incorporate a value-aware objective using an “advantage-weighted” training mechanism. This theoretically sound approach prioritizes actions that lead to higher returns, guiding the generative process beyond mere imitation towards optimal behavior.
Also Read:
- Boosting Diffusion Language Model Performance with Memory-Efficient Reinforcement Learning
- Robots Surpass Human Teachers: A New Learning Method for Enhanced Robotic Performance
State-of-the-Art Performance
The empirical results on the D4RL benchmarks demonstrate GTP’s effectiveness. In behavior cloning settings (where the policy only imitates data), GTP-BC achieved state-of-the-art performance in 11 out of 15 tasks, significantly outperforming prior generative approaches, especially in complex AntMaze tasks that require long-horizon planning and multi-modal behaviors. This highlights the inherent expressiveness of the GTP architecture.
When evaluated as a full actor-critic algorithm, GTP set a new state-of-the-art for generative policies in offline RL. It achieved the highest average returns on both Gym and AntMaze tasks, even reaching a perfect score on the antmaze-umaze task. These results confirm that the principled, advantage-weighted learning objective successfully leverages the critic’s signal to guide the powerful generative policy towards robust and effective policy improvement.
Ablation studies further validated the importance of the two key adaptations. Removing the score approximation led to substantially longer training times and weaker performance, while replacing the variational guidance with a linear Q-learning term resulted in highly brittle training and divergence. This underscores the stability and efficiency provided by GTP’s design choices.
This work represents a significant step forward in offline reinforcement learning, offering a policy class that effectively balances expressiveness and computational efficiency. While inference is fast, reducing the substantial training time of this model class remains an important avenue for future research. You can read the full paper here.
