Discovering Hierarchical Achievements in Reinforcement Learning via Contrastive Learning

Summary

This paper focuses on the problem of sequential decision-making within a hierarchical framework, where tasks exhibit a hierarchical decomposition structure, and the agent does not possess any prior knowledge of the task dependency graph. In contrast to prior hierarchical approaches that directly model dependencies and utilize two-level policies for task resolution, it investigates an even more demanding scenario. Specifically, the agent lacks information about the unlocked achievements within each episode. To tackle this challenge, the paper employs PPO as its backbone RL algorithm. Interestingly, it is found that the representations learned by PPO possess some ability to predict the next locked achievement, albeit with limited confidence. To further enhance this predictive capability, the paper proposes two contrastive loss mechanisms as representation learning objectives. These two mechanisms help the representation to be able to predict the agent’s next unlocked achievement (intra-trajectory contrastive loss) and also learn a representation of the achievement so that it is not capturing any environment-specific or spurious features (cross-trajectory matching).

Strengths

1. The author proposes a novel self-supervised loss within the context of hierarchical decision-making, which could be easily integrated into the existing RL algorithm (PPO). 2. Good empirical performance boost compared with the previous strongest model-based RL algorithms (Dreamer-v3) on the Crafter environment.

Weaknesses

My concerns is mostly around how relaxing the assumption of the agent's knowledge of its unlocked achievement during an episode. In particular, whether it is realistic in real applications, and how much does this affect the overall performance (in other words comparison to a baseline method that does assume such prior knowledge). Please see my detailed comments in the Questions section.

Questions

Here are my questions about the paper. Given the limited amount of time for rebuttal, I understand that it may not be possible to address all of the points with comprehensive additional experiments. But I would be happy to raise my score if my following concerns are adequately addressed. 1. **Assumptions of the problem setting**: The paper introduces an additional level of difficulty by assuming that the agent does not possess knowledge of the specific achievements unlocked throughout each episode. This assumption diverges from previous model-based and hierarchical approaches, which rely on such information. It would be helpful if the authors could provide a real-world example or scenario where relaxing this assumption becomes necessary. This would aid in understanding the practical applicability of such an assumption. From my perspective, this assumption appears weak, as it is difficult to envision an agent successfully completing an episode without awareness of the achievements attained during the process. 2. **Backbone RL algorithm**: The paper extensively discusses the effectiveness of the Proximal Policy Optimization (PPO) algorithm as the backbone for addressing tasks with hierarchical structures. The authors demonstrate how the representations learned by the PPO agent exhibit predictive abilities for the next achievement, as outlined in sections 3.1 and 3.2. However, it would be valuable to clarify whether these findings are specific to policy gradient methods or if they can be extended to value-based approaches as well. Understanding if the proposed contrastive loss mechanisms can be applied to value-based reinforcement learning algorithms would contribute to a more comprehensive evaluation of their potential. 3. **Baseline Comparison**: It would be insightful to compare the performance of the two proposed contrastive losses with an "oracle" objective, where the agent does possess prior knowledge of its unlocked achievements. For instance, for intra-trajectory achievement prediction, a simple 22-way classification could be employed as the loss, predicting the next achievement based on the current state-action pair. Similarly, for the cross-trajectory objective, minimizing the distance between representations of the same achievements could be considered. Such a comparison would facilitate a better understanding of the impact of relaxing the assumption regarding the agent's knowledge of achievements on the overall algorithm performance.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

This paper introduces achievement distillation, a representation learning method that is combined with PPO to obtain state-of-the-art results on the 2D crafter benchmark. First, the authors demonstrate that with simple hyper-parameter tweaks, the performance of vanilla PPO can be greatly improved. Next, detail the three components of achievement distillation which leverage achievement labels $g$. Intra-trajectory achievement prediction uses a contrastive objective to maximize the similarity between state action pairs and the next achievement. Cross-trajectory achievement matching maximizes the similarity between state-action pairs of the same achievement across episodes using optimal transport. Finally, they use the achievement representations as memory by concatenting the last achievement representation to the policy and value inputs. This results in a method that achieves high performance on the Crafter benchmark, especially for difficult to reach achievements.

Strengths

The paper was generally easy to follow. The experiments improving the performance of vanilla PPO were exciting! It’s great to see better tuned baselines. The improvements in PPO also naturally flowed into the introduction of achievement distillation. Achievement distillation is an interesting form of representation learning, and to my knowledge, is novel. The authors conducted experiments only on the Crafter environment. The results in this benchmark are compelling as it is extremely challenging, and as the authors highlight, achievement distillation can perform better on very hard to achieve tasks. The authors additionally include ablation studies by subtracting components of their method. I found the ideas in the paper to be easy to follow and straightforward in a good way.

Weaknesses

The approach appears slightly overfit to the chosen Crafter benchmark, and the authors do not test achievement distillation in any other settings. Though the results on Crafter are compelling, it would be great if the authors could provide more examples and discussion of when this could be applied to other environments, and if so, how. The comparison to baselines is a bit mis-leading, as other approaches, like DreamerV3, do not make use of the achievement labels necessary for achievement distillation. That being said, even PPO outperforms these baselines! I also found the log-scale on the axes of Figure 6 to be a bit confusing – is a success rate < 0.01% really significant? I generally recommend acceptance of this work, primary drawbacks being its limited applicability to broader settings and evaluation on only one environment.

Questions

Major: Since crafter is partially observed, do you use a recurrent network? Minor: Line 141: outperforms dreamer: for the same number of environment steps? I would expected model-based to have higher sample efficiency, though it may have lower asymptotic performance.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The work is heavily designed for the crafter benchmark. While this is impressive, it would be good if the authors could include a discussion of where achievement distillation would be useful beyond just this setting. Where is it inapplicable? What are the challenges in implementing this?

Summary

This work introduces achievement distillation, a model-free RL method designed to discover achievements without the need for explicit long-term planning components. The proposed method comprises three primary components: two self-supervised tasks and a memory component. The self-supervised tasks, namely Intra-trajectory achievement prediction and cross-trajectory achievement matching, guide the encoder to predict the next achievement to be unlocked using a contrastive learning loss function with different objectives. The memory component is formed by concatenating the latent state representation from the encoder with the action and the representation of the previous achievement. These components are integrated into the PPO algorithm through two alternating training phases involving policy training and auxiliary self-supervised tasks. The effectiveness of achievement distillation is evaluated in the Crafter environment, where it demonstrates significant performance improvements over strong baselines. Additionally, the authors conducted analyses on model sizes, representations, and the contribution of individual components.

Strengths

1. The primary strength of the paper lies in the significant results it presents. The main results in Table 1 and Figure 5 provide clear evidence that the proposed achievement distillation method outperforms the baseline methods to a substantial degree. Furthermore, the results in Figure 6 demonstrate that achievement distillation achieves success in several achievements that none of the baselines can accomplish (e.g., making an iron sword), which is impressive. 2. The authors have done a great job in implementing various relevant strong baselines. Reproducing these baselines are non-trivial given the complexity of them. 3. While the concept of self-supervised auxiliary tasks and the used self-supervised losses are not novel ideas, the authors have managed to execute them very well, resulting in an agent with remarkable performance. 4. The paper is well-written, and the ideas are effectively presented and justified.

Weaknesses

1. The primary weakness of this work lies in the limited range of environments in which the proposed method is evaluated. Since the authors only tested it on the Crafter environment, it remains uncertain whether their method would be effective in different settings and whether it has avoided overfitting to the Crafter environment. 2. The ablation study appears to be somewhat superficial. Although the contribution of cross-trajectory achievement matching is evident, the significance of cross-trajectory achieving matching and memory is not as convincing. The authors should provide additional evidence to support the role of these components. Please refer to the next section for suggestions on how to improve this aspect.

Questions

1. I would encourage the authors to address the main limitation of their work discussed above by conducting further experimentation in various environments, such as MiniHack and maybe ProcGen. 2. It would greatly enhance the understanding of the contribution of cross-trajectory achieving matching and memory if the authors could conduct additional research. For instance, including the individual success rates for all achievements (Figure 6) while ablating these components could be a valuable option. 3. It would be interesting if the authors could present the results in Table 1 and Figure 5 for a larger range of environment steps (e.g., 5M, 10M). This would clarify whether this method is only more data efficient or remains superior with more steps. 4. In Figure 6, why the Individual success rate is not shown for MuZero+SPR? 5. Do you think your method would work with RL algorithms other than PPO?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have adequately addressed the limitations.

Summary

This paper proposes a contrastive learning approach for representation learning in the problem of hierarchical achievement discovery. The proposed method leverages previous contrastive learning loss and combines it with optimal transport. Empirical results show that the learned representation could improve PPO in the Crafter environment.

Strengths

1. It is interesting that self-supervised representation learning can improve PPO, which significantly outperforms model-based approaches, even regarding sample efficiency. 2. Experimental results show the strong performance of the proposed method in the Crafter environment. 3. Presentation: this paper is generally well-organized and easy to follow.

Weaknesses

1. Domain knowledge: unlike baselines, the proposed method utilizes additional important knowledge, i.e., identifying when a new achievement is unlocked (which can be easily inferred from the observed rewards). This other knowledge can easily explain why the achievement prediction of the proposed method is much better than PPO. 2. Generalization: this paper only shows its results in the Crafter environment. It is highly recommended to conduct experiments in other environments to show its generality. 3. Novelty: though it is interesting to use optimal transport for matching achievements, the proposed method is a natural application of contrastive learning, which is not quite novel.

Questions

1. This paper defines MDPs with hierarchical achievements. Has prior work studied this model? If so, please add references. 2. It is interesting to see PPO is more sample-efficient than model-based methods in this paper’s experiments. Except for implementation practice improvement of PPO, any other insights for this result?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

This paper has discussed its limitation.

Thank you for your work!

I would like to thank the authors for continuing to improve their work. I find the new experiments across domains and with value-based algorithms convincing. I am consequently improving my score. I additionally think it would be of great value to the community if the authors could detail how they exactly implement the necessary achievement counting component of their method for new domains.

Thank you!

I appreciate the additional experiments conducted by the authors, and they have adequately addressed my concerns. I would love to see the work accepted and shared with the community at large.