SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Seungyong Moon $^{1}$ , Junyoung Yeom $^{1}$ , Bumsoo Park $^{2}$ , Hyun Oh Song $^{1*}$
$^{1}$ Seoul National University,  $^{2}$ KRAFTON
{symoon11, yeomjy, hyunoh}@mllab.snu.ac.kr
bumsoo . park96@krafton . com
Discovering achievements with a hierarchical structure in procedurally generated environments presents a significant challenge. This requires an agent to possess a broad range of abilities, including generalization and long-term reasoning. Many prior methods have been built upon model-based or hierarchical approaches, with the belief that an explicit module for long-term planning would be advantageous for learning hierarchical dependencies. However, these methods demand an excessive number of environment interactions or large model sizes, limiting their practicality. In this work, we demonstrate that proximal policy optimization (PPO), a simple yet versatile model-free algorithm, outperforms previous methods when optimized with recent implementation practices. Moreover, we find that the PPO agent can predict the next achievement to be unlocked to some extent, albeit with limited confidence. Based on this observation, we introduce a novel contrastive learning method, called achievement distillation, which strengthens the agent's ability to predict the next achievement. Our method exhibits a strong capacity for discovering hierarchical achievements and shows state-of-the-art performance on the challenging Crafter environment in a sample-efficient manner while utilizing fewer model parameters.
Deep reinforcement learning (RL) has recently achieved remarkable successes in solving challenging decision-making problems, including video games, board games, and robotic controls [31, 42, 18, 36]. However, these advancements are often restricted to a single deterministic environment with a narrow set of tasks. To successfully deploy RL agents in real-world scenarios, which are constantly changing and open-ended, they should generalize well to new unseen situations and acquire reusable skills for solving increasingly complex tasks via long-term reasoning. Unfortunately, many existing algorithms exhibit limitations in learning these abilities and tend to memorize action sequences rather than truly understand the underlying structures of the environments [28, 24].
To assess the abilities of agents in generalization and long-term reasoning, we focus on the problem of discovering hierarchical achievements in procedurally generated environments with high-dimensional image observations. In each episode, an agent navigates a previously unseen environment and receives a sparse reward upon accomplishing a novel subtask labeled as an achievement. Importantly, each achievement is semantically meaningful and can be reused to complete more complex achievements. Such a setting inherently demands strong generalization and long-term reasoning from the agent.
Previous work on this problem has mainly relied on model-based or hierarchical approaches, which involve explicit modules for long-term planning. Model-based methods employ a latent world model that predicts future states and rewards for learning long-term dependencies [20, 21, 1, 45]. While
these methods have shown effectiveness in discovering hierarchical achievements, particularly in procedurally generated environments, they are constructed with large model sizes and often require substantial exploratory data, which limits their practicality. Hierarchical methods aim to reconstruct the dependencies between achievements as a graph and employ a high-level planner on the graph to direct a low-level controller toward the next achievement to be unlocked [43, 10, 49]. However, these methods rely on prior knowledge of achievements (e.g., the number of achievements), which is impractical in open-world scenarios where the exact number of achievements cannot be predetermined. Additionally, they necessitate a significant number of offline expert data to reconstruct the graph.
To address these issues, we begin by exploring the potential of proximal policy optimization (PPO), a simple and flexible model-free algorithm, in discovering hierarchical achievements [39]. Surprisingly, PPO outperforms previous model-based and hierarchical methods by adopting recent implementing practices. Furthermore, upon analyzing the latent representations of the PPO agent, we observe that it has a certain degree of predictive ability regarding the next achievement, albeit with high uncertainty.
Based on this observation, we propose a novel self-supervised learning method alongside RL training, named achievement distillation. Our method periodically distills relevant information on achievements from episodes collected during policy updates to the encoder via contrastive learning [34]. Specifically, we maximize the similarity in the latent space between state-action pairs and the corresponding next achievements within a single episode. Additionally, by leveraging the uniform achievement structure across all environments, we maximize the similarity in the latent space between achievements from two different episodes, matching them using optimal transport [5]. This learning can be seamlessly integrated into PPO by introducing an auxiliary training phase. Our method demonstrates state-of-the-art performance in discovering hierarchical achievements on the challenging Crafter benchmark, unlocking all 22 achievements with a budget of 1M environment steps while utilizing only  $4\%$  of the model parameters compared to the previous state-of-the-art method [19].
2.1 Markov decision processes with hierarchical achievements
We formalize the problem using Markov decision processes (MDPs) with hierarchical achievements [49]. Let  $\mathbb{M}$  represent a collection of such MDPs. Each environment  $\mathcal{M}_i\in \mathbb{M}$  is defined by a tuple  $(S_i,\mathcal{A},\mathcal{G},p,r,\rho_i,\gamma)$ . Here,  $S_{i}\subset S$  is the image observation space, which has visual variations across different environments,  $\mathcal{A}$  is the action space,  $\mathcal{G}$  is the achievement graph with a hierarchical structure,  $p:S\times \mathcal{A}\to \mathcal{P}(\mathcal{S})$  is the transition probability function,  $r:S\times \mathcal{A}\times \mathcal{S}\rightarrow \mathbb{R}$  is the achievement reward function,  $\rho_{i}\in \mathcal{P}(S_{i})$  is the initial state distribution, and  $\gamma \in [0,1]$  is the discount factor.
The achievement graph  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  is a directed acyclic graph, where each vertex  $v\in \mathcal{V}$  represents an achievement and each edge  $(u,v)\in \mathcal{E}$  indicates that achievement  $v$  has a dependency on achievement  $u$ . To unlock achievement  $v$ , all of its ancestors (i.e., achievements in the path from the root to  $v$ ) must also be unlocked. When an agent unlocks a new achievement, it receives an achievement reward of 1. Note that each achievement can be accomplished multiple times within a single episode, but the agent will only receive a reward when unlocking it for the first time. Specifically, let  $b\in \{0,1\}^{|\mathcal{V}|}$  be a binary vector indicating which achievements have been unlocked and  $c:S\times \mathcal{A}\times \mathcal{S}\to \mathcal{V}\cup \{\emptyset \}$  be a function determining whether a transition tuple results in the completion of an achievement. Then, the achievement reward function is defined as
$$
r (s _ {t}, a _ {t}, s _ {t + 1}) = \left\{ \begin{array}{l} 1 \text {i f} \exists v _ {i} \in \mathcal {V}: b [ i ] = 0, c (s _ {t}, a _ {t}, s _ {t + 1}) = v _ {i} \\ 0 \text {o t h e r w i s e .} \end{array} \right.
$$
This reward structure provides an incentive for the agent to explore the environments and discover a new achievement, rather than repeatedly accomplishing the same achievements.
We assume that the agent has no prior knowledge of the achievement graph, including the number of achievements and their dependencies. Additionally, the agent has no direct access to information about which achievements have been unlocked. Instead, the agent must infer this information indirectly from the reward signal it receives. Given this situation, our objective is to learn a generalizable policy  $\pi : S \to \mathcal{P}(\mathcal{A})$  that maximizes the expected return (i.e., unlocks as many achievements as possible) across all environments of  $\mathbb{M}$ .
(a) Agent's observation
(b) Achievement graph with a hierarchical structure.
Figure 1: Overview of Crafter. Each environment is procedurally generated and partially observable. The objective is to unlock as many achievements as possible within a single episode.
2.2 Crafter environment
We primarily utilize the Crafter environment as a benchmark to assess the capabilities of an agent in solving MDPs with hierarchical achievements [19]. Crafter is an open-world survival game with 2D visual inputs, drawing inspiration from the popular 3D game Minecraft [17]. This is optimized for research purposes, with fast and straightforward environment interactions and clear evaluation metrics. The game consists of procedurally generated environments with varying world map layouts, terrain types, resource placements, and enemy spawn locations, each of which is uniquely determined by an integer seed. An agent can only observe its immediate surroundings as depicted in Figure 1a, which makes Crafter partially observable and thus challenging. To survive, the agent must acquire a variety of skills, including exploring the world map, gathering resources, building tools, placing objects, and defending against enemies. The game features a set of 22 hierarchical achievements that the agent can unlock by completing specific prerequisites, as illustrated in Figure 1b. For instance, to make a wood pickaxe, the agent needs to collect wood, place a table, and stand nearby. This achievement structure is designed to require the agent to learn and utilize a wide range of skills to accomplish increasingly challenging achievements, such as crafting iron tools and collecting diamonds.
2.3 Proximal policy optimization
PPO is one of the most successful model-free policy gradient algorithms due to its simplicity and effectiveness [39]. PPO learns a policy  $\pi_{\theta} : S \to \mathcal{P}(A)$  and a value function  $V_{\theta} : S \to \mathbb{R}$ , which are parameterized by neural networks. During training, PPO first collects a new episodes  $\mathcal{T}$  using the policy  $\pi_{\theta_{\mathrm{old}}}$  immediately prior to the update step. Subsequently, PPO updates the policy network using these episodes for several epochs to maximize the clipped surrogate policy objectives given by
$$
J _ {\pi} (\theta) = \mathbb {E} _ {(s _ {t}, a _ {t}) \sim \mathcal {T}} \left[ \min  \left(\frac {\pi_ {\theta} (a _ {t} \mid s _ {t})}{\pi_ {\theta_ {\mathrm {o l d}}} (a _ {t} \mid s _ {t})} \hat {A} _ {t}, \operatorname {c l i p} \left(\frac {\pi_ {\theta} (a _ {t} \mid s _ {t})}{\pi_ {\theta_ {\mathrm {o l d}}} (a _ {t} \mid s _ {t})}, 1 - \epsilon , 1 + \epsilon\right) \hat {A} _ {t}\right) \right],
$$
where  $\hat{A}_t$  is the estimated advantage computed by generalized advantage estimate (GAE) [38]. PPO simultaneously updates the value network to minimize the value objective given by
$$
J _ {V} (\theta) = \mathbb {E} _ {s _ {t} \sim \mathcal {T}} \left[ \frac {1}{2} \left(V _ {\theta} (s _ {t}) - \hat {V} _ {t}\right) ^ {2} \right],
$$
where  $\hat{V}_t = \hat{A}_t + V_{\theta_{\mathrm{old}}}(s_t)$  is the bootstrapped value function target.
In image-based RL, it is common practice to optimize the policy and value networks using a shared network architecture [13, 48]. An image observation is first passed through a convolutional encoder  $\phi_{\theta}: S \to \mathbb{R}^h$  to extract a state representation, which is then fed into linear heads to compute the policy and value function. Sharing state representations between the policy and value networks is crucial to improving the performance of agents in high-dimensional state spaces. However, relying solely on policy and value optimization to train the encoder can lead to suboptimal state representations, particularly in procedurally generated environments [9]. To address this issue, recent studies introduce an auxiliary training phase alongside the policy and value optimization that trains the encoder with auxiliary value or self-supervised objectives [9, 32].
3.1 PPO is a strong baseline for hierarchical achievements
Despite being a simple and ubiquitous algorithm, PPO is less utilized than model-based or hierarchical approaches for solving MDPs with hierarchical achievements. This is due to the fact that PPO does not have an explicit component for long-term planning or reasoning, which is believed to be essential for solving hierarchical tasks. However, a recent study has shown that a PPO-based algorithm is also successful in solving hierarchical achievement on the Minecraft environment, albeit with the aid of pre-training on human video data [4].
Based on this observation, we first investigate the effectiveness of PPO in solving hierarchical achievements on Crafter without pretraining. We adopt the recent implementation practices proposed in Andrychowicz et al. [2], Baker et al. [4]. Concretely, we modify the default ResNet architecture in IMPALA as follows [14]:
- Network size: We increase the channel size from [16, 32, 32] to [64, 128, 128] and the hidden size from 256 to 1024.
- Layer normalization: We add layer normalization before each dense or convolutional layer [3].
- Value normalization: We keep a moving average for the mean and standard deviation of the value function targets and update the value network to predict the normalized targets.
We train the modified PPO on Crafter for 1M environment steps and evaluate the success rates for unlocking achievements (please refer to Section 5.1 for the evaluation). Figure 2 shows that the slight modification in implementing PPO significantly improves the performance, increasing the score from 8.17 to 15.60. Notably, this outperforms the current state-of-the-art DreamerV3, which achieves a score of 14.77.
Figure 2: Score curves of PPO.
3.2 Representation analysis of PPO
Since Crafter environments are procedurally generated, simply memorizing successful episodes is insufficient for achieving high performance. We hypothesize that the PPO agent acquires knowledge beyond mere memorization of action sequences, possibly including information about achievements. To validate this, we analyze the learned latent representations of the encoder, as inspired by Wijmans et al. [46]. Specifically, we collect a batch of episodes using an expert policy and subsample a set of states for training. For each state  $s$  in the training set, we freeze its latent representation  $\phi_{\theta}(s)$  from the encoder. Subsequently, we train a linear classifier using this representation as input to predict the very next achievement unlocked in the episode containing  $s$ . Finally, we evaluate the classification accuracy and the prediction confidence of the ground-truth labels on a held-out test set. The detailed experimental settings are provided in Appendix A.
Surprisingly, PPO achieves a nontrivial accuracy of  $44.9\%$  in the 22-way classification. However, Figure 3 shows that the prediction outputs lack confidence with a median value of 0.240. This suggests that the learned representations of the PPO encoder are not strongly correlated with the next achievement to be unlocked and the agent may struggle to generate optimal action sequences towards a specific goal.
This finding warrants providing additional guidance to the encoder for predicting the next achievements with high confidence. However, since the agent has no access to the achievement labels, it is challenging to guide the agent in a supervised fashion. Therefore, it is necessary to explore alternative approaches to guide the agent toward predicting the next achievements.
Figure 3: Histogram for the confidence of next achievement prediction.
Figure 4: Illustration of achievement distillation.
4 Contrastive learning for achievement distillation
In this section, we introduce a new self-supervised learning method that works alongside RL training to guide the encoder in predicting the next achievement to be unlocked. This approach distills relevant information about discovered achievements from episodes collected during multiple policy updates into the encoder via contrastive learning. This method consists of two key components:
- Intra-trajectory achievement prediction: Within an episode, this maximizes the similarity in the latent space between a state-action pair and its corresponding next achievement.
- Cross-trajectory achievement matching: Between episodes, this maximizes the similarity in the latent space for matched achievements.
For ease of notation, we denote the sequence of unlocked achievements within an episode as  $(g_i)_{i=1}^m$  and their corresponding timesteps as  $(t_i)_{i=1}^m$ , where each achievement  $g_i$  is defined by a transition tuple  $(s_{t_i}, a_{t_i}, s_{t_i+1})$ . For each timestep  $t$ , we represent the very next achievement as  $g_t^+ = g_u$ , where  $u = \min\{i \mid t \leq t_i\}$ , and the very previous achievement as  $g_t^- = g_l$ , where  $l = \max\{i \mid t > t_i\}$ .
4.1 Intra-trajectory achievement prediction
Given a state-action pair  $(s_t, a_t)$  and its corresponding next achievement  $g_t^+$  within an episode  $\tau$ , we train the encoder  $\phi_\theta$  to produce similar representations for them through contrastive learning [34, 22]. Specifically, we regard  $g_t^+$  as the anchor and  $(s_t, a_t)$  as the positive. We also randomly sample another state-action pair  $(s_{t'}, a_{t'})$  from the same episode to serve as the negative. Subsequently, we obtain the normalized representations of the anchor, positive, and negative, denoted as  $\nu_\theta(g_t^+)$ ,  $\psi_\theta(s_t, a_t)$ , and  $\psi_\theta(s_{t'}, a_{t'})$ , respectively. Finally, we minimize the following contrastive loss to maximize the cosine similarity between the anchor and positive representations while minimizing the cosine similarity between the anchor and negative representations:
$$
L _ {\mathrm {p r e d}} (\theta) = - \mathbb {E} _ {(s _ {t}, a _ {t}) \sim \tau \atop (s _ {t ^ {\prime}}, a _ {t ^ {\prime}}) \sim \tau} \left[ \log \left(\frac {\exp (\psi_ {\theta} (s _ {t} , a _ {t}) ^ {\top} \nu_ {\theta} (g _ {t} ^ {+}) / \lambda)}{\exp (\psi_ {\theta} (s _ {t} , a _ {t}) ^ {\top} \nu_ {\theta} (g _ {t} ^ {+}) / \lambda) + \exp (\psi_ {\theta} (s _ {t ^ {\prime}} , a _ {t ^ {\prime}}) ^ {\top} \nu_ {\theta} (g _ {t} ^ {+}) / \lambda)}\right) \right],
$$
where  $\lambda > 0$  is the temperature parameter.
To obtain the state-action representation  $\psi_{\theta}(s_t,a_t)$ , we calculate the latent representation of the state  $\phi_{\theta}(s_t)$  from the encoder and concatenate it with the action  $a_{t}$  using a FiLM layer [35]. The resulting vector is then passed through an MLP layer and normalized. To obtain the achievement representation  $\nu_{\theta}(g_t^{+})$ , we simply calculate the residual of the latent representations of the two consecutive states from the encoder and normalize it, as motivated by Nair et al. [33].
While this contrastive objective encourages the encoder to predict the next achievement in the latent space, it can potentially lead to distortions in the policy and value networks due to the changes in the encoder. To address this, we jointly minimize the following regularizers to preserve the outputs of the policy and value networks, following the practice in Moon et al. [32]:
$$
R _ {\pi} (\theta) = \mathbb {E} _ {s _ {t} \sim \tau} \left[ D _ {\mathrm {K L}} \left(\pi_ {\theta_ {\mathrm {o l d}}} (\cdot | s _ {t}) \parallel \pi_ {\theta} (\cdot | s _ {t})\right) \right], R _ {V} (\theta) = \mathbb {E} _ {s _ {t} \sim \tau} \left[ \frac {1}{2} \left(V _ {\theta} (s _ {t}) - V _ {\theta_ {\mathrm {o l d}}} (s _ {t})\right) ^ {2} \right],
$$
where  $\pi_{\theta_{\mathrm{old}}}$  and  $V_{\theta_{\mathrm{old}}}$  are the policy and value networks immediately prior to the contrastive learning, respectively and  $D_{\mathrm{KL}}$  denotes the KL divergence.
4.2 Cross-trajectory achievement matching
Since Crafter environments are procedurally generated, the achievement representations learned solely from intra-trajectory information may include environment-specific features that limit generalization. To obtain better achievement representations, we leverage the common achievement structure shared across all episodes.
We first match the sequences of unlocked achievements from two different episodes in an unsupervised fashion. Given two achievement sequences  $\mathbf{g} = (g_i)_{i=1}^m$  and  $\mathbf{g}' = (g_j')_{j=1}^n$ , we define the cost matrix  $M \in \mathbb{R}^{m \times n}$  as the cosine distance between the achievement representations:
$$
M _ {i j} = 1 - \nu_ {\theta} (g _ {i}) ^ {\top} \nu_ {\theta} (g _ {j} ^ {\prime}).
$$
Subsequently, we regard these two sequences as discrete uniform distributions and compute a soft-matching  $T \in \mathbb{R}^{m \times n}$  between them using partial optimal transport, which can be solved by
$$
\begin{array}{l} T = \underset {T \geq 0} {\arg \min } \left\langle T, M \right\rangle + \alpha \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {n} T _ {i j} \log T _ {i j} \\ \mathrm {s u b j e c t t o} T \mathbf {1} \leq \mathbf {1}, T ^ {\top} \mathbf {1} \leq \mathbf {1}, \mathbf {1} ^ {\top} T ^ {\top} \mathbf {1} = \min  \{m, n \}, \\ \end{array}
$$
where  $\alpha > 0$  is the entropic regularization parameter [5]. Here, we set the total amount of probability mass to be transported to the minimum length of the two sequences for simplicity. However, some unlocked achievements in one sequence may not exist in the other sequence, and therefore should not be transported. In this case, the optimal amount of probability mass to be transported should be less than the minimum length. To address this, we compute the conservative hard matching  $T^{*}$  from  $T$  by thresholding the probabilities less than 0.5 (i.e.,  $T^{*} = \mathbb{1}[T > 0.5]$ ). Note that this also encourages each achievement to be matched to at most one other achievement. We provide examples of matching results in Appendix B.
We train the encoder to produce similar representations for the matched achievements according to  $T^{*}$  through contrastive learning. Specifically, suppose that the  $i$ th achievement of the source sequence  $g_{i}$  is matched with the  $k$ th achievement of the target sequence  $g_{k}^{\prime}$ . Then, we consider  $g_{i}$  as the anchor and  $g_{k}^{\prime}$  as the positive. We also randomly sample another achievement  $g_{j}^{\prime}$  from the target sequence to serve as the negative. Subsequently, we obtain the normalized representations of these achievements. Finally, we minimize the following contrastive loss to maximize the cosine similarity between the anchor and positive representations while minimizing the cosine similarity between the anchor and negative representations:
$$
L _ {\mathrm {m a t c h}} (\theta) = - \mathbb {E} _ {g _ {i} \sim \mathbf {g}, g _ {j} ^ {\prime} \sim \mathbf {g} ^ {\prime}} \left[ \log \left(\frac {\exp (\nu_ {\theta} (g _ {i}) ^ {\top} \nu_ {\theta} (g _ {k} ^ {\prime}) / \lambda)}{\exp (\nu_ {\theta} (g _ {i}) ^ {\top} \nu_ {\theta} (g _ {k} ^ {\prime}) / \lambda) + \exp (\nu_ {\theta} (g _ {i}) ^ {\top} \nu_ {\theta} (g _ {j} ^ {\prime}) / \lambda)}\right) \right],
$$
As in Section 4.1, we jointly minimize the policy and value regularizers to prevent distortions.
4.3 Achievement representation as memory
We further utilize the achievement representations learned in Sections 4.1 and 4.2 as memory for the policy and value networks. Specifically, given a state  $s_t$  and its corresponding previous achievement  $g_t^-$ , we concatenate the latent state representation  $\phi_\theta(s_t)$  with the previous achievement representation  $\nu_\theta(g_t^-)$ . The resulting vector is then fed into the policy and value heads to output an action distribution and value estimate, respectively.