Generative Trajectory Stitching through Diffusion Composition
Authors: Yunhao Luo
Utkarsh Aashu Mishra
Yilun Du
Danfei Xu
Venue: NeurIPS 2025 spotlight ()
Keywords: Trajectory Stitching
Diffusion Models
Compositionality
Planning
Primary area: reinforcement_learning
TL;DR: We illustrate how stitch trajectories together with generative models.

Abstract: Effective trajectory stitching for long-horizon planning is a significant challenge in robotic decision-making. While diffusion models have shown promise in planning, they are limited to solving tasks similar to those seen in their training data. We propose CompDiffuser, a novel generative approach that can solve new tasks by learning to compositionally stitch together shorter trajectory chunks from previously seen tasks. Our key insight is modeling the trajectory distribution by subdividing it into overlapping chunks and learning their conditional relationships through a single bidirectional diffusion model. This allows information to propagate between segments during generation, ensuring physically consistent connections. We conduct experiments on benchmark tasks of various difficulties, covering different environment sizes, agent state dimension, trajectory types, training data quality, and show that CompDiffuser significantly outperforms existing methods.

Generative Trajectory Stitching through Diffusion Composition
Yunhao Luo'
Utkarsh A. Mishra
Yilun Dua Danfei Xu lif
Georgia Tech Harvard University
Abstract
Effective trajectory stitching for long-horizon planning is a significant challenge robotic decision-making. While diffusion models have shown promise in planning, they are limited to solving tasks similar to those seen in their training data. We propose CompDiffuser, novel generative approach that can solve new tasks by learning to compositionally stitch together shorter trajectory chunks from pre- viously seen tasks. Our key insight modeling the trajectory distribution by subdividing it into overlapping chunks and learning their conditional relation- ships through single bidirectional diffusion model. This allows information propagate between segments during generation, ensuring physically consistent connections. We conduct experiments on benchmark tasks of various difficulties, covering different environment sizes, agent state dimension, trajectory types, train- ing data quality, and show that CompDiffuser significantly outperforms existing methods, Project website at http://coom--diffuerrgihhiii
1 Introduction
Generative models have demonstrated remarkable capabilities in modeling complex distributions across domains like images, videos, and shapes. robot planning, these models offer a promising approach by modeling distributions over plan sequences, which allows amortizing the computational cost of traditional search and optimization methods. This effectively transforms planning into sampling likely solutions given start and goal conditions. Recent works like Diffuser [26] and Decision Diffuser [1] have shown how diffusion models can learn to generate entire plans for long- horizon robotics tasks. However, exhaustively modeling joint distributions over entire plan sequences all possible start and goal states remains extremely inefeeifffffiient, as requires collecting long-horizon plan data covering all possible combinations of initial states and goals.
The concept of trajectory stitching [88] from Reinforcement Learning literature [62] presents potential solution by combining chunks of different trajectories to create new, potentially better policies. The methods work by identifying high-reward trajectory chunks and stitching them together at states where they overlap or are similar enough, creating composite trajectories that can inform better policy learning. This effectively enables compositional generalization since collecting long consecutive trajectories is costly, and these short chunks can be flexibly assembled to complete new tasks. The key challenge lies finding appropriate stitching points where trajectories can be combined while maintaining dynamic consistency and feasibility. Our goal is to enable generative planners to solve long-horizon tasks without requiring their ability to generate physically feasible, goal- -irected plans.
training data, while retaining
We propose a novel diffusion-based approach, Compositional Diffuser (CompDiffuser), that enables effective trajectory stitching through goal-conditioned causal trajectory generation.
Equual advising.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Our key insight is that we can model the trajec- tory distribution compositionally by subdividing it into distributions of overlapping chunks and learning their conditional relationships. Rather than learning separate models for each chunk, we train a single diffusion model that can gen-
STEC TC
erate trajectory chunks conditioned on neighbor- ing chunks' states (Figure 1). This allows infor- mation to propagate bidirectionally during the reverse diffusion process as illustrated in Fig- ure 2: each chunk's generation is influenced by both past and future chunks. This architec- ture naturally enables both parallel generation of chunks and causal autoregressive generation, each with different trade-offs in computational cost and planning quality.
Monolatbic Model (Prior Methods) Compositional (Our Method)
Model
Figure I: Compositional Trajectory Generation. Com- pDiffuser enables generative trajectory stitching through diffusion composition. Left: Monolithic generative plan- ner fails to generalize to tasks of longer horizon and collapses to the maze center. Right: Our method suc- cessfully navigates the ant agent from start to goal by compositionally stitching together shorter trajectories.
We conduct extensive experiments across bench- mark tasks of varying difficulty levels, including different environment sizes (from simple U-mazes complex giant mazes), agent state dimensions (from 2D point agents to 50D humanoid robots), trajectory types (from maze navigation trajectories to ball dribbling trajectories), and training data quality (from clean demonstrations noisy explo- ration data). Our results demonstrate that CompDiffuser significantly outperforms multiple imitation learning and offline reinforcement learning baselines across all settings. We show that our approach can effectively solve long-horizon tasks while maintaining plan feasibility and goal-reaching behavior. We validate the importance of our key technical components including the bidirectional conditioning mechanism, the autoregressive sampling process, and the flexible replanning capability.
In summary, the key contributions of this work are:
A noisy-sample conditioned diffusion planning framework that enables learning compositional trajectory distributions by decomposing the trajectory generation procedure into sequence of segments each generated by separate diffusion denoising process.
A compositional goal- coul-condiiioned trajectory planning method that uses bidirectional information propagation during denoising to maintain physical consistency between trajectory chunks.
A set of empirical results showing significant improvements over existing methods across multi-
ple trajectory stitching benchmarks, with detailed analysis of model capabilities and limitations.
2
Related Work
Diffusion Models For Planning. Many works have studied the applications of diffusion models [58,
24] for generative planning [26, 1, 54, 72, 22, 64, 42, 86, 7J. Diffusion planning has been widely applied various fields, such as motion planning [5, 43], procedure planning [67], task planning [76, 16], autonomous driving [75, 37, 68], reasoning [79], and reward learning [49]. Many techniques have also been combined with diffusion planning, including hierarchical planning [35, 8, 45], self-evolving
planner [36], preference alignment [10], tree branch-pruning [17], refining [31], replanning [84],
uncertainty-awar planning [60], equivariance [4]. However, these works are usually constrained to
plan within similar horizons as training data. Our work instead proposes a compositional diffusion planning approach that generalizes to much longer horizon via generative trajectory stitching.
Trajectory Stitching. A flurry of works have explored the trajectory stitching problem given offline data. One typical category of solution is based on data augmentation or goal relabeling along with various techniques, such as generative models [33, 41, 28, 30, 34, 40], model based approaches [6, 32, 85, 23], and clustering [21]. Other supervised learning based methods, such as sequence modeling [9, 25, 73, 87, 71], latent space learning [80, 66, 83], and learning with dynamic programming [18, 27, 74], have also demonstrated some extent of stitching capability. In this work, we propose different trajectory stitching approach based on generative modeling, where the model only learns from plain short trajectory segments while is able to directly perform goal-conditioned trajectory stitching through test-time compositional generation.

---

Figure 2: Illustrating the Trajectory Stitching Process. Given an unseen start (blue circle) and goal (green star), CompDiffuser generates a long-bo oug-horizon plan by progressively denoising three trajectory chunks in parallel, with each chunk conditioning its neighbors ensure smooth transitions.
Compositional Generative Models. Compositional Generative Models [13, 11, 19, 12, 46, 3, 50, 63] are widely studied in various domains, including visual content generation [39, 11, 77, 82, 59, 57], human motion generation [56, 61, 81], traffic generation [38], robotic planning [78, 2, 433, and policy learning [69, 53]. Most existing work on compositionality focuses on sampling under the conjunction of several given conditions. While several studies [48, 47] have explored sequential compositional models, they are restricted planning within given skeleton and hence unable to generalize to longer sequences or new tasks. In contrast, we propose a compositional planning framework that scales to generating much longer sequences and completing new tasks via directly stitching short trajectory segments, without relying on pre-defined task-dependent skeletons.
3 Planning through Compositional Trajectory Generation
We aim to develop a generative planning framework that can generate long-horizon trajectories by composing multiple modular generative models. Our method, Compositional Diffuser (CompDif- fuser), trains a single diffusion model on short-horizon trajectories. At inference time, given a start and goal, CompDiffuser runs parallel instances of this model to generate sequence of overlapping trajectory segments, coordinating their denoising processes to ensure they smoothly connect into a coherent long-horizon plan. This approach enables stitch together short-horizon subsequences of training trajectories to form novel long-horizon trajectory plans.
3.1 Compositional Trajectory Modeling
Given planning problem, consisting of a start state qs and goal state gg. we formulate planning as sampling a trajectory т from the probability distribution
[s 1:T 2221 ~ po(T|qe: 9g).
(1)
where ş I:T corresponds to future states to reach the goal state 9 and a I.T corresponds to set of future actions. To implement this sampling procedure, prior work [1, 26] learns generative model P(T) directly over previous trajectories in the environment. However, since the generative model is trained model the density of previously seen trajectories, is restricted generating plans with
start and goal that are similar to those seen in the past.
In this paper, propose to model the generative model over trajectories po(rlgs, 9g) composition- ally [12], where we subdivide trajectory 7 into a set K overlapping sub-chunks т (Figure 1). We then represent the trajectory distribution
po pT(Tlss: pi (Tilgs, T2) pT(TR|TK- 4n) II (T Tk41).
In the above expression, each trajectory chunk is only dependent nearby trajectory chunks Th-1 and Tk+1- This allows pe (rlge: 9,) to generate trajectory plans that significantly depart from previously seen trajectories, long intermediate trajectory chunks have been seen. Overall, the goal of our method is to enable long-horizon planning without long-horizon training data.
3.2 Training Compositional Trajectory Models
One approach to represent the composed probability distribution in Equation 2 is to directly learn separate generative models to represent each conditional probability distribution. However, sampling from the composed distribution is challenging, each individual trajectory chunk depends on the

---

Parallel Sampling x P(22 | (g)
Autoregressive Sampling (
p 1 12)
prra (2.
p 1% (2)
p(t (g,Ta)
p(rs (2:9g) xe
<
r2
x
r2
Denoising
12
z
ri
r
Cross Info Exchange
Denoising Flow
Start State qs
Goal State Qg
Figure 3: Compositional Trajectory Planning: Parallel Sampling and Autoregressive Sampling. We present an illustrative examplc of sampling three trajectories 71:3 with the proposed compositional sampling methods. Dashed lines represent cross trajectory information exchange between adjacent trajectories and black lines represent the denoising flow of each trajectory. In parallel sampling. T1.3 can be denoised concurrenty;; while in autoregressive sampling, denoising тı depends on the previous trajectory Th-1. e.g., the denoising of 72 depends on (as shown in the blue horizontal dashed arrows). Additionally, start state q, and goal state 4, conditioning applied to the trajectories in the two ends, and ra, which enables goal-conditioned planning. Trajectories Tha will be merged to form longer plan Teamp after the full diffusion denoising process.
values of neighboring chunks. As a result, to sample from the composed distribution, one would need a blocked Gibbs sampling procedure, where the value of each individual trajectory chunk is iteratively sampled given decoded values of neighboring trajectory chunks. This sampling procedure slow, and passing information across chunks to form consistent plan is challenging.
Information Propagation Through Noisy-Sample Conditioning, We propose a more efficient approach that addresses these challenges by leveraging the progressive denoising process of diffusion models. The key challenge in composing trajectories is ensuring feasible transitions between each pair of neighboring chunks, i.e., maintaining physical constraints and dynamic consistency connection points where segments overlap. Our key insight is that we can achieve this by having trajectory segments guide each other's generation during the diffusion process: one segment takes shape through denoising, it helps shape its neighbors into compatible configurations.
We implement this insight using diffusion model that generates trajectory chunks conditioning on their neighbors' noisy samples, Given a dataset D of trajectories we train denoising network eæ to learn the trajectory distribution pa (T Th- 1, Tk++1) i) with the training objective
Cabr D..D [||e 6o (T t T/-11 11+1)
where k identifies trajectory segment, is the noise level, and o represents segment k corrupted with noise level t. Crucially, when denoising each segment, the network conditions on noisy versions of neighboring segments 1 T at the same noise level. This allows each segment to influence its neighbors' denoising process, ensuring their final configurations are dynamically compatible. In addition, further training the network to condition on - can enable autoregressive compositional sampling, which we will discuss in Section 3.3. In practice we only need to condition on the small overlapping regions between consecutive trajectories, making the generation process efficient while maintaining consistency across connection points.
Representing Initial States and Goals. In addition, we further train the same denoising network to represent the distributions (T (T T2) and Po 11.1:99)).. This corresponds to the training objective
Caan E
[|le flle-eo(Th) ( qs, =
with an analogous objective for the conditioned goal state 4g- We train the same denoising network eo with both conditioning. Please see Appendix E for implementation details. We provide an overview of our proposed training strategy in Algorithm 1.
3.3 Compositional Trajectory Planning
Our compositional framework enables flexible sampling strategies for generating long-horizon plans with Equation 2. The basic sampling process starts by initializing each trajectory chunk тı with Gaussian noise. Then, through iterative denoising, each chunk is denoised while being conditioned

---

Algorithm 1 Training CompDifluser
Algorithm 2 Autoregressive Trajectory Sampling
Require: train dataset D, diffusion denoiser model
Models: trained diffusion denoiser model
co(r",t & st_cond, aad_cond), number of training steps N, diffusion timestep T
co(T, st_ st__ood, goo.d.o.ooo... g_cod))
2: Input: start state ge goal state dp. number composed trajectories K
2: for i= i=1- ++1+++0 N do
3: 70 D, sample clean trajectory from dataset
3: Initialize K trajectories TI-K ~ (0,
4: - + divide 7 to K overlapping chunks
4: for f= do
Demoise 71 conditioned q, and r
5: TiK f add noise to -P ~~ T, following DDPM
6: rf - _eo(Tf.& | 9s Tf)
6: Objective for noisy chunk condition
7: Denoise intermediate trajectories to 7k-1
7: Sake l|e-<e(TÉ_- 1/ Tk+ 12, k [2, K 1]
8: for == K do
8: Objective start end state condition
9: Cau || <0( qe. 12 1
-
10: end for
10: Lend = l
- )))l 2 = [[11]
11: Denoiise тk conditioned on rk and do
11: Cau = Lahe + Lean + fend
12: rh =o(kk + xk
12: Backprop to update es (.) using Call
13: end for
13 end for
14: Teump = Merge the denoised trajectories Ti.
14: return co()
15: return Toomp
on its neighbors. This structure allows for different ways of coordinating the denoising process across trajectory chunks, each offering different tradeoffs between information propagation and computational efficiency. We present two sampling schemes (illustrated in Figure 3):
Parallel Sampling. Our first sampling approach conditions denoising on the values of the noisy adjacent trajectory chunks from the previous denoising timestep, where the update rule is
T
--1 = a' - se(TETTk-I, + B*E), {~N(0,1) 1),
(5)
where a' and 3 are diffusion specific hyperparameters. This approach allows us to run denoising on each trajectory chunk in parallel, each denoising update only requires the values of the adjacent trajectory chunks at a previous noise level. However, information propagation between the values of adjacent trajectory chunks is limited each denoising timestep, as each trajectory chunk is denoised independently of the denoising updates of other trajectory chunks.
Autoregressive Sampling. To better couple the values of adjacent trajectory chunks, propose to denoise each trajectory chunk autoregressively dependent on the values of neighboring chunks at each denoising timestep. In particular, we iteratively denoise each trajectory rfzk starting from the ri, and condition the denoising of r the previously decoded chunk - the current noise level t and the future chunk r/41 the previous noise level t, giving us the sampling equation
k -1-ad
1-1 --1= a' a'( co c T + B)E), {~N(0,11) 1).
(6)
This sequential generation process enables stronger coordination among the chunks since each chunk is conditioned on the less noisy version of its previous chunk. However, it requires generating chunks one at time rather than simultaneously, making it computationally less efficient than parallel sampling. We compare the two sampling schemes in Table 6, where empirically find autoregressive sampling leads to improved performance. Additionally, we provide sampling time comparison in Table 10. We use this autoregressive sampling procedure throughout the experiments in the paper and illustrate pseudocode for sampling in Algorithm 2. Given such final set of generated chunks Ti:K, we then merge the chunks together to construct a final trajectory Toump by applying exponential trajectory blending to areas where subchunks тę overlap (See Appendix E.2 for implementation details).
Experiments
In this section, objective is to (1) validate that our method enforces coherent trajectory stitching on multiple benchmarks, with varying state space dimensions, task design, and training data collection policies (2) understand how planning with higher state dimensions, varying numbers of composed trajectories, different sampling schemes, and replanning affect the performance of the proposed method. Additional results are provided in Appendix and C with failure analysis in Appendix D.
Baselines. We compare CompDiffuser with three categories of existing methods: (1) for genera- tive planning methods, we include Decision Diffuser (DD) [1] for monolithic trajectory sampling
5

---

Env
Size
RvS
RvS (SA)
RvS (GA)
DT
DT (SA)
DT (GA)
DD
GSC
Ours
U-Maze
17=7
9745
76=5
17:5
65:4
5414
Ozo
100
100
PoimtMaze [21]
Medium
1+2
55+3
21=3
20+2
55+3
62+2
30+1
93+1
100=a
Large
314
3815
31_5
22.12
3512
3915
01.
9912
100_0
Eny
Type
Size
GCBC
GCIVI
GCIQL
QRL
CRL
HIQL
GSC
Ours
Medium
23
70
21 49
80 +-1
0.1
74 46
100
100
PointMaze
[51]
stitch
Large
7
12 :6
31 #2
84 115
#0
13 16
100=0
100
Giant
0 +0
+0
+0
50 488 +0
+0
29+3
68-3
Table 1: Quantitative Results on PointMaze Stitching Datasets in Ghugare et al. [21] and OGBench [51]. We compare CompDiffuser to baselines of multiple categories, including diffusion. data augmentation, and offline reinforcement learning. SA and GA stand for state augmentation and goal augmentation respectively, as described in Ghugare et [21]. Our results are averaged over seeds and standard deviations are shown after the + sign.
and Generative Skill Chaining (GSC) [48] for compositional sampling; (2) for data augmentation based methods, we include stitching specific data augmentation [21] with RvS [14] and Decision Transformer (DT) [9]; (3) for offline reinforcement learning methods, we include goal-conditioned behavioral cloning (GCBC) [44, 20]. goal-conditioned implicit V-learning (GCIVL) and Q-learnin (GCIQL) [29]. Quasimetric RL (QRL) [70]. Contrastive RL (CRL) [15], and Hierarchical implicit Q-learning (HIQL) [52]. See Appendix F for more details of baselines.
Evaluation Setup. For each environment, we report the success rate over all evaluation episodes. where the success criterion is that the agent or target object is close to the goal within small threshold. We evaluate all methods with random seeds for each experiment and report the mean and standard deviation. Specifically. in Ghugare et al. [21] datasets, we evaluate on 2 tasks in U-Maze, 6 tasks in Medium, and 7 tasks in Large with 10 episodes per task; OGBench [51], we evaluate on 5 tasks in each environment with 20 episodes per task. Each task is introduced in the respective papers and is defined by a base start and goal state that require trajectory stitching to complete. A random noise is added to the base start and goal state for each evaluation episode.
4.1 PointMaze
We present experiment results on two types of trajectory-stitchig datasets in point maze environments, which features different dataset collection strategies. Our method is trained on short trajectories of x-y positions of the point agent and is able to directly generate much longer trajectories from start and goal by composing multiple trajectories (detailed numbers of the composed trajectories in each experiment are provided in Appendix Table 13). Note that our method only requires training one model and we can use the proposed compositional inference method to construct coherent long-horizon trajectories.
PointMaze [21]. Following the original framework, the training data are curated by dividing each environment (here, maze) into several small regions, and feasible trajectories constrained within their respective regions are sampled. Possible stitching between segments is facilitated by a small overlap (one block) between different regions, which can used to join trajectories across regions. More dataset details are provided in Appendix A. We present the quantitative results on these datasets in Table 1, where we compare to coul-conditioned behavior cloning methods trained with data augmentation, Decision Diffuser, and GSC. Most baselines perform suboptimally, likely because they are unable to autonomously identify the small overlapping regions needed for trajectory stitching. In contrast, CompDiffusee successfully resolves all tasks across various maze sizes, demonstrating its ability to stitch trajectories even when the connecting regions are small.
PointMaze in OGBench [51]. In this particular setup, each trajectory in the training dataset is constrained to navigate no more than four blocks in the environment. The start and goal of each trajectory can be sampled from the entire environment provided that the travel distance between the start and goal is within four blocks. These trajectories are much shorter than the ones required for a feasible plan between a given start and goal at inference. More details about these datasets are provided in Appendix A.2. We present the quantitative results Table 1, where we compared our method to GSC and multiple offline RL baselines following the benchmark established in [51]. We observe that while GSC is able to perform on par with CompDiffuser in Medi um and Large mazes, struggles as the planning horizon further increases, illustrated in the qualitative results in Figure 4.

---

Env
Type
Size
GCBC
GCIVL
GCIQL
QRL
CRL
HIQL
GSC
Ours
Medium
45 =ll
44 is
29
59
53 16
94
9722
9612
stitch
Large
3
18
18
11 +2
67
66#2
86+2
antmaze
Giant
=0
0
40
21 =2
2011
6513
Medium Large
2
19
13
1 11
37 =10
90
81+2
explore
100
10 1:3
0
2113
27:1
Medium
29
12 1:2
12
18
71
96
92:1
9111
humanoid
stitch
Large
4:1
1:0
31 =3
70
72.3
maze
Giant
0 =0
0 =0:
0 40
#0
0
10
12
12
5+1
67:4
Table 2: Quantitative Results on AntMaze and HumanoidMaze in OGBench, We benchmark method on the 5 test-time tasks defined in OGBench with 20 episodes per task. Our results averaged over 5 seeds and standard deviations are shown after the + sign.
DD
GSC
Ours
Zoom-in Figure 5: Qualitative Results of Planning in High Dimension on OGBench AntMaze Large. Original plan sub-sampled for clearer view. Our method is able synthesize plans of valid dynam- ics while reaching the goal position (bottom right). Note that our method is only trained on trajectory segments of much shorter length.
Top-down Plan Overview
Figure 4: Qualitative Comparison of DD, GSC and Comp- Diffuser on OGBench PointMaze Giant. Effective bidi- rection information propagation enables CompDiffuser successfully synthesize trajectories from (bottom left) to goal (upper right), while other methods generate o.o.d tra- jectories that are disconnected with the start/goal or passing through walls. See Figure 10 for per-segment visualization.
CompDiffuser significantly outperforms all baselines in Giant maze, demonstrating its efficacy in complex environments. Additional results are provided in Appendix C.I, C.2, and C.3.
4.2 High Dimension Tasks
We present the evaluation results of our method on various trajectory stitching tasks involving higher-dimensional state spaces within OGBench: AntMaze, HumanoidMaze, and AntSoccer.
AntMaze and HumanoidMazee We conduct maze navigation experiments of multiple agents, ant and humanoid, using the pre-collected Stitch datasets provided in OGBench. The data collection strategy identical to OGBench PointMaze, where each episode constrained to travel at most blocks, while at inference, successful plan requires the agent travel up 30 blocks. We com- pare our method with the offline RL benchmark established in [51] along with the best-performing liffusion-based compositional stitching baseline GSC, shown in Table 2. Both GSC and CompDif- fuser generate plans in a planar x-y space while the agent follows the plans with a learned inverse dynamics model. We observe a pattern similar to the results in PointMaze. where CompDiffuser can consistently give high success rates as the planning horizon and complexity increase while other baselines start collapse. To complete the study, also conduct experiments where the full agent state is used for planning instead of the x-y space and provide the results Section 4.3.
AntMaze with Low-Quality Data. We evaluate CompDiffuser OGBench AntMaze with different data collection strategy, Explore. These datasets consist of extremely low-quality yet high-coverage data, where the data collection policy contains large amount of action noise and will randomly re-sample a new moving direction after every 10 steps (see Figure 9 in Appendix for qualitative examples). Hence, each demonstration episode typically oscillates within only 2-3 blocks. Our planner needs to learn from these clustered trajectories to construct coherent plans that reach goals in large spatial distances. We present the success rate of each method in Table 2.
AntSoccer. The AntSoccer environment OGBench requires the ant agent to move a soccer ball to designated goal in the environment, different from maze tasks that require the agent itself to

---

GSC
Ours
GSC
Ours
Env
Size
GCBC
GCIVL
GCIQL
QRL
CRL
HIQ
(4D)
(4D)
(17D)
(17D)
antsoccer
arena
34 #4
21 =3
5 +2
=1 2 +1
23
41+4
55
6543
69=3
stitch
medium
2 11
1 10
00
0.
0 10
8
12
512
1311
12.12
17.1
Table 3: Quantitative Results on OGBench AntSoccer Stitch. We evaluate two generative planners with different planning state dimensions: 4D planner that operates the x-y positions of the ant and the ball, and a 17D planner that additionally generates the 13 joint positions of the ant agent.
Figure 6: Qualitative Plan generated by CompDiffuser in OGBench AntSoccer Medium Stitch. The initial position of the ant is shown by the blue circle and the goal is individual trajectory Ti. separately (as shown from left to right) and mark the ball with a yellow circle for clearer view. These generated trajectories will be merged to form a long-horizon plan Teamp- Note that our model is only trained on two different types of short trajectories: 1) the ant moves without dribbling the hall; 2) the ant moves while dribbling the ball. At test time, the proposed compositional sampling method can stitch these two types of trajectories and construct end-to-end plans of longer horizon to solve the more difficult task first navigate to the ball from the side and then dribble ball to the goal position.
the ball to the pink circle. We plot each
reach the goal. OGBench provides two categories of trajectories in AntSoccer Stitch datasets: (1) navigates in the maze without the ball and (2) ant moves and dribbles the ball in the maze. The planner must stitch trajectories of these two skills to complete the goal-reaching objective, because the ant must reach the ball first and then dribble it to the given goal location (see Figure 6 and Figure 14).
We present experimental results in two distinct maze configurations, Arena and Medium, in Table relative to the benchmarking provided in OGBench. Specifically, as an ablation study, we evaluate with two planner state space configurations: (1) a 4D planner consisting of the x-y location of the ant and the ball: (2) a 17D planner consisting of the x-y location of the ant and the ball along with all the joint positions of the ant. We observe that our compositional method outperforms all the baselines in both configurations. Notably, the 17D variant demonstrates slightly higher success rates, likely due to the ability of joint positions to provide more fine-grained information for ball dribbling.
4.3 Ablation Studies
Planning in High Dimension Space. We report experiment results where CompDiffuser synthesizes trajectorics in state space of higher dimension. We compare the success rates of our method planning in dimension of 2D, 15D, and 29D in Table 4, and present qualitative plans in Figure 5 and Figure 13. Specifically, the planner operates in the x-y space for 2D, x-y with the joint positions for 15D, x-y with both the joint positions and velocities for 29D. Similar to Section 4.2, train MLP inverse dynamics model that takes in the current observation and a goal of the planner dimension to predict an action for the agent execute.
Our method performs consistently and achieves near-optimal success rates across all planning dimensions in AntMaze Medium. In AntMaze Large and Giant, the success rates decrease when planning in 15D and 29D, which is probably due to the increasing complexity of the trajectory modeling, since the joint positions and velocities of moving ant are highly dynamic and the tasks require planning hundreds of consecutive future states to reach the goal.
Different Numbers of Composed Trajectories. We study the effect of varying the numbers of trajectories to be composed K. To better study the planning performance with respect to K, we use the challenging PointMaze-On MaatMMaz--Gintt Ciant-Stiton in OGBench as the testbed. As shown in Figure 7, our method obtains consistent performance when composing 7 to 12 trajectories. We observe that the optimal K is around 9 and 10, which also corresponds to natural path length to reach the goal. Qualitatively, decreasing K will result sparser trajectory while increasing K will cause the final trajectory traveling back and forth to consume the redundant states (See Figure 10).
8

---

Ours
Ours
Ours
PointMaze Giant Stitch
AniMaze Giant Stitch
Size
HIQL
GSC CompDifliser
(2D)
(15D)
(29D)
70
Medium
94+1
96=2
9540
9712
60
Large
6745
8612
6645
662s
50
Giant
2142
65+3
41+3
28+4
40
Table Quantitative Results of Different Planning Dimensions on AntMaze Stitch Datasets in OGBench. Our method compo- sitionally constructs feasible plans that reach long-distance goals while modeling complex vironment dynamics, such as the positions and velocities of the agent's joints.
of Composed Trajectories
Inverse Dyanmies Model Horizon
Figure 7: Quantitative Results of Different Setup. Left:
success rate versus different numbers of composed trajecto- ries in OGBench PointMaze Giant (w/o replan). Right: success rate versus different inverse dynamics model horizons in OGBench AntMaze Ciant.
Ours w/o
Ours w/
Env
Size
QRL
HIQL
Env
Size
Replan
Parallel
AR
Replan
Replan
Medium
80a12
74=6
100
100
Large /
100
100=0
PointMaze
Large
84+15
13-6
100
100
PointMaze
Giant
x
45#1
53+6
Giant
50+s
O40
53=6
68
Giant
V
6611
6813
Medium
59
9411
9212
9612
Large
84.5
8612
AntMaze
Large
18+2
67=5
76=2
8622
AntMaze
Giant
X
18=4
27+4
Giant
0=o
21-2
27:4
65+33
Giant
V
48
6543
Table 5: Quantitative Results of CompDiffuser with and without Replanning. We report the success rates on OGBench PointMaze and AntMaze Stitch datasets, For w/o replan, CompDiffuser only synthesizes one tra- jectory and executes the trajectory in close-loop man- ner; for w/ replan, CompDiffuser will synthesize new trajectory if the agent loses track the current plan.
Table 6: Quantitative Comparison of two Sampling Schemes: Parallel and Autoregressive (AR). Paral- lel sampling performs on par with AR sampling in the casier Large mazc, while AR sampling can generate trajectorics higher quality when constructing longer plans (i.c., composing more trajectories), likely due to its causal denoising strategy.
Replanning with CompDiffuser. Our method can also flexibly replan during rollout, which enables the agent to recover from failure, such as when the agent fails to track the planned trajectory due to sub-optimal inverse dynamic actions. In practice, we replan if the distance between the agent's current observation and the synthesized subgoal is larger than a threshold. Please see Appendix E.3 for details. In Table 5, we present ablation studies of CompDiffuser with and without replanning in PointMaze and AntMaze Stitch datasets in OGBench. CompDiffuser outperforms the best performing baselines QRL and HIQL even without replanning that w/ and w/o replanning yield similar performance in maze size Medium um and Large, while offering significant performance boost in the more complex Giant maze.
out of tasks. We also observe
Parallel vs. Autoregressive Sampling. We compare the performance of the two proposed com- positional sampling schemes, parallel and autoregressive, as presented in Table 6. Autoregressive sampling consistently outperforms the parallel sampling across various tasks in plan quality, showing that the causal information flow, where each trajectory chunk is conditioned (less noisy) version of previous chunk, leads to more coherent and physically consistent transitions between chunks. We include additional discussion and sampling time comparison in Appendix B.3.
the already-denoised
5 Discussion and Conclusion
Limitations. Stitching short trajectories to solve for unseen longer- plans is challenging problem. While our method excels in our empirical evaluation, there still remain a number of future venues of research. When composing large sequence of trajectories, the error accumulation in the long chain of bidirectional information propagation might lead to infeasible plans. This issue can be potentially mitigated by using domain/task specific rejection sampling techniques to select from multiple candidate plans or leveraging more expensive MCMC sampling. Moreover, the optimal number of the test-time composed chunks K is task-dependent. Our future research will explore ways to automatically identify suitable number chunks K, potentially by incrementally increasing the number of chunks dependent on the quality of the generated plan.

---

Conclusion. We introduce CompDiffuser, a generative trajectory stitching method that leverages the compositionality of diffusion models. We introduce condiiiond score function formulation that helps in performing autoregressive sampling of multiple short-horizon trajectory diffusion models and eventually stitching them to form a longer-horizon goal-conditioned trajectory. Our method demonstrates effective trajectory stitching capabilities as evident from the extensive experiments on tasks of various difficulty, including different environment sizes, planning state dimensions, trajectory types, and training data quality.
Acknowledgments
We would like to thank Zilai Zeng for helpful early discussion trajectory stitching and feedback of the manuscript.
References
[1] Ajay Anurag, Du Yilun, Gupta Abhi, Tenenbaum Joshua B, Jaakkola Tommi S, Agrawal Pulkit, Is Conditional Generative Modeling all you need for Decision Making? The Eleventh International Conference on Learning Representations. 2022. 1, 2, 3, 5, 33,
[2] Ajay Anarag, Han Seungwook, Du Yilun, Li Shuang, Gupta Abhi, Jaakkola Tommi S., Tenenbaum Joshua B., Kaelbling Leslie Pack, Srivastava Akash, Agrawal Pulkit. Compositional Founda- tion Models for Hicrarchical Planning Thirty-seventh Conference on Neural Information Processing Systems. 2023.
[3] Bradley Arwen, Nakkiran Preetum, Berthelot David, Thornton James, Susskind Joshua M. Mechanisms of Projective Composition of Diffusion Models // arXiv preprint arXiv:2502.00549. 2025.
[4] Brehmer Johann, Bose Joey, De Haan Pim, Cohen Taco S. EDGI: Equivariant diffusion for planning with embodied agents // Advances in Neural Information Processing Systems. 2023. 36. 63818 8-88844. 2
[5] Carvalho Joao, Le An Baierl Mark, Koert Dorothea, Peters Jan. Motion planning diffusion: Learning and planning of robot motions with diffusion models // 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). 2023. 1916-1923.
[6] Char Ian, Mehta Viraj. Villaftor Adam, Dolan John M. Schneider Jeff. Bats: Best action trajectory stitching arXiv preprint arXiv:2204.11 12026. 2022.
[7] Chen Boyuan, Martí Monsó Diego, Du Yilun, Simchowitz Max, Tedrake Russ, Sitzmann Vincent. Diffusion foreing: Next-token prediction meets full-sequence diffusion // Advances Neural Information Processing Systems. 2024. 37. 24081-24125, 2
[8] Chen Chang, Deng Fei, Kawaguchi Kenji, Gulcehre Caglar, Ahn Sungjin. Simple hierarchical planning with diffusion // arXiv preprint N0V:24001026444 2024. 2
[9] Chen Lili, Lu Kevin, Rajeswaran Aravind, Lee Kimin, Grover Aditya, Laskin Misha, Abbeel Pieter, Srinivas Aravind, Mordatch Igor. Decision transformer: Reinforcement learning via sequence modeling // Advances neural information processing systems. 2021. 34. 15084 15097. 2, 6, 35
[10] Dong Zibin, Yuan Yifu, HAO Jianye, Ni Fei, Mu Yao, ZHENG YAN, Hu Yujing. Lv Tangjie, Fan Changjie. Hu Zhipeng. AlignDiff: Aligning Diverse Human Preferences via Behavior- Customisable Diffusion Model // The Twelfth International Conference on Learning Represen- tations. 2024.3 2
[11] Yilun, Durkan Conor, Strudel Robin, Tenenbaum Joshua B, Dieleman Sander, Fergus Rob, Juscha, Doucet Arnaud, Grathwohl Will Sussman. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme International conference machine learning. 2023, 8489-8510.
10

---

[12] Du Yilun, Kaelbling Leslie. Compositional generative modeling: A single model is not all you need arXiv preprint arXiv:2402.01103. 2024.
[13] Du Yilun, Li Shuang, Mordatch Igor. Compositional visual generation with energy based models
|/ Advances in Neural Information Processing Systems. 2020. 33, 6637-6647.
[14] Emmons Scott, Eysenbach Benjamin, Kostrikov Ilya, Levine Sergey. Rvs: What is essential for offline via supervised learning? // arXiv preprint arXiv:2 2021. 6, 35
[15] Eysenbach Benjamin, Zhang Tianjun, Levine Sergey, Salakhutdmov Russ R. Contrastive learning goal-conditioned reinforcement learning // Advances in Neural Information Processing Systems. 2022. 35. 35603-35620. 6, 35
[16] Fang Xiaolin, Garreli Caelan Reed, Eppner Clemens, Lozano-Pérez Tomás, Kaelbling Leslie Puck, Fox Dieter. Dimsam: Diffusion models samplers for task and motion planning under partial observability // 2024 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). 2024. 1412-1419.
[17] Feng Lang, Gu Pengjie, Bo, Pan Gang. Resisting Stochastic Risks Diffusion Planners with the Trajectory Aggregation Tree // Forty-first International Conference on Machine Learning. 2024.7
[18] Gao Chen-Xiao, Wu Chenyang, Cao Mingjun, Kong Rui, Zhang Zongzhang, Yu Yang. Act: Empowering decision transformer with dynamic programming via advantage conditioning / Proceedings of the AAAI Conference on Artificial Intelligence. 38, 11. 2024. 12127-12135.
[19] Garipov Timur, De Peuter Sebastiaan, Yang Ge, Garg Vikas, Kaski Samuel, Jaakkola Tommi. Compositional sculpting of iterative generative processes // arXiv preprint arXiv:2309.16115. 2023.
[20] Ghosh Dibya, Gupta Abhishek, Reddy Ashwin, Fu Justin, Devin Coline, Eysenbach Benjamin, Levine Sergey. Learning to reach goals via iterated supervised learning // arXiv preprint arXiv:1912.06088. 2019. 6, 35
[21] Ghugare Raj, Geist Matthieu, Berseth Glen, Eysenbach Benjamin. Closing the Gap between TD Learning and Supervised Learning-A A Generalisation Point of View. // The Twelfth International Conference on Learning Representations. 2024. 2, 6, 23, 24, 34, 35
[22] He Haoran, Bai Chenjia. Xu Kang. Yang Zhuoran, Zhang Weinan Wang Dong, Zhao Bin, Xuelong. Diffusion model an effective planner and data synthesizer for multi-task reinforce- ment learning // Advances in neural information processing systems. 2023. 36. 64896-64917.
2
[23] Hepburn Charles Alexander, Montana Giovanni. Model-based Trajectory Stitching for Improved Offline Reinforcement Learning // 3rd Offline RI. Workshop: Offline RI. as a "Launchpad"" 2022.
[24] Ho Jonathan, Jain Ajay, Abbeel Pieter. Denoising diffusion probabilistic models // Advances in neural information processing systems. 2020. 33. 6840-6851. 2
[25] Huang Xingshuai, Wu Di, Boulet Benoit. DRDT3: Diffusion-Refiei Decision Test-Time Training Model // arXiv preprint arXiv:2501.061 2025.
[26] Janner Michael, Du Yilun, Tenenbaum Joshua, Levine Sergey. Planning with Diffussonn for Flexible Behavior Synthesis // International Conference on Machine Learning. 2022. 9902-9915.
1, 2,
[27] Kim Jeonghye, Lee Suyoung. Kim Woojun, Sung Youngchul. Adaptive Q-Aid for Conditional Supervised Learning in Offline Reinforcement Learning // Advances in Neural Information Processing Systems. 2024. 37. 87104-87135.7
[28] Kim Sungyoon, Choi Yunseon, Matsanaga Duiki E, Kim Kee-Eung, Stitching Sub-trajectories with Conditional Diffusion Model for Goal-Conditioned Offline RI Proceedings of the AAAI Conference Artificial Intelligence. 38, 12. 2024. 13160-13167.

---

[29] Kostrikov Ilya, Nair Ashvin, Levine Sergey. Offline reinforcement learning with implicit q-learning // arXiv preprint 1rXiv:2110.0616..
2021. 3,
[30] Lee Jaewoo, Yun Sujin, Yun Taeyoung, Park Jinkyoo. GTA: Generative Trajectory Augmentation with Guidance for Offline Reinforcement Learning // The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.
[31] Lee Kyowoon, Kim Seongun, Choi Jaesik. Refining diffusion planner for reliable behavior syn- thesis by automatic detection of infeasible plans // Advances in Neural Information Processing Systems. 2024. 36. 2
[32] Lei Xing, Zhang Xuetao, Wang Donglin. MGDA: Model-based Goal Data Augmentation for Offline Goal-conditioned Weighted Supervised Learning // arvi preprint arXiv:2412.11410. 2024.
[33] Li Guanghe, Shan Yixiang. Zhu Zhengbang, Long Ting, Zhang Weinan. DiffStitch: Boost- ing Offline Reinforcement Learning with Diffusion-based Trajectory Stitching // Forty-first International Conference on Machine Learning. 2024. 2
[34] Li Shangzhe, Zhang Xinhua. Augmenting Offline Reinforcement Learning with State-only Interactions // arXiv preprint arXiv:2402.00807 2024. 2
[35] Li Wenhao, Wang Xiangfeng, Jin Bo, Zha Hongyuan. Hicrarchical diffusion for offline decision making // International Conference on Machine Learning. 2023. 20035-20064. 2
[36] Liung Zhixuan, Mu Yoo, Ding Mingyu, Ni Fei, Tomizuka Masayoshi, Luo Ping. AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners International Conference on Machine Learning. 2023. 10725-207455 2
[37] Liao Bencheng, Chen Shaoyu, Yin Haoran, Jiang Bo, Wang Cheng, Yan Sixu, Zhang Xinbang, Li Xiangyu, Zhang Ying, Zhang Qian, others. DiffusionDrive: Truncated Diffusion Model for End-to-End Autonomous Driving // arXiv preprint arXiv:2411.151: 2024. 2
[38] Lin Haohong, Huang Xin, Phan-Minh Tung, Hayden David S, Zhang Huan, Zhao Ding, Srini- vasa Siddhartha, Wolff Eric M, Chen Hongge. Causal Composition Diffusion Model for Closed-loop Traffic Generation // arXiv preprint M1X1:22412.1790000 2024. 3
[39] Liu Nan, Li Shuang, Du Yilun, Torralba Antonio, Tenenbaun Joshua B. Compositional visual generation with composable diffusion models // European Conference on Computer Vision. 2022. 423. 3
[40] Liu Zhihong, Qian Long, Liu Zeyang, Wan Lipeng, Chen Xingyu, Lan Xuguang. Enhancing Decision Transformer with Diffusion-Bare Based Trajectory Branch Generation // arXiv preprint arXiv:2411.11321 2024. 2
[41] Lu Cong, Ball Philip, Teh Yee Whye, Parker-THolder Jack. Synthetic experience replay // Advances in Neural Information Processing Systems. 2024. 36. 2
[42] Lu Haofei, Han Dongqi, Shen Yifei, Li Dongsheng. What Makes Good Diffusion Planner for Decision Making? // The Thirteenth International Conference on Learning Representations. 2025.
[43] Luo Yunhao, Sun Chen, Tenenbaum Joshua B, Du Yilun. Potential Based Diffusion Motion Planning // International Conference on Machine Learning. 2024. 33486-33510. 2,3
[44] Lynch Corey, Khansari Mohi, Xiao Ted, Kumar Vikash, Tompson Jonathan, Levine Sergey, Sermanet Pierre. Learning latent plans from play // Conference on robot learning. 2020. 1113-1132, 3,
[45] Ma Xiao, Patidar Sumit, Haughton lain, James Stephen. Hierarchical Diffusion Policy for Kinematics-Aware Basti Robotic Manipulation // Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition. 2024. 18081-18090.1 2
[46] Mahajan Divyat, Pezeshki Mohammad, Mitliagkas loannis, Ahuja Kartik, Vincent Pascul. Compositional Risk Minimization // arXiv preprint arXiv:2410.06303. 2024. 3
12

---

[47] Mishra Utkarsh Aashu. Chen Yongxin, Xu Danfei. Generative Factor Chaining: Coordinated Manipulation with Diffusion-basee based Factor Graph // 8th Annual Conference on Robot Learning. 2024. 3
[48] Mishra Utkarsh Aashu, Xue Shangjie, Chen Yongxin, Xu Danfei. Generative Skill Chaining: Long- Horizon Skill Planning with Diffusion Models // 7th Annual Conference on Robot Learning. 2023. 3, 6, 35
[49] Nuti Felipe, Franzmeyer Tim, Henriques João Extracting reward functions from diffusion models // Advances in Neural Information Processing Systems. 2024. 36. 2
[50] Okawa Maya, Lubana Ekdeep S. Dick Robert, Tanaka Hidenori. Compositional abilities emerge multiplicatively: Exploring diffusion models on a synthetic task // Advances in Neural Information Processing Systems. 2024. 36. 3
[51] Park Seohong, Frans Kevin, Eysenbach Benjamin, Levine Sergey. OGBench: Benchmarking Offline Goal-Conditioned RL // International Conference on Learning Representations (ICLR). 2025. 6, 7, 24, 35
[52] Park Seohong, Ghosh Dibya, Eysenbach Benjamin, Levine Sergey. Higl: Offlie goal- conditioned rl with latent states as actions Advances in Neural Information Processing Systems. 2024. 36. 6, 35
[53] Patil Omkar, Sah Anant, Gopalan Nakul. Composing Diffusion Policies for Few-shot Learning of Movement Trajectories // arXiv preprint arXiv:2410.17477 2024. 3
[54] Pearce Tim, Rashid Tabish, Kanervisto Anssi, Bignell Dave, Sun Mingfei, Georgescu Raluca, Macua Sergio Valcarcel, Tan Shan Zheng, Momennejad Ida, Hofmann Katja, others. Imitating human behaviour with diffusion models // The Eleventh International Conference on Learning Representations. 2023. 2
[55] Peebles William, Xie Saining. Scalable diffusion models with transformers // Proceedings of the IEEE/CVF International Conference Computer Vision. 2023. 4195-4205. 33
[56] Shafir Yonatan, Tevet Guy, Kapon Roy, Bermano Amit H. Human motion diffusion as generative prior // arXiv preprint arXiv:2303.01418. 2023.
[57] Skreta Marta, Atanackovic Lazar, Bose Avishek Joey, Tong Alexander, Neklyudov Kirill. The Superposition of Diffusion Models Using the Itô Density Estimator arXiv preprint arXiv:2412.1776.. 2024. 3
[58] Sohl-Dickstein Jascha, Weiss Eric, Maheswaranathan Niru, Ganguli Surya. Deep unsupervised learning using nonequilibrium thermodynamics // International conference on machine learning. 2015.
2
[59] Su Jocelin, Liu Nan, Wang Yanbo, Tenenbaum Joshua Du Yilun. Compositional image decomposition with diffusion models // arXiv preprint arXiv:2406, 19298. 2024. 3
[60] Sun Jiankai, Jiang Yigi. Qiu Jianing, Nobel Parth, Kochenderfer Mykel J, Schwager Mac. Con- formal prediction for uncertainty-awant planning with diffusion dynamics model // Advances in Neural Information Processing Systems. 2024. 36. 2
[61] Sun Shanlin, De Araujo Gabriel, Xu Jiaqi, Zhou Shenghan, Zhang Hanwen, Huang Ziheng, You Chenyu, Xie Xiaohui. CoMA: Compositional Human Motion Generation with Multi-moda. Agents Il arXiv preprint arXiv:2412.077300 2024.
[62] Sutton Richard Barto Andrew G. Reinforcement learning: An introduction. 2018.
[63] Thornton James, Bethune Louis, Zhang Ruixiang, Bradley Arwen, Nakkiran Preetum, Zhai Shuangfei. Composition and Control with Distilled Energy Diffusion Models and Sequential Monte Carlo // arXiv preprint arXiv:2502. 12786. 2025. 3
[64] Ubukata Toshihide, Li Jialong, Tei Kenji. Diffusion model for planning: A systematic literature review // arXiv preprint 11Xiv:2408.10266 2024. 2
13

---

[65] Vaswani Ashish, Shazeer Noam, Parmar Niki, Uszkoreit Jakob. Jones Llion, Gomez Aidan N, Kaiser Lukasz, Polosukhin Illia. Attention is all you need // Advances in neural information processing systems. 2017. 30. 33
[66] Venkatraman Siddarth, Khaitan Shivesh, Akella Ravi Tej, Dolan John, Schneider Jeff. Berseth Glen. Reasoning with latent diffusion offline reinforcement learning // arXiv preprint Siiv:23000065999 2023. 2
[67] Wang Hanlin, Wi Yilu, Guo Sheng. Wang Limin. Pdpp: Projected diffusion for procedure planning in instructional videos // Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023. 14836- 14836-14845. 2
[68] Wang Junming, Zhang Xingyu, Xing Zebin, Gu Songen, Guo Xiaoyang, Hu Yang, Song Ziying, Zhang Qian, Long Xiaoxiao, Yin Wei. HE-Drive: Human-Like End-to-End Driving with Vision Language Models // arXiv preprint arXiv:2410.05051. 2024.
[69] Wang Lirui, Zhao Jialiang, Du Yilun, Adelson Edward H, Tedrake Russ. Poco: Policy com- position from and for heterogeneous robot learning // arXiv preprint arXiv:2402.0251 2024.
[70] Wang Tongzhou, Torralba Antonio, Isola Phillip, Zhang Amy. Optimal goal-reaching reinforce- ment learning via quasimetric learning // International Conference on Machine Learning. 2023. 36411-36430) 6, 35
[71] Wang Yuanfu, Yang Chao, Wen Ying, Liu Yu, Qiao Yu. Critic-guided decision transformer for offline reinforcement learning // Proceedings of the AAAI Conference on Artificial Intelligence. 38, 14. 2024. 15706-15714.2
[72] Wang Zhendong. Hunt Jonathan J. Zhou Mingyuan. Diffusion policies as an expressive policy class for offline reinforcement learning // arXiv preprint arXiv:2208.06119: 2022. 2
[73] Wu Yueh-Hua, Wang Xiaolong. Hamaya Masashi. Elastic decision transformer // Advances in neural information processing systems. 2023. 36. 11832-15550. 2
[74] Yamagata Taku, Khalil Ahmed, Santos-Rodriguee Raul. Q-learning decision transformer: Lever- aging dynamic programming for conditional sequence modelling offline rl // International Conference on Machine Learning. 2023. 38989-39007.7 2
[75] Yang Brian, Su Huangyuan, Gkanatsios Nikolaos, Ke Tsung- Stang-Wei, Jain Ayush, Schneider Jeff. Fragkiadaki Katerina. Diffusion-es: Gradient-fres free planning with diffusion for autonomous driving and zero-shot instruction following // arXiv preprint 240..069990005559 2024. 2
[76] Yang Cheng-Fu, Fu, Xu Haoyang, Wu Te-Lin, Gao Xiaofeng, Chang Kai-Wei, Gao Feng. Planning as In-Painting: A Diffusion- Embodied Task Planning Framework for Environments under Uncertainty // arXiv preprint arXiv:2312.00 12.11997 2023. 2
[77] Yang Mengjiao. Du Yilun, Dai Bo, Schuurmans Dale, Tenenbaum Joshua B, Abbeel Pieter. Probabilistic adaptation of text-to-video models // arxiv preprint arXiv:2306.01872 2023. 3
[78] Yang Zhutian, Mao Jiayuan, Du Yilun, Wu Jiajun, Tenenbaum Joshua B, Lozano-Pérez Tomás, Kaelbling Leslie Pack. Compositional diffusion-based continuous constraint solvers // arXiv preprint arXiv:2309.0000 2023. 3
[79] Ye Jiacheng, Gao Jiahui, Gong Shansan, Zheng Lin, Jiang Xin, Li Zhenguo, Kong Lingpeng. Beyond autoregression: Discrete diffusion for complex reasoning and planning /| arXiv preprint arXiv:2410.141 14157. 2024. 2
[80] Zeng Zilai, Zhang Ce. Wang Shijie, Sun Chen. coal-conditioned predictive coding for offline reinforcement learning // Advances in Neural Information Processing Systems. 2023. 36. 25528-25548. 2
[81] Zhang Jianrong, Fan Hehe, Yong Yi. EnergyMoGen: Compositional Human Motion Generation with Energy-Based Diffusion Model in Latent Space arXiv preprint arXiv:2412. 14..2241..14706. 2024.
3
14

---

[82] Zhang Qinsheng. Song Jiaming, Huang Xun, Chen Yongxin, Liu Ming- Yu. Diffcollage: Parallel generation of large content with diffusion models 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). 2023. 101 80-1880111...
[83] Zhang Ziqi, Xu Jingzehua, Liu Jinxin, Zhuang Zifeng, Wang Donglin, Liu Miao, Zhang Shuai. Context-former Stitching via latent conditioned sequence modeling // arXiv preprint arXiv:2401.164552 2024. 2
[84] Zhou Siyuan, Du Yilun, Zhang Shun, Xu Mengdi, Shen Yikang, Xiao Wei, Yeung Dit-Yan, Gan Chuang, Adaptive online replanning with diffusion models // Advances in Neural Information Processing Systems. 2024. 36.
[85] Zhou Zhaoyi, Zhu Chuning, Zhou Runlong, Cui Qiwen, Gupta Abhishek, Du Simon Shaolei. Free from Bellman Completeness: Trajectory Stitching via Model-based Return- cosuur--ondiiooeed Supervised Learning // The Twelfth International Conference on Learning Representations. 2024. 2
[86] Zhu Zhengbang, Liu Minghuan, Mao Liyuan, Kang Bingyi, Xu Minkai, Yu Yong. Ermon Stefano, Zhang Weinan. Madiff: Offline multi-agent learning with diffusion models // Advances in Neural Information Processing Systems. 2024. 37. 4177-4206. 2
[87] Zhuang Zifeng, Peng Dengyun, Liu Jimxin, Zhang Ziqi, Wang Donglin. Reinformer: Max-Return Sequence Modeling for Offline RL // International Conference on Machine Learning. 2024. 62707-62722. 2
[88] Ziebart Brian D. Maas Andrew Bagnell J Andrew, Dey Anind others Maximum entropy inverse reinforcement learning, // Aaai. 8. 2008. 1433-1438.
15

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly and accurately outlined our scope and contribution in the Abstract and the Introduction Section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We provided description of limitations of the work in Section 5. We also provide additional analysis to the proposed method, such as the factors that might influence the performance of method (Appendix D) and computation efficiency (Appendix B.3).
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer:
[NA]
Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We have described our experiment settings in Section 4 and Appendix E. In addition, we will release our code upon acceptance.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
17

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We will release our code implementation and corresponding instructions upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / httppsscccppbccccccceecce geides/ Codeeubbm ssi aasonnooicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The training and evaluation details included in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: As shown in our quantitative results, we report the mean and standard deviation based on 5 random seeds for each of our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
18

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide description of the computation platform of our experiments along with resource requirements in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: http///neurippp hwtpp/cblll/cccccccceeeeeeeeeeee
Answer: [Yes]
Justification: This paper conforms with the NeurIPS Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We believe the societal impact of this work should be similar to a generic machine learning paper. No other issues feel must be specifically highlighted.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
19

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The assets used in paper are properly cited and related URLs are also provided.
Justification:
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. cow datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The assets in paper will be the corresponding code implementation of the proposed method. We will release the code along with documentation upon acceptance.
Guidelines:
The answer NA means that the paper does not release assets. Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: This research does not involve LLMs any important, original, or non- standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (attps://neul) htt.r/ppsccc/oooffeeeee////2255 for what should or should not be described.
22

---

Contents
Introduction
2 Related Work
3 Planning through Compositional Trajectory Generation
3.1
Compositional Trajectory Modeling
3.2
Training Compositional Trajectory Models
3.3
Compositional Trajectory Planning
Experiments
4.1
PointMaze
4.2 High Dimension Tasks
4.3 Ablation Studies
5
Discussion and Conclusion
Environment and Stitching Datasets
24
A.I Stitching Datasets in Ghugare et al. [2]]
24
A.2 OGBench Datasets
24
Additional Quantitative Results
26
B.1 Number Composed Trajectories
26
B.2 Inverse Dynamics Model
................................
B.3 Sampling Time Comparison: Parallel vs. Autoregressive
26
B.4 Starting Direction of Autoregressive Compositional Sampling
27
Additional Qualitative Results
28
C.I Composing Different Numbers of Trajectories
28
C,2 Diverse Trajectory Morphology
28
C33 Planning Results on Different Tasks
29
C.4 Compositional Planning in High Dimensional Space
29
C.5
Compositional Planning in AntSoccer Arena
30
Failure Mode Analysis
30
D.I Infeasible Generated Plan
30
D.2
Suboptimal Inverse Dynamics Model
31
D.3
Suboptimal Number of Composed Trajectories K
31
E
Implementation Details
32
E.1 Our Conditional Diffusson Model
32
E.2
Trajectory Merging
33
E.3 Replanning
34
F
Baselines
35
23

---

Appendix
In this appendix, we first introduce the evaluation environments and stitching datasets in Section Next, we provide additional quantitative results in Appendix B and additional qualitative results in Appendix C. We provide failure mode analysis in Section D. Following this, we provide implemen- tation details in Section E, including model architecture, training and evaluation setups, trajectory merging, and replanning. Lastly, we introduce notable baselines in Section
A Environment and Stitching Datasets
In this paper, we directly evaluate our method on public stitching datasets introduced in two recent papers Ghugare et al. [21] and OGBench [51]. In this section, we provide detailed descriptions of each dataset along with qualitative examples of trajectories in these datasets.
A.1 Stitching Datasets in Ghugare al. [21]
This paper [21] divides each evaluation environment into several small regions and each demonstration trajectory the training datasets can only navigate within a specific region. There is small overlap (one block) between each region, which can be used to stitch trajectories across regions. Therefore, to complete test-time goals, the agent needs conduct effective reasoning based on the given start state and goal state and identify the corresponding overlap joints. The division of regions is visualized in the original paper [21]. We use the environments and datasets from their official implementation release at https://github. https://gttubb.om/RjjCnni com/RajGhugare wwsstichinggss--omiinattttte sew.rraiallegeeaarttttoooooo....
A.2 OGBench Datasets
OGBench is comprehensive benchmark designed for offline goal-conditioned RL. Since our focus is evaluate the trajectory stitching ability of CompDiffuser, use the Stitch and Explore dataset types in OGBench.
In Stitch datasets, trajectories are constrained to navigate no more than 4 blocks in the environment. The start and goal state of each trajectory can be sampled from the entire environment provided that the travel distance between the start and goal is within 4 blocks. Qualitative examples of the trajectories the Stitch dataset are shown Figure 8.
AntMaze Medium Stitch
AntMaze Large Stitch
AntMaze Giant Stitch
Figure Trajectory Examples OGBench AntMaze Stitch Datasets. Each Trajectory limited travel most blocks for dataset type Stitch, while at inference, the distance between the start and goal can be up to 30 in the Giant Maze.
In Explore datasets, trajectories are of extremely low-quality though high-coverage. The data collection policy contains large amount of action noise and will randomly re-sample new moving direction after every 10 steps. Hence, each demonstration trajectory the training dataset typically moves within only 2-3 blocks due to the random moving direction. These datasets might be even more challenging due to the large noisy and cluster-like trajectory pattern. Qualitative examples of the trajectories in the Explore dataset are shown in Figure 9,

---

AntMaze Medium Explore
AntMaze Large Explore
Figure 9: Trajectory Examples in OGBench AntMaze Explore Datasets. Trajectories in Explore datasets are of extremely low-quality though high-coverage The data collection policy contains large amount of action noise and will randomly re-sample new moving direction after every 10 steps.
We show the environment names, corresponding datasets, and the maximum environment steps for each evaluation episode in Table 7. All methods are trained on the OGBench public re- lease datasets. We slightly increase the environment steps for some environments due to task difficulty (e.g., Giant Maze). For these environments, we follow the implementation in https://github.ccm https:/github.com/sohongggg wwtps//oohogggaaoooobooooooooo environment steps; for other environments, we directly adopt the reported success rates in the original paper.
and rerun all baselines with the increased maximum
Environment
Type
Size
Dataset Name
Env Steps
Medium
pointmaze- centmmaze-medimma stithhh--ttttttt
1000
pointmaze
stitch
Large
pointmaze-lage largesss--gg--ttch----
1000
Giant
gointmaze-giant-ao -iant-stitch---
1000
Medium
1000
AntMaze
Stitch
Large
hatmazz-larg--sticchh
2000
Giant
antmaye-giant-stiich-yy
2000
Arena
antsoccer-arrena--stih---o
5000
AntSoccer
stitch
Medium
antsocer-meeiiu--ttich--
5000
Medium
5000
HumanoidMaze
Stitch
Large
humanoid-large larraiiid-argge-sth---vo stitch-vo
5000
Giant
humanoid saamnoddgiian--stich---
8000
Table 7: Our Evaluation Environments and Datasets in OGBench.
25

---

B Additional Quantitative Results
In this section, we provide additional quantitative experiment results. In Section B.I, we study how varying the number of composed trajectories K affects the performance of our method. In Section B.2, we investigate the effect of inverse dynamic models of different horizons. Next, we provide a sampling time comparison of Decision Diffuser and the proposed parallel and autoregressive sampling schemes in Section B.3. Following this, in Section B.4, we analyze how the proposed autoregressive sampling scheme performs when using different denoising starting directions.
B.1 Number of Composed Trajectories
In Table 8, we compare CompDiffuser with GSC over composing different numbers of trajectories (results are also shown in Figure 7). Our method performs steadily when composing different numbers of trajectories while GSC collapses.
PointMaze Giant Stitch
# Comp
7
8
9
10
11
12
GSC
21:1
21
15:2
2+1
O
0=
CompDiffuser
51
5326
5514
5622
5026
50+3
Table 8: Quantitative Results over Different Numbers of Composed Trajectories. We report success rates (w/o replanning) of composing 7 12 trajectories in the OGBench PointMaze Giant Stitch dataset. CompDiffuser can consistently construct feasible trajectories over various numbers of composed trajectories while GSC gradually collapses.
B.2 Inverse Dynamics Model
In this section, we evaluate our planner with inverse dynamics models of different horizons (results also shown in Figure 7). Specifically, use an MLP implement the inverse dynamics model, which takes as input the start and goal state and outputs an action. We use the same training dataset (as to train the planner) train the corresponding inverse dynamics model. As shown in Table our method performs steadily across different inverse dynamics model horizons, showing that the subgoals generated by our planners are of high feasibility and are robust to various inverse dynamics models' configurations.
Env
Type
Size
8
12
16
20
24
28
Large
7714
8612
8042
8513
7622
7713
AntMaze
Stitch
Giant
61+s
65+3
68+4
63+a
60
63+t
Table 9: Quantitative Results of CompDiffuser with Inverse Dynamics Models of Different Horizons. We present the success rates of CompDiffuser with 6 different inverse dynamics model horizons. both environments, Large and Giant, our method performs consistently across all configurations, showing that the synthesized plans adhere to the transition dynamics and easy to follow.
B.3 Sampling Time Comparison: Parallel vs. Autoregressive
In Table 10, we compare the diffusion denoising sampling time of three methods: (1) monolithic model method, Decision Diffuser (DD), where we directly sample a long trajectory with the same horizon the proposed compositional sampling method; (2) parallel compositional sampling, as shown in the left of Figure 3, where we denoise all trajectories T in one batch; (3) autoregressive compositional sampling, as shown in the right of Figure 3, where we sequentially denoise each Tk-
We observe that both parallel and AR sampling methods require more sampling time than DD probably due to (I) the the simpler and smaller denoiser network of DD; (2) time for condition
26

---

encoding (our method will first encode the noisy adjacent trajectories 7k-1 and 7k+1 and feed the resulting latents to the denoiser Eω) and (3) the overhead of trajectory merging.
In addition, in our parallel sampling scheme, we stack all trajectories Ti-K to one batch and feed it to the denoiser network co. While it indeed requires only one model forward, the batch size increases implicitly. which is probably the major reason that the sampling time of Ours (Parallel) does not proportionally decrease as the number of composed trajectories K, in comparison Ours (AR).
Env
Type
Size
DD (Monolithic)
Ours (Parallel) Ours (AR)
Medium
0.23
1.02
1.54
PointMaze
Stitch
Large
0.23
1.67
3.39
Giant
0.23
2.73
5.00
Table 10: Quantitative Comparison of Sampling Time We report the time for sampling one (com- positional) trajectory in three different PointMaze environments using one Nvidia L40S GPU (unit: second). The reported results averaged over 20 sampling. In Parallel and AR (Autoregressive) mode, we use the proposed compositional sampling scheme as shown in Figure 3. Specifically, we compose 3, 6, and trajectories in Maze Medium, Large, and Giant, respectively. In Decision Dif- fuser (DD), we directly sample one trajectory with identical length as the compositional counterparts.
B.4 Starting Direction of Autoregressive Compositional Sampling
In the proposed autoregressive sampling scheme described in Figure 3, inside each diffusion denoising timestep, the denoising starts from and sequentially proceeds to тк (from left right), which denote as forward passing. Another implementation variant is to denoise in the reverse order, that is, first denoise Tк and sequentially proceed to (from right left), which denote backward passing.
We provide a quantitative comparison of these two starting directions in Table 11. We observe that these two methods yield similar performance, demonstrating that either autoregressive sampling direction can enable effective information propagation and exchange.
Env
Type
Size
Ours (Forward)
Ours (Backward)
Medium
100
100
PointMaze
Stitch
Large
100.0
100
Giant
5514
5612
Table 11: Quantitative Comparison of Forward and Backward Information Propagation. We
study the effect of the starting direction of the autoregressive sampling, from Tı vs. from Tк. To directly study the trajectory quality, we report the results w/o replanning for both methods. Either Forward Backward achieves similar performance, suggesting that our sampling method robust to different sampling configurations.
27

---

C Additional Qualitative Results
In this section, we present additional qualitative results of CompDiffuser. Videos and interactive demos are provided at our project website (see Abstract section for the link).
C.1 Composing Different Numbers of Trajectories
We present qualitative results of composing 8 to 11 trajectories in OGBench PointMaze Giant Stitch environments in Figure 10. We compositionally sample multiple trajectories to construct long-hone long-horizon plan where the given start is at the bottom-left corner and the given goal is at the top-right corner For clearer view, we present the results before applying trajectory merging (i.c., we show each individual trajectory of Ti:k) and different colors to highlight different trajectories.
K=8
K=9
K=10
K=11
Figure 10: Composing Different Numbers of Trajectories Inference. Our method is only trained 4 blocks, while able compositionally generate coherent long
short trajectory segments that travel trajectories given the start (cirele) and goal (star). When smaller and barely sufficient to reach the goal (c.g., 8), the longth of the overlapping part between segments decrcases so as to extend the travel distance of the overall compositional plan. In contrast, if given larger K (e.g., 11), some parts of the compositional plan the extra length.
might travel back and forth
C.2 Diverse Trajectory Morphology
The proposed compositional sampling method allows direct generalization to long-horizon planning tasks at test time through noisy-sample conditioning and bidirectional information propagation design. Meanwhile, this sampling approach also preserves the multi-modal nature of the diffusion model, enabling a diverse range trajectory morphology. As shown Figure 1 11, given a similar start and goal pair, our method can construct trajectories that reach the goal via various possible paths. With such multi-modal flexibility, the proposed sampling process can further customized with specific preferences by integrating additional test-time steering techniques.
Trajectory
Trajectory
Trajectory
Trajectory
Figure 11: Diverse Trajectory Morphology. We present four trajectories with similar start and goal in OGBench PointMaze Giant Stitch. trajectories are generated CompDiffuser, composing trajectories. CompDiffuser preserves multi-modal nature of diffusion models and able flexibly sample trajectories diverse morphology.
28

---

C.3 Planning Results on Different Tasks.
In Section C.2 above, present multiple trajectories of Task 1 in OGBench PointMaze Giant. In this section, we present qualitative results of the following Task 2 to Task 5, as Figure 12. We set the number of composed trajectories K to 9 in all tasks. The state state is shown by the black circle and the goal state is shown by the black star. Our method successfully constructs feasible plans for various sarr-gooal configurations across different spatial distances.
+C
+
Task
Task 3
Task 4
Task 5
Figure 12: Different Tasks in OGBench PointMaze Giant Stitch. We present qualitative plans by CompDif- fuser OGBench Task 2-5 (See Figure 11 for results of Task 1). We the number of composed trajectories all the tasks above. black circle denotes start state and black star denotes the goal state. In task and 3, our method effectively constructs long horizon plans that reach the goals on the opposite environment (despite heing trained only on trajectories that travel at most blocks). In comparison, Task and 5 feature relatively smaller spatial gap between the start and goal, thus requiring shorter planning horizon. Our method still generalizes these tasks by generating plans that traverse additional distance or leveraging back-and-forte movements consume the extra plan length.
of
C.4 Compositional Planning in High Dimensional Space
In this section, we present additional high dimensional trajectories generated by CompDiffuser in OGBench AntMaze Large Stitch environment. Similar to other experiments in the paper, the models are trained on the corresponding OGBench public release datasets.
AntMaze Large Stitch 29D. We train CompDiffuser on the state space of x-y position along with the ant's joint positions and velocities, resulting 29D planning task. Note that CompDiffuser is only trained on short trajectory segments (we set its horizon to 160), while at test time, we compose trajectories to directly construct trajectory plans horizon 584. present additional qualitative results in Figure 13. Note that we uniformly sub-sample the length of the trajectory to 50 for clearer view. Corresponding quantitative results are reported Table 4.
Figure 13: Compositional Planning in 29D OGBench AntMaze Large Stitch Tasks. Original trajectory
plans are much denser and we uniformly sub-samplc 50 states from the original generated trajectories for better view. Five trajectories are composed as shown by different colors: the blue one indicates the first trajectory tI and the purple one indicates the fifth trajectory 75.
29

---

C.5 Compositional Planning in AntSoccer Arena
We provide additional qualitative results in OGBench AntSoccer Arena Stitch environment in Figure 14. In this task, the ant is initialized to the location of the blue circle and is tasked to move the ball the goal location indicated by the pink circle. Hence, the ant needs first reach the ball from the far side of the environment and dribble the ball to the goal.
However, such long-horizon trajectories (the ant reaches the ball and then dribbles the ball the goal) do not exist in the training dataset. The training dataset only contains two distinct types of trajectories: (1) the moves the environment without the ball, (2) the ant moves while dribbling the ball. Therefore, the planner needs to generalize and stitch in zero-shot manner constructing an end-to-end trajectory that first approaches the ball and then dribbles the ball to the goal.
We train CompDiffuser on the state space of the ant's x-y position and joint positions along with the x-y position of the ball, resulting in 17D planning task. Similar Figure 6, present each individual trajectory TIKK in Figure 14 for clearer view. Corresponding quantitative results are provided in Table 3.
SORLGE
Init State
T1
T3
T4
Figure 14: Compositional Planning on OGBench AntSoccer Arena Stitch. We present the initial state for above. The start position of the ant is shown by the blue circle
planning and each individual trajectories (bottom right) and the goal is to move the ball to the pink circle (upper right). We compositionally sample trajectories (as shown from left to right), which then be merged form long-horizon plan. The ball highlighted with a small yellow cirele. Our compositional sampling method effectively stitches two different types of trajectories and generalizes to more difficult tasks unseen in the training data.
D Failure Mode Analysis
In this section, we present several failure cases of our method along with analyses and qualitative examples of these failure modes. Our proposed framework consists of two components: generative planner and an inverse dynamics model. We outline and discuss several failure modes below.
D.1 Infeasible Generated Plan
As the number of composed trajectories K increases, our method may show inconsistencies in the trajectory plan, especially when planning in high dimensional state space. We observe that although the start and goal conditioning can consistently ensure that the composed plan begins at the initial state and terminates at the provided goal, the intermediate trajectories may include implausible transitions, such trajectories that jump walls, making the overall plan infeasible. However, empirically these failures usually occur at considerably higher K as compared to baseline methods (e.g., GSC, Table and Table 2).
In Figure 15, we present a qualitative example of infeasible plan in OGBench AntMaze Giant Stitch when planning in 29D space (ant's x-y position, joint positions, and joint velocity) and composing 9 trajectories. The original white walls are rendered transparent for better view of infeasible states. Though the generated compositional plan is mostly valid, the second trajectory T2. as shown yellow, is infeasible as it passes through walls. This is probably due to the suboptimal coordination among trajectories in the compositional sampling process.
As illustrated the figure, certain trajectories (e.g., the neighboring blue, green, and red ones) barely proceeds, moving only 2-3 blocks. To bridge the resulting spatial gap, the intermediate yellow trajectory must span much longer distance. However, the training data does not contain such long Long-horizon trajectories, making it difficult for single trajectory to extend to such length, which
30

---

finally results in an o.o.d . sample that goes through walls. We could potentially mitigate this issue with rejection or guided sampling techniques select feasible candidate plans.
D.2 Suboptimal Inverse Dynamics Model
In our experiments, we observe that even if CompDiffuser synthesizes feasible trajectory for the agent to follow, the agent might not successfully reach the goal due to the error by the inverse dynamics model. For example, the agent might bump into walls due unstable locomotion get stuck in some local region, yet the synthesized plan collision-fre free and coherent.
In implementation, we use a simple MLP to parameterize the inverse dynamics model and train it with regression MSE loss. We believe that more optimized model architecture or specific finetuning might further boost the performance of the inverse dynamics model, hence boosting the overall performance of our method. We deem that out of the scope of this work.
In Figure 16, we present qualitative example of failure due to the inverse dynamics model in OGBench AntMaze Giant Stitch. The planning is in 15D space (ant's x--y position and joint positions) and composes 3 trajectories. While we observe that the synthesized plan is feasible and successfully reaches the goal. the environment execution rollout fails during the yellow trajectory. In this instance. the ant agent fails to execute the right turn, loses track of subsequent subgoals, and becomes trapped in local region. To address this issue, incorporating effective replanning strategies could help the agent recover (since the new plan will start from the current trapped state). In addition, we deem that employing more robust or specialized inverse dynamics models may further mitigate such failure scenarios.
Infeasible Plan: Passing through Walls
Feasible Plan
Rollout: Suboptimal Actions
Figure 15: Failure Mode: Infeasible Plan. We increase the transpareney of the original white walls for better view of infeasible states. The start state is at the top left corner and the goal state is at the bottom right corner. Nine trajectories are composed, highlighted by dif- ferent colors. The second trajectory 72 (shown in yellow) infeasible and passes through walls. This probably due to the suboptimal coordination between trajectories in the com- positional sampling process: the neighboring blue and green trajectories barely progress, leaving a long o.o.d gap for the yellow trajec- tory to fill.
Figure 16: Failure Mode: Suboptimal Inverse Dynamics Actions. We present a failure case of planning in 15D space (ant's x-y position and joint positions) in AntMaze Medium um Stitch. Left: compositional plan where three trajectories (represented in blue, yellow, and green) are composed. We samlee the plan to 36 states for visualization (the original plan dense with ing enviromment execution rollout of the compositional plan. Though the synthesized plan is valid and successfully reaches the goal, the inverse dynamics model may fail, as illustrated on the right which gets stuck in right turn and is unable proceed. Further incorporating some effective replanning strategies or employing more robust inverse dynamics models could potentially mitigate these failure scenarios.
300 states). Right: the correspond-
D.3 Suboptimal Number of Composed Trajectories K
In our current implementation, the number of composed trajectories K needs to manually specified at test time. As shown by the quantitative results Table 8 and qualitative results Figure 10, relatively larger K does not significantly affect the model performance as the extra plan length will consumed by staying or circulating within certain valid regions in the environment. However, an aggressively smaller might lead to failure plans since the overall planning horizon becomes insufficient to cover the substantial spatial gap between the start and the goal.
31

---

In this section, we provide qualitative examples of infeasible plans due to impractical K. In Figure 17, we show each individual trajectory Tk generated by our compositional sampling method when K ranges from 3 to 6. Note that feasible horizon to reach the goal from the start (bottom left) to the goal (upper right) requires composing at least trajectories, i.e., K' 8.
Therefore, while the generated trajectories can begin at the start state and terminate at the goal state, the intermediate segments become disconnected since there are too few trajectories to bridge the significant spatial gap (given that the training data contains only short trajectories). Nonetheless, we observe that the overall flow of the trajectories is directed toward the goal, and as increases, the plan's 's structure gradually becomes more feasible.
T
+
K=3
K-4
K=5
K=6
Figure 17: Failure Model: Suboptimal Number of Composed Trajectories K. Our method may generate infeasible plans if K, the number of composed trajectories, set significantly below the required minimum. For example, in the planning task illustrated above, the start state located in the bottom left corner while the goal state at upper right comer. feasible plan for task typically requires composing at least trajectories (i.c., = 8). We present qualitative examples of plans when K is smaller than the minimum threshold. Though the first and last trajectory segments correctly attach to the start and goal states, the intermediate trajectories are disconnected due the insufficient plan length. However, the overall plan still progresses towards the goal, which may provide useful guidance signals the agent.
E Implementation Details
Software: The computation platform is installed with Ubuntu 20.04.6, Python 3.9.20, PyTorch 2.5.0.
Hardware:
We use I NVIDIA GPU for each experiment, A GPU with 24GB memory is sufficient to train our models and it takes 1-2 days to train model using a recent mid-level NVIDIA GPU.
E.1 Our Conditional Diffusion Model
In this section, we introduce detailed implementation of our conditional diffusion model <o(7',4 ||st__ooddd end_cond).
Model Inputs and Outputs. Our diffusion model takes as input the noisy trajectory r', diffusion noise timestep f, and the start condition st_ st_cond and end condition end_cond for 7" st__sssce can be either noisy chunk start state qs and end_ _ond_ can be either noisy chunk or goal state qg- We the predicting 70 formulation to implement this network, i.e., the network directly predicts the clean sample ?
Implementation of Noisy Chunk Conditioning, As described in Section 3, we only train one diffusion model Eg that is able to condition on both the start state g., goal state and noisy chunk Th-11 T such that in test time construct long-houd horizon plans. This unified one model design eliminate the need for manual task subdivision (across multiple models) reliance predefined planning skeleton, thereby enabling holistic end-to-end planning.
use the proposed compositional sampling approach to
In practice, we can implement the noisy neighbor conditioning Late simply using the overlapping part between the two chunks instead of full chunk T conditioning, because the overlapping part is sufficient to ensure the connectivity between two adjacent trajectories. Such design further simplifies the training procedure: instead of dividing a training trajectory 7 to K chunks, we only need to sample noisy sub-chunk from itself.
32

---

Specifically, assume T' is the noisy trajectory to be denoised and f is another independent noisy version of also noisy timestep We denote the length of h and the length of the overlapped part as ha then we can use + ha] and 1'hh h, h) as the noisy sub-chunks for st, utcond and end_cond during training, respectively. Hence, when inference, the end_ cond for Ti can be set to 1 ho] (the front chunk of the second trajectory 72) and the st_cond for can be set to #5 ho] (the tail chunk of the first trajectory TT).
Model Architecture. For planning in 2D x-y space, we follow Decision Diffuser [1] and use conditional U-Net as the denoiser net- work. For planning in high dimension state space, we use DiT [55] based transformer [65] as the denoiser network bttps://gthuu..cm/fcccboorrssrr com//ccceookressrch/ii....
(https://github.o
Training Pipeline. We provide detailed hyperparameters for training our model on PointMaze Giant Stitch environment in Table 12. We do not apply any hyperparameter search learning rate scheduler. Please refer to our codebase for more implementation details.
Hyperparameters
Value
Horizon
160
Diffusion Time Step
512
Probability of Condition Dropout
0.2
Iterations
1.2M
Batch Size
128
Optimizer
Adam
Learning Rate
2e-4
U-Net Base Dim
128
U-Net Encoder Dims
(128, 256, 512, 1024)
Table 12: Hyperparameters for Training on PointMaze Giant tch environment.
Inference Pipeline. In Table 13, we present the single model horizon (length of individual rT) and the inference-time number of composed trajectories K corresponding to the reported results in Table 1, Table 2 and Table 3.
For each evaluation problem, we generate B samples in a batch and we use a simple heuristic to select one sample as the output plan. Specifically, we compute the L2-distance of each overlapping parts in the generated trajectory segments TI:K. The one with the smallest average distance will be adopted as the output plan, the sense that small distance in the overlapping parts indicates better coherency between adjacent trajectories. we deem that developing some more advanced inference-ting methods with CompDiffuser may be an interesting future research direction, such as probability or density based plan selection methods or compositional sampling with flexible preference steering.
E.2 Trajectory Merging
As introduced in Method Section 3 and Algorithm 2, the generated trajectories FIKK are mutually overlapped and we merge these K trajectories to form a long-horizon compositional trajectory Teomp- In this section, we describe the implementation of the exponential trajectory blending technique which we use for merging.
We directly leverage the classic exponential trajectory blending formulation. For simplicity. let Ti and 72 denote the trajectories to blend, 7i [t] denote the t-th state in Tı, fsum and fend denote the start and end index of the region to apply blending. Note that, in practice, we only blend the overlapped part between two adjacent trajectories. We provide the equation for exponential blending below,
-8(
E-Saen
Feed tan
eB
Toomp [t] w())= [1[t] (1 w(t) ro(t). where w(t)
-8
We set across all our experiments. In practice, various other trajectory blending techniques can also be directly applied, such as cosine blending and linear blending.
33

---

Environment
Type
Size
Single Model Horizon
of Composed Trajectories
U-maze
40
PointMaze [21]
Medium
144
5
Large
192
Medium
160
pointmaze
Stitch
Large
160
Giant
160
8
Medium
160
3
AntMaze
Stitch
Large
160
Giant
160
9
Medium
192
AntMaze
Explore
Large
192
6
Arena
160
5
AntSoccer
Stitch
Medium
160
Medium
336
HumanoidMaze
Stitch
Large
336
6
Giant
336
11
Table 13: Number of Composed Trajectories for Each Evaluation Environment. Our diffusion
models are trained with short horizon listed in the Single Model Horizon column. In test time, compositionally generate multiple such short trajectories to enable trajectory stitching and construct plans of much longer horizon.
E.3 Replanning
In this section, we describe the detailed implementation of replanning. While our method is designed directly generate an end-to-end trajectory from the given start state the goal state, replanning can be performed at any given timesteps during rollout. Specifically, we initiate replanning if the agent loses track of the current subgoal, i.e., the L2 distance between the agent and the subgoal is larger than threshold.
In larger maze, the required number of composed trajectories, denoted K, is usually large due to the distance between the start state and the goal state. However, if we keep replanning with a similar large K even when the agent is already close to the goal, the generated trajectory might travel back and forth to consume the unnecessary intermediate length (see (2) and (3) Figure 12), thus delaying the agent's progress toward the goal.
To address this, we use receding scheme for K to encourage faster convergence to the goal. Let K denote the number of composed trajectories of the current plan (the plan that the agent is following), heomp denote the length of the current plan, hęee denote the length the current plan that the agent already executes in the environment. The number of composed trajectories used for replanning Keeplan given by
max(/exe 8, 0)
Keellnn ceil
hoomp
where 8 is haper-parameter that controls the convergence speed to the goal, for example, K Kyeptan will decrease faster if is set to a small (or negative) number while Keeptan will decrease very slowly if 8 is large positive number.
34

---

Baselines
We compare our approach with a wide variety of baselines, including culfussion-asse trajectory planning algorithms, data augmentation based stitching algorithms, and goal-conditioned offline RL algorithms.
Particularly, we include the following methods:
For generative planning methods, we include Decision Diffuser (DD) [1] for monolithic trajec- tory sampling and Generative Skill Chaining (GSC) [48] for trajectory stitching:
For data augmentation based methods, we include stitching specific data augmentation [21]
with RvS [14] and Decision Transformer (DT) [9]. Since the evaluation setting is identical, we directly adopt the reported numbers in the original paper [21].
For offline reinforcement learning methods, we include goal-conditioned behavioral cloning (GCBC) [44, 20J, goal-conditioned implicit V-learning (GCIVL) and Q-learning (GCIQL) [29], Quasimetric RL (QRL) [70], Contrastive RL (CRL) [15], and Hierarchical implicit Q-learning (HIQL) [52]. For these baselines, we follow the implementation setup established by OG-
Bench [51] throughout our experiments.
We describe the implementation details of a few notable ones below.
Generative Skill Chaining (GSC) [48]. GSC recent diffusion model-based skill stitching method. We directly adopt its original score-averaging based stitching algorithm and apply it in our tasks. Specifically, when composing K trajectories Tı:K, GSC averages the scores of the overlapping segments between adjacent trajectories (in total K 1 overlapping segments this case) prior to each denoising timestep. For fair comparison, we re-use the same diffusion denoiser networks in CompDiffuser for every GSC experiment.
Hierarchical Implicit Q-Learning (HIQL) [52]. HIQL is recently proposed Q-learning based method that employs a hierarchical framework for training goal-conditioned RL agents. learns a goal conditioned value function and uses it to learn feature representations, high-level policy, and low-level policy. We follow the original implementation of the method OGBench [51].
Goal-Conditioned Behavioral Cloning (GCBC) [44]. GCBC is classic imitation learning based method. In our experiments, GCBC trains MLP that takes input the observation state and future goal state from the same offline trajectory and outputs corresponding action for the agent. use the same implementation as OGBench [51].
35


Reviews:
[{"review_id": "t95kfVL5fd", "reviewer": "Reviewer_yhYf", "summary": "Briefly summarize the paper and its contributions. This is not the place to critique the paper; the authors should generally agree with a well-written summary. This is also not the place to paste the abstract—please provide the summary in your own understanding after reading.\n\nThis paper considers the problem of building compositionally generalizable diffusion-based planners, i.e. models that can diffuse long sequences of actions (and/or states), without being bound to a fixed plan length. While prior work on diffusion-based planners (e.g. Diffuser or Decision Diffuser) did allow for generating multiple fixed-size plans in sequence, they struggle to generalize to sequences of much longer sequences than ones used in training. One of the challenges is ensuring a smooth and physically feasible transition between denoised sub-plans.\n\nThis work proposes CompDiffuser, which addresses this by denoising several sub-plans at the same time, with each sub-plan being conditioned on its adjacent sub-plans. In addition, the sub-plans are chosen to be overlapping. This enables the model to sample sub-plans that can actually be stitched together more easily.\n\nConcretely, the trajectories are factorized as:  \n$$\np_\\theta(\\tau \\mid q_s, q_g) \\;\\propto\\;\np_1\\bigl(\\tau_1 \\mid q_s, \\tau_2\\bigr)\\,\np_K\\bigl(\\tau_K \\mid \\tau_{K-1}, q_g\\bigr)\\,\n\\prod_{k=2}^{K-1}\np_k\\bigl(\\tau_k \\mid \\tau_{k-1}, \\tau_{k+1}\\bigr).\n$$\n\nWhere $q_s$ and $q_g$ are the start and goal states, respectively.\n\nThese conditional distributions are modeled using a joint diffusion model. At training time, the noise prediction network (in the DDPM sense) for sub-trajectory $\\tau_k$ at noise level $t$ (denoted $\\tau_k^t$) takes as input $t, \\tau_k^t, \\tau_{k-1}^t$ and $\\tau_{k+1}^t$, i.e. it is conditioned on the adjacent sub-trajectories at the same noise level. There is also a loss term for conditioning the first and last plans on the start and goal states, respectively. The fact that adjacent chunks influence each other during sampling aims to enable the model to make the sub-plans compatible “by construction”, allowing them to be easily stitched together at the end of sampling.\n\nThe authors propose two alternatives for sampling:\n\n- **Parallel sampling**: the noise predictor at $\\tau_k^t$ conditions on $\\tau_{k-1}^t$ and $\\tau_{k+1}^t$, as in training. This means that the samples at noise level $t-1$ only depend on samples at noise level $t$, meaning the denoising steps can be run in parallel.  \n- **Autoregressive sampling**: the noise predictor at $\\tau_k^t$ conditions on $\\tau_{k-1}^{t-1}$ and $\\tau_{k+1}^t$ (i.e. the sample for the previous trajectory is fed in at noise level $t-1$). Now one needs to do each denoising step sequentially (across time) at each noise level. The authors claim this enables stronger coordination across sub-plans, and find that autoregressive sampling has better empirical performance.\n\nExperiments focus on validating the empirical performance of CompDiffuser and studying how its performance varies with respect to state dimensionality, number of composed trajectories, sampling schemes and replanning.\n\nEvaluation covers PointMaze stitching datasets and OGBench locomotion tasks (AntMaze, HumanoidMaze, AntSoccer), comparing CompDiffuser against diffusion planners (Decision Diffuser, GSC), sequence-modeling methods (RvS, DT), and offline RL algorithms (GCBC, GCIVL, GCIQL, QRL, CRL, HIQL). The training data for these environments is constrained to shorter trajectories, as opposed to whole plans, producing a setting where models need to compositionally generalize in order to solve the task. CompDiffuser consistently outperforms these baselines on PointMaze and high-dimensional stitching tasks, demonstrating its ability to maintain feasibility and goal reachability under varied conditions.\n\nAblations show autoregressive sampling yields higher success than parallel sampling; replanning offers robustness gains primarily in the most complex (Giant) mazes; and performance remains stable across segment counts and inverse-dynamics configurations but degrades as planning dimensionality increases.", "questions": "- I think one interesting statistic about planning methods is the size of the largest problem they can solve reliably. In the setting of PointMaze environments, this could be measured along the lines of “how large does the maze have to be for the average success rate to drop below 95%”. 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? We see that from the Large maze to the Giant maze in OGBench there is a significant drop; I wonder at what problem size the degradation starts.  \n- On a similar vein, 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?  \n- A more general instance of the question above: what are the main variables with respect to which the generalization ability of the method scales? (Potential variables could be amount of data, ratio of training trajectory length to evaluation trajectory length, computation budget for denoising during the forward pass, replanning budget, etc).  \n- In appendix, the authors mention that CompDiffuser’s autoregressive performance is robust to whether one starts from the starting position (hence working “forward” to the goal position) or from the goal position (hence walking “backward” to the starting position). However, the locomotion problem here is in a sense very symmetric, in the sense that a valid plan from point A to point B can be inverted to obtain a valid plan from B to A. 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”?  \n  - For an example of such an environment, consider a graph with a root node $s$, which is always a starting node, a set $I$ of intermediate nodes, and a set $T$ of target nodes. Each target node is connected to a single intermediate node. Intermediate nodes can be connected among themselves arbitrarily. $s$ is connected all intermediate nodes. Given a target position $t$, one immediately knows the intermediate node it is connected to (because it is unique), and from there one can easily construct a path from $s$ to $t$. However, starting from $s$, there are $|I|$ possibilities for the first move. If $I$ is very large, it seems intuitive to me that solving the problem “backward” should be significantly easier given finite memory capacity compared to solving it “forward”. Please let me know if I’m missing something with this example!  \n- 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).  \n- Recent work by Chen et al. (2024) introduced Diffusion Forcing, and had strong long rollout results extending to complex environments such as Minecraft. Do you see any complementarity in your approaches to e.g. enable stronger temporal generalization of decision making diffusion models in complex control settings? What would an ultimate, ideal solution for this more general problem in complex environments look like in your view, and what are the main unsolved problems for making it happen?  \n- 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?  \n- On a similar vein, would one expect CompDiffuser to inherit similar global guidance properties from full-trajectory diffusion methods like Diffuser or Decision Diffuser? My intuition is it would, since sub-plans condition on each other during sampling.", "limitations": "The authors address limitations in the conclusion.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**  \n- Clear problem motivation: the authors clearly outline their goal, which is to address length generalization in decision making diffusion models.  \n- Experimental results indicate successful completion of research goal and extensively study relevant variables in the method’s implementation and performance. In particular, they demonstrate that CompDiffuser does indeed show much better length generalization than its full-sequence diffusion counterparts.  \n- Thorough and high-quality exposition that makes it very clear what the proposed algorithmic approach is, and relevant design choices (e.g. sampling approach).\n\n**Weaknesses:**  \n- Overall I consider this to be a good paper. On the algorithmic side, the three main novel points in their approach are: (1) subdividing large trajectories into chunks when denoising; (2) choosing the chunks to be overlapping; and (3) conditioning each chunk on its previous and following chunk. While ablations consider several aspects of the proposed setup (e.g. sensitivity to replanning, sampling approach, etc), I believe it would be constructive to include ablations of at least points (2) and (3).  \n  - Namely, these ablations would be: (2) testing a version of CompDiffuser where the chunks don’t overlap, and (3) testing a version of CompDiffuser that conditions only on the previous trajectory, and not on both adjacent ones. While I can see reasons why design decisions (2) and (3) should be beneficial, it’s not clear to me that e.g. the method would completely collapse or be mis-specified without them. These ablations would help understand the value of each algorithmic design decision made by the authors.  \n  - Also, ablation (3) above would bear some resemblance to Diffusion Forcing (Chen et al. 2023). It seems to me the paper would be strengthened by including a comparison between CompDiffuser and Diffusion Forcing in maze planning settings, since this is another recent diffusion decision making model which had good results relating to length generalization (e.g. infinite rollouts in Minecraft).", "was_revised": true, "final_justification": "Following my observations in **Strengths** and the author's replies (particularly regarding ablations), I believe this work opens a lot of interesting directions for further work in compositional planning, and I maintain my recommendation of acceptance.", "rebuttal": "We thank Reviewer yhYf for the positive and careful review. Please see our response to each listed concern below.\n\n> 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.\n\n\nThank you for your valuable suggestions.\n\nWe 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.\n    \n| Env | Only Previous | Both (Ours) |\n| - | - | - | \n| Medium | 100±0   | 100±0 | \n| Large | 94±3  |  100±0 | \n| Giant | 17±3 |   68±3 | \n\n\nDue 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.\n\n\n\n\n\n\n> W2. It seems to me the paper would be strengthened by including a comparison between CompDiffuser and Diffusion Forcing in maze planning settings.\n\nThank you for the valuable suggestions.\n\nWe 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.\n\n\n\n\n\n> 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? \n\n\nGiven 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.\n\nSpecifically, 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.\n\n\n\n> 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?\n\n\n\nThank you for your comments. \nWe 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.\n- We can leverage a plan verifier to select a promising plan out of a batch of candidate plans.\n- Inference-time searching [1] can be further integrated in the sampling process to enhance the performance.\n- 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.)\n- Increased replanning budget. Allowing more replanning steps may help the agent recover from failures, such as when getting stuck by walls.\n\nPlease also see the response to Q3 below for additional strategies.\n\n\n\n> Q3. what are the main variables with respect to which the generalization ability of the method scales? \n\n\nThank you for your insightful question. We believe that the following factors effect the generalization ability of the model:\n- 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).\n- 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.\n- 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.\n- 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).\n\n\n\n\n> 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”?\n\n\n\nIn 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$). \nThough 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.\n    \n    \n\n\n\n> 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).\n\n\nThank 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. \n\nFor 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. \nIn 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.\n\n\nIn 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. \n\n\n\n\n> Q6. Complementarity between  Diffusion Forcing and the proposed method.\n\nThank 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\n\nTowards 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).\n\n\n\n> 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?\n\n\nWe 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. \nAs 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.\n\n\n\n> Q8. Would one expect CompDiffuser to inherit similar global guidance properties from full-trajectory diffusion methods like Diffuser or Decision Diffuser?\n\n\nWe believe that CompDiffuser is able to inherit the global guidance properties from full-trajectory diffusion methods.\n\nThough 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.\n\n\n\n\n[1] Inference-time Scaling of Diffusion Models through Classical Search"}, {"review_id": "dJOshniDzS", "reviewer": "Reviewer_a2a3", "summary": "This article introduces a method for generative planning using stitching. Their method denoises short segments conditioned on the noisy states of their neighbors, to form overall long coherent trajectories. They introduce two flavors of this method, one of which only conditions on the previous noise level, and another which takes the current noise level for one neighbor. The paper presents extensive evaluation on maze-type benchmarks as bundled in the OGBench benchmark and others, and demonstrate an advantage over various baselines.", "questions": "In the algorithm, it states that \"clean trajectories\" are sampled from the dataset, while in the experiment description it is stated that the experment is done on benchmarks with short segments of trajectories. Are they full trajectories or short segments? If they are short segments, how are these connected together with their neighbors, as is required for the training algorithm? I.e., to train the diffusion model, the two neighbors on either side are required inputs to the model.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strenghts:\n- The paper is nice to read and the method appears elegant and fits the problem well\n- There is an extensive number of baselines and experiments presented in the paper\n- Effective learning based only on chunks that generalizes to longer horizons is an important task\nWeaknesses:\n- While the baselines are plentiful and there are a lot of experiments, it is not made entirely clear how appropriate the baselines actually are, and while there are a lot of experiments, there is not a lot of variation in task type. I will address each of these concern in a bit more detail. \nBaselines:\n- In appendix B.3, they compare against \"monolithic\", which is much faster, they say because it is a much smaller architecture. This seems like a strange comparison: it is not apples to apples, and therefore not very useful. As far as I see it, their parralel \"chunked\" attention is effectively monolithic attention with masking of all but three chunks, so, it should be possible for their method to be at least as fast or faster in like for like scenarios. This would indicate the overall model for the monolithic case is smaller, making the comparison invalid. \n- Continuing on this same topic, the parameter count / model size is not really well presented, but seems like an important thing to report since they give such strong performance against baselines. 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. It is more useful to the reader if they can compare like for like. 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.\n- 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. \n- 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?\nExperiments:\n- 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.", "was_revised": true, "final_justification": "The authors have adequately answered the comments and concerns of my review, but the rebuttal and other reviews have not moved my score into the next bracket.", "rebuttal": "We thank Reviewer a2a3 for the positive and careful review. Please see our response to each listed concern below.\n\n\n#### Baselines:\n\n> 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.\n\n\n\nThank 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.\n\n\n**Model Size.** \nFor the result shown in Table 1, the denoiser model of DD is indeed smaller than ours. \nHowever, in our implementation, we already matched the architecture, depth, and feature dimension of the main denoiser network for DD and ours.\nThe 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). \n\nPlease see the parameter count of the networks below:\n\nPointMaze from [1]\n| Env | DD | Ours |\n| - | - | - |\n| U-Maze |  35.7 M  | 59.8 M | \n| Medium |  35.7 M  | 61.1 M | \n| Large |   35.7 M  |  61.4 M | \n    \nOGBench PointMaze (networks used for reporting the sampling time in Table on Page 20)\n| Env | DD | Ours |\n| - | - | - |\n| Medium |  63.8 M  | 92.9 M | \n| Large |   63.8 M  |  92.9 M | \n| Giant |   63.8 M  | 92.9  M | \n\nFor 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.\n    \n    \n**Inference Time.**\nWe 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:\n    (1) DD's denoiser network is indeed smaller;\n    (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.\n    (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.\n    (4) The overhead of trajectory merging. The construction of the final plan involves merging multiple trajectory segments, which might introduce further overhead.\n\n\n    \nAdditionally, 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.\n\n\nIn 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.\n| Env | DD (35.7M) | DD (63.8M) | Ours (~60M) |\n| - | - | - | - |\n| U-Maze |  0±0  | 0±0 | 100±0 |\n| Medium |  30±1  | 16±0 | 100±0 |\n| Large |   0±0  | 28±0 | 100±0 |\n\n    \nIn 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.\n\n| Env |DD (63.8M) | DD (99.0M) | Ours |\n| - | - | - | - |\n| Medium |  80±0  | 55±1 | 100±0 |\n| Large |  40±0 | 39±0 | 100±0 |\n| Giant |   0±0  | 0±0  | 53±6  |\n\nIn 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.\n\n[1] Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View.\n    \n\n\n    \n\n\n\n> 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.\n\n\n\nThank you for the valuable suggestions.\nWe 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.)\n\nThe hyperparameters of our method is actually very similar across different experiments.\nIn fact, the hyperparameters (e.g., training dynamics, model architectures) is identical for OGBench PointMaze and AntMaze as shown in Table 12. \n\nSince 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:\n\n|Hyperparameters | Value |\n| - | - |\n|Horizon | 336 |\n|Diffusion Time Step | 1000 |\n|Probability of Condition Dropout | 0.2 |\n|Iterations | 1.2M |\n|Batch Size | 192 |\n|Optimizer | Adam |\n|Learning Rate | 2e-4 |\n|U-Net Base Dim | 128 |\n|U-Net Encoder Dims | (128, 256, 512, 1024) |\n\n\n\n\n\n> 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? \n\n\n\nYes. 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.\n    \nMany 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.\nWhile 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.\n\n\n\n\n#### Experiments:\n\n\n> 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.\n\n\n\n\nThank 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. \n\n| Env | GCBC | GCIVL | GCIQL | QRL | CRL | HIQL | Ours |\n| - | - | - | - | - | - | - | - |\n| `cube-single-play-v0` | 6 ± 2 | 53 ± 4 |68 ± 6 | 5 ± 1 | 19 ± 2 | 15 ± 3 | 92±4 |\n\n\n\n\n\n    \n\n> 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?\n\n\nIn 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). \nThis '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).\n\nThen 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."}, {"review_id": "GjKzQUSQZm", "reviewer": "Reviewer_AQzW", "summary": "The paper presents a novel approach to planning for long-horizon tasks, called CompDiffuser. This method involves breaking down trajectories into shorter segments with overlapping parts, and then training diffusion models to learn the distribution of consecutive states conditioned on the start and end states of neighboring segments.\n\nAt test time, given a starting state and a goal state, CompDiffuser generates a plan by combining these segments in an autoregressive or parallel manner, resulting in a physically plausible plan. The plan is then executed using a learned inverse dynamic policy.\n\nThe paper provides empirical results, primarily focusing on navigation tasks with various embodiments (such as pointmass, ant, and humanoid) and different mazes. These results are compared to other state-of-the-art methods in goal-conditioned offline RL and diffusion-based approaches", "questions": "see comments above", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n\n- The paper is well-structured, clearly motivated, and presents a straightforward methodology.\n- The proposed CompDiffuser algorithm demonstrates good performance in stitching trajectories, outperforming baselines in most domains, particularly in longer horizon tasks (e.g., maze giant). Visualizations confirm that the generated plans are physically plausible.\n- The paper provides comprehensive ablation studies on various aspects, including chunk size, state dimensionality, parallel vs. autoregressive sampling, and replanning procedures.\n\n\nWeaknesses:\n\n- 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? \n- 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.\n\n- The trajectory merging approach appears to assume linear interpolation between produced trajectories. However, this might not be suitable for cases where two states are separated by obstacles (e.g., walls). Clarification on this aspect is necessary.\n\n- Although the method seems generalizable, 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) ? \n- 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 ? \n- the method seems to rely on high quality data, expert demonstration. when trained on navigate data, the method performance decreases", "was_revised": true, "final_justification": "the authors address my concerns, I raised my score", "rebuttal": "We thank Reviewer AQzW for the positive and careful review. Please see our response to the listed concerns below.\n\n> 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?\n\n\nThank 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.\n\nPlease see the additional implementation details below:\n\n**Architecture details:**\nFor 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.\n    \nFor 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.\n\n\n**Network capacity** (total number of parameters needed):\nFor experiment results for Table 1, we have the parameter count as follows:\nPointMaze from [1]\n| Env | DD | Ours |\n| - | - | - |\n| U-Maze |  35.7 M  | 59.8 M | \n| Medium |  35.7 M  | 61.1 M | \n| Large |   35.7 M  |  61.4 M | \n    \nOGBench PointMaze\n| Env | DD | Ours |\n| - | - | - |\n| Medium |  63.8 M  | 92.9 M | \n| Large |   63.8 M  |  92.9 M | \n| Giant |   63.8 M  | 92.9  M | \n\n    \nOGBench AntMaze\n| Env | Ours |\n| - | - |\n| Medium  | 92.9 M | \n| Large   |  92.9 M | \n| Giant   | 92.9  M | \n    \nOGBench HumanoidMaze\n| Env | Ours |\n| - | - |\n| Medium  | 95.8 M | \n| Large   | 95.8 M | \n| Giant   | 95.8  M | \n\n\n**Data Normalization:**\nWe use the standard min-max normalization to normalize the training data to range (-1, 1).\n\n**loss weighting:**\nWe do not use any weighting in the diffusion loss.\n\n    \n\n\n\n\n> 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.\n\n\n\nGenerative 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. \n\nMathematically 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\n\n$$\\epsilon_k[:H_o] = (\\epsilon_{k-1}[-H_o:] + \\epsilon_k[:H_o])/2$$\n$$\\epsilon_k[-H_o:] = (\\epsilon_{k}[-H_o:] + \\epsilon_{k+1}[:H_o])/2$$\n\n    \nBoth GSC and CompDiffuser can enable cross trajectory information exchange during denoising, hence promoting stitching capability. \nGSC 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).    \n\nIn 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.\n\n    \n\n\n> 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.\n\n\nWe do not use trajectory merging to avoid obstacles, and CompDiffuser as a planner is tasked to sample valid trajectories that satisfy the obstacle contraints.\n\nSpecifically, we propose CompDiffuser which stitches trajectories by generating a set of feasible and consecutive segments $\\tau_{1:N}$. \nThen, 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.\n\nAs 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}$.\n\nAdditionally, 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.  \n\n    \n\n\n\n> 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) ?\n\nThank you for the insightful suggestions.\n\nIn 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].\n\nDuring the rebuttal period, due to the time and computation constraints, we conduct additional experiments on the `cube-single-play-v0` env.\nSince 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. \nThat is, the dataset contains 1000 trajectories of length 1000 and we train CompDiffuser only on short segments of length 48. \nIn 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. \nSpecifically, 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.\n\n\nPlease see the results below. Results of the baselines are from the OGBench paper. \n\n    \n| Env | GCBC | GCIVL | GCIQL | QRL | CRL | HIQL | Ours |\n| - | - | - | - | - | - | - | - |\n| `cube-single-play-v0` | 6 ± 2 | 53 ± 4 |68 ± 6 | 5 ± 1 | 19 ± 2 | 15 ± 3 | 92±4 |\n\n\n\n\n\n\n\n\n> 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?\n\n\nThe 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. \n\nIn 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). \nQuantitatively, 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. \n\nFor further reference, we conduct additional experiments on OGBench `pointMaze-stitch` environments, where CompDiffuser predicts full state, both x-y positions and velocities (4D). \nPlease see the success rate comparison below.\n\n| Env | qpos | qpos+qvel |\n| - | - | - | \n| `pointmaze-medium-stitch-v0` | 100±0 | 100±0  |\n| `pointmaze-large-stitch-v0` | 100±0  | 100±0 |\n| `pointmaze-giant-stitch-v0` |  68±3   | 72±4 |\n    \n    \nFor 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. \nWe 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).\n\n\n\n    \n\n\n\n\n\n \n> W6. The method seems to rely on high quality data, expert demonstration. when trained on navigate data, the method performance decreases.\n\n\nIn our experiments, we evaluated CompDiffuser on various data qualities and the results demonstrated that CompDiffuser is robust to data quality.\n\nAs in Table 2, our method is able to outperform all the baselines on extremely low-quality datasets: the OGBench *Explore* datasets. \n(The data quality of these datasets is much lower than the OGBench *navigate* datasets.)\nThe 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). \n\nOur 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.\n\n\n[1] Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View."}, {"review_id": "C7ILoSWLR7", "reviewer": "Reviewer_qv69", "summary": "The authors present CompDiffuser a method for planning trajectories in a sampling driven generative manner based on diffusion. Long and potentially complex trajectories are decomposed into several shorter shorter trajectory segments to solve the overall motion planning problem. The conditional relationships between consecutive trajectory chunks are then learned using a bidirectional diffusion model. The proposed method enables both parallel generation of trajectory chunks and causal autoregressive generation and is shown to significantly outperform several relevant baselines.\nThe authors present an experimental evaluation of their approach both against multiple baselines and for several scenarios in simulation. Experimental simulation settings include 2d maze navigation problems and higher dimensional maze navigation settings with ant and humanoid embodiments where the method is shown to significantly outperform existing baselines.", "questions": "- 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?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The paper is clearly written and provides a good overview of the proposed method. \n- The supplementary material is exemplary and extensive and covers important aspects such as the failure mode analysis.\n- The paper tackles a crucial current shortfall of diffusion based methods in motion planning and appears to make significant advances over baseline state of the art methods.\n\nWeaknesses:\n\n- 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 and a real world implementation in the robotics context would also be interesting to see in future work.", "was_revised": true, "final_justification": "Following the authors' rebuttal I do maintain my accept recommendation for this paper which I believe proposes a valuable new method with the potential to have impact in motion planning for robotics. I do maintain my suggestion to the authors to consider additional more challenging motion planning problems in the experimental evaluation, but believe that the existing experimental evaluation is sufficient to accept the work to this conference.", "rebuttal": "We thank Reviewer qv69 for the positive and careful review. Please see our response to the listed concerns below.\n\n> 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.\n\n    \nThank you for your valuable comments.\n\nIn our experimental setup, we include tasks of varying difficulty levels — from relatively simple point maze navigation tasks to very challenging tasks such as AntSoccer.\n\nFor 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.\n\n\n\n> W1-2. A real world implementation in the robotics context would also be interesting to see in future work.\n\n\nThank you for the valuable suggestions.\n\nThe 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.\n\nSince 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.\n    \nAdditionally, 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. \n\nSpecifically, the dataset contains 1000 trajectories of length 1000 and we train CompDiffuser only on short segments of length 48. \nIn 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.\n\n\n\nPlease see the results below. Results of the baselines are from the OGbench paper. \n\n    \n| Env | GCBC | GCIVL | GCIQL | QRL | CRL | HIQL | Ours |\n| - | - | - | - | - | - | - | - |\n| `cube-single-play-v0` | 6±2 | 53±4 | 68±6 | 5±1 | 19±2 | 15±3 | 92±4 |\n    \n\n\n\n> 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?\n\n\nThank you for the insightful comments.\n\nWe 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.  \n\nOur 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.\n\nTo 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). \nWe can then denoise the whole graph in a similar fashion to Algorithm 2.\nSpecifically, 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."}]

Author rebuttal:

Meta-review: This paper introduces CompDiffuser, a novel diffusion-based composition method for generative planning. The paper is well-written, addresses an important problem, and presents a method that is clearly motivated and supported by extensive and successful results. It received unanimous positive reviews. During the rebuttal and discussion phases, the authors effectively addressed reviewers’ questions by providing new experimental results in a robotic manipulation environment, additional ablation studies on single-way conditioning at inference, a fair comparison with the Decision Diffuser baseline using matched model sizes, and further experiments on joint planning over both qpos and qvel. Overall, this is a strong and interesting paper that merits acceptance.
Decision: Accept (spotlight)