SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mineui Hong, Minjae Kang, and Songhwai Oh
Department of Electrical and Computer Engineering and ASRI
mineui.hong@rllab.snu.ac.kr, minjae.kang@rllab.snu.ac.kr, songhwai@snu.ac.kr
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.
Developing a general-purpose agent that can handle a variety of decision-making problems is a long-standing challenge in the field of artificial intelligence. In recent years, there has been a 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 of 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 sub-trajectories from the undirected multi-task data, the majority of 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 to predict future states and actions conditioned on the goal state or the sum of rewards along the trajectory. In particular, Janner et 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 sub-trajectories from offline data and generating feasible trajectories under various
Figure 1: Training process of DTAMP. (Left) DTAMP learns the latent milestones through a 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 of constraints (e.g., maximizing rewards, approaching to a given goal state). As a result, DDPM based models are able to flexibly plan trajectories and show promising results on 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 vision-based tasks.
In this paper, we present a method named Diffused Task-Agnostic Milestone Planner (DTAMP) which aims to generate a sequence of milestones using the diffusion model to have an agent to accomplish a given task by following them. By doing so, DTAMP divides a long-horizon problem into shorter goal-reaching problems, and address them through goal-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, we propose a method to train an encoder along with a goal-conditioned actor and critic (Figure 1, Left). By doing so, we expect the encoder to extract the control-relevant features and to embed them into milestones. Then, we present a method to train the milestone planner to predict milestones for a 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 to train a diffusion model to reconstruct them. (Figure 1, Right). Furthermore, we also propose a method to guarantee that the milestone planner predicts the shortest path to the goal, based on the classifier-free diffusion guidance technique [17].
The main strengths of 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-based planner enables flexible planning for a variety of tasks and takes advantage on multi-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 once at the beginning of an 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-art 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 to offline RL methods.
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 a goal state [28, 41, 35]. Accordingly, we consider the goal-reaching
problem in our work and handle the problem through a goal-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  $\mathcal{A}$ , a transition model  $p(s'|s, a)$ , an initial state distribution  $\rho_0$ , and a goal space  $\mathcal{G}$ . We assume that there exists a mapping  $f: S \to \mathcal{G}$ , to determine whether the current state  $s \in S$  has reached a given goal  $g \in \mathcal{G}$  by determining whether  $f(s)$  equals to  $g$ . We further assume that there exists a set of previously collected data  $\mathcal{D}$  which consists of the trajectories sampled by a behavior policy  $\mu(a|s)$ . Then, we define goal-conditioned imitation learning by training a goal-conditioned actor  $\pi_{\psi}: S \times \mathcal{G} \to \mathcal{A}$  to predict the most likely action to take for the goal  $g \in \mathcal{G}$ :
$$
\operatorname {m a x i m i z e} _ {\psi} \log \mu (a | s, g), \text {f o r} s \in \mathcal {S} \text {a n d} a = \pi_ {\psi} (s, g). \tag {1}
$$
We note that  $\mu(a|s, g) = \frac{p_{\mu}(g|s, a)\mu(a|s)}{p_{\mu}(g|s)}$ , where  $p_{\mu}(g|s, a)$  and  $p_{\mu}(g|s)$  are the finite horizon occupancy measures defined by,
$$
p _ {\mu} (g | s, a) := \sum_ {s ^ {+} \in \{s \in \mathcal {S} | f (s) = g \}} p _ {\mu} (s ^ {+} | s, a) \text {a n d} p _ {\mu} (g | s) := \sum_ {a \in \mathcal {A}} p _ {\mu} (g | s, a) \mu (a | s),
$$
$$
\text {w h e r e}, p _ {\mu} \left(s ^ {+} \mid s, a\right) = \mathbb {E} _ {\mu \left(a _ {t} \mid s _ {t}\right), p \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right)} \left[ \frac {1}{T} \sum_ {t = 1} ^ {T} P \left(s _ {t} = s ^ {+} \mid s _ {0} = s, a _ {0} = a\right) \right]. \tag {2}
$$
Then the objective in (1) can be reformulated as a loss function defined by,
$$
J (s, g; \psi) := - \log p _ {\mu} (g | s, a) - \log \mu (a | s), \text {f o r} a = \pi_ {\psi} (s, g). \tag {3}
$$
We further note that the presented framework is highly related to the existing offline RL methods. Intuitively, the first term  $\log p_{\mu}(g|s,a)$  in (15) is similar to the  $Q$ -functions used in the studies that address goal-conditioned RL problems [6, 7], and the second term  $\log p_{\mu}(a|s)$  works as a regularization to prevent the policy from taking out-of-distribution 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(\mathbf{x})$  by iteratively denoising noises sampled from  $\mathcal{N}(\mathbf{0},\mathbf{I})$ . The diffusion models consist of the diffusion process which iteratively corrupts the original data  $\mathbf{x} \sim q(\mathbf{x})$  into a white noise, and the denoising process which reconstructs the data distribution  $q(\mathbf{x})$  from the corrupted data. First, in the diffusion process, a sequence of corrupting data  $\mathbf{x}^{1:N} = \{\mathbf{x}^1,\mathbf{x}^2,\dots ,\mathbf{x}^N\}$  is sampled from a Markov chain defined by  $q(\mathbf{x}^{n + 1}|\mathbf{x}^n) = \mathcal{N}(\sqrt{1 - \beta_n}\mathbf{x}^n,\beta_n\mathbf{I})$ , where  $\beta_{n}\in (0,1)$  is a variance schedule and  $n = 0,1,\dots ,N$  indicates each diffusion timestep. Then the denoising process aims to learn a denoising model  $p_{\theta}(\mathbf{x}^{n - 1}|\mathbf{x}^n)$ , to sequentially denoise the noisy data  $x^{N}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$  into the original data  $x^0$ .
The denoising model  $p_{\theta}(\mathbf{x}^{n - 1}|\mathbf{x}^n)$  can be trained by minimizing the following variational bound:
$$
\mathbb {E} _ {q} \left[ D _ {\mathrm {K L}} \left(q \left(\mathbf {x} _ {N} \mid \mathbf {x} _ {0}\right) \| p \left(\mathbf {x} _ {N}\right)\right) + \sum_ {n = 2} ^ {N} D _ {\mathrm {K L}} \left(q \left(\mathbf {x} ^ {n - 1} \mid \mathbf {x} ^ {n}, \mathbf {x} ^ {0}\right) \| p _ {\theta} \left(\mathbf {x} ^ {n - 1} \mid \mathbf {x} ^ {n}\right)\right) - \log p _ {\theta} \left(\mathbf {x} ^ {0} \mid \mathbf {x} ^ {1}\right) \right]. \tag {4}
$$
However, the variational bound (4) is often reformulated to a simpler rescaled loss function using a function approximator  $\epsilon_{\theta}$  to predict the noise injected into the original data.
$$
J \left(\mathbf {x} ^ {0}, n; \theta\right) := \mathbb {E} _ {\epsilon \sim \mathcal {N} (\mathbf {0}, \mathbf {I})} \left[ \| \epsilon - \epsilon_ {\theta} \left(\sqrt {\bar {\alpha} _ {n}} \mathbf {x} ^ {0} + \sqrt {1 - \bar {\alpha} _ {n}} \epsilon , n\right) \| ^ {2} \right], \text {w h e r e} \bar {\alpha} _ {n} = \prod_ {m = 0} ^ {n} (1 - \beta_ {m}). \tag {5}
$$
The surrogate loss function (5) is known to reduce the variance of the loss and improve the training stability [18].
3 Diffused task-agnostic milestone planner
In this section, we introduce diffused task-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 a sequential decision-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.1 Goal-conditioned imitation learning from offline data
We first introduce our approach to train the goal-conditioned actor  $\pi_{\psi}$ , using a discriminative function  $D_{\phi}: S \times \mathcal{A} \times \mathcal{G} \to [0,1]$  as the critic function to estimate the log-likelihood  $\log p_{\mu}(g|s,a)$  in (15). Here, we assume that a goal is designated by a single state  $s_{\mathrm{goal}} \in S$  which can be a raw sensory observation or an image. In order to extract the underlying goal representation  $g \in \mathcal{G}$  from the given goal state, we also train the encoder  $f_{\omega}: S \to \mathcal{G}$  alongside the actor and critic, where  $\mathcal{G}$  is a 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.
To estimate the log-likelihood  $\log p_{\mu}(g|s,a)$  in (15), we first train the discriminative critic  $D_{\phi}$  to minimize the following cross-entropy loss:
$$
J _ {\text {c r i t i c}} (s, a, s ^ {+}, s ^ {-}; \phi , \omega) := - \log D _ {\phi} (s, a, f _ {\omega} (s ^ {+})) - \log \left(1 - D _ {\phi} (s, a, f _ {\omega} (s ^ {-}))\right), \tag {6}
$$
where,  $s \sim p_{\mu}(s)$ ,  $a \sim \mu(a|s)$ ,  $s^{+} \sim p_{\mu}(s^{+}|s,a)$ ,  $s^{-} \sim p_{\mu}(s^{-})$ .
We note that  $p_{\mu}(s)$  is a prior distribution over the state space computed by  $\sum_{s_0 \in S} p_{\mu}(s | s_0) \rho_0(s_0)$ , which is estimated by a uniform distribution over the all data points of the offline data. Intuitively, for a given state-action pair  $(s, a)$ , the positive sample  $s^+$  is sampled from the future states comes after  $(s, a)$ , and the negative sample  $s^-$  is sampled uniformly across the entire data distribution. As a result, 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.
Proposition 3.1. The optimal critic function  $D_{\phi}^{*}$  that minimizes the cross-entropy loss (16) satisfies the following equation for a given goal  $g = f_{\omega}(s_{\text{goal}})$ .
$$
D _ {\phi} ^ {*} (s, a, g) = \frac {p _ {\mu} (g | s , a)}{p _ {\mu} (g | s , a) + p _ {\mu} (g)}. \tag {7}
$$
Proof. Proof of Proposition A.1 is presented in Appendix A.
Using Proposition A.1, we estimate  $p_{\mu}(g|s,a)$  for a given goal  $g \in \mathcal{G}$  with the trained critic function.
$$
p _ {\mu} (g | s, a) = p _ {\mu} (g) \frac {D _ {\phi} (s , a , g)}{1 - D _ {\phi} (s , a , g)}. \tag {8}
$$
Now, the remaining part is to estimate  $\log p_{\mu}(a|s)$  in (15) to train the policy  $\pi_{\psi}$ . However, although some previous studies try to estimate the marginal likelihood using conditional generative models [11] or 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 p_{\mu}(a|s)$  in the objective (15) into a 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 is defined as the following loss function:
$$
J _ {\text {a c t o r}} (s, \tilde {a}, s ^ {+}; \psi , \omega) := - \lambda \log \frac {D _ {\phi} (s , \tilde {a} , f _ {\omega} (s ^ {+}))}{1 - D _ {\phi} (s , \tilde {a} , f _ {\omega} (s ^ {+}))} + \| \tilde {a} - a \| ^ {2}, \tag {9}
$$
where,  $s \sim p_{\mu}(s)$ ,  $a \sim \mu(a|s)$ ,  $s^{+} \sim p_{\mu}(s^{+}|s,a)$ ,  $\tilde{a} = \pi_{\psi}(s,f_{\omega}(s^{+}))$ .
where  $\lambda$  is a hyperparameter which balances the weight between the maximizing log-likelihood  $\log p_{\mu}(f_{\omega}(s^{+})|s,\tilde{a})$  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 is trained only for the positive goals, it has its limitation in inferring a suitable action for a negative goal that has never been reached from the current state in the data collection process. To overcome this limitation, we further propose a diffusion-based milestone planner which predicts milestones to guide an agent to reach the distant negative goal state, in the next section.
Algorithm 1: Sequential decision-making using DTAMP
1: Predict milestones  $g_{1:K} \sim p_{\theta}(g_{1:K}^{0}|g_{1:K}^{N}, \triangle_{\text{target}}, s_0, s_{\text{goal}})$ , through denoising process (11).  
2: Set the index  $k$  for targeting a milestone:  $k \gets 1$ .  
3: Set internal timestep  $\tau \gets 0$ .  
4: for each environment timestep  $t$  do  
5: Take an environment step with an action  $a = \pi_{\psi}(s_t, g_k)$ .  
6: Increase internal timestep  $\tau \gets \tau + 1$ .  
7: if  $\| f_{\omega}(s_{t+1}) - g_k \|_2^2 < \delta$  or  $\tau > \tau_{\text{lim}}$  then  
8: Switch the target milestone to the next:  $k \gets k + 1$ .  
9: Reset internal timestep  $\tau \gets 0$ .  
10: end if  
11: end for
3.2 Diffusion model as a milestone planner
We now introduce our method to train the diffusion model as a milestone planner to 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  $(s_0,s_{\mathrm{goal}})$  and  $K$  -intermediate states  $\{s_{\triangle},\dots ,s_{K\triangle}\}$  at interval  $\triangle$  , where  $s_{(K + 1)\Delta} = s_{\mathrm{goal}}$  . Then the sampled intermediate states are encoded into a series of latent milestones  $g_{1:K}\coloneqq \{g_1,g_2,\dots ,g_K\}$  , such that  $g_{k} = f_{\omega}(s_{k\triangle})$  for  $k = 1,2,\dots ,K$  . Then, the diffusion model is trained to reconstruct  $g_{1:K}$  by minimizing the following loss function:
$$
J _ {\text {d i f f u s i o n}} \left(\boldsymbol {g} _ {1: K}, n, s _ {0}, s _ {\text {g o a l}}; \theta\right) := \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I})} \left[ \| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} _ {\theta} \left(\sqrt {\bar {\alpha} _ {n}} \boldsymbol {g} _ {1: K} + \sqrt {1 - \bar {\alpha} _ {n}} \boldsymbol {\epsilon}, n | s _ {0}, s _ {\text {g o a l}}\right) \| ^ {2} \right]. \tag {10}
$$
Then, the trained milestone planner can predict milestones for a given pair of initial and goal state  $(s_0,s_{\mathrm{goal}})$  by sequentially denoising a series of noisy milestones  $\pmb{g}_{1:K}^{N}\sim \mathcal{N}(\pmb {0},\pmb {I})$
$$
\boldsymbol {g} _ {1: K} ^ {n - 1} \sim \mathcal {N} \left(\sqrt {\frac {\bar {\alpha} _ {n - 1}}{\bar {\alpha} _ {n}}} \left(\boldsymbol {g} _ {1: K} ^ {n} - \sqrt {1 - \bar {\alpha} _ {n}} \hat {\boldsymbol {\epsilon}} _ {n}\right), \sqrt {1 - \bar {\alpha} _ {n - 1}} \boldsymbol {I}\right), \tag {11}
$$
$$
\text {f o r} \hat {\boldsymbol {\epsilon}} _ {n} = \boldsymbol {\epsilon} _ {\theta} \left(\boldsymbol {g} _ {1: K} ^ {n}, n | s _ {0}, s _ {\text {g o a l}}\right) \text {a n d} n = N, N - 1, \dots , 1.
$$
3.3 Minimum temporal distance diffusion guidance
The denoising process (11) makes it possible to 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 are interested in minimizing the length of trajectories represented by milestones, minimizing the temporal distances between the pairs of successive milestones is a natural choice. Therefore, we modify the diffusion model to be conditioned on the temporal interval  $\triangle$  of the sampled sequence of intermediate states  $\{s_{\triangle}, \dots, s_{K\triangle}\}$ . The loss function (5) is then extended to the following unconditional loss function  $J_{\mathrm{uncond}}$  and conditional loss function  $J_{\mathrm{cond}}$ :
$$
J _ {\text {u n c o n d}} \left(\boldsymbol {g} _ {1: K}, n; \theta\right) = \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (\mathbf {0}, \mathbf {I})} \left[ \| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} _ {\theta} \left(\sqrt {\bar {\alpha} _ {n}} \boldsymbol {g} _ {1: K} + \sqrt {1 - \bar {\alpha} _ {n}} \boldsymbol {\epsilon}, \varnothing , n\right) \| ^ {2} \right], \tag {12}
$$
$$
J _ {\text {c o n d}} \left(\boldsymbol {g} _ {1: K}, \bigtriangleup , n; \theta\right) = \mathbb {E} _ {\epsilon \sim \mathcal {N} (\mathbf {0}, \mathbf {I})} \left[ \| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} _ {\theta} \left(\sqrt {\bar {\alpha} _ {n}} \boldsymbol {g} _ {1: K} + \sqrt {1 - \bar {\alpha} _ {n}} \boldsymbol {\epsilon}, \bigtriangleup , n\right) \| ^ {2} \right].
$$
Then, the shortest path planning is done with the guided prediction of  $\hat{\epsilon}_n$  in (11):
$$
\hat {\epsilon} _ {n} = \epsilon_ {\theta} \left(\boldsymbol {g} _ {0: K} ^ {n}, \varnothing , n\right) + \beta \left[ \epsilon_ {\theta} \left(\boldsymbol {g} _ {0: K} ^ {n}, \triangle_ {\text {t a r g e t}}, n\right) - \epsilon_ {\theta} \left(\boldsymbol {g} _ {0: K} ^ {n}, \varnothing , n\right) \right], \tag {13}
$$
where  $\triangle_{\mathrm{target}}$  denotes the target temporal distance between milestones which is set to be relatively smaller than the maximum value  $\triangle_{\mathrm{max}}$ , and  $\beta$  is a scalar value that determines the weight of the guidance. Please note that the conditions on the pair of initial and goal state  $(s_0, s_{\mathrm{goal}})$  are omitted in (12) and (13) for the legibility.
3.4 Sequential decision-making using DTAMP
Using the learned milestone planner, we can predict a series of milestones for a given goal state, and let the agent to follow the milestones using the learned actor function. We note that the goal-conditioned actor works as a feedback controller which can adapt to stochastic transitions. This
Environment	CQL	IQL	ContRL	DT	TT(+IQL)	DD	DTAMP
antmaze	medium-play	61.2	71.2	72.6	0.0	0.0 (81.9)	0.0	93.3±0.94
medium-diverse	53.7	70.0	71.5	0.0	0.0 (85.7)	0.0	88.7±3.86
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±8.83
Average	36.4	57.1	61.7	0.0	0.0 (74.5)	0.0	85.0±4.84
kitchen	mixed	52.4	51.0	-	-	-	65.0	74.4±1.39
partial	50.1	46.3	-	-	-	57.0	63.4±8.80
Average	51.3	48.7	-	-	-	61.0	68.9±6.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 at 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  $\delta$  and make the switch if the distance between the current state and the targeted milestone measured in the latent space is less than the threshold. In addition, to make the agent recover from failure to reach a milestone, we also set an internal time limit  $\tau_{\mathrm{lim}}$  and allow the agent to move on to the next milestone when it fails to reach the current one within the time limit. Algorithm 1 summarizes our method for sequential decision-making.