Technical Summary: BiTrajDiff

Problem Statement

Offline Reinforcement Learning (RL) aims to learn effective policies from static, pre-collected datasets without further environment interaction. A primary challenge in this domain is the Out-of-Distribution (OOD) problem, where the learned policy encounters state-action pairs not present in the logged data, leading to extrapolation errors and value overestimation. To mitigate this, existing methods often enforce conservative constraints, which inadvertently restrict learning to the structural biases of the dataset, limiting behavioral diversity and generalization.

Recent Data Augmentation (DA) approaches attempt to address this by using generative models (particularly diffusion models) to synthesize new transitions. However, current DA methods predominantly operate within a single rollout paradigm (either forward-future or backward-history) conditioned on in-dataset states. As a result, these methods tend to preserve the original trajectory-level connectivity of the dataset. They generate local variations within existing behavior patterns but fail to recover connections between distinct, disconnected behavior patterns (e.g., bridging two separate rooms in a navigation task where no cross-door trajectories exist in the original data).

Methodology

The authors propose BiTrajDiff (Bidirectional Trajectory Diffusion), a novel DA framework designed to explicitly recover global trajectory-level connectivity by synthesizing trajectories that bridge previously disconnected behavior patterns.

Core Framework

BiTrajDiff decomposes trajectory synthesis into two independent yet complementary diffusion processes:

1. Forward-Future Diffusion: Models the generation of future state sequences starting from an intermediate state.
2. Backward-History Diffusion: Models the reconstruction of past state sequences ending at the same intermediate state.

Both processes are conditioned on the same intermediate anchor states ($s_t$) sampled from the offline dataset, along with a target cumulative return signal.

Key Components

1. Bidirectional Diffusion Training:

   • Two separate diffusion models ($\epsilon_{\theta_f}$ for forward, $\epsilon_{\theta_b}$ for backward) are trained on the offline dataset.
   • The forward model generates trajectories $\{s_t, \dots, s_{t+H-1}\}$ conditioned on the initial state $s_t$ and target return $R_f$.
   • The backward model generates trajectories $\{s_{t-H+1}, \dots, s_t\}$ conditioned on the terminal state $s_t$ and target return $R_b$.
   • Training utilizes Classifier-Free Guidance (CFG) to handle conditional and unconditional generation, with the cumulative return serving as the primary condition.

2. Bidirectional Trajectory Generation & Stitching:

   • During inference, an anchor state $s_t$ is sampled from the dataset.
   • The forward model generates a future segment, and the backward model generates a history segment.
   • These segments are stitched at the shared anchor state $s_t$ to form a complete bidirectional trajectory: $\hat{\tau}_s(s_t) = \{ \hat{s}_{t-H+1}, \dots, \hat{s}_{t-1}, s_t, \hat{s}_{t+1}, \dots, \hat{s}_{t+H-1} \}$.
   • This stitching mechanism allows the synthesis of trajectories that connect states which were previously unreachable or disconnected in the original dataset.

3. Trajectory Completion:

   • Since the diffusion models generate state sequences, an Inverse Dynamics Model (IDM) is used to predict actions ($a_t$) given state transitions ($s_t, s_{t+1}$).
   • A Reward Model (RM) predicts rewards ($r_t$) for the state-action pairs.
   • This completes the trajectory with action and reward signals, forming a full transition tuple.

4. Bidirectional Trajectory Filtering:
   To ensure the quality of the augmented data, a two-stage filtering mechanism is applied:

   • OOD Filter: Uses an Isolation Forest to detect and remove trajectories that deviate significantly from the original data distribution (novelty detection).
   • Greedy Filter: Selects the top trajectories based on the highest cumulative reward, ensuring the augmented data guides the policy toward high-performing behaviors.

Key Contributions

1. Identification of Limitations: The authors identify that single-rollout DA methods preserve original trajectory-level connectivity, failing to explore connections between distinct behavior patterns.
2. Bidirectional Synthesis: The proposal of BiTrajDiff, which explicitly decomposes generation into forward and backward processes conditioned on shared anchors, enabling the synthesis of trajectories that bridge disconnected behavior modes.
3. Empirical Validation: Extensive experiments on the D4RL benchmark demonstrate that BiTrajDiff consistently outperforms advanced DA baselines (including Synther, DiffStitch, and RTDiff) across various offline RL backbones (IQL, TD3BC, CQL, DT).

Experimental Results

The paper evaluates BiTrajDiff on the D4RL benchmark, covering locomotion, navigation (Maze2D, AntMaze), and manipulation (Kitchen) tasks.

• Performance Gains: BiTrajDiff achieves the highest average returns across most tasks and algorithms. For instance, in locomotion tasks, it consistently outperforms single-direction baselines. In sparse-reward environments like Maze and Kitchen, where data quality is critical, BiTrajDiff shows significant improvements in stability and performance.
• Ablation Studies:
  • Directionality: BiTrajDiff (bidirectional) consistently outperforms both forward-only and backward-only variants, validating the necessity of stitching both past and future segments to recover global connectivity.
  • Filtering: The combination of OOD and Greedy filters is shown to be superior to using either alone or no filters, balancing diversity with reliability.
  • Augmentation Ratio: A ratio of 30% augmented data to original data ($\sigma=30\%$) yields optimal stability and performance; excessive augmentation (100%) introduces noise that degrades performance.
• n-step TD Estimator: BiTrajDiff maintains its performance advantage over baselines even as the n-step horizon increases, suggesting the generated trajectories are of higher quality and more consistent with the underlying dynamics than those from other DA methods.
• Diversity vs. Accuracy: Visualization metrics (Dynamic Error and L2 Distance) show that BiTrajDiff generates trajectories with significantly higher diversity (larger L2 distance from original data) while maintaining competitive dynamic accuracy, unlike single-direction methods which tend to overfit the original distribution.

Significance and Claims

The paper claims that BiTrajDiff addresses a fundamental limitation in offline RL data augmentation: the inability of single-process rollouts to bridge disconnected behavior patterns. By explicitly modeling and stitching forward-future and backward-history trajectories around shared anchor states, BiTrajDiff recovers global trajectory-level connectivity absent from the original dataset.

The authors assert that this approach allows offline RL algorithms to learn from a more diverse set of behavioral patterns without sacrificing physical consistency, leading to substantial performance improvements across a wide range of tasks and algorithms. The work is presented as a step forward in overcoming the structural biases of conservative offline learning through generative data augmentation. Future work is noted to explore bidirectional generation in multi-task settings with broader behavioral patterns.