Offline Reinforcement Learning with Generative Trajectory Policies
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.
Lay Summary
Many AI systems learn by trial and error, but in real-world settings this can be expensive, risky, or impossible. For example, a robot, autonomous system, or control algorithm may need to learn from past records rather than repeatedly testing new actions in the real world. This paper studies how to build better decision-making systems from such fixed datasets. A key challenge is that good behavior can be diverse: there may be many reasonable actions in the same situation. Existing AI methods can model this diversity, but some are slow to use, while faster methods often lose quality. We propose a new approach called Generative Trajectory Policies, which learns not just a single shortcut from randomness to an action, but the full path by which good actions are generated. This makes the policy both expressive and efficient. Our results show that this approach performs strongly on standard offline decision-making benchmarks, especially on challenging navigation tasks. In the long term, this could help build more capable AI systems that learn from existing data while reducing the need for costly or unsafe real-world exploration.