SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section we empirically verify that DTAMP is able to handle long-horizon, 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 DTAMP on the most challenging CALVIN benchmark [29], and show that DTAMP achieves the state-of-the-art performance. Finally, we also present an ablation study to investigate how DTAMP's performance is 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 a 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 baselines (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, we also compare DTAMP to a variant of TT which utilizes a value function trained by IQL for using the beam search algorithm  $(\mathbf{TT} + \mathbf{IQL})$ , as proposed by Janner et al. [20].
The results in Table 1 show that DTAMP 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 initial
Environment	IQL+HER	ContRL	TT+IQL	DTAMP
antmaze multi-goal	medium-play	34.7 (-51.3%)	58.7 (-19.1%)	68.9 (-15.9%)	89.3±3.86 (-4.3%)
medium-diverse	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.0±2.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±5.86 (-17.3%)
Table 2: The table shows averaged scores on the multi-goal setting of Antmaze environments. The values in the parentheses indicate performance difference between single-goal and multi-goal settings in percentage.
(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  $x$ -axis indicates the inference time in logarithmic scale. The inference times are measured using an NVIDIA GeForce 3060 Ti graphics card.
state<sup>1</sup> 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 with multi-goal setting to show that DTAMP can flexibly plan the milestones according to the 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-goal setting. We note that DTAMP and ContRL already learn goal-conditioned actors and do not require to be modified for the 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 the 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 to predict the trajectory to reach a new goal that was 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 in 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. It 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
Figure 3: Images on the left show a pair of initial state and goal state, which implies three different tasks. The series of images on the right represents the milestones planned by DTAMP. The image of each milestone is reconstructed using the trained decoder.
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 require the agent to sequentially accomplish multiple tasks. Figure 3 shows an example that three different tasks ('place pink block in drawer', 'move slider left', and 'turn on led') are implied by a single goal image. We utilize the dataset provided by Mees et 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 in 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 to DTAMP in this experiment, as 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 skill-based policy as done by Lynch et al. [28] and Rosete-Beas et al. [35], to let the agent predict actions based on the current and previous observations. More detailed discussion of how modifications affect the performance of DTAMP will be presented in the ablation study.
In this experiment, our method is compared against an offline RL baseline combined with goal relabelling (CQL+HER: conservative q-learning [26] with hindsight experience replay [2]) and hierarchical RL methods (LMP: latent motor plans [28], RIL: relay imitation learning [13], TACO-RL: task-agnostic offline reinforcement learning [35]). Furthermore, we also evaluate DTAMP+Replanning which allows the agent to re-plan the rest of milestones at every time it reaches a milestone. The results in Table 3 show that our approach achieves the state-of-the-art performance in the CALVIN environment against existing methods by a significant margin. We also would like note that only DTAMP and DTAMP+Replanning show the meaningful success rates on the hardest setting that a goal image implies three different tasks in a row. The results indicate that our approach enables planning on image-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 be 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 to verify that the minimum temporal distance guidance (13) enables DTAMP to plan the shortest path for a given goal. We demonstrate DTAMP on a simplified open Antmaze environment with varying target temporal distance  $\triangle_{\mathrm{target}}$ . In order to train DTAMP, one million steps of offline data was collected by a 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
Number of tasks	CQL +HER	LMP	RIL	TACO-RL	DTAMP	DTAMP +Replanning
1	58.0	61.4	66.2	80.3	82.4±2.87	87.4±1.43
2	2.4	2.7	13.3	27.0	52.4±1.64	62.6±3.57
3	0.0	0.0	0.0	0.4	29.7±6.80	38.1±5.78
Average	20.1	21.3	26.5	35.9	54.8±4.37	62.7±3.60
Figure 4: The figure visualizes the milestones planned under different conditions (red: unconditioned, green:  $\triangle_{\mathrm{target}} / \triangle_{\mathrm{max}} = 0.1$ , blue:  $\triangle_{\mathrm{target}} / \triangle_{\mathrm{max}} = 1.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.
△target/△max	Success rate (%)	Timesteps to goal
0.1	97.7	178
±1.77	±3.46
1.0	96.3	310
±1.89	±10.9
Uncond	96.7	249
±1.70	±18.8
Table 4: The table shows averaged success rates and the number of timesteps taken to reach the goal under different conditions.
unconditional model tends to 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 not affect the success rate.
Ablation study in image-based planning In order to apply DTAMP on the image-based planning problems, we made two modifications on DTAMP: 1) Training a 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, we evaluate three variations of DTAMP in the CALVIN benchmark (DTAMP-Recon. Loss: the encoder is trained without reconstruction loss, DTAMP-GCIL Loss: the encoder is trained with only reconstruction loss without goal-conditioned imitation learning loss, DTAMP-LMP: the agent does not utilize skill-based policy). The results shown in 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 a large performance degradation.
Diffusion models [39, 18, 38] have received active attention recently. There exists a number of applications of diffusion model, such as 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 a given task are proposed by Janner et al. [21] and Ajay et 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 a method to train a cost function alongside with a diffusion model, and to utilize it to select a proper grasping angle. It is worth mentioning that our approach is the first work that utilizes diffusion model to plan trajectories in a 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 reinforcement learning (RL) 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-distribution actions [10, 11], which makes it challenging to apply RL on 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], or model-based reinforcement learning [22]. In addition, applying offline RL on multi-task RL problems is showing promising results in recent years [35], as it is able to extract various skills from offline data and compose them to accomplish a new challenging task. Accordingly, our work also considers the offline setting that leverages the previously collected data, and mainly focus on designing a method to stitch the separated trajectories in the offline data to accomplish a given task.
Cannot solve a 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 to 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 to 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 a wider area of robotic tasks.
This paper presents a method to plan a 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 a goal-conditioned imitation learning manner. Our method also utilizes a denoising diffusion model to generate an 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 it possible to plan the milestones adapting to various tasks, and to manage long-term tasks without using the unstable bootstrapping method.
This work was supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2019-0-01190, [SW Star Lab] Robot Learning: Efficient, Safe, and Socially-Acceptable Machine Learning).
[49] Minjong Yoo, Sangwoo Cho, and Honguk Woo. Skills regularized task decomposition for multitask offline reinforcement learning. In Advances in Neural Information Processing Systems, Long Beach, US, Dec 2022.
[50] Daquan Zhou, Weimin Wang, Hanshu Yan, Weiwei Lv, Yizhe Zhu, and Jiashi Feng. MagicVideo: Efficient video generation with latent diffusion models. arXiv preprint arXiv:2211.11018, Nov 2022.
Recap We consider a goal-conditioned Markov decision process defined by a state space  $S$ , an action space  $\mathcal{A}$ , a transition model  $p(s'|s,a)$ , and a goal space  $\mathcal{G}$ . We further consider an encoder  $f_{\omega}: S \to \mathcal{G}$  which maps a goal state  $s_{\mathrm{goal}} \in S$  into a goal vector  $g \in \mathcal{G}$ . Goal-conditioned imitation learning is defined by training a goal-conditioned actor  $\pi_{\psi}: S \times \mathcal{G} \to \mathcal{A}$  to predict the most likely action that a behavior policy  $\mu(a|s)$  will take for a given 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 S \text {a n d} a = \pi_ {\psi} (s, g). \tag {14}
$$
Using Bayes' rule,  $\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, the objective in (14) can be reformulated as:
$$
J (s, g; \psi) := - \log p _ {\mu} (g | s, a) - \log \mu (a | s), \text {f o r} a = \pi_ {\psi} (s, g). \tag {15}
$$
In order to estimate the log-likelihood  $\log p_{\mu}(g|s,a)$ , we train a discriminative critic function  $D_{\phi}: S \times \mathcal{A} \times \mathcal{G} \to [0,1]$  to distinguish between states that are more likely to come after the current state-action pair, by minimizing the following loss function:
$$
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 {16}
$$
$$
\text {w h e r e ,} s \sim p _ {\mu} (s), a \sim \mu (a | s), s ^ {+} \sim p _ {\mu} (s ^ {+} | s, a), s ^ {-} \sim p _ {\mu} (s ^ {-}).
$$
We now aim to prove the following proposition:
Proposition A.1. The optimal critic function  $D_{\phi}^{*}$  that minimizes the cross-entropy loss (16) satisfies the following equation for a given goal 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 {17}
$$
Proof. We can rewrite the loss function (16) as,
$$
J (s, a; \phi , \omega) = - \sum_ {\tilde {g} \in \tilde {\mathcal {G}} _ {\omega}} \left[ p _ {\mu} (\tilde {g} | s, a) \log D _ {\phi} (s, a, \tilde {g}) + p _ {\mu} (\tilde {g}) \log \left(1 - D _ {\phi} (s, a, \tilde {g})\right) \right], \tag {18}
$$
where,  $\tilde{\mathcal{G}}_{\omega} = \{f_{\omega}(s)|s\in S\}$ . Then, by taking derivative of (18) with respect to  $D_{\phi}(s,a,g)$ ,
$$
\frac {\partial J (s , a ; \phi , \omega)}{\partial D _ {\omega} (s , a , g)} = - \frac {p _ {\mu} (g | s , a)}{D _ {\omega} (s , a , g)} + \frac {p _ {\mu} (g)}{1 - D _ {\omega} (s , a , g)}. \tag {19}
$$
At  $D_{\phi}(s,a,g) = D_{\phi}^{*}(s,a,g)$ , the derivative in (19) becomes 0. Therefore,
$$
D _ {\phi} ^ {*} (s, a, g) = \frac {p _ {\mu} (g | s , a)}{p _ {\mu} (g | s , a) + p _ {\mu} (g)}.
$$
B Experiment setting
Multi-goal Antmaze experiments To evaluate DTAMP and baseline methods on multi-goal setting, we predefined three different target positions as shown in Figure 5. During the evaluation, a goal was randomly sampled from the three target positions for each rollout. The models were trained using the same data (antmaze-medium-play, antmaze-medium-diverse, antmaze-large-play, antmaze-large-diverse) that was used for single-goal Antmaze experiments.
CALVIN experiments We used the same data split and tasks used in Rosete-Beas et al. [35] to train and evaluate the model in the CALVIN experiment. We set time horizon of a rollout to 300 timesteps when a goal image implies one or two tasks, and 450 timesteps when a goal image implies three tasks.
(a) Antmaze-medium
Figure 5: Each position marked by an ant illustrates a predefined target position.
(b) Antmaze-large
C Source of baseline performance
D4RL experiments The performances of CQL [26], IQL [25], and ContRL [7] are taken from their corresponding papers. The performance of DT [3] in Antmaze environments is taken from Janner et al. [20]. The performance of 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<sup>2</sup>, 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], it is inaccurate as only 15 rollouts were used for evaluation. For more accurate evaluation, TT and IQL are trained with three different random seeds using the code<sup>3</sup> 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<sup>4</sup> 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 a goal-conditioned policy, ContRL is trained with three different random seeds using the author's code<sup>5</sup> without modification. As also TT does not require a 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-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 a goal image, we trained each model using the code provided by Rosete-Beas et al. [35], and evaluated each model with 1,000 rollouts.
Encoder The encoder consists of two neural networks one for each actor and critic. Each neural network has two hidden fully-connected layers with a size of 512, and an 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 of 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].
$^{2}$ https://github.com/anuragajay/decision-diffuser/tree/main/code
<sup>3</sup>The author thankfully provided a private github repository
<sup>4</sup>https://github.com/ikostrikov/implicit_q_learning
<sup>5</sup>https://github.com/google-research/google-research/tree/master/contrastive_r1
<sup>6</sup>https://github.com/ErickRosete/tacorl
Diffusion model The diffusion model for planning milestones is implemented based on the code provided by Ajay et al. [1], which consists of a temporal U-Net architecture with six residual 1-D convolutional blocks. To let the diffusion model predict milestones  $g_{1:K}$  based on an initial state  $s_0$  and goal state  $s_{\mathrm{goal}}$ , we fixed  $g_0 = f_\omega(s_0)$  and  $g_{K+1} = f_\omega(s_{\mathrm{goal}})$  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 \times 200 \times 3$  into  $128 \times 128 \times 3$ . 2) Perform stochastic image shifts of 0-6 pixels. 3) Perform color jitter transform augmentation with a contrast of 0.1, a brightness of 0.1 and hue of 0.02. 4)Normalize value of each pixel to let it falls between -1 and 1.