Diffused Task-Agnostic Milestone Planner
Authors: Mineui Hong
Minjae Kang
Songhwai Oh
Venue: NeurIPS 2023 poster (2023)
Keywords: Multi-task decision-making
Offline reinforcement learning
Planning
Diffusion model
TL;DR: This paper propose a method to plan latent milestones with diffusion model for multi-task, long-horizon decision-making.

Abstract: Addressing decision-making problems using sequence modeling to predict future trajectories shows promising results in recent years.
In this paper, we take a step further to leverage the sequence predictive method in wider areas such as long-term planning, vision-based control, and multi-task decision-making.
To this end, we propose a method to utilize a diffusion-based generative sequence model to plan a series of milestones in a latent space and to have an agent to follow the milestones to accomplish a given task.
The proposed method can learn control-relevant, low-dimensional latent representations of milestones, which makes it possible to efficiently perform long-term planning and vision-based control.
Furthermore, our approach exploits generation flexibility of the diffusion model, which makes it possible to plan diverse trajectories for multi-task decision-making.
We demonstrate the proposed method across offline reinforcement learning (RL) benchmarks and an visual manipulation environment.
The results show that our approach outperforms offline RL methods in solving long-horizon, sparse-reward tasks and multi-task problems,
while also achieving the state-of-the-art performance on the most challenging vision-based manipulation benchmark.

Diffused Task-Agnostic Milestone Planner
Mineui Hong, Minjae Kang, and Songhwai Oh
Department Electrical and Computer Engineering and ASRI Seoul National University
mineui. hon.rrilon@@@llb............ snu anu.ac.ke ac.kr, minjae. kang@rllab snu anu.ac.kr, songhwai@smu. ac kr
Abstract
Addressing decision-making problems using sequence modeling to predict future trajectories shows promising results in recent years. In this paper, we take a step further to leverage the sequence predictive method in wider areas such as long-term planning, vision-based control, and multi-task decision-makinn this end, we propose a method to utilize a diffusion- bassdd generative sequence model plan series of milestones in latent space and to have agent follow the milestones to accomplish a given task. The proposed method can learn control-relevant, low-dimensional latent representations of milestones, which makes it possible to efficiently perform long-term planning and vision-based control. Furthermore, our approach exploits generation flexibility of the diffusion model, which makes it possible to plan diverse trajectories for multi-task decision-making, We demonstrate the proposed method across offline reinforcement learning (RL) benchmarks and an visual manipulation environment. The results show that our approach outperforms offline RL methods in solving long- long-horizon, sparse-reward tasks and multi-task problems, while also achieving the state-of-the-aue performance on the most challenging vision-based manipulation benchmark.
To
Introduction
Developing general-purpose agent that can handle variety of decision-maker problems is long-standing challenge the field of artificial intelligence. In recent years, there has been growing interest in using offline reinforcement learning (RL) to tackle this problem [48, 35, 49]. One of the main advantages of offline RL is its ability to leverage the data comprised multiple sub-tasks by stitching together the relevant transitions to perform a new and more complex task [8]. This property of offline RL allows an agent to learn diverse behaviors demonstrated by a human or other robots performing similar tasks. In order to stitch the task-relevant uub-trajectories from the undirected multi task data, the majority the offline RL methods rely on temporal difference learning which measures the value of the policy by bootstrapping the approximated value functions. However, despite the successes in the online setting which allows an agent to actively collect data [30, 37, 14), the bootstrapping method often causes instability in offline learning. To address this problem, the offline RL methods require special treatments such as using batch-constrained policies [11, 32, 25] and penalizing the value of out-of-distribution samples [26]. Nevertheless, these methods often show unstable results depending on the how hyperparameters are tuned [9, 20].
To overcome the instability and tuning complexity of the offline RL methods, bootstrapping free methods that directly predict the future trajectories are also presented in recent studies [3, 20, 21, 1]. These methods, called sequence modeling methods, utilize the generative sequence models predict future states and actions conditioned on the goal state or the sum of rewards along the trajectory. In particular, Janner al. [21] and Ajay et al. [1] propose the methods to predict future trajectories using the denosing diffusion probabilistic models (DDPM), which have strengths in stitching the task relevant -rajrrjjcttoiies from offline data and generating feasible trajectories under various
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Offline D Dataset
Offline Dataset
tangled Injectory
Interval à- Uniff) Aauu))
Encoder
Encoder
Encoder
(optional)
Decoder
Milestuss
1o....
Critic
maximize jog (n )) log (1 DgG6..08
Diffuee::
Actor: maximize, logpo(a]5..8)) for mpbrno")
minimizeg E "rir
|
Figure Training process of DTAMP. (Left) DTAMP learns the latent milestones through goal-conditioned imitation learning manner. (Right) The diffusion model is trained to reconstruct the sequences of milestones sampled from the trajectories in the offline data.
combinations constraints (e.g., maximizing rewards, approaching to a given goal state). As result, DDPM based models are able to flexibly plan trajectories and show promising results multi-task learning. However, the inherent high computational cost of the denoising process limits its application to wider areas, such as real-time control and ussoonbbased tasks.
In this paper, we present method named Diffused Task-A Lask-Agnostic Milestone Planner (DTAMP) which aims to generate sequence of milestones using the diffusion model to have an agent to accomplish a given task by following them. By doing so, DTAMP divides long-horizon problem into shorter goal-reaching problems, and address them through soal-conditioned imitation learning which does not require the unstable bootstrapping method and has minimal hyperparameters to be tuned. Here, we note that a milestone represents a latent vector encoding a state (or an image), which guides an agent to reach the corresponding state. To learn the meaningful milestone representations, propose method to train encoder along with Left). By doing so, we expect the encoder to extract the control-relevant features and to embed them into milestones. Then, present method to train the milestone planner predict milestones for given goal. To this end, we first sample a series of intermediate states at intervals from the trajectory approaching the goal, and encode them into milestones train diffusion model to reconstruct them. (Figure 1, Right). Furthermore, we also propose a method to guarantee that the milestone planner predicts the shortest path the goal, based on the classifier-free diffusion guidance technique [17].
condiionddt actor and critic (Figure 1,
The main strengths DTAMP are summarized as follows: 1) The encoder trained through the proposed method encodes control-relevant features into compact latent representations, which enables efficient planning even if the observations are high-dimensional. 2) The diffusion planner enables flexible planning for variety of tasks and takes advantage multi-tass task decision-making problems. 3) Using the trained actor as a feedback controller makes an agent robust against the environment stochasticity and makes it possible to plan milestones only episode which largely reduces inference time compared to the existing sequence modeling methods. The proposed method is first evaluated on the goal-oriented environments in the D4RL benchmark [8]. Especially, the results on Antmaze environments show that DTAMP outperforms offline RL methods in long-horizon tasks without using the bootstrapping method. Finally, DTAMP is demonstrated on the CALVIN benchmark [29], which is the most challenging image-based, multi-task environment, and shows the state-of-the-an performance achieving 1.7 times higher success rate compared to the second-best method. The results show that our approach can broaden the area of application of sequence modeling, while being superior offline RL methods.
the beginning of an
Preliminaries
2.1 Goal-conditioned imitation learning
Multi-task decision-making problems are often treated as a goal-reaching problem, assuming that each task can be specified by goal state [28, 41, 35). Accordingly, we consider the goal-reaching

---

problem in our work and handle the problem through coal-conditioned imitation learning manner. We first formulate the problem with a goal-conditioned Markov decision process defined by a state space S, an action space A, a transition model p(s'|s. 1s1 a), an initial state distribution Po. and goal space G. We assume that there exists mapping f > 9. to determine whether the current 9 by determining whether f(s) equals to g. We further assume that there exists a set of previously collected data D which consists of the trajectories sampled by behavior policy p(a)s). Then, we define goal-conditioned imitation learning by training goal-conditioned actor Tv S x G - A to predict the most likely action to take for the goal g E g: maximizes log (ogu((ass.g) g), for S and (T(5.)) g).
state 8 E S has reached a given goal
(1)
We note that plals, g) = pancy measures defined by,
(a(ols,a)u(ale) where Pu(g|s, a) and Pu(gls) are the finite horizon occu- Pa(a(s)
Pu(g)s,a)=
Σ
Pu(st|s, and pulgls) = Σ Pulgls, a))(als))
(+E)))
(2)
1
s, =0)
where, Pu(s* [s, a) ( T
P(s st a,
Then the objective in (1) can be reformulated as loss function defined by,
J(s, gi v) == - (gls a) - log p(a))),, for a = (8.
We further note that the presented framework is highly related to the existing offline RL methods. Intuitively, the first term logp, (gll,, a) (15) is similar to the Q-functions used the studies that address goal-conditioned RL problems [6. 7], and the second term log Po (als) works as regularization to prevent the policy from taking -fd-triiiittittoionn actions [32, 11, 25].
2.2 Denosing diffusion probabilistic models
A denoising diffusion probabilistic model (DDPM) [18] is a generative model which learns to reconstruct the data distribution q(x) by iteratively denoising noises sampled from (N I). The diffusion models consist of the diffusion process which iteratively corrupts the original data x ~ q(x) into white noise, and the denoising process which reconstructs the data distribution q(x) from the corrupted data. First, in the diffusion process, a sequence of corrupting data x x (V11- Bn B,
{xt x2, ...., x" } is sampled from Markov chain defined by q(x"+1 where B, E (0, 1) is variance schedule and n = 0, 1.... ,N indicates each diffusion timestep. Then the denoising process aims to learn a denoising model pe(x"-1[xx" [x""), sequentially denoise the noisy data x ~ NO, I) into the original data x
The denoising model po (x(") can be trained by minimizing the following variational bound:
Eng DKL (g(xN (x(xx\x=)llx(xx))) + Σ DKL (g(x"= x²)
log po (x |x))
(4)
However, the variational bound (4) is often reformulated to a simpler rescaled loss function using function approximator eg to predict the noise injected into the original data.
m
J(x", nz == ( [|le
+ vI - a, n) where ā. = II (1 Bm). (5)
m==
The surrogate loss function (5) is known to reduce the variance of the loss and improve the training stability [18].
3 Diffused task aask-agnostic milestone planner
In this section, we introduce diffused aask-agnostic milestone planner (DTAMP). We first present the goal-conditioned imitation learning method for simultaneously training the encoder, the actor, and the critic functions. Secondly, the method to train the diffusion model to plan milestones is presented. In addition, we also propose a diffusion guidance method to let the milestone planner predict the shortest path to reach a given goal. Finally, we present sequential recision-making algorithm combining the goal-conditioned actor and the milestone planner. Note that more detailed explanation about our implementation of DTAMP is presented in Appendix D.
3

---

3.1 Goal-conditioned imitation learning from offline data
We first introduce our approach to train the goal-conditioned actor Tah using a discriminative function Da 8 A x G [0, 1] as the critic function to estimate the log-likelihood log Py) Here, we assume that goal is designated by single state Sgeal E & which can be a raw sensory observation an image. In order to extract the underlying goal representation 9 E & from the given goal state, we also train the encoder fw S - G alongside the actor and critic, where G is latent goal space. Note that the more discussion about how the encoder is trained through imitation learning and visualization of the learned goal space are presented in Appendix G.
in (15).
To estimate the log- legllikeihhodd log (g a) in (15), first train the discriminative critic D& to minimize the following cross-entropy loss:
Jentie (s; a, stt s : w) = log Do (o(")) log (1 Do a, fuls 1)), where, s ~ Ps(s), a ~ p(als), st ~ pu(st(|s & ~ Ps(s)).
(6)
We note that Py. (s) is prior distribution over the state space computed by Enyes Pu(s 50 Po( which is estimated by uniform distribution over the all data points of the offline data. Intuitively, for a given state-action pair (s, a), the positive sample st is sampled from the future states comes after (s, a), and the negative sample we expect the critic to be trained to distinguish between states that are more likely to come after the current state-action pair. This intuition can be formally formulated with the following proposition.
is sampled uniformly across the entire data distribution. As result,
Proposition 3.1. The optimal critic function D that minimizes the cross- cross-entropy loss (16) satisfies the following equation for given goal g = fa (Sgoa).
Pa 8, a) D)(s,a.g) ( a) (g)
Proof. Proof of Proposition A.1 A.1. is presented in Appendix A.
Using Proposition A.1, estimate Pr, (gls,, a) for given goal g with the trained critic function.
De(s,a,g) = (s.....g)
(8)
Now, the remaining part estimate log (a(ss) in (15) train the policy тỤ. However, although some previous studies try to estimate the marginal likelihood using conditional generative models [11] density estimation methods [47, 24], they require additional computational cost while not showing the significant improvement in performance [9]. Therefore, we circumvent this issue by replacing log Py (als) in the objective (15) into simple behavior-cloning regularization term. This regularization has been shown to be an effective way to stabilize offline policy learning requiring minimal hyperparameter tuning [9]. Consequently, the training objective for the policy function
defined as the following loss function:
Do(s,ū, u, Su(st)) + l|ā al?, Jacior (8, a 8 : w)) = -Alog 1 (s, ā fo fs(
where, s s~ ps)s,,
p(als), s ~ ps a), ā= = T T (s+))
where is a hyperparameter which balances the weight between the maximizing log-likelle likeliild log (fs [s, ā) and the regularization.
Please note that training the actor and the critic through the proposed goal-conditioned imitation learning method does not require temporal difference learning (or bootstrapping). Meanwhile, as the actor trained only for the positive goals, it has limitation inferring suitable action for negative goal that has never been reached from the current state in the data collection process. To overcome this limitation, we further propose diffusion-based milestone planner which predicts milestones to guide an agent to reach the distant negative goal state, in the next section.

---

end for Algorithm Sequential decision-making using DTAMP
1: Predict milestones 91.K ~ po Po gi Aurget $o Sgual), through denoising process (11).
2: Set the index for targeting milestone: 1.1.
3: Set internal timestep т + 0.
4: for each environment timestep do
5: Take an environment step with action (8t:
6: Increase internal timestep +- T +
7: if
(8++1) - 9s 22 < 8 or >> Tim then
Switch the target milestone to the next: à + + 9: Reset internal timestep +- 0.
8:
10: end
11:
3.2 Diffusion model as a milestone planner
We introduce our method train the diffusion model milestone planner guide the actor to reach the distant goal state. Intuitively. this can be done by letting the diffusion model to predict intermediate sub-goals that should pass. To this end, we first sample from offline data a pair of initial and goal state (So, Sgeal) and K K-intermediiae states {sa..... SKA] } at interval A, where 8(K+)) &gual- Then the sampled intermediate states are encoded into a series of latent milestones gl:K := [g1 92. 9KK). such that 9ı fw (SkA) for k 1, K. Then, the diffusion model is trained to reconstruct g by minimizing the following loss function:
Jaiffusion (gI:K, 80, ggaal 0) ((,.) [||e ea(vōngix VI- n n|/0, 8gma) ||1
(10)
Then, the trained milestone planner can predict milestones for given pair of initial and goal state (so, Sgea)) by sequentially denoising a series of noisy milestones g ~ (0, I):
gi.k N (NVV ( (gl. v1 anEn), v1 an-
on
(11)
for én (gT.K) n|50, Sgoal) and n N, -1,.,,,...
1,
3.3 Minimum temporal distance diffuusio guidance
The denoising process (11) makes it possible generate milestones for the given goal state, but does not guarantee that the planned trajectory is the shortest path. To address this issue we make an additional modification using classifier-free diffusion guidance [17]. Since we interested in minimizing the length of trajectories represented by milestones, minimizing the temporal distances between the pairs successive milestones is natural choice. Therefore, we modify the diffusion model to be conditioned on the temporal interval A of the sampled sequence of intermediate states {SA,"* AKA. ^A) The loss function (5) is then extended to the following unconditional loss function Juncond and conditional loss function Jeona:
Jumooma [ 0) (e-A((..) [|le - voloangii + VT- - 11
(12)
Jeond ( (1:K A, n; 0) (c-A((o)) [|le - ca(Vāngi. + vI a, A, n)] Then, the shortest path planning done with the guided prediction of é, in (11):
En E (g K. Ø, + [Ea (gü.k Amreet n) Ee E K Ø, n)],
(13)
where Aurget denotes the target temporal distance between milestones which is set to be relatively smaller than the maximum value Amax and / is a scalar value that determines the weight of the guidance. Please note that the conditions on the pair of initial and goal state (50, Sgoal) are omitted in (12) and (13) for the legibility.
3.4 Sequential decision-making using DTTAP
Using the learned milestone planner, we can predict series of milestones for given goal state, and let the agent to follow the milestones using the learned actor function. We note that the goal- conditioned actor works a feedback controller which adapt to stochastic transitions. This
5

---

Environment
CQL
IQL
ContRL
DT
TT(+IQL)
DD
DTAMP
medium-plas -lay
61.2
71.2
72.6
0.0
0.0 (81.9)
0.0
93.3
antmaze
medium-diverse
53.7
70.0
71.5
0.0
(85.7)
0.0
88.713386
large-play
15.8
39.6
48.6
0.0
0.0 (64.8)
0.0
80.0:3.27
large-diverse
14.9
47.5
54.1
0.0
0.0 (65.7)
0.0
78.0
Average
36.4
57.1
61.7
0.0
0.0 (74.5)
0.0
85.0
kitchen
mixed
52.4
51.0
65.0
74.4 1.39
partial
50.1
46.3
57.0
63.4
Average
51.3
48.7
61.0
68.916.30
Table 1: The table shows averaged scores on D4RL benchmarks. DTAMP was trained with three different random seeds, and evaluated with 100 rollouts for each random seed. The source of the baseline performance is presented in Appendix C.
property allows to perform the time-consuming denoising process only once the beginning of an episode. In contrast, the other sequence modeling methods [20, 21, 1] have to predict future trajectories at every timestep. which largely increases inference time. The remaining important question is how to decide whether the agent has reached the current milestone and move on to the next one. We address this issue by setting a threshold and make the switch if the distance between the current state and the targeted milestone measured in the latent space is less then the threshold. In addition, to make the agent recover from failure to reach milestone, we also set an internal time limit Tim and allow the agent to move to the next milestone when it fails to reach the current one within the time limit, Algorithm summarizes our method for sequential decision making.
decision-making.
Experiments
In this section we empirically verify that DTAMP is able to handle honggooizoo,, multi-task decision- making problems. We first evaluate the proposed method on the goal-oriented environments in the D4RL benchmark [8]. and show that our approach outperforms the other methods on the long-horizon, sparse-reward tasks. Then, we evaluate DTTAMP on the most challenging CALVIN benchmark [29], and show that DTAMP achieves the state-of-the- performance. Finally, we also present an ablation study to investigate how DTAMP's performance influenced by the minimum time distance guidance method (13) and the various other design choices of DTAMP. Please note that the detailed experiment settings are presented in Appendix B, and additional experiments on the D4RL locomotion tasks and stochastic environment are presented in Appendix H.
4.1 Offline reinforcement learning
D4RL benchmarks We first evaluate DTAMP on the goal-oriented environments in the D4RL benchmark (Antmaze and Kitchen environments) to verify that DTAMP can handle long-horizon tasks without bootstrapping the value function. These tasks are known to be especially challenging to be solved as the useful reward signals only appear when the agent reaches the goal states that are far apart from starting states. The proposed method is compared against offline RL basclines (CQL: conservative q-learning [26], IQL: implicit q-learning [25]. ContRL: contrastive reinforcement learning [7]), and the sequence modeling methods (DT: decision transformer [3], TT: trajectory transformer [20], DD: decision diffuser [1]). In addition, also compare DTAMP variant of TT
which utilizes a value function trained by IQL for using the beam search algorithm (TT+IQL), as proposed by Janner al. [20].
The results in Table show that DTTAP outperforms the baseline methods on the every environment. Especially, the other sequence modeling methods (DT, TT, and DD) that do not use temporal difference learning fail to show any meaningful result on the Antmaze environments. The reason for their poor performance is that the goal state of the Antmaze environment is too far from the imitial

---

Environment
IQL+HER
ContRL
TT+IQ
DTAMP
medium play
34.7 (-31.3%)
58.7 (-19.1%)
68.9 (-15.9%)
89.313.86
(-4.3%)
antmaze
multi-goul
medium-divers
50.7 (-27.6%)
57.3 (-19.9%)
75.6 (-11.8%)
76.7=4.50 (-13.5%)
large-play
46.0 (+16.2%)
24.7 (-49.2%)
28.9 (-55.4%)
62.012.94 (-22.5%)
large-diverse
42.0 (-11.6%)
26.7 (-50.6%)
23.3 (-64.5%)
53.3+9.67 (-31.7%)
Average
43.4 (-24.0%)
41.9 (-32.2%)
49.2 (-34.0%)
70.3 (-17.3%)
Table The table shows averaged scores on the multi-goal setting of Antmaze environments. The values in the parenthesises indicate performance difference between single-goal and multi-goal settings in percentage.
80
Avereged score
70
60
werrgee
COL
50
IGL
CDL
ContAL
IDL
40
DTAMP
50
DTAMP
10
10
10
10"
10
20
10
10"
Inference time
Inference time
(a) Antmaze environments
(b) Kitchen environments
Figure 2: The plots compare the inference times and the average scores of the different algorithms on the D4RL benchmarks. The r-axis indicates the inference time in logarithmic scale. The inference times are measured using an NVIDIA GeForce 3060 graphics card.
state' to accurately predict the whole trajectory to the goal. Meanwhile, DTAMP can efficiently plan a long-horizon trajectory by generating temporally sparse milestones.
Multi-goal Antmaze experiment In addition, we demonstrate DTAMP on Antmaze environments flexibly plan the milestones according to the
with multi-goal setting show that DTAMP varying goals. In this experiment, a goal for each rollout is randomly sampled from three different predefined target positions including the goals that were not reached from the initial state during the data-collection (see Appendix B for more details). We compare DTAMP against the three baseline methods (IQL [25], ContRL [7], and TT+IQL [20]), which achieve high scores in the single -oal setting. We note that DTAMP and ContRL already learn coll-conditioned actors and do require to be modified for the multi multi-goal setting. Meanwhile, we modify IQL and TT+IQL by adding hindsight experience replay (HER) [2]. In the multi-goal setting, DTAMP shows the least performance degradation compared to the other methods (Table 2). Especially, TT+IQL shows largest performance degradation compared to the results evaluated on the single-goal setting. This result indicates that the transformer architecture [45] used in TT struggles predict the trajectory to reach new goal that not seen in the training time, while DTAMP adaptively plan the milestones by stitching the relevant sub-trajectories.
Comparison of inference time In order to evaluate the computational efficiency of the proposed method, we further compare DTAMP against the other methods terms of inference time it takes to predict a single step of action. The comparative result in Figure 2 shows that our approach largely reduces the computational cost compared to the other sequence modeling methods, while achieving the highest average score. In particular, the inference time of DTAMP is about 1,000 times faster than TT and about 300 times faster than DD. is because using the actor as a feedback controller makes the agent robust on the stochasticity of the environment, and allows to perform the time-consuming sequence generation process only once at the beginning of an episode. In contrast, the other sequence
'about 500 timesteps apart

---

Initial state
Goal
Time
Planned milestones
Implicd tasks: (place pink block in drawer, mave slider left, "uum an led)
Images on the left show a pair of initial state and goal state, which implies three different tasks. The series of images the right represents the milestones planned by DTAMP. The image of each milestone is reconstructed using the trained decoder.
Figure 3:
modeling methods have to predict the future trajectories for every timestep which results in high computational cost.
4.2 Image-based planning
We now evaluate the proposed method on the CALVIN benchmark [29], which is the most challenging multi-task visual manipulation task. In this experiment, the observations and goals are designated by images, and reaching a goal image may requires the agent to sequentially accomplish multiple tasks. Figure shows an example that three different tasks ('place pink block in drawer', 'move slider left', and 'turn on led') are implied by single goal image. We utilize the dataset provided by Mees al. [29]. which consists of the trajectories performing a variety of manipulation tasks (e.g.. picking up a block, opening the drawer, pushing the button to turn on the led, etc.) collected by a human teleoperating the robotic arm the simulation. To investigate how the performance of DTAMP and other baselines varies with the difficulty of the task, we group the pairs of initial and goal states into three groups according to the number of tasks they imply, and evaluate the models on each group separately.
We note that two modifications are added DTAMP in this experiment, the CALVIN environment has the partially observable characteristic. First, we additionally train a decoder to reconstruct image observations from the encoded latent milestone vectors in order to let the encoder capture the essential visual features. Second, we utilize
policy as done by Lynch et al. [28] and Rosete-Bear al. [35], let the agent predict actions based on the current and previous observations. More detailed discussion of how modifications affect the performance of DTTAP will be presented in the ablation study.
In this experiment, our method is compared against an offline RL baseline combined with goal rela- belling (CQL+HER: conservative q-learning [26] with hindsight experience replay [2]) and hierarchi- cal RL methods (LMP: latent motor plans [28]. RIL: relay imitation learning [13]. TACO-RL: task- agnostic offline reinforcement learning [35]). Furthermore, we also evaluate Replanneplanning which allows the agent to re-plan the rest of milestones at every time it reaches a milestone. The results in Table shows that our approach achieves the sfate-of-the-ar performance in the CALVIN environment against existing methods by significant margin. We also would like note that only DTAMP and DTAMP-Replanning show the meaningful success rates on the hardest setting that goal image implies three different tasks in row. The results indicate that our approach enables planning on mmage-based environments, while also outperforming the other hierarchical RL methods that utilize skill-based policy (LMP and TACO-RL) or generate sub-goals (RIL). Furthermore, we also see that the performance of DTAMP can further improved by allowing agents to re-plan milestones during an episode.
4.3 Ablation study
Minimum temporal distance guidance We present an ablation study verify that the minimum temporal distance guidance (13) enables DTAMP to plan the shortest path for a given goal. We demonstrate DTAMP on simplified open Antmaze environment with varying target temporal distance Augett In order to train DTAMP, one million steps of offline data was collected by pretrained locomotion policy. Then, the unconditional version (10) and conditional version (12) of DTAMP were trained using the collected data. The trajectories planned by the unconditional model and conditional model are show in Figure 4. As expected, the guidance method enables DTAMP to plan trajectories of varying lengths according to the given target temporal distances, while the
8

---

Number of tasks
CQL +HER
LMP
RIL
TACO-RL
DTAMP
DTAMP
+Replanning
1
58.0
61.4
66.2
80.3
82,412.87
87.4
2
2.4
2.7
13.3
27.0
52.4::1.64
62.6:13.57
3
0.0
0.0
0.0
0.4
29.7#
38.1=5.78
Average
20.1
21.3
26.5
35.9
54.8.14.77
62.7+13.0
Table 3: The table shows averaged success rates on the CALVIN benchmark in percentage. DTAMP is trained with three random seeds, and the each model is evaluated with 1000 rollouts.
Lared Amx
Success rate
Timesteps to goal
(%)
0.1
97.7
178
1..7
#3.46
1.0
96.3
310
+1.89
110
Uncond
96.7
249
+1.70
+18.8
Figure 4: The figure visualizes the milestones
Table 4: The table shows averaged success
planned under different conditions (red: un- conditioned, green: Durget max = 0.1, blue:
rates and the number of timesteps taken reach the goal under different conditions.
Aarge Amax 1.0)
unconditional model tends plan sub-optimal trajectories. In addition, Table 4 shows that the guidance method is able to reduce the number of timesteps required to reach the goal, while it does affect the success rate.
Ablation study in image-based planning In order to apply DTAMP on the image-based planning problems, made two modifications DTAMP: 1) Training decoder to reconstruct images. 2) Utilizing skill-based policy. In this experiment we further investigate how each modification affects to the performance of DTAMP. To this end, evaluate three variations of DTAMP in the CALVIN benchmark (DTAMP Recon. Loss: the encoder is trained without reconstruction loss, DTAMP GCIIL Loss: the encoder is trained with only reconstruction loss without concicndit imitation learning loss DTAMP -TMMP-LMP: the agent does not utilize skill-based policy). The results shown Table 5 indicate that using both reconstruction loss and GCIL loss when training the encoder improves DTAMP's performance the most. We also find that ablating skill-based policy causes large performance degradation.
Related Work
Diffusion models [39, 18, 38] have received active attention recently. There exists a number of applications of diffusion model, such photo-realistic image generation [5, 34], text-guided image generation [36, 31, 33), and video generation [19, 15, 16, 50]. While the majority of work related to diffusion model is focused on image generation and manipulation, there is also growing interest in applying diffusion model on robotic tasks or reinforcement learning. The methods to plan a trajectory with diffusion models to perform given task are proposed by Janner et al. [21] and Ajay al. [1]. In the meantime, Wang et al. [46] and Chi et al. [4] present the methods to utilize diffusion models for predicting actions. Furthermore, Urain et al. [42] proposes method train cost function alongside with diffusion model, and to utilize it to select a proper grasping angle. It is worth mentioning that approach is the first work that utilizes diffusion model to plan trajectories in learnable latent space and handle visuomotor tasks, to the best of our knowledge.
Offline reinforcement learning has also been widely studied to improve data efficiency of rein- forcement learning (RI.) algorithms [30, 37, 44, 27, 14]. The most of RL. methods utilize temporal difference method to estimate the value of current action considering over the further timesteps. How-

---

Number of tasks
DTAMP
DTAMP Recon. Loss
DTAMP GCIL Loss
DTAMP
LMP
1
82.4
80.2
86.7
55,9
2
52.4
49.4
43.3
24.3
3
29.7
23.0
17.2
15.2
Average
54.8
50.9
49.1
31.8
Table 5: The results of ablation study on the CALVIN benchmark. The values represent success rates in percentage.
ever, bootstrapping the estimated values often causes overestimating the value of out-of out-of-distribution actions [10, 11]. which makes it challenging to apply RL fixed offline data. To tackle this issue, Fujimoto et al. [11] proposes batch constrained policy learning which restricts the policy to output the actions within the data. There also various methods to stabilize the offline RL are proposed, by learning policy through advantage weighted regression [32], augmenting q-learning objective [26, 25], model-based reinforcement learning [22]. In addition, applying offline RL on multi-task RI problems is showing promising results recent years [35], as it is able to extract various skills from offline data and compose them accomplish a new challenging task. Accordingly, our work also considers the offline setting that leverages the previously collected data, and mainly focus on designing method stitch the separated trajectories the offline data to accomplich given task.
Limitations
Cannot solve new task not included in dataset. In this work, we considered multi-task decision- making problems assuming that the offline data contains the segments of trajectories for performing the tasks we are interested in. However, there may also be cases where an agent needs to find new skills to solve new tasks that were not included in the previously collected data. Exploring the environment using the pretrained model to find new skills and continually learning perform more challenging tasks would be an interesting topic for the future work.
Efficient re-planning strategy is needed. Although the proposed method leverages faster inference time compared to other generative model-based planning methods by allowing the agent to plan milestones only once, our empirical results show that having the agent re-plan the trajectories during episodes can improve performance. Furthermore, in order to apply DTAMP on the problems that their environments are constantly changing, the agent is required re-plan the trajectories according to the changed environment. In this context, an efficient re-planning strategy to determine when and how to modify the planned path must be further studied to utilize DTAMP for wider area robotic tasks.
7 Conclusion
This paper presents a method to plan sequence of latent sub-goals, which is named as milestones, for a given target state to be reached. The proposed method trains an encoder to extract milestone representations from observations, and an actor to predict actions to follow the milestones, through goal-conditioned imitation learning manner. Our method also utilizes a denosing diffusion model to generate array of milestones conditioned on a target state. Furthermore, the minimum temporal distance diffusion guidance method is also presented in this paper, to make the proposed method plan the shortest path to reach the target state. Experimental results support that the proposed method makes possible to plan the milestones adapting to various tasks, and to manage long-tern teng-term tasks without using the unstable bootstrapping method.
Acknowledgements
This work was supported by Institute of Information & Communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (No. 2019-0-01199 [SW Star Lab] Robot Learning: Efficient, Safe, and Socially-Aceeptale Machine Learning).
10

---

References
[1] Anurag Ajay, Abhi Gupta Yilun Du Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making. arXiv preprint 11X1522211.15657, Dec 2022.
[2] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In Advances in Neural Information Processing Systems, Long Beach, US, Dec 2017.
[3] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, Virtual conference, Dec 2021.
[4] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchtiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), Daegu, KR, Jul 2023.
[5] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems, Virutal Conference, Dee 2021.
[6] Benjamin Eysenbach, Sergey Levine, and Ruslan Salakhutdinov. Replacing rewards with examples: Example-based policy search via recursive classification. In Advances in Neural Information Processing Systems, Virtual Conference, Dec 2021.
[7] Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Ruslan Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. In Advances in Neural Information Processing Systems, New Orleans, US, Dec 2022.
[8] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072/99 Apr 2020.
[9] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Advances in Neural Information Processing Systems, Virtual conference, Dec 2021.
[10] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximiation error in actor-critic methods. Proceedings of International Conference on Machine Learning. Stockholm, SE, Jul 2018
[11] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of the International Conference on Machine Learning, Long Beach, US, Jun 2019.
[12] Zhaohan Guo, Shantanu Thakoor, Miruna Pislar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean- Bastien Grill, Yunhao Tang, et al. Byol-cxplore: Exploration by bootstrapped prediction. In Advances in neural information processing systems, New Orleans, US, Dec 2022.
[13] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving horizoorizon tasks via imitation and reinforcement learning. In Proceedings of the Conference on Robot Learning, Osaka, JP, Oct 2019.
[14] Tuomas Haarnoja, Schoon Ha, Aurick Zhou, Jie Tan, George Tucker, and Sergey Levine. Learning to walk via deep reinforcement learning. Proceedings of the Robotics: Science and Systems Conference, Messe Freiburg, DE, Jun 2019.
[15] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Dietrich Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. In Advances in Neural Information Process- ing Systems, Virtual Conference, Dec 2022.
[16) Yingqing He, Tianyu Yang, Yong Zhang, Ying Shan, and Qifeng Chen. Latent video diffusion models for high-fidelity long video generation. arXiv preprint arXiv:2211.13221, Mar 2023.

---

[17] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. 12X15 Jul 2022.
arXiv preprint
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, Virtual conference, Dec 2020.
[19] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458, Apr 2022.
[20] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems, Virtual conference, Dec 2021.
[21] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In Proceedings of the International Conference on Machine Learning. Baltimore, US, Jul 2022.
[22] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. MOReL: Model-based offline reinforcement learning. In Advances in Neural Information Processing Systems, Virtual Conference, Dec 2020.
[23] Diederik P. Kingma and Jimmy Lei Ba. Adam: A method for stochastic optimization. In Proceedings of International Conference on Learning Representations, San Diego, US, May 2015.
[24] Ilya Kostrikov, Jonathan Tompson, Rob Fergus, and Ofir Nachum. Offline reinforcement learning with Fisher divergence critic regularization. In Proceedings of the International Conference on Machine Learning. Virtual Conference, Jul 2021.
[25] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In Proceedings of the International Conference on Learning Representations, Virtual conference, Apr 2022.
[26] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances Neural Information Processing Systems, Virtual conference, Dec 2020.
[27] Timothy P. Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Proceedings of International Conference on Learning Representations, San Juan, PR, May 2016.
[28] Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Proceedings of the Conference on Robot Leurning, Osaka, JP, Oct 2019.
[29] Oier Mees, Lukas Hermann, Erick Rosete-Beas and Wolfram Burgard. CALVIN: A benchmark for ainuuag-connitioong policy learning for long- horizon robot manipulation tasks. IEEE Robotics and Automution Letters, 7(3):7327-7334, Jul 2022.
[30] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 13Xiv:
Dec 2013.
[31] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. GLIDE: Towards photorealistic image generation and editing with text-guided diffusion models. In Proceedings of the International Conference on Machine Learning, Baltimore, US, Jul 2022.
[32] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off off-policy reinforcement learning. arXiv preprint arXiv: 1910000777 Oct 2019.
12

---

[33] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with CLIP latents. arXiv preprint arXiv:2204.06/25, Apr 2022.
[34] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceddings of the IEEE/CVF Conference Computer Vision and Patlern Recognition, New Orleans, US, Jun 2022.
[35] Erick Rosete-Beas, Oier Mees, Gabriel Kalweit, Joschka Boedecker, and Wolfram Burgard. Latent plans for task agnostic offline reinforcement learning, In Proceedings of the Conference on Robot Learning, Auckland, NZ, Dec 2022.
[36] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho. David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems, New Orleans, US, Dec 2022.
[37] David Silver, Aja Huang, Christopher Maddison, Arthur Guez, Laurent Sifre. George Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529:484 489, Jan 2016.
[38] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Proceedings of the International Conference on Learning Representations, Virtual Conference,
May 2021.
[39] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, Vancouver, CA, Dec 2019.
[40] Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal plan- ning networks. In Proceedings of International Conference on Machine Learning, Stockholm, SE, Jul 2018.
[41] Stephen Tian, Suraj Nair, Frederik Ebert, Sudeep Dasari, Benjamin Eysenbach, Cehlsea Finn, and Sergey Levine. Model-based visual planning with self-supervised functional distances. In Proceedings of the International Conference on Learning Representations, Virtual Conference, May 2021.
[42] Julen Urain, Niklas Fune, Jan Peters, and Georgia Chalvatzaki. SEf3)-DiffusiioFiitdd: Learning smooth cost functions for joint grasp and motion optimization through diffusion. arXiv preprint 2rXiv:2209.03855. Mar 2023.
[43] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9:2579-2666 2605, Nov 2008.
[44] Hado van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. Proceedings of AAAI Conference on Artificial Intelligence, Phoenix, US Feb 2016.
[45] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances Neural Information Processing Systems, Long Beach, US Dec 2017.
[46] Zhendong Wang, Jonathan J. Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.0611 0679,, Aug 2022.
[47] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 1911. Nov 2019.
[48] Ming Yin and Yu-Xiang Wang. Optimal uniform OPE and model-based offline reinforcement learning in home-nomogeneous, reward-free and task-agnostic settings. In Advances Neural Information Processing Systems, Virtual Conference, Dec 2021.
13

---

[49] Minjong Yoo, Sangwoo Cho, and Honguk Woo. Skills regularized task decomposition for multi- task offline reinforcement learning. Advances Neural Information Processing Systems, Long Beach, US, Dec 2022.
[50] Daquan Zhou, Weimin Wang, Hanshu Yan, Weiwei Lv, Yizhe Zhu, and Jiashi Feng. MagicViden Efficient video generation with latent diffusion models. arXiv preprint arXiv22211.1001, Nov 2022.
14

---

Appendix
A Proof of Proposition 1
Recap We consider a goal-conditioned Markov decision process defined by state space S, an action space A, transition model a)), and goal space G. We further consider an encoder fu S - G which maps goal state Sgaal E S into goal vector learning is defined by training a goal-conditioned actor Te. :
G. Goal-conditioned imitation - A to predict the most likely action that behavior policy μ(a(s) will take for given goal y E G:
maximizes log u(aas, g), for s E S and a = (s,
(14)
Using Bayes' rule, p(a|s, g) = P=(ols..)a(all)) where Pu(gls, a) and Pu(g]s) are the finite horizon Pa(gla)
occupancy measures, the objective in (14) can be reformulated as:
J( == (ols, a) - log p(als), for = Ts
(15)
In order to estimate the log- log-likelihood log (g]8, a), we train discriminative critic function Da S x A G + [0, 1] to distinguish between states that are more likely to come after the current state-action pair, by minimizing the following loss function:
Jenic (8, a, sts", :=: w) == - log Do fu(st)) - -log (1 - Do(s, a, fulss)))))) where, ~ (s(s), a ~ p(als), st ~ pa a). & ~ pp((s).
(16)
We now aim prove the following proposition:
Proposition A.1. The optimal critic function D that minimizes the cross-entropy loss (16) satisfies the following equation for given goal goal g = fw (ggoat).
Pa Pu(g s,a)
D7 a,(s, (n a) Pg)
(17)
Proof. We can rewrite the loss function (16) as.
((,,a;
= - Σ
(āls, a) log D Do(s,a, + Py (r) log (1 a Ø))]
(18)
BEC..
where, G. = [s)(s) E S}. Then, by taking derivative of (18) with respect to Da(s,a,0))
(Jss,a;0,00 Pp(gls,a) a) pp(g)
(19)
(D, Ds,,,,, 1- (w, a,))
At Do(s,a,9) = (s, o) the derivative in (19) becomes 0. Therefore,
Pa(g)s,a) D2(s, Pa(ols, + (g(
B Experiment setting
Multi-goal Antmaze experiments To evaluate DTAMP and baseline methods multi-goal setting, we predefined three different target positions as shown in Figure During the evaluation, a goal was randomly sampled from the three target positions for each rollout. The models were trained us- ing the same data
antmaze-mes metizee-ediun www.-m-diiiiiiiiee wu-tuum-iiii dinerseeee that was used for single-goal Antmaze experiments.
aatmaze-largo-diveeree
CALVIN experiments We used the same data split and tasks used Rosete-Bean et al. [35] to train and evaluate the model in the CALVIN experiment. We set time horizon of rollout to 300 timesteps when goal image implies one or two tasks, and 450 timesteps when goal image implies three tasks.
15

---

(a) Antmaze-medium
(b) Antmazo-large
Figure 5: Each position marked by an ant illustrates a predefined target position.
C Source of baseline performance
D4RL experiments The performances CQL [26], IQL. [25], and ContRI. [7] are taken from their corresponding papers. The performance of DT [3] in Antmaze environments is taken from Janner DD [1] in Kitchen environment is taken from its corresponding paper. To evaluate DD on Antmaze environments, DD is trained with three different random seeds using the author's ' code", and each trained model is evaluated with 30 rollouts (90 rollouts in total) for each task. While the performance of TT+IQL is presented in Janner et al. [20].
et al. [20]. The performance
inaccurate only 15 rollouts were used for evaluation. For more accurate evaluation, TT and IQL are trained with three different random seeds using the code' provided by the author, and each model is evaluated with 30 rollouts (90 rollouts in total) for each task.
Multi-goal Antmaze experiments IQL+HER is implemented based on the author's code by concatenating two-dimensional goal position on the original state vector. We train IQL+HER with three different random seeds and evaluate each model with 150 rollouts. As ContRL. already learns goal-conditioned policy, ContRL is trained with three different random seeds using the author's ' code" without modification. As also TT does require goal-specific training, we evaluate TT+IQL using the model trained for single-goal experiments while using the value functions trained by IQL+HER.
CALVIN experiments The performances of CQL+HER, LMP [28], RIL [13], and TACO-RL [35] are taken from Rosete-Bei Beas et al. [35] for the cases that one or two tasks are implied by a goal image. For the case that three tasks are implied by goal image, we trained each model using the code provided by Rosete-Beat et al. [35], and evaluated each model with 1, 1,000 rollouts.
D Implementation details
Encoder The encoder consists two neural networks one for each actor and critic. Each neural network has two hidden fully-connected layers with size of 512, and output layer. In addition, we normalized each output of the two neural networks, to provide consistent signal-to-noise ratio to the diffusion model during training process.
Actor and critic Each actor and critic has two hidden fully-connected layers with a size 512, and an output layer. We train five independent critic networks and use the smallest of the values predicted by critics to train actors, as done in Eysenbach et al. [7].
http://wwwnniiiiiiiiiiiidiisiiie 'The author thankfully provided a private github repository
https://github.c
"httss//gihuh.co//ooglllo/ "htps://github.comm http://hhhbb.om////ccceeeeet//no
16

---

Diffusion model The diffusion model for planning milestones is implemented based on the code provided by Ajay al. [1], which consists of temporal U-Net architecture with six residual 1-D convolutional blocks. To let the diffusion model predict milestones 91:K based on an initial state so and goal state Sgaal, we fixed go = fa (80) and gK+1 (8gen)) during the denoising process, as done in Janner et al. [21].
Image preprocessing Our image preprocessing method used for CALVIN experiment consists of three steps: 1) Resize RGB images size of 200 x 200 x 3 into 128 x 128 x 3. 2) Perform stochastic image shifts of 0-6 pixels. 3) Perform color jitter transform augmentation with a contrast of 0.1, brightness of 0.1 and hue of 0.02. 4) Normalize value of each pixel to let it falls between -1 and 1.
Visual perception For CALVIN experiments, utilize visual perception network consists of three convolutional layers and a spatial softmax layer. It is the same as the perception network used in Lynch et al. [28] and Posete-Beas al. [35].
Skill encoder and decoder The skill encoder and decoder are trained using the code provided by Rosete-Beas al. [35], and have the same architecture as LMP and TACO-RL.
E
Training details
D4RL experiments We train the encoder, goal-conditioned actor, critic and diffusion model simultaneously using a unified loss function: Junified := Jactor + Jeritic aiifiiuusion where the coefficient is fixed to 0.001 for all the experiments. We train the models for 2.5M training steps for Antmaze environments and 1.0M steps for Kitchen environments. We utilize NVIDIA Geforce RTX 3060 Ti graphics card for training our models, taking approximately 14 hours per 1.0M training steps. We use Adam optimizer [23], with learning rate of 0.0001.
CALVIN experiments We train our model for 1.0M training steps using the same graphics card used for D4RL experiments, taking approximately 40 hours per 1.0M training steps. We use the same optimizer and learning rate used in D4RL experiments.
F Hyperparameter setting
In this section, we describe hyperparameter details.
We set dimension of goal space 9 to 16 for Antmaze environments, and 32 for Kitchen and CALVIN environments.
We use the number of milestones environments, and K = 14 with maximum interval A = 16 for Kitchen and CALVIN environments.
30 with maximum interval Amax 32 for Antmaze
We set À in Equation (9) in the main paper to be adaptive scale of estimated log- likelihood log p(g]s; a), by setting it as X
के. Et..n.se log p(ols,a) as done in Fujimoto and Gu [9], where B denotes a mini-batch. We use X of 2.5 for Antmaze environments, and 0.05 for Kitchen and CALVIN environments.
We use diffusion timestep N of 300, diffusion guidance coefficient (3 of 0.5, and target temporal distance Aurget of 0.5Amax for the all experiments.
We use threshold 5 of 0.1 to determine whether the agent has reached the targeted milestone.
G Visualization of the learned latent goal space
In this section, we discuss about the ability of DTAMP to learn latent goal space that captures the dynamics of an environment. We train the encoder along with the actor and critic, through goal-conditioned imitation learning. By doing so, we make the encoder capture useful information from states, to let the actor and critic distinguish the states and actions that can lead an agent to the goal. There are also similar methods that have been studied through prior work, such as learning forward or inverse dynamics model in a latent space [40, 12] or using contrastive learning based on state occupancy measures [7]. Figure 6 visualizes an example of latent goal space learned by DTAMP
17

---

1.0
0.66
0.2
Figure The figure shows 1-SNE [43] embedding of the learned latent space. Each data point represents state in kitchen-mixed-vo data. The black x's indicate milestones planned by DTAMP, and the black triangles indicate three states sampled from an arbitrary trajectory. Each image was created by visualizing the data point closest to the corresponding milestone. The colors represent state-values estimated using IQL.
using kit siithen-mixed-ee data. The figure shows that the trained encoder can capture the dynamics of environment, and also let the milestone planner successfully predict milestones the learned latent space.
H Additional experiments
H.1 D4RL locomotion tasks
In this experiment, aim verify that DTAMP also deal with the general reinforcement learning problems in which the objectives of the tasks are specified by the reward functions. To this end, we evaluate DTAMP D4RL locomotion tasks (Halfcheetah, Hopper, and Walker2d). In order to let the milestone planner to predict a trajectory that maximizes the sum of rewards, we made the classifier-fre guidance method be conditioned on sum of rewards as done in Ajay al. [1]. The comparative result shown in Table 6 shows that DTAMP achieves a marginally higher average score compared the baselines, which indicates that the proposed method also can be used for general reinforcement learning problems. We would like to note that the D4RL. locomotion tasks provide dense rewards and can be performed by planning relatively short trajectories compared to antmaze tasks (to predict more than 100) timesteps forward does not affect much on the performance on the locomotion tasks). This explains why DTAMP does not show more significant performance improvement over Diffuser and DD in these tasks.
H.2 Robustness against environment stochasticity
In order demonstrate the robustness against environment stocahsticity, we further evaluate DTAMP size2d-umaze. environment of the D4RL benchmark, while adding stochasticity to the
18

---

Environment
CQL
IQL
DT
TT
Diffuser
DD
DTAMP
halfchectah-
42.6
47.4
42.6
46.9
44.2
49.1
47.3
halfchectah- medium-replas
36.6
44.2
36.6
41.9
42.2
39.3
42.6
91.6
86.7
86.8
95.0
79.8
90.6
88.2
hopper-medium
52.9
66.3
67.6
61.1
58.5
79.3
80.7
hupper-medium-rella
18.1
94.7
82.7
91.5
96.8
100.0
100.0
hopper-medium medium- expert
52.5
91.5
107.6
110.0
107.2
111.8
109.4
waldir2d-medium medium
75.3
78.3
74.0
79.0
79.7
82.5
82.7
walker2d-ment meddimm-repplae
26.0
73.9
66.6
82.6
61.2
75.0
79.5
(elierdddmdiumm exxpert
107.5
109.6
108.1
101.9
108.4
108.8
108.2
Average
77.6
77.0
74.7
78.9
75.3
81.8
82.1
Table 6: The table compares the performance of DTAMP evaluated on D4RL locomotion tasks with baselines. The result shows that DTAMP achieves the highest average score, which indicates that the proposed method also able to handle general reinforcement learning problems as well.
Diffuser
DD
DTAMP
Ref. max score
0.0
61.9
92.6
97.6
80.5
0.3
29.6
97.0
99.0
72.3 99.4
0.5
18.5
70.5
64.3
Table Normalized average score with different levels of stochasticity.
transition model. We also compare DTAMP against Diffuser [21]. which predicts future states and actions using diffusion model, and Decision Diffuser (DD) [1], which utilizes an inverse dynamics model to predict actions from the states planned by diffusion model. In this experiment we let Diffuser and DD predict the whole trajectory only once the beginning of each rollout done in DTAMP.
The environment either executes a random action a ~ Unif((A) with probability p or executes the action given by the agent with probability (1 p). Three sets of offline data were collected separately according to the choice of p. and each set of data consists of one million transitions collected by waypoint controller. We also note that the control frequency reduced threefold by repeating the same action three times in this experiment. By doing so, we reduce the time horizon of the environment (300 timesteps to 100 timesteps) to let Diffuser and DD predict the whole trajectory only once at the beginning of each rollout.
Setting We create three different levels of stochasticity by choice of p
The normalized score of each model is shown in Table 7. The column named Ref. max score (reference max score) shows the average score of the waypoint controller, which represents the performance of an optimal agent. We note that even if an agent predicts optimal actions, stochasticity of environment causes inherent performance degradation. In order to neglect this effect and only compare the robustness of each model, the scores shown in Table are normalized respect to the reference max scores (i.e., normalized score 100 X
Result
average score ) ToT. maN score
The result shows that DTAMP is robust against the environment stochasticity and does not show performance degradation in the terms of normalized score. On the other hand, Diffuser shows greater performance degradation as stochasticity increased. DD shows performance degradation when the stochasticity is small (p- 0.3) while showing performance degradation when the stochasiticity is large (p 0.5). The result indicates that approach using goal-conditioned actor makes DTAMP robust against environment stochasticity and allows to plan milestones only once for each rollout, while the other diffusion based planners do not.
19


Reviews:
[{"review_id": "wAPSCx8UWG", "reviewer": "Reviewer_Haes", "summary": "The paper proposes uses a diffusion-based generative model to plan a sequence of milestones in a latent space and have the agent follow this latent plan to accomplish a given task. The authors show  results on AntMaze, and more importantly on the CALVIN Benchmark to show the effectiveness of their method for image-based planning.  ", "questions": "- For an exhaustive one-to-one comparisons with existing methods, it’d be good to show experiments for HalfCheetah, Hopper, and Walker 2D as done in both Janner et al, and Ajay et al. Additionally, these aforementioned papers also include block stacking experiments which aren’t included in this manuscript.\n\n- I only partially understood Fig 2 in Supplementary. How did the authors generate corresponding images for certain milestones (X) in the planned trajectory? Are these images corresponding to the state the agent reached when it was \"closest\" to the milestone? ", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The direction of using diffusion-model based generative models for planning is definitely quite interesting. I also agree with the authors that generating the plan in a learned latent space is a promising direction.\n- I also thought that the use of classifier-free diffusion guidance to encourage the trajectory follow the shortest path was pretty neat.\n- I appreciate the authors presenting experiments on a partially observed environment like CALVIN! Showing that their method can come up with a reasonable latent plan for partially observed environment is promising.\n- I also liked that the authors tried visualising the learned latent space. (Fig 2. in Supp). The visualisation shows how the proposed approach, can plan successfully in the learned latent space to sequentially accomplish multiple tasks in the environment.", "weaknesses": "- In L175, the authors mention that their approaches can only do the diffusion step at the beginning once, compared to existing methods which have to predict future trajectories at every step. Can the authors describe how can the policy recover from failure to reach a `milestone' in their case?  If the agent's current state and the desired state is never closer than the given threshold, then according to the Alg. 1, the agent will never recover from this.\n- Can the authors comment on the performance of their approach will work when both the environment, and the state space is more complex and partially observed? Something that comes to mind is ImageNav in indoor environments [1] [2]. These environment seem more challenging since trajectories are often longer, and predicting intermediate milestones is much harder than CALVIN like benchmarks.\n- For Robustness against environment stochasticity experiments in the supplementary, isn’t the comparison to DD, and Diffuser a bit unfair? The experiments compare to DD and Diffusion by only allowing them to only predict the whole trajectory at once. But inference time related issues aside, DD and Diffuser model which predict the future trajectory at each timestep should be more robust to stochasticity in the environment than the methods which predict the whole trajectory at the beginning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for valuable comments, including acknowledging the contribution of our work and pointing out unclear explanations. \nWe present the responses to the reviewer's concerns and questions below.\n\n**Q. Can the authors describe how can the policy recover from failure to reach a `milestone' in their case?**\n\nA. To avoid the mentioned issue, we set a time limit and let an agent to move on to the next milestone when it fails to reach the current one within the time limit.\nWe found that the performance of DTAMP is not significantly affected by the small changes in the setting of time-limit and setting the time limit of about twice the maximum interval $\\triangle_{max}$ was sufficient for our experiments.\nWe agree that there was a lack of explanation, and we will include it in the final version of our paper.\n\n**Q. Can the authors comment on the performance of their approach will work when both the environment, and the state space is more complex and partially observed? (e.g. ImageNav)?**\n\nA. Experimental results in antmaze environments show that DTAMP can handle long-horizon tasks that require an agent to plan paths of up to a thousand of timesteps.\nTherefore, DTAMP would be able to successfully plan the trajectories for the environments that have even longer time horizon than CALVIN environment.\nIn our opinion, learning the milestone representation of images obtained from changing camera poses will be a major challenge for applying DTAMP to visual navigation domain such as ImageNav environment.\nMeanwhile, there are existing studies proposing promising ways to learn latent representation that capture 3D scenes by encoding radiance fields and utilize it for indoor navigation [1][2].\nCombined with these methods, DTAMP will also be able to address the challenging visual navigation problems as well.\n\n[1] Bautista et. al., \"GAUDI: A Neural Architect for Immersive 3D Scene Generation\", NeurIPS 2022  \n[2] Kwon et. al., \"Renderable Neural Radience Map for Visual Navigation\", CVPR 2023\n\n**Q. For Robustness against environment stochasticity experiments in the supplementary, isn’t the comparison to DD, and Diffuser a bit unfair?**\n\nA. We would like to note that the experiment of Table 1 in our supplementary material was performed to explain why DTAMP allows an agent to predict the trajectory only once whereas Diffuser and DD cannot.\nTo that end, we think the performed experiment sufficiently supports our claim:\nthe learned goal-conditioned actor can adapt to stochastic transitions and allows us to perform the time-consuming denoising process only once at the beginning of an episode (in Section 3.4 L174).\nIn addition, the result in Table 2 in our supplementary material demonstrates that allowing DTAMP to replan milestones during an episode makes it more robust to stochasticity and enhances its performance.\nHowever, we understand the confusion and will clarify the purpose of the experiment the revised version of our manuscript.\n\n**Q. For an exhaustive one-to-one comparisons with existing methods, it’d be good to show experiments for HalfCheetah, Hopper, and Walker2D.**\n\nA. Thank you for your suggestion to improve the impact of our work.\nTo reflect the reviewer's comment, we conducted further experiments on the mentioned tasks by modifying the diffusion guidance method to maximize the sum of rewards rather than to minimize temporal distance between milestones.\nThe attached PDF file presents the performance of DTAMP evaluated in the D4RL locomotion tasks, which achieves a marginally higher average score compared to the baselines.\nWe would like to note that the main purpose of our approach is to address long-horizon, sparse-reward problems and image-based tasks which the existing diffusion-based sequence modeling methods (Diffuser and Decision Diffuser) cannot handle.\nHowever, the D4RL locomotion tasks provide dense rewards and can be performed by planning relatively short trajectories compared to antmaze tasks (to predict more than 100 timesteps forward does not affect much on the performance on the locomotion tasks).\nThis explains why DTAMP does not show a more significant performance improvement over Decision Diffuser in these tasks.\nMeanwhile, we would like to emphasize that the greater contribution of our paper is to broaden the field where generative flexibility of diffusion models can be exploited, rather than performing better in the problems already covered by the existing diffusion-based models.  \n\n**Q. How did the authors generate corresponding images for certain milestones (X) in the planned trajectory?**\n\nA. To visualize each milestone, we selected a state in the \"kitchen-mixed-v0\" dataset, which is the closest to each milestone in the latent space.\nWe thank the reviewer for pointing out unclear explanation, and we will refer to it for revising the paper."}, {"review_id": "tAXV0WzmjT", "reviewer": "Reviewer_g1BQ", "summary": "The paper studies the problem of latent-space planning for sequential decision-making tasks using generative diffusion models. Similarly to works such as Diffuser, they train an endpoint-conditionned generative model to generate sub-goals along a path starting at a given state and reaching a certain goal. Differently to Diffuser, they perform the generation in a latent space pre-trained with imitation learning. The method achieves good results on D4RL environments as well as on a pixel-based planning task.", "questions": "* Can you elaborate on the distinction of your DTAMP method vs the Diffuser algorithm?\n* If possible, can you show a comparison against the method from “Contrastive Learning As a Reinforcement Learning Algorithm”?\n* Can you elaborate on the discrepancy in results between Decision Diffuser paper and your Table 2?\n* How was the noise schedule chosen for various experiments? If you perform state-based denoising, even in a latent space, I would imagine its structure to be somewhat different to the latent space of a pixel-based tasks.", "limitations": "Main limitation is that the empirical contribution and evaluation is focused on goal-conditionned tasks, which reduces to maximizing the 0-1 sparse reward function, as opposed to the general “sum-of-trajectory-rewards” objective.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-written and, for someone who is not expert in planning or diffusion models, can be very easy to understand.\n* The proposed approach makes sense in the current context, i.e. planning to visit subgoals in long-horizon tasks. I appreciate the approach and specifically the contribution related to this line of work, as it reduces advances in the field of planning and/or control to advances in the field of generative models (i.e. better diffusion models -> better goal trajectories -> better performance to some extent).\n* Experimental results show that the proposed method has competitive performance with baselines and latest SOTA, both on state-based and pixel-based domains.", "weaknesses": "* I am not sure how exactly the approach differs from Diffuser (Janner et al, 2022), at least in the second phase when generating subgoals? From what I understand, the difference is in the first phase, where a latent space with desirable features for control is learned.\n* Based on Figure 4 of the Decision Diffuser paper, the method achieves high returns on the D4RL locomotion task, while results from your Table 2 paper seem to indicate otherwise. Can you explain this discrepancy in results?\n* I am curious how the method performs against recent generative-model-based model-free algorithms such as the  contrastive value estimation line of work (Contrastive Learning As a Reinforcement Learning Algorithm by Eysenbach et al. 2023 and Contrastive Value Learning: Implicit Models for Simple Offline RL by Mazoure et al. 2023)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable reviews including the positive comment on legibility.\nWe present the responses to the reviewer's concerns and questions below.\n\n**Q. Can you elaborate on the distinction of your DTAMP method vs the Diffuser algorithm?**\n\nA. The main distinctive features of DTAMP compared against Diffuser are learning latent representation of subgoals (we named it milestones) and generating temporally sparse milestones instead of predicting the every states and actions connecting the current state and the goal.\nIt is an intuitive idea and easy to be built on an implementation of Diffuser, while providing large benefits in terms of reducing computational costs and making it possible to perform long-horizon and partially observable tasks which Diffuser cannot handle.\nFurthermore, proposing a diffusion guidance method to ensure the planned trajectory to be the shortest path to the goal is another contribution of our work, while Diffuser does not guarantee the planned trajectory to be the shortest.\n\n**Q. Can you show a comparison against the method from “Contrastive Learning As a Reinforcement Learning Algorithm”?**\n\nA. We already compared DTAMP against the method from the mentioned paper [1].\nThe algorithm is indicated by \"ContRL\" in Table 2 and Table 3 of our manuscript.\nWhile ContRL also utilizes a bootstrapping-free method to train a critic function, it shows limited performance in long-horizon tasks such as Antmaze environments.\nOn the other hand, our approach to generate a series of milestones divides a long-term problem into short-term problems and let an agent solve them more easily without using the bootstrapping method.\n\n[1] Eysenbach et. al., \"Contrastive Learning as a Goal-Conditioned Reinforcement Learning\", NeurIPS 2022\n\n**Q. Can you elaborate on the discrepancy in results between Decision Diffuser paper and your Table 2?**\n\nA. The performance of the Decision Diffuser shown in Figure 4 of its paper [2] was evaluated on D4RL locomotion tasks (halfcheetah, hopper, and walker2d), and the results in Table 2 of our paper were evaluated on D4RL antmaze tasks.\nThe D4RL locomotion tasks provide dense rewards (depending on the robot's velocity and posture at each time step), while the antmaze tasks provide a reward of +1 only when the robot reaches the goal.\nWhen dense rewards are provided, an agent can achieve high performance by planning relatively short trajectories (planning trajectories of about 100 timesteps is sufficient to accomplish the locomotion tasks), while antmaze tasks require an agent to plan trajectories of longer than 500 timesteps.\nAs a result, Decision Diffuser exhibits poor performance in antmaze tasks due to its limited ability to predict long-term trajectories.\n\n[2] Ajay et. al., \"Is Conditional Generative Modeling All You Need for Decision-Making?\", ICLR 2023\n\n**Q. How was the noise schedule chosen for various experiments?**\n\nA. We utilized the same noise schedule (cosine schedule) used by Diffuser and Decision Diffuser for fair comparison, and the same scheduling method was used for demonstrating DTAMP on image-based environments.\nIt is worth mentioning that the output of the encoder is normalized in our algorithm, so that the signal-to-noise ratio is consistent across various domains, which reduces the burden of choosing a different scheduling method for each different environment.\nWe agree that there was a lack of explanation about how we designed the diffusion model, and we will include it in the final version of our paper.\n\n**Q. Main limitation is that the empirical contribution and evaluation is focused on goal-conditioned tasks.**\n\nA. We conducted further experiments on the D4RL locomotion tasks (Halfcheetah, Hopper, and Walker2d) which provide dense rewards and should be handled by maximizing \"sum-of-trajectory-rewards\" objective. The additional experiments were done by modifying the diffusion guidance method to maximize the sum of rewards rather than to minimize temporal distance between milestones.\nThe attached PDF file presents the performance of DTAMP evaluated in the D4RL locomotion tasks, which achieves a marginally higher average score compared to the baselines.\nThis result indicates that DTAMP can also handle the general reinforcement learning problems as well with a small modification."}, {"review_id": "XHOfoICUUD", "reviewer": "Reviewer_d8Sq", "summary": "This paper uses sequence modeling method in applications like long-term planning, vision-based control, and multi-task decision-making. They formulate a novel method which uses diffusion-based generative sequence model to plan a series of milestones in a latent space and to have an agent to follow the milestones to get the task done. Their method can learn control-relevant, low-dimensional latent representations of milestones that makes it possible to efficiently perform long-term planning and vision-based control.\n\nThey train an encoder to extract milestone representations from observations, and an actor to predict actions to follow the milestones, using a goal-conditioned imitation learning fashion. Their method uses a denosing diffusion model to generate an array of milestones conditioned on a target state. The minimum temporal distance diffusion guidance method is used in this paper to make the proposed method plan the shortest path to reach the target state. ", "questions": "1) Did the authors re-implement all the baselines mentioned in Table2 for the comparison?\n\n2) why and how did authors take the decisions to built the diffusion model how they did?\n", "limitations": " 1) The proposed method might not work on the goal-reaching tasks and authors have not used it for the task not having goal reaching objective.\n2) Their method cannot solve a new task not included in dataset.\n\nAuthors can also elaborate on societal impact of their work. ", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The encoder trained using the method proposed, encodes control-relevant features into unique latent representations, which lets efficient planning even if the observations are high-dimensional. \n2) The diffusion-based planner enables flexible planning for a variety of tasks and takes advantage on multi-task decision-making problems. \n3) Using the trained actor as a feedback controller makes an agent robust against the environment stochasticity, and makes it possible to plan milestones only once at the beginning of an episode which largely reduces inference time compared to the existing sequence modelling methods. DTAMP has lower inference time compared to other baselines.\n4) DTAMP outperforms the baseline methods on the every environment.", "weaknesses": "1) Language of the paper could be improved. Some of the words are repetitive and there are quite some typos. for eg: \"letting the the diffusion \" and many other - at other places.\n2) DTAMP: the ablation study on how the design of diffusion model was done is missing.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments including summary of our contributions and pointing out typos we missed.\nWe present the responses to the reviewer's concerns and questions below.\n\n**Q. Language of the paper could be improved.**\n\nA. Thank you for your valuable comment for improving the quality of our paper, and we will refer to it for revision.\n\n**Q. Did the authors re-implement all the baselines mentioned in Table2 for the comparison?**\n\nA. We did not re-implement the baseline algorithms.\nThe performances of baselines are mostly taken from their original paper, except we evaluated Diffuser, Decision Diffuser and Trajectory Transformer on Antmaze environments using the code provided by the authors of their papers.\nFor a more detailed description of the source of baseline performance, please refer to Section C (Source of baseline performance) of the supplementary material.\n\n**Q. Why and how did authors take the decisions to built the diffusion model how they did?**\n\nA. We used a network architecture of the same diffusion model as Diffuser and Decision Diffuser for fair comparison.\nWe made this decision to verify that our approach can broaden the area of application of diffusion-based sequence modeling methods without changing the architecture of the diffusion model. \nFor a more detailed description of implementation details, please refer to Section D (Implementation details) of the supplementary material.\n\n**Q. The proposed method might not work on the task not having goal reaching objective.**\n\nA. We conducted further experiments on the D4RL locomotion tasks (Halfcheetah, Hopper, and Walker2d) which provide dense rewards and should be handled by maximizing \"sum-of-trajectory-rewards\" objective. The additional experiments were done by modifying the diffusion guidance method to maximize the sum of rewards rather than to minimize temporal distance between milestones. The attached PDF file presents the performance of DTAMP evaluated in the D4RL locomotion tasks, which achieves a marginally higher average score compared to the baselines. This result indicates that DTAMP can also handle the general reinforcement learning problems as well with a small modification.\n\n**Q. Authors can also elaborate on societal impact of their work.**\n\nA. The proposed method can be applied to a wide range of robotic tasks that require long-term planning in our society, e.g. autonomous driving, indoor navigation.\nIn addition, the advantages of our method of performing multiple tasks with a single agent can lead to cost savings in industrial robot development."}, {"review_id": "dv1uJJU2V0", "reviewer": "Reviewer_8zhN", "summary": "This paper introduces a hierarchical architecture named DTAMP to solve sequential decision-making problems. Specifically, the high-level part of the architecture is realized by a diffusion model to decompose the long-term goal into several short-term milestones. Then, the low-level part makes basic decisions at fixed intervals until the milestone is reached, and the policy is trained by goal-conditioned imitation learning. The results show that the proposed approach can achieve state-of-the-art performance on D4RL and CALVIN benchmarks.", "questions": "1. The authors state that the trained actor is limited in inferring a suitable action for a distant goal if the offline data does not contain a trajectory connecting the current state and the goal state at the end of Section 3.1. Then, when the amount of data decreases, will the performance of the proposed approach decrease significantly?\n2. In this paper, the number of planned intervals K is a fixed value. Will different values of K have a significant impact on performance, and is there a situation where the milestone cannot be reached after K intervals?\n", "limitations": "See above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The method combines the diffusion model with hierarchical reinforcement learning, which can achieve state-of-the-art performance on multiple benchmarks.\n2. The intermediate latent goals, i.e., milestones obtained by the proposed approach, can be represented not only by traditional discrete labels but also by images, as shown in Figure 4.", "weaknesses": "1. The biggest concern would be that the training process is not clearly explained. The authors do not state whether they trained the diffusion model and imitation learning separately or synchronously. If they're trained separately, how to obtain the labeled data of milestones? The authors should supplement enough information.\n2. In Table I, the target interval value denotes the target temporal distance between milestones, which is set to be relatively smaller than the maximum value. Then, there is no restriction on the target interval value when the ratio is 1.0, which should be the same as in the unconditioned experiment. The authors should make a further clarification on the increased timesteps when the ratio is 1.0.\n3. The structural relationship of Figure 1. is quite confusing. According to my understanding, the diffuser (as shown in the right subfigure) should be performed first, and then the actor and critic (as shown in the left subfigure) should be performed to obtain the trajectory. It is recommended to arrange the subfigures from left to right and establish the relationship between the two sub-figures.\n4. There are some minor mistakes in this paper, such as typos in Line 122, 176, and 'inference time' in Figure 3.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for valuable comments including summary of our contributions and pointing out some of our explanations that was not clear enough.\nWe present the responses to the reviewer's concerns and questions below.\n\n**Q. The biggest concern would be that the training process is not clearly explained.**\n\nA. We trained the diffusion model and imitation learning synchronously by adding up the losses as, \n$J_\\text{unified}=J_\\text{actor} + J_\\text{critic} + \\alpha J_{diffusion}$,\nwhere the coefficient $\\alpha$ is fixed to 0.001 for all the experiments.\nWe would like to mention that it is presented in Section E (Training details) of our supplementary material.\n\n**Q. The authors should make a further clarification on the increased timesteps when the ratio is 1.0.**\n\nA. We would like to note that setting a ratio of 1.0 leads to planning the longest path by making the temporal distance between milestones the maximum value $\\triangle_\\text{max}$.\nOn the other hand, in the case of planning without condition of target temporal distance, the generated path can have various lengths from the shortest length to the longest length.\nAs a result, the trajectories planned under the condition of $\\triangle_\\text{target}/\\triangle_\\text{max}=1$ shows longer average path-length than the trajectories planned unconditionally.\nWe will clarify how the setting of target temporal distance affects the length of generated trajectory in the revised version of our manuscript, to avoid the possible confusion.\n\n**Q. The structural relationship of Figure 1. is quite confusing.**\n\nA. We wanted to illustrate the training process of DTAMP in Figure 1.\nWe thought that we should explain how to learn milestone representations first, and then explain how to train the diffusion model to plan the learned milestones.\nHowever, we understand the confusion and we will rearrange and revise the figure for the final version of our paper to more clearly describe how the proposed method works.\n\n**Q. There are some minor mistakes in this paper, such as typos in Line 122, 176, and 'inference time' in Figure 3.**\n\nA. We appreciate for pointing out the typos we missed. \nIt will help a lot to revise our manuscript.\n\n**Q.  When the amount of data decreases, will the performance of the proposed approach decrease significantly?**\n\nA. No, the statement at the end of Section 3.1 does not mean that the performance of DTAMP will decrease significantly when the amount of data decreases. \nIt means that even if the trained actor may suffer to predict a suitable action for a distant goal, the proposed diffusion-based milestone planner provides milestones that can guide the actor, so that enables the actor to eventually reach the distant goal by following the relatively close milestones.\nWe understand the confusion and will clarify how the proposed milestone planner enhances sample efficiency for the long-horizon, sparse-reward tasks in the revised version of our manuscript.\n\n**Q. Will different values of K have a significant impact on performance, and is there a situation where the milestone cannot be reached after K intervals?**\n\nA. We agree that there is a lack of explanation for how changes in the number of milestones affect performance.\nTo answer to the reviewer's question, we constructed additional experiments to determine how changes in the number of milestones affect performance. The results of the additional experiments are presented in the attached PDF file. The results show that using the small number of milestones leads to marginal degradation of performance. This is because the temporal distance between milestones increases as the number of milestones decreases, which results to increase the burden of the goal-conditioned actor to reach distant subgoals. On the other hand, if the number of milestones exceeds a certain threshold (30 milestones for antmaze-medium task), the change in the number does not affect performance of DTAMP.\nFrom this empirical analysis, we can conclude the issue mentioned by the reviewer can be resolved by setting the sufficient number of milestones that covers the time horizon of the targeted environment."}, {"review_id": "S8Uc9PtCEU", "reviewer": "Reviewer_M7TA", "summary": "This paper extends diffusion-based latent milestone planners to long-term planning, vision control, and multi-task settings. Specifically, an encoder is trained to project observations into the latent space. The authors employ goal-conditioned imitation learning to train both the encoder and the prior goal-conditioned policy. After that, a diffusion probabilistic model is trained to model the milestone trajectories.\n\nIn order to generate the shortest path leading to the goal state more efficiently, the authors take into account the temporal distance between consecutive milestones in addition to the initial state and goal state. At the low level, the prior goal-conditioned policy is used to sample atomic actions, guiding the agent toward the given milestone.\n\nExperiments conducted on selected tasks from the Deep RL for Real Life (D4RL) and CALVIN benchmarks demonstrate superior performance over the baseline models. ", "questions": "- Why are there only three different goals in the multi-goal setting? Isn't the initial state and goal state assigned manually and the goal-conditioned plans are sampled with inpainting?", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper is well-organized. They proposed shortest-path guidance method is effective.", "weaknesses": "\n- I think the current experiments are not enough. Could the authors also provide results on the D4RL MuJoCo locomotion tasks? It would strengthen the paper if better performance could also be achieved there. \n- I'm wondering why the PointMaze tasks were not selected for comparison with the results from the Diffuser[1].\n- Can the authors ablate on the choice of image encoder in the image-based planning to verify that joint train a goal-conditioned actor and critic is necessary?\n\n\n\n[1]: Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In Proceedings of the International Conference on Machine Learning, Baltimore, US, Jul 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments including acknowledging our contribution on proposing an effective method to generate shortest path and suggestions to strengthen our paper. \nWe present the responses to the reviewer's concerns and questions below.\n\n**Q. Could the authors also provide results on the D4RL MuJoCo locomotion tasks?**\n\nA. Thank you for your suggestion to improve the impact of our work.\nTo reflect the reviewer's comment, we conducted further experiments on the mentioned tasks by modifying the diffusion guidance method to maximize the sum of rewards rather than to minimize temporal distance between milestones.\nThe attached PDF file presents the performance of DTAMP evaluated in the D4RL locomotion tasks, which achieves a marginally higher average score compared to the baselines.\nWe would like to note that the main purpose of our approach is to address long-horizon, sparse-reward problems and image-based tasks which the existing diffusion-based sequence modeling methods (Diffuser and Decision Diffuser) cannot handle.\nHowever, the D4RL locomotion tasks provide dense rewards and can be performed by planning relatively short trajectories compared to antmaze tasks (to predict more than 100 timesteps forward does not affect much on the performance on the locomotion tasks).\nThis explains why DTAMP does not show a more significant performance improvement over Decision Diffuser in these tasks.\nMeanwhile, we would like to emphasize that the greater contribution of our paper is to broaden the field where generative flexibility of diffusion models can be exploited, rather than performing better in the problems already covered by the existing diffusion-based models.  \n\n**Q. I'm wondering why the PointMaze tasks were not selected for comparison with the results from the Diffuser.**\n\nA. We have conducted experiment on a Pointmaze environment (U-maze env.) and presented the result in our supplementary material (see Section H of the supplementary material).\nIn this experiment, we added three different levels of stochasticity to the system dynamics to demonstrate DTAMP's robustness against environment stochasticity.\nAs a result, DTAMP achieved the score close to the maxium possible value at all three levels of stochasticity.\nThe reason why we did not further evaluated the proposed method on the other variations of Pointmaze environment (medium and large maze) is that we believe that the experimental results on Antmaze tasks (in Table 2 and Table 3 of our paper) are sufficient to show that DTAMP can plan trajectories for various goal positions in maze of various sizes.\nThis is because the Pointmaze tasks can be solved much more easily than the Antmaze tasks as the Pointmaze tasks take less timesteps to reach their goals and have a smaller action space, making point robots easier to control than ant robots.\nWe would like to further mention that the papers of Contrastive RL [1], Decision Diffuser [2], and Trajectory Transformer [3] also do not provide the performance of their algorithms on the mentioned environments.\nThis fact makes it challenging to fairly compare DTAMP with baselines on Pointmaze environments, as the environment specific setting of hyper-parameters might be needed to accurately evaluate the baseline methods.\n\n[1] Eysenbach et. al., \"Contrastive Learning as a Goal-Conditioned Reinforcement Learning\", NeurIPS 2022  \n[2] Ajay et. al., \"Is Conditional Generative Modeling All You Need for Decision-Making?\", ICLR 2023  \n[3] Janner et. al., \"Offline Reinforcement Learning as One Big Sequence Modeling Problem\", NeurIPS 2021\n\n**Q. Can the authors ablate on the choice of image encoder in the image-based planning?**\n\nA. Thank you for your suggestion to clarify our contribution on proposing a method to learn latent milestone representation.\nTo reflect the reviewer's comment, we conducted an additional experiment of DTAMP using a variational autoencoder (VAE) as an image encoder instead of training the encoder jointly with actor and critic.\nThe result of the ablation study is presented in the attached PDF file.\nOur empirical analysis shows that the representations learned by VAE cannot capture dynamical distances (distance in terms of how far apart are states in timesteps),\nwhich results in prediction of infeasible trajectories by DTAMP.\n\n**Q. Why are there only three different goals in the multi-goal setting?**\n\nA. We found that increasing the number of different goals increases the variance of success rates and requires more rollouts for accurate evaluation.\nEspecially, computational burden of Trajectory Transformer makes it difficult to simulate a sufficient number of rollouts (about an hour per episode).\nTherefore, we selected three goals that can include as diverse paths as possible on the maze."}, {"review_id": "DtT4M3CBE0", "reviewer": "Reviewer_xaPC", "summary": "This paper proposed a novel method to solve long-horizon, sparse-reward tasks and multi-task problems, which outperforms offline RL methods on many benchmarks.", "questions": "Mentioned in weakness.", "limitations": "No obvious societal impact.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Provide an elegant and general idea to solve sparse-reward problem -- which is hard for RL-based methods.\n- The paper is well-written and easy to be understood.\n- Both theoretical proof and empirical details are provided, which make the claims in the paper persuasive. \n- Good results.", "weaknesses": "- How much do the sampling methods matter? In the paper, the number of sampled points are fixed. If the authors can provide more ablation results with varying number of sampling points or use other sampling methods, the results will be more instructive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments including the positive comment on the soundness of our paper and the suggestion of ablation study to make our manuscript more instructive.\nWe present the responses to the reviewer's concerns and questions below.\n\n**Q. If the authors can provide more ablation results with varying number of sampling points or use other sampling methods, the results will be more instructive.**\n\nThank you for your suggestion of ablation study on the number of sampled milestones to make the results more instructive.\nTo reflect the reviewer's comments, we constructed additional experiments to determine how changes in the number of milestones affect performance.\nThe results of the additional experiments are presented in the attached PDF file.\nThe results show that using the small number of milestones leads to marginal degradation of performance.\nThis is because the temporal distance between milestones increases as the number of milestones decreases, which results to increase the burden of the goal-conditioned actor to reach distant subgoals.\nOn the other hand, if the number of milestones exceeds a certain threshold (30 milestones for antmaze-medium task), the change in the number does not affect performance of DTAMP.\n\n**Q. No obvious societal impact.**\n\nA. The proposed method can be applied to a wide range of robotic tasks that require long-term planning in our society, e.g. autonomous driving, indoor navigation.\nIn addition, the advantages of our method of performing multiple tasks with a single agent can lead to cost savings in industrial robot development."}]

Author rebuttal:
Dear all the reviewers.

We thank the reviewers for acknowledging the contributions of our work and for making constructive comments to improve the submitted manuscript.
We are pleased that reviewers find out the major strengths of the proposed method, which can be summarized as follows:
* Planning trajectories in the learned latent space using a diffusion model is interesting and promising way to solve partially observable problems.
* The proposed diffusion guidance method to predict the shortest path to a given goal is effective and neat.
* The proposed framework to use a goal-conditioned actor makes an agent robust against the environment stochasticity and largely reduces inference time.
* The proposed method shows state-of-the-art performance on the long-horizon, sparse-reward tasks and an image-based, multi-task benchmark.
* The paper is well-written and easy-to-understand.

We also find that the majority of the reviewers' comments can be summarized in two folds:
* Additional experiments would make the paper more instructive (demonstration on the locomotion tasks, ablation studies to investigate how the number of milestones affects the performance, etc.).
  - We constructed several additional experiments to reflect the reviewers' suggestions, and present the results in the attached PDF file. We hope that the extended empirical analysis will enrich our paper and resolve the reviewers' concerns.
* There are a few unclear explanations and typos that need to be revised.
  - We are sincerely appreciate for pointing out the shortcomings of our presentation. We will actively reflect the advice for revising the manuscript.

Finally, we thank again the reviewers for their valuable reviews and hope that our responses will be sufficient answers to their concerns and questions.

Meta-review: The paper is concerns with using diffusion-based models to generate multiple future milestones that are subsequently used for planning in control tasks.

The paper has received 1 x Strong Accept, 3 x Weak Accept, 2 x Borderline Accept.

All reviewers are exciting in seeing the application of diffusion models to control tasks and appreciate the exact formulation the authors propose — predicting long horizon trajectories in latent space. They find the approach novel (closest one is Diffuser but the current work does operate in latent space which is appreciated), and effective (the reviewers appreciate the results presented on several benchmarks, both more simple as well as pixel-based ones; also ones with partially observable spaces). The approach is also considered timely and promising research area (diffusion models for planning). Hence, the AC in consultation with SAC accepts the paper to Neurips.

The authors are encouraged to add the clarifications from the rebuttal (a more clear explanation of differences to Diffuser, explanation of the results are asked, clarification of the training procedure, additional D4RL results).
Decision: Accept (poster)