State-Covering Trajectory Stitching for Diffusion Planners

Thank you for the detailed review and constructive feedback on our work. We especially appreciate your insightful question regarding the generalizability of SCoTS beyond the controlled OGBench setting. Please find our detailed responses below.

• “How robust is SCoTS to noisy or low-quality datasets beyond the controlled OGBench settings, e.g., occupancy maps?”

To evaluate the robustness of SCoTS on noisy and low-quality data beyond the controlled OGBench settings, we conducted additional experiments using a mobile robot navigation task in realistic, large-scale environments (a 35m x 39m corridor and a 50m x 110m building), following the setup from HRL [1].

We simulated low-quality data by generating short, disconnected segments. Specifically, trajectories were collected by randomly selecting start and goal locations within short ranges (5m) using HRL [1]. To handle the noisy raw 2D LiDAR measurements (512 rays over 360°), we converted these scans into 64x64 occupancy grid images. These images were then encoded using a pre-trained $\beta$-VAE into a compact 8-dimensional latent representation, providing structured, low-dimensional state inputs for the agent. Subsequently, the SCoTS stitching method was applied to augment this fragmented dataset.

The table below shows the goal-reaching success rate (over 100 episodes) of HIQL when trained with and without our SCoTS data augmentation.

These results clearly indicate that SCoTS effectively enhances performance, even with highly fragmented and noisy datasets, demonstrating its robustness and practical applicability to robot navigation tasks.

• “Can SCoTS be extended or adapted for online learning or continual RL settings?”

Although our work primarily addresses the offline setting, the SCoTS framework can naturally be extended to online or continual RL settings. One straightforward approach would involve periodically incorporating new trajectories collected by the agent into the training dataset. This evolving dataset could then be used to periodically retrain both the temporal distance-preserving embedding $\phi$ and the diffusion-based stitcher $p_{\theta}^{\text{stitcher}}$. This strategy aligns with successful online data augmentation approaches demonstrated in prior work [2]. We believe further extending the idea behind SCoTS to other setting is an exciting future research direction.

• “Computational overhead”

We acknowledge that SCoTS introduces additional computational overhead, primarily due to the training of the diffusion-based stitcher and the trajectory augmentation process, which together roughly equate to 50% of the training time required for the final diffusion planner.

However, we believe this overhead is justified by substantial performance gains. For instance, on challenging tasks such as antmaze-explore-giant, SCoTS improved the hierarchical diffusion planner's success rate by up to seven times. Such considerable improvement in generalization, achieved through a single offline augmentation phase, is highly beneficial, particularly when compared to the significantly higher costs associated with collecting a diverse set of high-quality, long-horizon data.

• “Assumption of Symmetric Temporal Embeddings”

As we discussed in our paper's limitations section, our temporal distance-preserving embedding assumes symmetric temporal embeddings. We believe this limitation can be resolved by leveraging universal quasimetric embedding as explored in [3] which allows the representation to more faithfully model irreversible dynamics, which we leave for future work.

• “Evaluation bias toward diffusion planners”

While our motivation and main narrative center around enhancing diffusion planners, we emphasize that our experimental results show substantial improvements for traditional offline goal-conditioned RL (GCRL) algorithms as well.

Specifically, our work is the first, to our knowledge, to enable GCRL methods to surpass 50% success rates on particularly challenging OGBench tasks (antmaze-stitch-giant and antmaze-explore-large). These results demonstrate that SCoTS serves as a general-purpose data augmentation framework that significantly improves performance across diverse offline RL algorithms, not merely diffusion-based approaches.

We believe these results strengthen rather than weaken our contribution.

• “Given the nature of navigation robotics tasks, SCoTS could benefit from blending diffusion and gradients from task objectives as in [4], improving collision-free performance.”

We thank the reviewer’s suggestion. We would like to clarify that our SCoTS framework and the suggested guidance mechanism are orthogonal and complementary. SCoTS is a data augmentation method that enriches the training data to provide the planner with a stronger, more diverse prior. In contrast, the guidance technique is an inference-time method that refines a generated plan to satisfy specific constraints like collision avoidance. We believe this combination is an interesting direction for future research and will add a discussion of this potential synergy in the final version of the paper.

References:

[1] Lee et al., Adaptive and Explainable Deployment of Navigation Skills via Hierarchical Deep Reinforcement Learning. ICRA, 2023.

[2] Lu et al., Synthetic Experience Replay. NeurIPS, 2023.

[3] Wang et al., Optimal goal-reaching reinforcement learning via quasimetric learning. ICML, 2023.

[4] Carvalho et al., Motion planning diffusion: Learning and planning of robot motions with diffusion models. IROS, 2023.

Thank you for the detailed review and constructive feedback on our work. We especially appreciate your helpful comments and suggestions for further improving the work. Please find our detailed responses below.

• “Can $\phi$ actually learns the embedding of optimal temporal distance?”

Learning temporal embedding $\phi$ is grounded in the well-established equivalence between the optimal temporal distance and the optimal goal-conditioned value function [1, 2, 3]: $V^{\*}(\boldsymbol{s},\boldsymbol{g}) = -d^{\*}(\boldsymbol{s},\boldsymbol{g})$

Here, $V^{\*}(\boldsymbol{s},\boldsymbol{g})$ is the maximum possible cumulative sum of rewards for setting where an agent receives a reward of $0$ if the $l_2$ distance between states $\boldsymbol{s}$ and $\boldsymbol{g}$ is within a small threshold $\delta_g$, and $-1$ otherwise. We leverage this by parameterizing our value function directly as a distance in the latent space: $V(\boldsymbol{s}, \boldsymbol{g}) \coloneqq -||\phi(\boldsymbol{s}) - \phi(\boldsymbol{g})||_2$

By training $\phi$ to minimize a standard IQL-based temporal difference (TD) error for this value function, we make $V(\boldsymbol{s}, \boldsymbol{g})$ to approximate the true optimal value $V^{\*}(\boldsymbol{s}, \boldsymbol{g})$ . Consequently, this implicitly forces our learned latent distance, $||\phi(\boldsymbol{s}) - \phi(\boldsymbol{g})||_2$ to approximate the optimal temporal distance, $d^{\*}(\boldsymbol{s},\boldsymbol{g})$. Empirically, this approximation proves effective, though we acknowledge it remains an approximation.

• “How will embedding error affect the subsequent process and final performance?”

Thank you for this important question. The robustness of our method to imperfections in the learned embedding is a core design consideration. We explicitly acknowledge that the embedding $\phi$ can be a biased approximation, and SCoTS incorporates two key mechanisms to mitigate the potential negative effects during the stitching process.

First, embedding errors can lead to the retrieval of candidate segments that are not dynamically consistent with the current trajectory. This can result in stitches with large dynamic inconsistency errors at the connection points. To address this, we employ our Diffusion-based Stitching Refinement step, which is specifically designed to smooth the transition between segments and ensure dynamic feasibility, as shown in Figure 7 of our Appendix.

To demonstrate its necessity, we ablate this component. The results show that removing refinement significantly degrades performance for both the diffusion planner (HD) and, more dramatically, the GCRL algorithm (HIQL), which is highly sensitive to the dynamic validity of transitions.

Second, geometric distortions in the latent space can make a straight-line exploration path suboptimal. Consequently, a straight line between two latent states, $\phi(s)$ and $\phi(s')$, may not correspond to a feasible path in the underlying MDP. Therefore, exploratory detours are often necessary to find a valid temporally extended path. This is precisely where the novelty score becomes critical. It encourages these exploratory detours by rewarding the selection of segments that lead to less-visited states. This allows the agent to navigate around the imperfections in the latent space, discovering more temporally extended paths.

The following ablation study on the novelty weight $\beta$ demonstrates this complementary effect. The results show that while directional stitching alone $(\beta=0)$ significantly outperforms the baseline, introducing a novelty score $(\beta>0)$ provides substantial further gains.

• “Formulation of TopKNeighbours”

The function $\text{TopKNeighbors}(\text{query}, \text{search space}, k)$ is a standard k-nearest neighbor search algorithm. In the context of our paper, its query is the latent embedding of the final state of the currently composed trajectory, $\phi(\text{end}(\boldsymbol{\tau}_{\text{comp}}))$, and the search space is the the set of latent embeddings of the initial states of all trajectories within the entire offline dataset, denoted by the shorthand $\phi(\mathcal{D})$.

The function operates by searching this space to find the $k$ vectors closest to the query, using the Euclidean distance in the learned latent space as the metric. It then returns the $k$ original trajectory segments from the dataset that correspond to these nearest initial states, which form the candidate set for the next stitching step. (Practical implementation can be found in our Appendix B.)

We appreciate the reviewer highlighting this ambiguity and will clarify this point carefully in the final version of the paper.

• “How can diffusion-based stitching refinement avoid to “exhibit minor dynamic inconsistencies or sub-optimal transitions”?

The key to avoiding inconsistencies lies in the fact that our method is generative, not simple concatenation. A naive approach would be to simply concatenate two segments, which could create an abrupt, or even dynamically impossible transition at the stitching point. Our diffusion-based stitcher avoids this by generating an entirely new bridging path. Because the stitcher is a diffusion model trained on the entire dataset of real trajectories, it learns a strong prior over the distribution of physically plausible state sequences. When conditioned on the boundary states ($\boldsymbol{s}_1$ and $\boldsymbol{s}_H$), it samples a trajectory from the learned distribution of valid trajectories.

• “Ablation study and sensitivity analysis.”

Thank you for your valuable feedback. Following the suggestion, we have conducted additional ablation and sensitivity analyses to better understand the impact of key hyperparameters and the necessity of our core components.

1. Sensitivity to Sub-trajectory Length $H$

The results show that on the stitch dataset, performance is relatively robust to the choice of $H$. However, on the explore dataset, shorter segments perform best. We hypothesize this is because the explore dataset consists of low-quality, noisy trajectories. Using shorter segments allows SCoTS to be more selective, finding and connecting the temporally-extended parts of trajectories.

2. Sensitivity to Number of Retrieved Segments $K$

Performance is robust as long as $K$ is not too small. A very small $K$ limits the diversity of candidate segments, hindering the effectiveness of our progress and novelty-based selection. While performance is stable for larger $K$, we expect an excessively large $K$ could eventually degrade performance by increasing the chance of retrieving dynamically inconsistent segments.

3. Sensitivity to Novelty Weight $\beta$

We respectfully refer the reviewer to our detailed response to second question.

4. Necessity of Core Components

They are discussed in the rebuttal for reviewer RD6p due to the space limitation.

“We are happy to see the results on other offline RL benchmarks” and “How does reward affect the agent in this method?”

While our paper focuses on tasks where extrinsic rewards are absent, the SCoTS framework is not limited to goal-conditioned tasks.

To clarify, reward-free refers to our stitching mechanism, which is guided by the intrinsic temporal structure rather than extrinsic task rewards. This allows the core data augmentation to be task-agnostic. To demonstrate its applicability to general, dense-reward offline RL, we conducted additional experiments on four standard MuJoCo locomotion benchmarks from D4RL, utilizing suboptimal medium and medium-replay datasets.

For a fair comparison, we used the same Diffuser architecture and planning horizon (32) as described in the original paper. We upsample the original data to 2M samples using the same SCoTS procedure, with sub-trajectory lengths set to 4. Additionally, to label these new trajectories with rewards, we additionally trained a reward model, $r_t = f_{\omega}(s_t,a_t)$, on the original offline dataset. The table below shows the normalized average returns over 50 planning seeds.

The results show that training on SCoTS-augmented data consistently improves the performance of the original Diffuser. Even without architectural or sampling strategy changes, by stitching sub-trajectories in a way that is temporally extended and exploratory, SCoTS generates novel behaviors (e.g., trajectories that travel further) that are not present in the original suboptimal dataset.

References:

[1] Kaelbling et al., . Learning to achieve goals. IJCAI, 1993.

[2] Wang et al., Optimal goal-reaching reinforcement learning via quasimetric learning. ICML, 2023.

[3] Park et al., Foundation Policies with Hilbert Representations. ICML, 2024.

Thank you for the detailed review and constructive feedback on our work. We especially appreciate your insightful question about the necessity of balancing directional progress with exploratory novelty in our trajectory stitching method, which is particularly helpful for clarifying our contribution. Please find our detailed responses below.

• “Contradiction between our Fig.1 and Diffuser original paper's Fig.3.b?”

Thanks for asking this valid question. The success shown in the original Diffuser paper's Figure 3.b was on a relatively simple toy experiments. In such settings, standard U-Net-based diffusion models can perform well. However, this architectural property becomes a critical limitation in the long-horizon and high-dimensional OOD settings we target. We hypothesize this is because the aggressive downsampling from pooling and strided convolutions reduces the bottleneck representation so drastically that even small kernels can access distant sequence information, a property that is helpful in simple settings but detrimental in complex ones where temporal details are lost.

This limitation has been recognized issue in the community, as acknowledged in prior works [2, 3]. For example, the appendix of [3] presents a similar analysis and empirically illustrates this failure mode (Figure 4 in Appendix). Therefore, we argue there is no contradiction. The original work demonstrated success in a constrained setting, while our work and Figure 1 address the failure of these models to generalize to more challenging scenarios.

• “Why balance directional progress with novelty?”

As illustrated in our paper's Figure 4, using only the progress score ($\beta=0$), which quantifies the alignment with a latent exploration direction, is effective at generating long, temporally-extended trajectories with clear directional distinctions.

However, relying solely on this directional guidance can be suboptimal. The learned temporal distance-preserving embedding $\phi$ is a powerful but imperfect approximation of the true environment topology. Due to embedding errors, the geometry of the latent space can be distorted. Consequently, a straight line between two latent states, $\phi(s)$ and $\phi(s')$, may not correspond to a feasible path in the underlying MDP. Therefore, exploratory detours are often necessary to find a valid temporally extended path.

This is precisely where the novelty score becomes critical. It encourages these exploratory detours by rewarding the selection of segments that lead to less-visited states. This allows the agent to navigate around the imperfections in the latent space, discovering more temporally extended paths.

To demonstrate this complementary effect quantitatively, we conducted an additional ablation study on 2 OGBench tasks. The table below compares the performance of the HD planner trained with SCoTS data generated using different novelty weights $\beta$ (5 seeds).

The results suggest that while directional stitching alone ($\beta=0$) provides a notable performance improvement, introducing and balancing it with the novelty score ($\beta>0$) yields substantial further gains. This effect is especially pronounced on the more complex Giant task.

• Applicability of SCoTS to general offline RL.

While our paper focuses on tasks where extrinsic rewards are absent, the SCoTS framework is not limited to goal-conditioned tasks.

To clarify, reward-free refers to our stitching mechanism, which is guided by the intrinsic temporal structure rather than extrinsic task rewards. This allows the core data augmentation to be task-agnostic. To demonstrate its applicability to general, dense-reward offline RL, we conducted additional experiments on four standard MuJoCo locomotion benchmarks from D4RL, utilizing suboptimal medium and medium-replay datasets.

For a fair comparison, we used the same Diffuser architecture and planning horizon (32) as described in the original paper. We upsample the original data to 2M samples using the same SCoTS procedure, with sub-trajectory lengths set to 4. Additionally, to label these new trajectories with rewards, we additionally trained a reward model, $r_t = f_{\omega}(s_t,a_t)$, on the original offline dataset. The table below shows the normalized average returns over 50 planning seeds.

The results show that training on SCoTS-augmented data significantly and consistently improves the performance of the original Diffuser. Even without architectural or sampling strategy changes, by stitching sub-trajectories in a way that is temporally extended and exploratory, SCoTS generates novel behaviors (e.g., trajectories that travel further) that are not present in the original suboptimal dataset.

• “Additional experiments to demonstrate the necessity of each component.”

We appreciate the reviewer's valuable suggestion. We have conducted additional ablation studies to empirically demonstrate the necessity of each key component of SCoTS.

To isolate the contribution of our learned latent space, we compare the performance of the HD planner trained on SCoTS-augmented data where the stitching process was guided by either our temporal distance-preserving embedding or raw state space.

The results clearly show that the temporal distance-preserving representation is crucial for effective stitching. This is especially true in high-dimensional environments like antmaze, where a small Euclidean distance in the raw state space (e.g., between two similar joint configurations) does not guarantee reachability. Even in pointmaze, where state-space distance is more intuitive, the learned compact latent space provides a better-structured representation for stitching temporally extended trajectories.

To demonstrate the importance of the refinement step, we compare the performance of both a diffusion planner (HD) and a GCRL algorithm (HIQL) trained on SCoTS data generated with and without the diffusion-based refinement.

These results reveal the critical importance of the refinement component. Without refinement, the connection points between stitched segments can suffer from large dynamic inconsistency errors, as illustrated in Figure 7 of our appendix. This performance degradation is particularly significant when training GCRL algorithms like HIQL, which are highly sensitive to the dynamic validity of the training transitions.

References:

[1] Janner et al., Planning with Diffusion for Flexible Behavior Synthesis. ICML, 2022.

[2] Chen et al., Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion. NeurIPS, 2024.

[3] Chen et al., Simple Hierarchical Planning with Diffusion. ICLR, 2024.

Thank you for the detailed review and constructive feedback on our work. We especially appreciate your insightful question comparing the benefits of our trajectory stitching approach against subgoal decomposition methods, which is particularly helpful for clarifying our contribution. Please find our detailed responses below.

• “Does SCoTS really bring substantial benefits compared to methods like MCTD [1] and HD [2]?”

Thank you for this important question. We would like to clarify the fundamental difference between our approach and prior methods such as MCTD [1] and HD [2].

The fundamental limitation of these methods is that their planning horizon is inherently constrained by the maximum length of trajectories seen during training. Consequently, they typically struggle with tasks requiring extensive stitching, such as those in the OGBench Stitch dataset, where individual trajectories are limited to 200 steps but successful task completion requires trajectories of 800+ steps.

In contrast, SCoTS systematically generates diverse, extended trajectories, effectively breaking the horizon limitation of existing planners. To demonstrate this empirically, we compare SCoTS with MCTD [1] and HD [2] on PointMaze environments with two dataset types:

• Navigate: Contains trajectories long enough to directly solve test tasks.
• Stitch: Contains short trajectories (max 4 cells), requiring extensive stitching.

As shown, methods like HD and MCTD struggle significantly when extensive trajectory stitching is required. In contrast, HD trained with SCoTS-augmented data achieves consistent success, demonstrating our method's clear benefit.

• “Are the candidate trajectory segments sampled from the same original trajectories, or are they drawn across different trajectories?”

To clarify, the candidate trajectory segments $\boldsymbol{\tau}_j$ sampled are drawn from across all different trajectories within the entire offline dataset $\mathcal{D}$, not from the same original trajectory.

This is explicitly handled by our search mechanism described in Section 3.2 and Equation (8): $\text{TopKNeighbors}(\phi(\text{end}(\boldsymbol{\tau}_{\text{comp}})),\phi(\mathcal{D}),k).$

Here, the search space $\phi(\mathcal{D})$ consists of the latent embeddings of the initial states of all trajectories in the offline dataset. Our method finds the nearest neighbors to the current endpoint $\text{end}(\boldsymbol{\tau}_{\text{comp}})$ from this global pool of candidates.

This cross-trajectory stitching is fundamental to effectiveness of SCoTS. By connecting segments from entirely different source trajectories, we can synthesize novel behavioral sequences that were never demonstrated together in the original data. This enables the generation of diverse, extended trajectories that explore previously unseen regions of the trajectory space, significantly enhancing the diffusion planner's generalization capabilities beyond the original data distribution.

We appreciate the reviewer highlighting this ambiguity and will clarify this point carefully in the final version of the paper.

• “How directional and exploratory trajectory stitching would help planning.”

The rationale behind our stitching strategy is inspired by recent successes in unsupervised reinforcement learning [3, 4, 5], demonstrating that state-covering behaviors significantly benefit downstream tasks.

Our work addresses a critical limitation in offline RL: datasets are often composed of disconnected, short-horizon, or narrow-distribution trajectories. Consequently, the strategy for how to stitch these trajectories becomes a crucial factor. For example, a naive stitching approach, such as randomly connecting segments with nearby endpoints, would inevitably generate cycles or static behaviors confined to local regions, failing to provide meaningful state coverage for downstream planning.

In contrast, SCoTS generates trajectories that are both temporally extended and diverse through a principled mechanism. We achieve temporal extension by guiding the stitching process along consistent directions within a learned temporal distance-preserving latent space. Diversity is ensured by assigning a unique, randomly sampled latent direction to guide the generation of each augmented trajectory. This principled strategy produces augmented datasets containing meaningful state-covering trajectories beneficial for training both diffusion planners and general GCRL methods.

• In the first paragraph of Section 3.2, there seems to be no difference between the unit vector $\boldsymbol{z}$ and the vector itself in $\boldsymbol{z}\gets \boldsymbol{z}/\|\boldsymbol{z}\|$.

Thank you for pointing out this. The expression $\boldsymbol{z}\sim\mathcal{N}(\mathbf{0},\mathbf{I})\,\boldsymbol{z}\gets \boldsymbol{z}/\|\boldsymbol{z}\|$ is intended to describe a two-step procedure: 1) a vector is first sampled from a standard Gaussian distribution, and 2) this vector is then normalized to become a unit vector. To enhance clarity in the final version, we will revise the text to explicitly separate these two steps,

• The state encoder $\phi(s)$ should take the state as input, but in Eq. (8), $\phi(\mathcal{D})$ appears to take the dataset $\mathcal{D}$ as input.

$\phi(\mathcal{D})$ is a concise representation for the set of latent embeddings of the initial states of all trajectories within the dataset $\mathcal{D}$. This set forms the search space for candidate segments. We will clarify this shorthand in the final version to improve readability.

References:

[1] Yoon et al., Monte Carlo Tree Diffusion for System 2 Planning. ICML, 2025.

[2] Chen et al., Simple Hierarchical Planning with Diffusion. ICLR, 2024.

[3] Park et al., Lipschitz-constrained Unsupervised Skill Discovery. ICLR, 2022.

[4] Park et al., METRA: Scalable Unsupervised RL with Metric-Aware Abstraction. ICLR, 2024.

[5] Bae et al., TLDR: Unsupervised Goal-Conditioned RL via Temporal Distance-Aware Representations. CoRL, 2024.

Dear Area Chair,

Thank you for your insightful question and for carefully reviewing our rebuttal. We appreciate you raising this important point regarding the generalizability of our trajectory augmentation method.

We believe the temporal-aware embedding employed by SCoTS is robust and generalizable to other domains. Specifically, even in high-dimensional state spaces, such as AntMaze (29 dimensions), leanred temporal-aware embedding successfully captures relevant temporal and spatial aspects beyond simple x-y coordinates. For instance, in AntMaze, the learned embedding effectively distinguishes states not only by spatial position but also by orientation and posture.

Nevertheless, we acknowledge that tasks involving precise object manipulation, such as controlling the cube in OGBench, may require additional considerations. For example, even minor discrepancies in the cube's position can imply large temporal distances, significantly limiting the availability of suitable stitching candidates. To handle this, we could consider a hindsight-inspired candidate generation step: when searching for candidate segments, we can artificially adjust cube states from other trajectories that fall within a small epsilon-distance from the current query trajectory's cube state. Specifically, we would relabel these nearby states to exactly match the cube position of the current query trajectory, effectively enlarging the candidate set and facilitating more robust trajectory stitching.

In summary, while the locomotion tasks provide a clear and intuitive validation, the principles of SCoTS are general. We believe that with a principled modification, such as the proposed hindsight-inspired state matching, the framework can be effectively extended to object-centric domains, which we leave for future work.

Best regards,

The Authors