BiTrajDiff: Bidirectional Trajectory Generation with Diffusion Models for Offline Reinforcement Learning
Abstract
Offline Reinforcement Learning (RL) relies on static datasets and often enforces conservative constraints to mitigate out-of-distribution errors, but this inevitably gives rise to learning dataset biases and limited behavioral generalization. Recent Data Augmentation (DA) methods leverage generative models to enrich offline data, yet they mainly operate within a single rollout paradigm and tend to preserve the original trajectory-level connectivity of the dataset. As a result, such methods often introduce local variations and fail to recover connections between distinct behavior patterns. In this paper, we propose Bidirectional Trajectory Diffusion (BiTrajDiff), a novel DA framework that explicitly addresses this limitation. BiTrajDiff decomposes trajectory synthesis into two independent diffusion processes that generate forward-future and backward-history segments conditioned on shared intermediate anchor states. By stitching the generated segments at these anchors, BiTrajDiff can synthesize trajectories that bridge disconnected behavior patterns and recover global trajectory-level connectivity absent from the original data. Extensive experiments demonstrate that BiTrajDiff consistently outperforms advanced DA methods across a range of offline RL backbones. Our code is available at https://github.com/Plankson/BiTrajDiff.
Lay Summary
Offline reinforcement learning often suffers from limited trajectory diversity and missing trajectory-level connectivity in static datasets, which restricts the ability of agents to compose long-horizon behaviors. We propose BiTrajDiff, a bidirectional trajectory diffusion framework that augments offline datasets by generating trajectory bridges in both forward and backward directions. Unlike prior methods that only model single-direction generation, BiTrajDiff explicitly recovers missing compositional connections between disconnected trajectories, enabling more effective trajectory stitching and long-horizon planning. Various experiments demonstrate that BiTrajDiff consistently improves multiple offline RL algorithms and achieves stronger performance than existing diffusion-based data augmentation and stitching methods.