Paper: Pre-Training Goal-based Models for Sample-Efficient Reinforcement Learning (ICLR-2024)

Reviewer comment (Novelty):
**Novelty**: PTGM draws on several ideas from existing work by Pertsch et al.[1]. It would be helpful to get a better understanding of the similarities and differences between the two methods to judge the novelty of the work.

Author response:
**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

Reviewer comment (Experiments):
**Section 3.1** says *‘We study tasks that provide binary rewards, offering a reward of +1 only upon reaching a non-trivial success state’*. However, the Minecraft results have a MineCLIP reward that appears to be dense. Since the majority of the results are in Minecraft, I would revise the above sentence.

Author response:
**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

Reviewer comment (Experiments):
The low-level policy is learned purely from the dataset. However, the trajectories in the dataset could be generated by sub-optimal agents which could result in a sub-optimal low-level goal-conditioned policy. The assumptions of the dataset and the properties of the low-level policy obtained from that should be clarified further. On the same note, I am curious to know if fine-tuning the low-level policy with RL was explored.

Author response:
**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

Reviewer comment (Presentation):
The caption in **Table 1** should be rephrased for clarity. In particular, *‘two rows’* is confusing as the actual data is contained in a single row.

Author response:
**Q4.** The caption in Table 1 should be rephrased for clarity.

**A4.** Thanks! We have revised this caption.

Reviewer comment (Evaluation):
It would be interesting to see the number of different (unique) goals output by the high-level policy during a trajectory. For example, for a trajectory of length 1000 with k=100, are all 10 goals output by the high-level policy different?

Author response:
**Q6.** About the number of unique goals output by the high-level policy during a trajectory. 

**A6.** We pick the last training checkpoints to test for 1. the number of unique goals picked per episode; 2. the probability that the high-level policy picks a goal that is different from the last selected goal. We test each checkpoint for 100 episodes and report average results in the Table below.
 
 Task | Kitchen | Log | Water | Cobblestone | Iron | Spider
 ------------- | ------------- | ------------- | ------------- | ------------- | -------------  | ------------- 
 Unique goals  |  13.71  |  1.42  |  2.30  |  1.00  |  1.12  |  1.59  
 Change probability | 0.57 |  0.11  |  0.32  |  0.00  |  0.01  |  0.34  

In Kitchen, Log, Water, and Spider tasks, the learned high-level policy in PTGM picks various goals in an episode and is possible to switch to a different goal at each step. In Kitchen, since the downstream task is explicitly defined with four sub-tasks, the high-level policy should pick a large number of unique goals on average.
In the Cobblestone task, the high-level policy can simply choose to dig down without switching to other goals, thus the policy converges to always picking a unique goal. The Iron task is similar to Cobblestone but is much more challenging in its long-horizon nature. Though it does not require switching for diverse goals, only PTGM finds a good policy that chooses to dig down for many high-level steps and achieves a high success rate on this task (Figure 2).

Reviewer comment (Novelty):
With reference to **2)** in weakness, what are the similarities and differences between PTGM and SPiRL [1]?

Author response:
**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

Reviewer comment (Writing):
Task and goal seem to be used interchangeably in the paper. For more clarity, it would be useful to clearly define the difference between the two terms in the *Problem Formulation* section.

Author response:
**Q7.** It would be useful to clearly define the difference between task and goal in the Problem Formulation section.

**A7.** Thanks for your suggestion. In our paper, (downstream) tasks refer to long-horizon tasks defined with MDP. Goals refer to the last states of short-term behaviors in the dataset, which can represent low-level skills. 
In problem formulation, we distinguish tasks and low-level behaviors and have not mentioned goals. We revised Section 3.2 to make the definition of goals clear.

Reviewer comment (Evaluation):
I don’t think the visualization in Figure 4 can support that the clustered goals are meaningful. You can always find images from different clusters representing different behaviors. What I would expect is that images clustered into the same group should present similar patterns.

Author response:
**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.

Reviewer comment (Theory):
Why does the clustering can produce meaningful goals? The clustering is done without any prior knowledge or inductive bias. If the image background is noisy, can the proposed clustering method work?

Author response:
**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.

Reviewer comment (Experiments):
Ablations are fairly limited and leave me with several unanswered questions that seem important to address given the technical contributions of the paper. Firstly, it is evident that too few clusters (10) fails to capture the diversity of the goal space, and that no clustering fails to learn at all. However, it is not clear to me what the breaking point would be in terms of number of clusters: would e.g. 5,000 clusters lead to similar or *more* diverse behaviors than 500 clusters, or would it collapse to similar performance as the no-clustering ablation? Similarly, it is not clear based on the ablation on the number of low-level steps that more than 100 steps (thus offloading more of the learning to behavior cloning rather than RL) would lead to worse behaviors. Given that the low-level policy is trained on a very large dataset, I imagine that behavior cloning over longer horizons, e.g., 1,000 steps, could still lead to very meaningful behaviors. Additionally, # of clusters and # of steps are highly dependent hyper-parameters given that they jointly balance how much of the behavior should be offloaded to the high-level RL policy vs. the low-level BC policy, but I didn't find any discussion or results that highlight this.

Author response:
**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.

Reviewer comment (Reproducibility):
The paper is lacking in terms of implementation details on the proposed method + the baselines. It would be helpful if the paper was more self-contained and described the overall architecture etc. in the appendix rather than simply referencing prior work.

Author response:
**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

Reviewer comment (Experiments):
I understand the motivation behind the KL weight and that neither alpha=0 (no prior) nor too large of a weight are desirable. However, it appears that the authors choose to train the high-level policy from scratch and only leverage the goal prior to guide exploration. Given that the goal prior and high-level policy share the same action space, why do the authors decide against initializing the high-level policy as the goal prior and simply finetuning it using the proposed objective (reward + KL)?

Author response:
**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.

Reviewer comment (Experiments):
I am left wondering how many of the design choices in the proposed framework are uniquely beneficial for MineDojo due to its enormous state and action space. For example, are discrete goals really necessary for the simpler Kitchen environment compared to providing the raw physical state as a continuous goal? The MineDojo results are surely impressive, but it would be useful to contextualize the proposed method and design choices more for other domains besides MineDojo, with or without additional experiments to back any claims.

Author response:
**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.