Pre-Training Goal-based Models for Sample-Efficient Reinforcement Learning

Q1. About the novelty of PTGM and the differences to SPiRL [1].

A1. While both SPiRL and our work study pre-training skills on task-agnostic datasets for RL, we find issues in extending SPiRL-like methods to challenging domains and propose quite different technical approaches.

We observe that SPiRL fails in domains that feature large-scale datasets (with large amounts of data and diverse behaviors) + challenging open-world environments (with high-dimensional observation and action spaces, require long action sequences to execute certain behaviors). The reasons why SPiRL fails are threefold:

• SPiRL models behaviors in compact Gaussian latent variables $p(z|s_t,a_{t:t+k})$ regularized with KL divergence, which cannot accurately distinguish diverse behaviors in the large-scale dataset.
• The low-level policy in SPiRL decodes multi-step actions at one state $\pi(a_{t:t+k}|s_t,z)$, which cannot accurately reconstruct the high-dimensional long action sequences in open-world domains. Thus, it fails to execute many complicated skills in Minecraft (e.g. breaking a log requires taking more than 10 attack actions repeatedly).
• In SPiRL, the high-level policy acts in the continuous skill space $z$, making downstream RL inefficient (especially for long-horizon tasks in high-dimensional environments).

To address these issues respectively, PTGM provides novel approaches that:

• It models behaviors with goals (states in the dataset).
• The goal-conditioned policy $\pi(a_t|s_t,s^g)$ learns one-step action prediction on each state. Such a low-level policy can perform diverse skills accurately as presented in [2].
• The high-level policy acts in the discretized goal space. To this end, the methods of goal space clustering and goal prior regularization are also our original contributions to implementing such a framework.

We agree that this discussion is helpful for readers to understand our work better and add it to Appendix B.2.

[1] Accelerating Reinforcement Learning with Learned Skill Priors, 2021

[2] STEVE-1: A Generative Model for Text-to-Behavior in Minecraft, 2023

Q2. Section 3.1 says ‘We study tasks that provide binary rewards, offering a reward of +1 only upon reaching a non-trivial success state’.

A2. Thanks for pointing out our writing mistakes! We revised the above sentence since we use a MineCLIP reward in Minecraft. We claim in Section 5.2 (the VPT-finetune. paragraph) that such reward cannot lead vanilla RL to task success. MineCLIP only provides a high-reward when the target items mentioned in the task description are close in front of the agent [3]. The exploration in RL is difficult when the target is unseen.

[3] CLIP4MC: An RL-Friendly Vision-Language Model for Minecraft, 2023

Q3. The trajectories in the dataset could be generated by sub-optimal agents. Is fine-tuning the low-level policy with RL explored?

A3. We assume that the task-agnostic dataset contains non-trivial behaviors generated by the agent (e.g. human players) while performing various tasks in the environment. Though trajectories in the dataset are sub-optimal for solving downstream tasks, the short-term behaviors $a_{t:t+k}$ in the dataset represent meaningful skills and can be combined sequentially to accomplish downstream tasks better. Recent work [4] also demonstrates that stitching short-term goals from sub-optimal data can lead to better performance. We revised the second paragraph in Section 3.1 to clarify it.

Finetuning to further improve the performance of the low-level policy is promising and remains unexplored in the literature of pre-training for RL. In future work, we could pre-train the goal-conditioned policy with offline RL methods and finetune it with online data collected in downstream tasks.

[4] Waypoint Transformer: Reinforcement Learning via Supervised Learning with Intermediate Targets, 2023

Q4. The caption in Table 1 should be rephrased for clarity.

A4. Thanks! We have revised this caption.

Q5. What are the different tasks used for the dataset collection in Minecraft?

A5. We actually use only one task-agnostic human gameplay dataset introduced in VPT [5] to pre-train models for all downstream tasks. This satisfies our problem formulation that low-level behaviors are learned from a large task-agnostic dataset and the high-level policy can reuse the pre-trained skills to learn arbitrary tasks. For downstream tasks, the behaviors of combatting spiders and harvesting water buckets rarely appears in the dataset. Our method can stitch and generalize the goals to accomplish these tasks.

[5] Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos, 2022

Q1. The baselines are not sufficient. About using SPiRL, VPT, and Steve-1 as baselines.

A1. Thanks for pointing out these related works! We have to emphasize that the main literature we study is "pre-training from data to accelerate RL", in which we find specific issues in learning skills from large task-agnostic datasets in open worlds (the third paragraph in Section 1) and propose the first goal-based approach PTGM to address this problem. While GCRL is related to our study, most works in GCRL are orthogonal to our work in the problem formulation (though we admit that techniques proposed in GCRL can be alternative implementations of some components in PTGM). TACO-RL[1] studies hierarchical approaches in fully offline RL; DGRL[2] studies goal representation learning for GCRL, which is not an issue considered in our study and can be used as an alternative method to provide goal representations for PTGM as discussed in Section 6. We added these related works in Section 2.

Nevertheless, we take your comments and implement an online RL variant for TACO-RL to fit our settings. We update the results in Figure 2. While TACO holds competitive performance in harvest water bucket, it fails on other Minecraft tasks. We discuss that TACO has the following drawbacks compared with PTGM:

• TACO provides a continuous high-level action space for downstream RL, which makes the high-level policy likely to output unseen latent skills and is harmful for sample-efficiency.
• TACO requires manually picking a state representing the task goal for each downstream task. While the goal state for each task is deterministic in simple domains experimented in TACO-RL and SPiRL, the goal image for a Minecraft task can be diverse (e.g., for Harvest Log, there are various kinds of trees in Minecraft and the agent has diverse view angles), making it difficult to learn a good skill prior to enhance high-level RL.

About the necessity of our selected baselines:

• Though SPiRL was proposed in 2021, it is a strong baseline in the field of "pre-training skills from task-agnostic data for RL". Recent papers [3,4] under this setting also adopt this baseline.
• VPT+RL is actually the strongest baseline in Minecraft which adopts the same setting of pre-training + RL finetuning. As reported in [5], it can solve the most challenging tasks in Minecraft with billions of environment steps.
• Steve-1 uses the same goal-conditioned controller as our method and can serve as a strong baseline for Minecraft tasks. While Steve-1 maps goals to language instructions to build a multi-task agent, our high-level policy learned with RL can overcome its limitations that 1. cannot switch goals to perform hierarchical long-horizon control; 2. fails on tasks with unseen instructions.

[1] Latent plans for task-agnostic offline reinforcement learning, 2022

[2] Discrete factorial representations as an abstraction for goal conditioned reinforcement learning, 2022

[3] Skill-Critic: Refining Learned Skills for Reinforcement Learning, 2023

[4] Subwords as Skills: Tokenization for Sparse-Reward Reinforcement Learning, 2023

[5] Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos, 2022

Q2. About cluster visualization in Figure 4.

A2. Thanks! We have modified Figure 4 to visualize images within the same cluster. The updated figure shows that images within the same cluster demonstrate goals of similar behaviors in Minecraft like tree-chopping or mining. We put more visualization results in Appendix F.

Q3. Why does the clustering can produce meaningful goals without inductive bias?

A3: As presented in Section 5.1, in Kitchen, we cluster on the environment states which contain physical positions of objects in the scene and are good representations to distinguish different goals; In Minecraft, we adopt the MineCLIP image embedding pre-trained from Internet-scale YouTube videos with captions, which may be able to capture the information of semantic items in Minecraft. While our study focuses on the skill pre-training + RL framework, learning good representations for goals (especially for image-based environments) can be complementary to our work, as we discussed in Section 6.

Q1. Ablations are fairly limited and leave several unanswered questions.

A1. Thanks for your valuable suggestions to improve our ablation studies! We added ablation results in different number of clusters and low-level steps in Section 5.3. Results for both Log and Spider are presented in Appendix E and we will add results in all other tasks later. We summarize our conclusions below. More discussions can be found in Section 5.3 and Appendix E.

For the choice of cluster numbers:

• With N=10, the agent fails to improve the task success rates, due to the low capacity of behaviors in the small goal space. But the task success rates are non-zero, due to the generalization ability of the low-level policy.
• With an increasing number of goal clusters, the performance of PTGM is robust, while the sample efficiency and stability slightly decrease due to the larger high-level action space.
• PTGM-no-cluster fails on all tasks, which is caused by ineffectiveness of RL acting in high-dimensional continuous action spaces. It can be viewed as the case where infinite number of goals are either in clusters or inexistent in the dataset.

For the choice of the number of low-level steps:

• With a small number of low-level steps (k=10), the behavior is offloaded more to the high-level policy, resulting in worse sample efficiency of RL.
• With an increasing number of low-level steps, PTGM shows better sample efficiency and success rates. k=100 is a good choice for all the Minecraft tasks, where the high-level controller is able to switch about 10 goals per episode.
• With k=500, PTGM converges to a lower performance compared with k=100 because the behavior is offloaded too much to the low-level policy. In this case, the high-level policy cannot switch many goals in an episode, making the task performance limited by the capability of the low-level controller conditioned on a single goal.

Q2. Lacking implementation details on the method and the baselines. Explanation about the VPT baseline in Kitchen.

A2. We enrich the implementation details of our method in Appendix A and the baselines in Appendix B. As explained in Section 5.2 (the "VPT-finetune." paragraph), the VPT baseline in Kitchen is named BC-finetune, which pre-trains a behavior-cloning low-level policy and finetunes it with RL. Its failure is related to the lack of temporal abstraction for RL and the forgetting issue during RL finetuning. Pertsch et al. [1] report the similar results of this baseline in Kitchen.

[1] Accelerating Reinforcement Learning with Learned Skill Priors, 2021

Q3. Why don't we initialize the high-level policy with the goal prior in RL?

A3. We conduct an ablation study that initializes the high-level policy with the goal prior model (PTGM-prior-init). As shown in Appendix E, Figure 11, PTGM-prior-init outperforms PTGM only on harvesting logs, but has worse performance in other Minecraft tasks. We find that since chopping trees is the most frequent behavior in the Minecraft Contractor dataset, the behavior-cloning models learned from such data can draw higher probabilities on harvesting logs than other behaviors. Thus, PTGM-prior-init is biased to the task of harvesting logs and fails to explore for other tasks in which the goal prior model draws low probabilities on the task-relevant goals. On the contrary, PTGM with random initialization of the high-level policy adopts a uniform goal prior for exploration, showing strong capabilities in solving out-of-distribution tasks. Since our problem formulation is to learn from task-agnostic data and finetune for arbitrary tasks, we believe that uniformly exploring in the goal space for RL is better.

Q4. Are discrete goals really necessary for the simpler Kitchen environment?

A4. We answer this question by implementing the action space of the high-level policy with the state space of Kitchen without clustering (PTGM-no-cluster). The results are presented in Appendix D, Figure 8. Though the dimension of the goal in Kitchen is relatively lower (21 dimensions), we find that PTGM-no-cluster fails in this task. We observe that, with the continuous action space, the high-level policy should learn to output valid goals that lay in the manifold of states from the dataset to make the low-level policy perform reasonable behaviors. On the contrary, for PTGM with discrete goal clusters, the output of the high-level policy is always a valid goal (in the clusters), making RL efficient. Though the KL reward encourages PTGM-no-cluster to output goals close to the goal prior, it cannot make the output accurate enough due to the Gaussian sampling of the action head.