Generative Trajectory Stitching through Diffusion Composition

We thank Reviewer qv69 for the positive and careful review. Please see our response to the listed concerns below.

W1-1. The current experimental evaluation is still only based on relatively limited simulation problems. It would be interesting to see examples of challenging planning problems where the method still fails.

Thank you for your valuable comments.

In our experimental setup, we include tasks of varying difficulty levels — from relatively simple point maze navigation tasks to very challenging tasks such as AntSoccer.

For example, in the PointMaze-Large-Stitch task, our method achieves a 100% success rate. In contrast, in the challenging AntSoccer-Medium-Stitch task, although our method outperforms all baselines, the absolute success rate remains relatively low (17%), highlighting significant room for future improvement.

W1-2. A real world implementation in the robotics context would also be interesting to see in future work.

Thank you for the valuable suggestions.

The focus of our work is to propose a core algorithm for trajectory stitching based on compositional diffusion models. Empirically, we have conducted extensive experiments on various simulation environments (e.g., different task objectives, data quality, different embodiments and environment sizes) and have shown the efficacy of the proposed method.

Since the proposed method makes no privileged knowledge about the simulation environments or tasks, it is applicable to a wide-range of real-world trajectory planning problems, which will be our future research focus.

Additionally, during the rebuttal phase, we conduct an experiment on an OGBench Cube task, 'cube-single-play-v0'. Since this dataset is of play data type and not designed for trajectry stitching, we segment each long play-style trajectory in the dataset to multiple short segments and train CompDiffuser on these short segments.

Specifically, the dataset contains 1000 trajectories of length 1000 and we train CompDiffuser only on short segments of length 48. In test time, since a plan of length 48 is not sufficient to complete a task, we compose three models and use the sampling scheme in Algorithm 2 to sample a plan of length 112 (i.e., $\tau_{\text{comp}}$). Similarly, we train an inverse dynamic model to output actions, such that the plan can be followed by the robot.

Please see the results below. Results of the baselines are from the OGbench paper.

Q1. Classical motion planning techniques in robotics such as probabilistic roadmaps/rapidly exploring random trees construct graphs or trees with nodes and connecting trajectory segments between sets of states. Could the proposed stitching techniques potentially be extended to construct not just a single multi-segment path, but a complete graph of segments to solve motion planning problems between sets of source/target locations as well?

Thank you for the insightful comments.

We agree that extending the proposed method to solve motion planning problems between sets of source/target locations is a promising research direction, and we will explore along this line in our future work.

Our current method generates a single long-horizon trajectory by compositionally denoise each segment $\tau_{1:N}$. These segments, namely $\tau_{1:N}$, are sequentially arranged in a line to form a continuous multi-segment path.

To extend this to a graph of segments, we can first organize segments $\tau_{1:M}$ in a graph like structure (each segment can be viewed as an edge and $M$ is the number of segments in the graph). We can then denoise the whole graph in a similar fashion to Algorithm 2. Specifically, if the segment starts from a source location, we use Line 6 of Algorithm 2; if it is an intermediate segment, we use Line 9; if the segment terminates at a target location, we use Line 12.

Dear Reviewer qv69,

We sincerely appreciate your thoughtful and detailed feedback, as well as your positive evaluation of our manuscript. Thank you once again for your time and effort in reviewing our work.

Best,
Paper Authors

We thank Reviewer AQzW for the positive and careful review. Please see our response to the listed concerns below.

W1: A significant weakness of the paper is the lack of detailed implementation specifics, such as: Architecture details, the way they condition on overlapping parts, Network capacity (total number of parameters needed), any data normalization techniques, any weighting used in the diffusion loss?

Thank you for your valuable comments. In our original manuscript, we described our implementation in Appendix E, including model inputs/outputs, noisy (overlapping) chunk conditioning, and model architectures, and we will make sure to add all the mentioned implementation specifics in our final manuscript.

Please see the additional implementation details below:

Architecture details: For planning in 2D x-y space, we follow Decision Diffuser and use a conditional U-Net as the denoiser network. The main modification we made it to concatenate the encoded latents of the noising neighbor trajectories to the conditional latents for denoising.

For planning in high dimension state space (Table 3 and Table 4), we use a DiT based transformer as the denoiser network. We stack 4 states to construct a transformer input token. (For example, there will be 40 tokens for a trajectory with 160 states.) The latent dimension and depth of the transformer are 768 and 16.

Network capacity (total number of parameters needed): For experiment results for Table 1, we have the parameter count as follows: PointMaze from [1]

OGBench PointMaze

OGBench AntMaze

OGBench HumanoidMaze

Data Normalization: We use the standard min-max normalization to normalize the training data to range (-1, 1).

loss weighting: We do not use any weighting in the diffusion loss.

W2. The usage of the generative skill chaining baseline is unclear; more details, equations, or explanations would be helpful to understand its application in this context.

Generative Skill Baseline (GSC) is a compositional inference method that averages the scores of the overlapping parts between two segments throughout the denoising process. In contrast, CompDiffuser's compositional inference directly conditions the denoising score function on the scores of the neighboring segments.

Mathematically at denoising step $t$, with score of the neighboring segments as $\epsilon_{k-1}$ and $\epsilon_{k+1}$, with an overlapping region length $H_o$, we have the score of the overlapping region as

$\epsilon_k[:H_o] = (\epsilon_{k-1}[-H_o:] + \epsilon_k[:H_o])/2$ $\epsilon_k[-H_o:] = (\epsilon_{k}[-H_o:] + \epsilon_{k+1}[:H_o])/2$

Both GSC and CompDiffuser can enable cross trajectory information exchange during denoising, hence promoting stitching capability. GSC enables such information exchange via simply replacing the score of the overlapping part between two segments by the average score of the corresponding parts from the two segments; while in CompDiffuser, we propose a more unified stitching framework based on the proposed autoregressive sampling method and bi-directional information passing (Algorithm 2).

In Appendix F of our original manuscript, we described how we apply GSC in our benchmark and we will make sure to include the detailed introduction above in our updated manuscript.

W3. Trajectory merging might not be suitable for cases where two states are separated by obstacles (e.g., walls). Clarification on this aspect is necessary.

We do not use trajectory merging to avoid obstacles, and CompDiffuser as a planner is tasked to sample valid trajectories that satisfy the obstacle contraints.

Specifically, we propose CompDiffuser which stitches trajectories by generating a set of feasible and consecutive segments $\tau_{1:N}$. Then, trajectory merging (detailed in Appendix E.2) is only used to merge these compositionally-generated segments $\tau_{1:N}$ and form a long-horizon trajectory $\tau_{comp}$ for the agent to execute.

As in our experiments (e.g., Table 1 and Figure 10/11), we also find that CompDiffuser is able to generate states that travel in valid areas and hence can result in a successful $\tau_{comp}$.

Additionally, in some cases when the planner does fail (two states are separated by obstacles), other techniques can also be employed to refine the generated trajectory, such as perform replanning or leverage some sampling-based planning method to refine the local collision region.

W4. the empirical investigation primarily focuses on navigation tasks and planning with xy positions, leaving questions about its applicability to other domains, such as manipulation or locomotion. Could you provide results in cube domains in OGbench for example (single and double) ?

Thank you for the insightful suggestions.

In this paper, we aim to study the trajectory stitching problem and hence, in our original manuscript, we focus on evaluating our method on stitching-specific benchmarks, such as various maze stitching tasks and ant soccer stitching tasks in OGBench, along with multiple stitching datasets from [1].

During the rebuttal period, due to the time and computation constraints, we conduct additional experiments on the cube-single-play-v0 env. Since this play env is not designed for trajectory stitching, we segment each long play-style trajectory in the dataset to multiple short segments, and train CompDiffuser on these short segments. That is, the dataset contains 1000 trajectories of length 1000 and we train CompDiffuser only on short segments of length 48. In test time, since a plan of length 48 is not sufficient to complete a task, the planner must stitch multiple trajectories to construct a successful plan. Specifically, we compose three models and use the sampling scheme in Algorithm 2 to sample a plan of length 112 (i.e., $\tau_{\text{comp}}$). Similarly, we train an inverse dynamic model to output actions, such that the plan can be followed by the robot.

Please see the results below. Results of the baselines are from the OGBench paper.

W5. Why in most of the results except in antmaze, the authors ignore the qvel part of the states, how bad the prediction would be if we predict the full state instead only the positions?

The primary objective of a trajectory stitching algorithm is to find sequence of subgoals that are reachable from the start state and allows us to progress towards the goal state. Thus, only the xy or xy with joint states is sufficient to determine a feasible subgoal.

In our original submission, we presented an ablation study on the prediction dimension in Table 4 where we compare the success rates of predicting xy (2D), xy with joint states (17D), and xy with joint states and joint velocities (full state, 29D). Quantitatively, we observe that the performance loss is minor in smaller mazes such as Medium and Large, but we do see a drop in success rate in the most challenging Giant maze, likely due to modeling complexity of the full state planning.

For further reference, we conduct additional experiments on OGBench pointMaze-stitch environments, where CompDiffuser predicts full state, both x-y positions and velocities (4D). Please see the success rate comparison below.

For the low-dimensional setup of pointmaze-giant-stitch-v0, we can see comparable performance when predicting full state (qpos+qvel) as compared to xy (qpos) positions only. This is in contrast with the findings from the experiment with antmaze-giant-stitch-v0 environment. We speculate that the decreased results on the antmaze is due to the insufficient neural network capacity (higher modeling complexity to plan in a larger state space) and limited training data (while planning in antmaze is more difficult, the size of the antmaze dataset is identical to the size of the pointmaze dataset).

W6. The method seems to rely on high quality data, expert demonstration. when trained on navigate data, the method performance decreases.

In our experiments, we evaluated CompDiffuser on various data qualities and the results demonstrated that CompDiffuser is robust to data quality.

As in Table 2, our method is able to outperform all the baselines on extremely low-quality datasets: the OGBench Explore datasets. (The data quality of these datasets is much lower than the OGBench navigate datasets.) The collection policy for the Explore datasets will randomly re-sample a new moving direction after every 10 steps and contain a large amount of action noise, making it very challenging to learn an effective policy (See Figure 9 for qualitative dataset examples).

Our method can also surpass baselines when the data are of relatively higher quality, for example, as in the OGBench Stitch datasets (see Table 1 and 2). These datasets are of relatively higher qualitiy than Explore, where the data collection policy is a low-level directional policy trained via SAC with a high-level waypoint controller.

[1] Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View.

We thank Reviewer a2a3 for the positive and careful review. Please see our response to each listed concern below.

Baselines:

W1&W2. The inference time in appendix B.3 would indicate the overall model for the monolithic case is smaller, making the comparison invalid. Continuing on this same topic, the parameter count / model size is not really well presented, This raises the question whether the authors simply train much bigger networks than their competitors and this is why they get such a big advantage. In Table 12 there are some hyperparameters reported, but it is not possible to understand overall parameter count from that, and it is not possible to understand how it compares to DD.

Thank you for your valuable comments. Please see our response to the model sizes and inference time below. In our final manuscript, we also will make sure to include all the mentioned implementation details.

Model Size. For the result shown in Table 1, the denoiser model of DD is indeed smaller than ours. However, in our implementation, we already matched the architecture, depth, and feature dimension of the main denoiser network for DD and ours. The additional parameters come from the neighbor trajectory encoding networks (which encode $\tau_{k-1}$ and $\tau_{k+1}$ to latents) and MLPs for handling the resulting condition features of higher dimension (since these two neighbor trajectory latents are also used to condition the denoising network).

Please see the parameter count of the networks below:

PointMaze from [1]

OGBench PointMaze (networks used for reporting the sampling time in Table on Page 20)

For our method, we use the same main denoiser network across these mazes. However, since the trajectory lengths are different for mazes of different sizes in PointMaze[1], the size of the encoder networks for neighbor trajectories (i.e. $\tau_{k-1}$, $\tau_{k+1}$) also varies, resulting in a minor difference in model sizes for our method. We will make sure to enclose detailed model sizes in our updated manuscript accordingly.

Inference Time. We agree that given a similar network, the parallel denoising scheme should yield a similar inference time as DD. We have also discussed the reasons for the time difference in Appendix B.3 (Line 664-671) in our original submission. We summarize the reasons below: (1) DD's denoiser network is indeed smaller; (2) Our method requires encoding two noisy trajectories into latent representations, which are then used to condition the main denoiser network. This encoding step cannot be performed in parallel with the forward pass of the denoiser, resulting in additional latency. (3) Parallel sampling increases the inference batch size. In the parallel sampling scheme, all trajectories $\tau_{1:N}$ are stacked to one batch. And as described in Appendix E.1 (Line 822-831), we generate $B$ (set to 10) samples in a batch and we compute the L2-distance of each overlapping parts to select the best plan (the one with the smallest distance). As a result, the batch size increases to $K \times B$, which will increase sampling time. (4) The overhead of trajectory merging. The construction of the final plan involves merging multiple trajectory segments, which might introduce further overhead.

Additionally, we further increase the parameters of DD and conduct additional experiments for fairer comparison. To better reflect the planning performance, we show the results of DD and Ours without replanning.

In the table below, we present results on PointMaze from [1], where we increase the base feature dimension from 96 to 128, such that DD's network is in similar size as ours.

In the table below, we present results on OGBench PointMaze Stitch, where we increase the base feature dimension from 128 to 160, such that DD's network is in similar size as ours.

In these datasets, while the two DD models achieve similar performance on Large and Giant mazes, the smaller DD model even yield better performance than larger model in the simple Medium maze. We deem that this is likely due to overfitting in the larger model.

[1] Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View.

W3. In the checklist, authors state that their experiments are reproducible since they "describe it in Appendix E and they release their code on acceptance". In appendix E, the authors only describe the setup partially (only for pointmaze giant), and refer the reader to "refer to our codebase for details", which is not provided, making it impossible to judge any of these details. The authors should have an appendix with experimental settings for each experiment, and list total number of hyperparameters for each model, as well as how it compares to baselines.

Thank you for the valuable suggestions. We will make sure to include an appendix section with complete experimental settings for each experiment in our updated manuscript. We also promise to release the code for reproducibility. (Due to the character limits, we are not able to present all the detailed experimental settings in the response.)

The hyperparameters of our method is actually very similar across different experiments. In fact, the hyperparameters (e.g., training dynamics, model architectures) is identical for OGBench PointMaze and AntMaze as shown in Table 12.

Since the trajectories in the humanoidmaze dataset is much longer (humanoidmaze 400 vs antmaze 200), we update several hyperparameters to fit the longer trajectories. Please see the table below:

W4. Something that is not made clear in the paper is whether baselines are trained on snippets as well, and then required to produce full trajectories at inference time? If this is the case that also seems like a comparison of limited use, since those methods were not developed to perform well on these tasks. Can the authors comment?

Yes. The baselines are also trained on the snippets (which are short trajectories with respect to the horizon of the test-time task). For each experiment, we use the same training dataset for all methods.

Many of our baselines do process trajectory stitching ability. For example, RL methods can inherit the stitching ability through their formulation such as value propagation; while GSC is specifically designed for stitching tasks. While monolithic model (Diffuser[1], Decision Diffuser[2]) is not specifically designed for stitching, it is also shown to exhibit the ability of length generalization, potentially due to the generalizability of diffusion model itself and the neural network (convolutional U-Net). Our goal to add DD as a baseline is to show that our compositional sampling method is able to offer significant advantage on stitching and length generalization compared to a monolithic model method.

Experiments:

W5. The experimental evaluation is limited to maze-like RL benchmarks, especially given the baselines such as GSC, it would be ineresting to see performance on such robot manipulation tasks as well.

Thank you for the valuable comments. We agree that extending the proposed method to robot manipulation tasks is an interesting direction. During the rebuttal period, we also conduct additional experiments on the OGBench cube-single-play-v0 environments and validate that our method can also work well on robot manipulation tasks. Please see the results below.

Q1. Are the sampled "clean trajectories" in Algorithm 1 full trajectories or short segments? If they are short segments, how are these connected together with their neighbors?

In Line 3 of Algorithm 1, we wrote 'sample clean trajectory from dataset'. The 'clean trajectory' here refer to a trajectory without any added noise (hence in diffusion noise level 0). This 'clean trajectory' can be considered as a 'full trajectory' in the dataset. However, to effectively evaluate stitching ability, the dataset only contains 'short' trajectories ('short' here is relative to the horizon of the test-time tasks).

Then we can obtain multiple overlapping short segments from the 'full trajectory' from the dataset using Line 4 in Algorithm 1. Specifically, to obtain the neighbors for training, we further divide the sampled short segment $\tau_0$ to $K$ overlapped pieces. Since these $K$ pieces are from one consecutive trajectory, we can obtain valid data to train the model. For clarity, we will update our manuscript regarding the training data accordingly.

The authors have adequately answered my questions and concerns, but I will maintain my score.

Dear Reviewer a2a3,

We sincerely appreciate your thoughtful and detailed feedback, as well as your positive evaluation of our manuscript. Thank you once again for your time and effort in reviewing our work.

Best,
Paper Authors

We thank Reviewer yhYf for the positive and careful review. Please see our response to each listed concern below.

W1. It would be constructive to include ablations of at least points (2) and (3), where (2) is testing a version of CompDiffuser where the chunks don’t overlap and (3) is testing a version of CompDiffuser that conditions only on the previous trajectory.

Thank you for your valuable suggestions.

We include the ablation studies of inference with single-way conditioning below. As shown, the proposed bi-directional information propagation design outperforms the single-way conditioning counterpart that conditions only on the previous trajectory. The improvement becomes more prominent as the size of the maze increases, demonstrating the effectiveness of our method.

Due to limited time, we are unable to provide the numbers for the no-overlapping design by the deadline of the rebuttal -- we are running them now and will make sure they are in the final manuscript.

W2. It seems to me the paper would be strengthened by including a comparison between CompDiffuser and Diffusion Forcing in maze planning settings.

Thank you for the valuable suggestions.

We are working on adding this baseline, but due to limited time and compute, we are unable to provide the numbers for the no-overlapping design by the deadline of the rebuttal. We will make sure to include this baseline in the final manuscript.

Q1. Do the authors have a sense of, in the setting of Table 1, for how large a horizon does CompDiffuser’s performance start degrading, if any?

Given that CompDiffuser achieves 100% in Large maze and 68% in Giant Maze, we speculate that CompDiffuser's success rate might start degrading if the planning horizon is between the horizon of the Large maze and Giant maze.

Specifically, we compose 8 models to obtain a plan of horizon 888 for PointMaze-Giant-Stitch and 5 models for a plan of horizon 572 to solve PointMaze-Large-Stitch. Hence, for this specific setup, the CompDiffuser's performance might start degrading when composing 6-7 models, i.e., a horizon around 700.

Q2. do the authors have a sense of how they would build a CompDiffuser to achieve 100% accuracy in the Giant maze? Is it a matter of simply collecting more data? Or is it potentially a fundamental limitation of the method due to accumulating errors throughout the autoregressive sampling process at longer horizons?

Thank you for your comments. We agree that the error accumulation when composing a large number of trajectories could negatively impact the performance. However, we believe that there are still multiple ways to enhance the performance of CompDiffuser.

• We can leverage a plan verifier to select a promising plan out of a batch of candidate plans.
• Inference-time searching [1] can be further integrated in the sampling process to enhance the performance.
• More training data can also enhance the performance, because it can provide more intersection/overlapping between different trajectories, promoting cross-trajectory stitching. (The size of the Giant maze dataset is identical to the size of the datasets of Large maze and Medium Maze. However, since the Giant maze is much larger, its data density is actually smaller compared to the Large/Medium maze datasets.)
• Increased replanning budget. Allowing more replanning steps may help the agent recover from failures, such as when getting stuck by walls.

Please also see the response to Q3 below for additional strategies.

Q3. what are the main variables with respect to which the generalization ability of the method scales?

Thank you for your insightful question. We believe that the following factors effect the generalization ability of the model:

• Amount/Quality of data: We hypothesize that the quality of domain coverage and existence of overlaps between training trajectories will improve the performance of CompDiffuser, since higher state-space coverage and more intersections between trajectories in the dataset can promote cross-trajectories combinatorial generalization (hence encouraging stiching ability).
• Ratio of training and eval trajectory length: Generally, we see that the performance starts to drop with increasing number of the composed trajectories. Hence, we believe that the performance of CompDiffuser can increase if it has access to longer training trajectories, as it will require less number of test-time composed trajectories.
• Computation Budget: In our current setup, we use DDIM 50 steps to sample trajectories and we found that increasing the inference denoising timestep does not further improve the performance. However, we believe that techniques such as (1) inference-time searching technique [1] can be incorporated during the sampling process and (2) using an additional plan verifier to select a more reliable path will be helpful.
• Replanning budget: Higher replanning budge can also enhance the performance of CompDiffuser. For instance, replanning can help the agent escape from suboptimal cases during a rollout (e.g., when the agent is stuck in a corner or collides into walls).

Q4. Do the authors expect the method to be robust to sampling direction in more general settings, even e.g. for problems that are more naturally solved by “working back from the goal position”?

In the given example, we agree that sampling "backward" would be a more efficient sampling direction since it would be easier to find the intermediate node when starting from $t$ due to the high probability density (only one edge connecting to $t$). Though our method might still be able to solve such asymmetric case via its bi-directional information passing design (as in Figure 3), we believe that manually or leveraging an additional model to decide the denoising direction according to the task structure would be beneficial as this might reduce the planning search space.

Q5. Do the authors expect there to be any major methodological hurdles in generalizing CompDiffuser to settings requiring processing of e.g. visual inputs? (or, more generally, non-ground-truth-state observations, POMDPs, etc).

Thank you for the valuable comments. As our method do not assume any specific knowledge about the tasks, it is generally applicable to any trajectory planning settings.

For instance, if the starting observation is given in images, our method can use the latents of the images to condition the diffusion denoiser network to generate a trajectory of states. In addition, our method can also directly plan in the image-space (i.e., synthesize a sequence of images instead of states). A simple implementation is to use some existing pre-trained models such as Cosmos or DINO to encode the images into latents and run CompDiffuser on these latents.

In the case of POMDPs, due to uncertainty about the underlying ground truth state of the environment, there may not be one correct ground truth plan we can follow to solve our desired task. In this setting, we can leverage the probabilistic nature of CompDiffuser and sample multiple candidates trajectories, and then chose the sequence of actions that maximizes the likelihood of satisfying the set of the trajectories.

Q6. Complementarity between Diffusion Forcing and the proposed method.

Thank you for your insightful question! We believe that our method can be combined with the ideas in Diffusion Forcing to solve more complex long-horizon tasks. In such long-horizon tasks, it may not make sense to construct a fully detailed plan to solve the entire task, as such a plan can be computationally intractable to generate. We can instead adopt the idea of Diffusion Forcing, and learn each generative component with variable levels of noise. This then enables us to construct plans where individual portions of trajectories are either fully realized or abstracted, depending on the varied noise levels of each component, allowing us to more effectively generate long horizon plans

Towards solving very complex, long-horizon problems -- we think that there are still open questions with respect to the right abstraction to plan over tasks (for instance, directly planning in continuous actions to make trip to Europe would not make sense), as well as how to form plans when the overall state of the world is unknown (for instance, making a plan to make dinner when you do not know the locations of pots or food).

Q7. Do the authors see a natural extension of their method to non-goal-directed settings, instead driven by e.g. reward functions or other success metrics?

We believe that extending our method to non-goal-directed settings is an interesting future research topic. For example, a simple way to adapt our method to cost or reward directed planning is that when we are compositionally sampling $N$ segments, instead of using a goal state to guide the denoising of the final segment $\tau_N$, we can use some specific cost function or reward value as guidance. As we are not doing a final-goal directed planning, such guidance can also be flexibly applied to any intermediate segment $\tau_n$ to achieve other specific behavior.

Q8. Would one expect CompDiffuser to inherit similar global guidance properties from full-trajectory diffusion methods like Diffuser or Decision Diffuser?

We believe that CompDiffuser is able to inherit the global guidance properties from full-trajectory diffusion methods.

Though we are sampling multiple segments $\tau_{1:N}$ concurrently as in Figure 3, each segment actually depends on the values of other segments. Hence, the sampling of the segment set $\tau_{1:N}$ can be viewed as a holistic sampling process, which enables the global guidance properties.

[1] Inference-time Scaling of Diffusion Models through Classical Search