STITCH-OPE: Trajectory Stitching with Guided Diffusion for Off-Policy Evaluation

We would like to thank the reviewer for their detailed comments. We appreciate that the reviewer found our empirical results versus a variety of recent model-based and model-free OPE methods convincing, including the detailed ablation studies testing various hyper-parameter choices of STITCH-OPE. Please find our responses to your questions below.

---

(Q1) What is the direct relationship between Figure 1 and Off-Policy Evaluation?

Figure 1 illustrates that several components and innovations in model-based reasoning are required to achieve robust model-based OPE in practice. This includes:

1. careful processing of the offline data set to ensure good sample efficiency
2. chunked prediction, composition and guidance during inference are crucial to handle policy distribution shift, which is central to all OPE methods, and
3. empirical reward and return estimation, which is the main task in OPE.

In more details:

1. Figure 1 (Part A) shows how the offline data is split into overlapping sub-trajectories. Critical to model-based OPE methods, this builds an offline training data set on which accurate short-horizon transition dynamics can be learned using less offline data.
2. Once these dynamics are learned using a chunked diffusion model, we effectively obtain a model that accurately captures how the state evolves under the behavior policy (Part B). A critical question of all OPE methods is how to tackle the distribution shift that arises between the behavior and target policy. Figure 1 (Part D) shows that this question can be effectively addressed by sub-trajectory composition and target policy guidance. These ideas work synergistically to account for the policy distribution shift, thus accurately approximating what would happen had we executed the target policy in the real world, without access to a simulator or the real world which is a key feature of OPE.
3. Finally, return estimation involves estimating the empirical reward model from the offline data (Part C), and together with accurate synthetic rollouts from our diffusion model, we can generate accurate estimates of the target policy value, which is the main goal of OPE.

---

(Q2) How does Stitch-OPE handle the OOD (out-of-distribution) issue? Intuitively, since Stitch-OPE trains a reward predictor using behavior data, wouldn't it lead to significant errors in evaluating the performance of the target policy if the trajectories generated by the target policy fall into OOD regions?

Great question! We thank the reviewer for raising a very important point about the necessary assumptions for successful implementation and the limitations of STITCH-OPE.

In summary, we expect STITCH-OPE to generate novel trajectories, but by leveraging the existing structure of the behavior data set through correct composition and stitching of the data, we believe the burden of predicting reward on OOD transitions is significantly reduced because stitching and composition do not involve generating OOD states. We now explain this point in greater detail:

Clarification on the Meaning of OOD: When we refer to out-of-distribution in the context of STITCH or other OPE methods, we mean that rollouts simulated from behavior and target policies might be extremely dissimilar (i.e. in terms of standard measure like L2 norm, KL-divergence), but at the level of individual transitions or chunks, they must still share some common structure. This is typically expressed in OPE literature as requiring that the support of the target policy's action distribution (at every individual decision time step) must be within the support of the behavior distribution, which is also required by STITCH-OPE (Section 3.4). Consequentially, in the normal application of STITCH-OPE, we do expect to generate entirely novel trajectories through effective composition and stitching, but individual transitions within these rollouts are still expected be connected to the behavior data.

Stitching Improves Reward Estimation in OOD Regions: We also argue that successful stitching is a necessary component of accurate reward estimation, since it allows the synthetic rollouts to better match the target rollout distribution, while also respecting the overlapping policy support assumption that is necessary for successful reward prediction (for instance, see row A of Table 1 in our paper).

Standard Model-Based RL Assumptions: We agree with the reviewer that estimation of immediate reward can be particularly challenging when it is nonsmooth or non-Lipschitz, i.e. discontinuous or pathological, even with successful stitching. However, our practical assumption is that the reward is both Markovian and sufficiently well-behaved that it can be approximated from the offline data sets in standard RL benchmarks. This assumption also underlies all model-based offline RL methods that leverage deep learning to approximate both the dynamics and the reward function for value estimation.

Accurate Dynamics Estimation is a Difficult Problem in RL: We also want to emphasize that learning the state dynamics accurately from purely offline data -- while also avoiding the error compounding problem over long horizons during evaluation -- is very often highly non-trivial in OPE. For instance, even if errors of estimated $P(s' | s, a)$ are moderately small at any given time step, they can quickly blow up over time rendering them ineffective over sufficiently long horizon. Meanwhile, reward estimation error typically does not compound, since rewards are predicted individually at the transition level and thus their error accumulates linearly in $T$. The recent paper of Jiang [1] discusses a number of practical challenges of dynamics estimation that are not present in reward estimation, including error compounding and challenges of accurately learning stochastic dynamics. It also shows generally that the error in value increases quadratically in terms of the dynamics error; this is asymptotically looser than the bound on the value error that grows linearly in terms of the reward error. Thus, we wish to re-emphasize that the main goal of STITCH-OPE is to estimate the dynamics accurately offline through composition and stitching, which tackles a crucial problem in model-based OPE that has been inadequately addressed by prior work.

Ablation Experiment to Study the Effects of Reward Estimation on D4RL: To determine the practical effect of reward estimation on the performance of STITCH-OPE on standard offline benchmarks, we ran STITCH-OPE on the Hopper-medium data set using both the true reward function and the empirically estimated reward function. The aggregated performance comparison is provided below:

Aggregated Reward Ablation on Hopper-medium

and the per-policy performance is also included:

Per-Target-Policy Reward Ablation on Hopper-medium

In summary, we found that the reward estimation error accumulated has a minimal effect on OPE performance measured using aggregate metrics. We also found that the error of the OPE estimation remains stable and relatively low across all target policies, including those that differ the most from the behavior policy ($\pi_0$, $\pi_9$) and result in the most significant OOD shift between target and behavior rollouts. This concludes that the prediction errors of the reward were not a significant concern in the Gym and D4RL benchmark suite. We once again thank the reviewer for raising this important point. We will clarify the aforementioned points in our revision and add the previously discussed ablation experiments to our appendix.

---

(Weakness 1) The baselines in Figure 2 and Table 2 are inconsistent.

We thank the reviewer for pointing out this typo and we will fix this in our revision.

---

1. Jiang, Nan. "A note on loss functions and error compounding in model-based reinforcement learning." arXiv, 2024.

We thank the reviewer for your insightful comments. We also agree that the problem of OPE is an important problem that has also not been fully addressed in the context of high-dimensional, long-horizon prediction. We thank the reviewer for acknowledging that our use of negative behavior guidance to address over-regularization and our OPE-specific theoretical bounds were not studied in prior diffusion work, and appreciate the feedback that our empirical results were convincing. Please find our answers to your questions below.

---

(Weakness 1) I'm not fully convinced by a claimed innovation/contribution, i.e., the "sub-trajectory stitching using a sliding window".

Thank you for pointing out the apparent similarity between STITCH-OPE and the receding-horizon approach in Diffusion Policy (cited as [7] in our paper). There are a number of key distinctions between STITCH and Diffusion Policies, most notably:

1. the difference in meaning of sliding windows between the two methods (action vs. state + action prediction),
2. the difference between training and inference procedures of the two methods, and
3. OPE-specific innovations in STITCH -- such as theoretical analysis showing exponential reduction in return variance when using sliding windows -- that are not discussed in the diffusion policy paper.

Distinct Objectives: Diffusion Policy and variants apply a sliding window to predict actions for online control, thus all windows are derived from a single demonstration or from the same trajectory. In contrast, the goal of STITCH-OPE is to produce accurate synthetic trajectories (states and actions) for off-policy evaluation that remain accurate over long horizon and in the presence of distribution shift. This is achieved by drawing short segments from different parts of the offline data set, allowing the diffusion model to optimally compose/stitch them end-to-end dependent on the target policy guidance.

Distinct Training/Inference Procedures: Diffusion Policy is designed for imitation learning by training on complete demonstration trajectories, slicing them only at inference time. Thus, there is no mechanism to recombine slices from different real trajectories. In contrast, STITCH-OPE trains the diffusion model on overlapping windows (sampled randomly) within demonstration trajectories. This allows us to build an explicit model of the length $w$ transition dynamics that can then generalize well across all possible starting states (see row B of Table 1) and even target policies (see row A of Table 1).

OPE Specific Innovations: The unique training and inference setting in STITCH-OPE enables:

1. bridging states across distinct offline trajectories,
2. applying positive guidance to correct distribution shift across policies and even policy classes,
3. applying negative guidance to prevent over-regularization, and
4. provide a theoretical framework for analysing OPE perforamnce (under weak assumptions) that guarantees exponential reductions in variance of the return estimate as compared to other trajectory diffusion models.

The combination of these innovations, none of which have been discussed by Diffusion Policy variants nor other prior work on trajectory diffusion, are critical to ensure reliable OPE performance across many long-horizon tasks (Figure 2, Table 2).

---

(Weakness 2) Does the authors try different time horizon under the same benchmark?

Thank you for raising this point. To address this, we ran an additional experiment to show that STITCH-OPE remains effective over a large range of evaluation horizons.

Specifically, we evaluated STITCH-OPE against all baselines on the Hopper-medium data set for various values of $T \in \lbrace 128, 256, 512, 768, 1000 \rbrace.$ All values have been normalized by the global minimum and maximum achieved via the aggregation of value estimates across all the horizons. We used a discount factor of $\gamma = 0.999$ to show the difference in performance between $T = 768$ and $T = 1000$:

Normalized Log RMSE

In summary, we found that:

1. STITCH-OPE performance degrades very slowly as $T$ increases, validating the theoretical bound of Theorem 3.3 empirically,
2. as predicted, model-based and importance sampling based methods degrade rather quickly as $T$ increases, whereas the performance of STITCH-OPE remains more stable across $T$, and
3. STITCH-OPE outperforms all other baselines according to MSE across all values of $T$, striking the optimal balance between fully auto-regressive and existing trajectory diffusion methods.

We will include results for all metrics (rank correlation and Regret@1) for this ablation experiment in our revised version.

---

(Question) Does that mean 768 seconds for trajectory (i.e., frequency is 1HZ), or the frequency is not 1HZ?

When we write $T = 768$, we refer to the number of decision epochs in each evaluation episode. Indeed, as you correctly pointed out, it would imply 768 seconds at a frequency of 1Hz.

We thank the reviewer for the detailed comments. We appreciate that the reviewer found value in our theoretical analysis, which indeed shows an exponential reduction in variance versus existing methods, and our thorough empirical evaluation and competitive OPE performance across all tasks and metrics.

---

(Q1) Could you detail the specific diffusion sampling algorithm (e.g. DDPM, DDIM) and hyperparameters?

Our implementation of the diffusion model builds on the PyTorch trajectory diffuser implementation by Janner et al. [7], which is a DDPM whose hyperparameters are described in Appendix C of that paper. We are thoroughly committed to reproducible research, and as such have provided key details regarding the diffusion model implementation in Appendix J in the supplementary material (and configuration files with detailed hyperparameters in our attached code base). We will describe all hyperparameters from [7] that we used that were omitted from Table 10 of our appendix in our revised version.

---

(Q2) The choice of $T = 768$ for D4RL (vs. the standard episode length 1000) appears arbitrary. What empirical or theoretical justification supports this?

Thank you for highlighting an important implementation detail of our work. In summary, the main reason for choosing $T = 768$ as the evaluation horizon for the generative baselines was that time steps beyond this contributed little variance to the cumulative return. We also ran an ablation over $T$ to demonstrate that STITCH-OPE works across different values of $T$.

Ablation Experiment over $T$: Since the goal of our paper is to long-horizon correctness of OPE, we conducted an additional ablation over multiple values of $T \in \{128, 256, 512, 768, 1000\}$ on Hopper-medium plus all baselines. We also set the discount factor to $\gamma = 0.999$, to better see the performance differences across methods for $T = 1000$.

Normalized Log RMSE

In summary:

1. STITCH-OPE performance degrades very slowly as $T$ increases, validating the theoretical bound of Theorem 3.3
2. MBR and IS methods degrade rather quickly as $T$ increases, whereas the performance of STITCH-OPE remains more stable across $T$, and
3. STITCH-OPE outperforms all other baselines according to MSE across all $T$, striking the optimal balance between fully auto-regressive and trajectory diffusion methods.

We will include results for all metrics (rank correlation and Regret@1) for this ablation in our revised version.

Practical Justification for $T = 768$: We agree that our choice of $T=768$ could benefit from more explicit justification.

1. An important distinction is that the truncation to $T=768$ applies specifically to generative methods such as STITCH-OPE and PGD, and not truncation of the data set. In contrast, non-generative baselines (e.g., FQE, DRE, DR, IS) use the full offline dataset without truncation, processing entire episodes up to 1000 steps -- our implementation reflects this to ensure fair comparison.

2. The primary rationale for $T=768$ in STITCH-OPE is tied to the discount factor $\gamma=0.99$, which we borrow from standard work in Offline RL and D4RL benchmarks [2]. With $\gamma=0.99$, rewards collected towards the end of the rollout contribute negligibly to the discounted return on these problems. Truncation to $T=768$ thus captures nearly all meaningful return variance across all three D4RL benchmark problems.

3. Additionally, 768 was selected as a convenient multiple for implementation efficiency, since it aligns with common batch sizes and diffusion step multiples in our codebase.

---

(Q3) Could STITCH-OPE's effectiveness be demonstrated using higher-performing target policies?

We would like to clarify a misunderstanding. Our target policies are not the SAC checkpoints from the D4RL medium datasets; they are trained online to span the full spectrum from near-random to near-optimal. Concretely, our best performing policies achieve the following raw returns, which are in par with the strongest results achieved by recent offline methods on the same tasks (medium-expert dataset).

Results are taken from the Robomimic documentation page [8].

An added benefit of training online is that these policies are produced with fresh environment interaction (not the behaviour dataset), thus they introduce a larger shift than re‑using policies trained on the same data, making the OPE problem harder and more realistic. In summary, our current protocol already evaluates STITCH‑OPE on high‑return, high‑shift targets while remaining fully compatible with the reviewer’s request for stronger offline policies.

---

(Q4) The model-based (MB) baseline appears outdated.

We thank the reviewer for suggesting Diffusion World Models (DWM) [1] as an alternative diffusion model-based approach. We have implemented DWM in our setup and included its results in the hopper-medium benchmark below, where it is clear that STITCH-OPE (and PGD) outperform DWM.

We observe that the log MSE values of DWM are worse than PGD [2], and significantly worse than STITCH-OPE which is still by far the best-performing method. This is intuitive, since DWM does not use policy guidance, and thus cannot adapt to the significant distribution shift present between the behavior and target policies.

More generally, we wish to point out that while there have been significant advancements in diffusion modeling over the last few years, we believe PGD [2] is the most representative and recent model-based baseline we found to date that is relevant for OPE. This is because addressing distribution mismatch between behavior and target policy trajectories is a highly non-trivial problem, which only STITCH-OPE (and partially PGD) are able to tackle effectively by using positive and negative policy guidance.

We will include the Diffusion World Models baseline in our evaluation (with accompanying details) and our related work section accordingly. We will also clarify why some recent diffusion work is not suited for OPE.

---

(Q4) Please confirm if other baselines (e.g., PGD, DRE, DR) represent the most current competitive methods.

We thank the reviewer for their concern over fair baselines and believe rigorous evaluation is crucial for advancing OPE. We believe our choice of baselines is up to date.

Single-Step Baselines: Within the class of single-step dynamics modeling, the state-of-the-art in OPE still fit a single neural network to predict the next state given the current state and action, and then do Q-learning with the model rollouts [4,5,6]. While models that are one-step and auto-regressive in predicting individual state components are an alternative, Fu et al. [3] shows that using a single MLP to predict the entire next state performs better.

Trajectory Generation: We believe that trajectory modeling is considerably more challenging in OPE, as it is largely unclear how to address the distribution shift between the behavior and the target policy. While many general diffusion modeling papers cited in our text could in principle model dynamics over multiple time steps, it is largely unclear how they can be adapted to address the distribution shift. PGD (published in 2024) [2] was the most recent paper and most representative member within this class of methods that is relevant to OPE, since it naturally corrects for the policy distribution shift using target policy guidance.

DRE, DR and Other Standard Baselines: The other baselines in our evaluation are considered standard in the OPE literature, as evidenced by the most recent surveys on OPE [4, 5, 6]. They are also considered representative baselines across a variety of diverse classes of OPE methods, i.e. visitation densities, Q-values, doubly-robust, etc. For these reasons, we felt that it would be unfair to exclude these standard baselines, despite that some of them were introduced less recently.

---

1. Ding, Zihan, et al. "Diffusion world model: Future modeling beyond step-by-step rollout for offline reinforcement learning." arXiv 2024.
2. Jackson, Matthew T., et al. "Policy-Guided Diffusion." RLC 2024.
3. Fu, Justin, et al. "Benchmarks for deep off-policy evaluation." ICLR 2021.
4. Kiyohara, Haruka, et al. "Towards Assessing and Benchmarking Risk-Return Tradeoff of Off-Policy Evaluation." ICLR 2024.
5. Uehara, Masatoshi et al. "A Review of Off-Policy Evaluation in Reinforcement Learning." arXiv 2022.
6. R. Figueiredo Prudencio et al., "A Survey on Offline Reinforcement Learning: Taxonomy, Review, and Open Problems," IEEE TNNLS, 2024.
7. Janner, Michael, et al. "Planning with diffusion for flexible behavior synthesis." arXiv 2022.
8. Mandlekar, Ajay, et al. "What Matters in Learning from Offline Human Demonstrations for Robot Manipulation", CoRL 2021.

We appreciate that the reviewer found our paper clear and easy to follow, our theoretical analysis rigorous and insightful, and our empirical evaluation thorough.

---

(Q1) Is the guidance applied at every step during the denoising process?

Yes, you are correct in that we apply guidance at every step of the denoising process.

To clarify this, we first want to point out an interesting feature of guided diffusion not discussed in the literature. Even with a constant guidance coefficient, the actual strength of the guidance across denoising steps is varied by the denoising schedule, i.e. $\alpha_k$. In Appendix A, a concrete illustration with the standard parameterization $\sigma_k^2 = \beta_k = 1 - \alpha_k$ shows the guidance strength is $\sqrt{\alpha_k}$.

Theoretical Justification: Although we agree that the guidance strength is larger at the initial denoising steps, we also want to clarify that this choice is necessary according to theory of classifier diffusion guidance [1]. It also parallels the standard implementation of Langevin dynamics where the ``learning rate'' is annealed over time to achieve (asymptotic) theoretical convergence guarantees.

Practical Justification: Annealing the guidance is also intuitive, since we want to take larger steps to speed up convergence when the sample is far from its true distribution at the early denoising steps. Meanwhile, during later denoising steps, it makes practical sense to take smaller steps to ensure convergence and avoid escaping out of the high-density region. Our empirical results did not require additional tricks to improve convergence (beyond guidance normalization), and worked well across diverse target policies and policy classes.

---

(Q2) Clarification of Diffusion Policy experiments

When we refer to approximating the score at $k=1$, we refer to the calculation of the target policy's score function when the target policy is a diffusion policy. Also, our diffusion policies output immediate/single-step actions.

Details about $k = 1$: Naturally, when the target policy is parameterized as a diffusion policy, the closed-form expression of its score function is no longer available, so we instead use the known fact that the score at denoising step $k$ is $\epsilon(a | s, k)/\sigma_k$. Since we learn $\epsilon( a|s,k)$ to match the score of the convolution of $\pi(a|s)$ with $\mathcal{N}(0, \sigma_k)$ we would ideally estimate the score function at the last denoising step $k = 0$. But this is practically impossible since $\sigma_0 = 0$ and so we use the estimate at $k = 1$. Most notably, these calculations are described in terms of the denoising process of the target policy, and not the guided denoising process of STITCH (for sampling trajectories) which is a separate process. Thus, once we derive the approximate score function of the diffusion policy, we use this score function to guide across all denoising steps within STITCH.

$k = 1$ as a Policy Approximation: As we show in Table 3, STITCH-OPE still provides competitive OPE performance when the target policy score function must be approximated as discussed above, which introduces some measurement errors into the guidance function (which is often the case in many real-world settings). It also outperforms baselines (such as MBR or DRE) that use the target policy only for sampling actions during OPE, which are technically unbiased samples from the exact target policy.

Output of the Diffusion Policy: All diffusion policies we use output single-step (i.e. immediate) actions. This follows exactly the offline diffusion policy literature [2], where single-step actions were also used in their practical implementation.

---

(Q3) Does the trajectory dataset need to contain rollouts that are similar to the behavior policy and target policies? I'm asking this because training generative models like diffusion models is much more data-hungry than other models.

We would like to thank the reviewer for raising this important point. We want to clarify that STITCH-OPE does not require that rollouts be similar between behavior and target policies. This is due to the following reasons:

Chunked Diffusion Training: One of the key technical novelties of STITCH-OPE is training on short overlapping windows sampled from the behavior data. By reducing the prediction horizon, we cut down on the diffusion model's training complexity and data requirements compared to training on full rollouts. More importantly, by observing the same state-action dynamics across different rollouts, they can be reused at different epochs, across multiple trajectories, and across target policies. This means that the behavior data set does not need to contain rollouts that are similar to the target policies, as long as the support of the target policy's immediate action distribution is within the support of the behavior policy -- a standard assumption in OPE literature.

Compositional Inference: At inference time, STITCH-OPE stitches together generated trajectory chunks to reconstruct full-horizon rollouts. Together with policy guidance, this process generates novel rollouts that were never observed in the behavior data. Crucially, these rollouts remain faithful to the target policy even when the target policy's action distribution is significantly different from the behavior policy, as long as the target policy stays within the support of the behavior policy. Please see row A of Table 1 for a visual example and Appendix C for a formal proof that the entropy of stitched trajectories is higher than the behavior policy.

Empirical performance: We validate these contributions on D4RL and demonstrate that STITCH can reliably predict target policy returns even when the target policy differs from the behavior policy (low MSE, high correlation and low regret aggregated across all target policies in Table 2). We also want to clarify that the construction of the target policies in our paper of varying ability from random to expert, as well as the inclusion of target diffusion policies in our work, makes for a considerably more challenging evaluation benchmark than prior OPE work.

---

(Q4) How sensitive is the algorithm to the selection of the behavior policy?

We found that STITCH-OPE is robust to the selection of the behavior policy. To check this, we ran STITCH-OPE and baselines on the medium-expert and expert data sets from the Hopper D4RL benchmark using default parameters ($T = 768, \gamma = 0.99$). The metrics are shown below:

Medium-Expert

Expert

STITCH-OPE outperforms all other baselines across all three metrics using both medium-expert and expert data sets, and performance remains stable across both data sets. This is surprising, given that we would expect the distribution shift for the expert data set to be greater than the medium data set (which contains more diverse rollouts) or the medium-expert data set (which is a mixture of the prior two). This suggests that STITCH-OPE remains effective at composing trajectories across a variety of behavior data sets generated by different behavior policies. Finally, we observe a steep decline in performance of the baselines on the medium-expert and expert data sets as compared to the medium data set (Table 2 in the paper), which suggests that prior OPE work is not robust to large distribution shifts between the behavior and target policies.

---

(Q5) For the visualization, Figures 9 and 11, what is the optimal policy?

The random and optimal policies correspond to the worst and best target policies within the 10 target policies evaluated. We selected the best and worst target policies to demonstrate that STITCH-OPE can generate accurate synthetic trajectories even when guided by target policies that are significantly different from each other and from the behavior policy. The trajectories labeled as STITCH Guided or PGD Guided are purely synthetic trajectories generated from the trained diffusion model with policy guidance, i.e. there is no simulator involved.

---

1. Dhariwal, Prafulla et al. "Diffusion models beat gans on image synthesis." NeurIPS, 2021.
2. Wang, Zhendong, et al. "Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning." ICLR, 2023.