METRA: Scalable Unsupervised RL with Metric-Aware Abstraction

We thank the reviewer for the thorough review and constructive feedback about this work. Below, we present new ablation study results and provide answers to the questions. We believe that these changes strengthen the paper, and welcome additional suggestions for further improving the work.

• “Could METRA be modified to model non-linear skills even in the latent space?”, Ablation study on approximated objectives

Yes, it is possible to modify METRA to model non-linear skills in the latent space by using the full $\psi(z)$ form (instead of simplifying it with $\psi(z) = z$ in Eq. 7), and we discuss this full version of METRA in detail in Appendix E.3. Below, we empirically compare METRA with the original version before the $\psi(z)=z$ simplification (Eq. 6). We test two different dimensionalities for $\psi$ ($D=2$ and $D=16$).

The table above shows the results (policy state coverage) on Ant and HalfCheetah (at epoch 10000, 4 seeds each). In these environments, while both the simplified version (with $\psi(z) = z$) and the full version (without setting $\psi(z) = z$) perform similarly, the setting with $D=16$ exhibits slightly slower learning due to the added complexity.

In general, we believe using the full version of METRA could lead to more diverse skills in highly complex environments, as there is less restriction on latent behaviors (though it could be more computationally expensive or relatively harder to optimize). Another interesting aspect of the full METRA objective is that it resembles contrastive learning (Appendix E.3). As such, combining the full version of METRA with contrastive learning techniques may further help scale up the idea of METRA to more complex or even real-world environments, which we leave for future work.

• “Are there any other metrics that are promising for their applications to the reinforcement learning setting in the METRA framework?”

As the reviewer pointed out, it is possible to combine our WDM objective (Eq. 3) with distance metrics other than temporal distances. For example, as mentioned in Sec 4.2, LSD (a previous unsupervised skill discovery method) can be viewed as a special case of WDM with the Euclidean distance metric (though this is not as universally applicable as the temporal distance metric). Another potentially interesting metric is the bisimulation metric [1, 2], which measures the similarity between states based on both environment dynamics and a reward function. Although this metric requires a pre-defined reward function(s), it may serve as a way to inject human priors about tasks into unsupervised learning of skills. We believe considering and comparing various distance metrics (or quasimetrics) is an interesting future research direction.

We thank the reviewer for the thorough review and constructive feedback about this work. Please find our answers to the questions below. We welcome additional suggestions for further improving the work.

• “If we apply METRA to a traditional RL setting with external rewards, how would the discovered skills help in finding task-specific policies?”

In Section 5, we have (already) shown two ways to utilize skills learned by METRA for downstream tasks. First, we train hierarchical high-level policies $\pi^h(z|s)$ that selects skills as temporally extended actions. In Fig. 6, we show that the high-level policies (trained on fixed METRA skills) effectively solve challenging downstream tasks defined by external rewards (e.g., jumping over multiple hurdles, navigating through multiple goals, etc.). Second, specifically for goal-reaching tasks, METRA provides a very simple way to directly select a skill in a zero-shot manner (Section 4.2). We empirically demonstrate this capability of METRA in Fig. 8, which shows that METRA achieves the best zero-shot goal-reaching performance.

Moreover, while we did not mention in the paper, we can also combine METRA with successor features by noting that the METRA reward $(\phi(s’) - \phi(s))^\top z$ can be interpreted as the inner product between the feature vector $\tilde \phi(s, a, s’) = \phi(s’) - \phi(s)$ and the task vector $z$. Based on this interpretation, we can compute via linear regression the most suitable skill $z$ for any arbitrary reward function in a zero-shot manner, similarly to successor feature methods. We are currently working on this as a follow-up to METRA.

• “Does METRA have the potential to be applied to environments with temporal dependencies (NetHack, MineDojo, ...)?”

METRA can be applied to any fully observable MDP in principle. However, METRA in its current form does not particularly deal with partially observable environments (e.g., Minecraft), like many prior works in unsupervised skill discovery [1, 2, 3, 4]. We believe this limitation can be resolved by considering observation histories, and combining METRA with world models that use recurrent structures (e.g., Dreamer) would be an exciting future research direction. We have clarified this limitation in the revised manuscript (Appendix A).

We thank the reviewer again for the helpful feedback and please let us know if there are any additional concerns or questions.

[1] Eysenbach et al., Diversity is all you need: Learning skills without a reward function (2019).

[2] Sharma et al., Dynamics-aware unsupervised discovery of skills (2020).

[3] Park et al., Lipschitz-constrained unsupervised skill discovery (2022).

[4] Laskin et al., Unsupervised reinforcement learning with contrastive intrinsic control (2022).

We thank the reviewer for the thorough review and constructive feedback about this work. Below, we describe how we have added new comparisons with DGPO and a new ablation study with different sizes of latent states. We believe that these changes strengthen the paper, and welcome additional suggestions for further improving the work.

• Additional baseline – DGPO

Thank you for suggesting the new baselines. We additionally compared METRA with DGPO [2] on the MuJoCo Ant and HalfCheetah environments. With the addition of this comparison, the paper compares to a total of $12$ different methods.

Since DGPO [2] (as well as RSPO [1]) was originally proposed for the quality diversity (QD) setting (i.e., learn diverse behaviors while maximizing rewards), unlike our unsupervised RL setting (i.e., learn diverse behaviors without rewards), we only use the intrinsic motivation part of DGPO for our experiments. Specifically, we use $\min_{z’ \neq z} \log{\frac{q(z|s’)}{q(z|s’) + q(z’|s’)}}$ as the reward function for DGPO.

The table above compares the performances of METRA, DIAYN, and DGPO on Ant and HalfCheetah (at epoch 10000, 4 seeds each). The results suggest that both DIAYN and DGPO struggle to cover the state space in the absence of supervision. This is an expected result, given that DGPO also considers a KL divergence as a divergence measure, which is a metric-agnostic quantity. As such, DGPO shares the same limitation as DIAYN of not necessarily encouraging exploration (Section 2). Hence, while DGPO succeeds in the QD setting with the guidance of rewards, it does not necessarily achieve strong performance in our unsupervised settings.

That being said, we believe combining existing KL-based intrinsic rewards, such as DGPO, with our temporal distance-based Wasserstein metric could potentially lead to another performant unsupervised RL algorithm, which we leave for future work. We have updated the new results in the draft (Appendix F.4).

• How does the dimensionality of the latent space affect the performance of METRA? Is there a trade-off between the dimensionality of the latent space and the complexity of the behaviors that can be learned?

As the reviewer pointed out, the size of the latent skill space affects the performance of METRA. Since METRA maximizes state coverage under the capacity of the latent space $\mathcal{Z}$, a too small skill latent space could lead to less diverse skills. To empirically demonstrate this, we evaluate METRA with 1-D, 2-D, and 4-D continuous skills as well as 2, 4, 8, 16, and 24 discrete skills on MuJoCo Ant and HalfCheetah.

The table above shows the goal-reaching performances (negative goal distance, the higher the better) of different skill latent spaces (at epoch 10000, 4 seeds each). In general, the diversity of skill (and thus goal-reaching performance) increases as the capacity of $\mathcal{Z}$ grows. However, using a too large skill space may slow down the training of the skill policy, as the agent needs to learn more behaviors. We have added qualitative results from these experiments to the revised draft (Appendix F.3).

We thank the reviewer for the thorough review and constructive feedback about this work. Please find our answers to the questions below. We welcome additional suggestions for further improving the work.

• Failure of DIAYN/LEXA in goal-conditioned tasks

As the reviewer pointed out, DIAYN and LEXA struggle with complex goal-conditioned tasks. This is mainly because they often fail to cover the state space, as shown in Figure 3 (for LEXA, please refer to the “P2E/Disag.” column, which is the exploration method LEXA uses). This is an expected result, given that DIAYN maximizes mutual information, which does not necessarily encourage exploration (Section 2), and LEXA (a pure exploration method) struggles in complex environments, where it is infeasible to completely cover every possible state or fully capture environment dynamics.

• Does the advantage of METRA mainly come from visual learning?, “Will the proposed method still outperform if we only consider state-based tasks?”

In our experiments, Ant and HalfCheetah are state-based environments (Fig. 4), in which METRA outperforms most previous unsupervised RL methods by significant margins (Fig. 5, Fig. 7). This suggests that the performance gain mainly comes from our objective, not from visual representation learning. In state-based environments, however, LSD (a previous Euclidean-distance maximizing skill discovery method) shows comparable performance to METRA. Yet, LSD fails to scale to pixel-based environments because Euclidean distances do not necessarily provide meaningful learning signals in pixels.

The fact that the performance gain of METRA mainly comes from its objective becomes further clearer if we contrast the objectives of DIAYN and METRA: DIAYN essentially minimizes $\\\|\phi(s) - z\\\|_2^2$ (Appendix E.1), while METRA maximizes $(\phi(s') - \phi(s))^\top z$ under the $\\\|\phi(s) - \phi(s')\\\|_2 \leq 1$ constraint (Eq. 8). They both have the exactly same (visual) encoder $\phi$ in their objectives, but the results are very different (Fig. 3). This is, again, because mutual information does not encourage the agent to cover the state space (Section 2).

• “In Figure 7, some methods are very unstable on the Kitchen tasks (variance of different seeds is large). Can the authors give any reason?”

We believe this is mainly because we use a queue (or policy) coverage metric for Kitchen, while we use a total coverage metric for the other environments (please refer to Appendix G.1 for details). The rationale behind this choice is that, as mentioned in Sec 5.3, most methods max out the total coverage metric in Kitchen. The queue coverage metric in Kitchen measures the number of achieved subtasks from the trajectories within a specific window. Exploration methods may exhibit high variances in this metric since they may visit different parts of the space on different iterations. That being said, we would like to note that METRA outperforms the best exploration method for Kitchen (P2E = LEXA) in a different goal-reaching metric by a statistically significant margin (Fig. 8).

We thank the reviewer again for the helpful feedback and please let us know if there are any additional concerns or questions.