Paper: METRA: Scalable Unsupervised RL with Metric-Aware Abstraction (ICLR-2024)

Reviewer comment (Experiments):
Does METRA have the potential to be applied to environments with temporal dependencies (NetHack, MineDojo, ...)?

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

Reviewer comment (Experiments):
If we apply METRA to a traditional RL setting with external rewards, how would the discovered skills help in finding task-specific policies?

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

Reviewer comment (Experiments):
While the paper presents results on a variety of environments, it would be beneficial to see how METRA performs on more complex environments such as Atari[1] or Google Research Football[2]. This would provide a more comprehensive evaluation of the method's scalability and effectiveness.

Author response:
* **Applying METRA to other types of environments such as Atari or Google Research Football**

Thank you for the suggestion. While we agree that evaluating METRA on more diverse benchmarks like Atari games would further strengthen the paper, we would like to note that we have included a large fraction of the most widely used benchmarks for unsupervised RL and unsupervised skill discovery [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], including both locomotion and manipulation environments. Please understand that adapting the current METRA implementation to Atari games or Football (with discrete action spaces) requires a considerable amount of time and work, as they require a very different RL backbone (other than SAC) due to the difference between continuous vs. discrete action spaces. (We also tried to find an existing unsupervised skill discovery implementation for Atari, but most prior methods do not evaluate on Atari games, and the only available [repository](https://github.com/mklissa/dceo) we found seems to require 5-7 days of training.) That being said, we believe further extending the idea behind METRA to other types of environments is an exciting future research direction.

Reviewer comment (Experiments):
The paper could benefit from a comparison with more diversity RL baselines, such as RSPO[3] and DGPO[4]. This would provide a more complete picture of how METRA compares to other state-of-the-art methods in the field.

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

| Environment (Metric) | DIAYN | DGPO | METRA (ours) |
|---|---|---|---|
| HalfCheetah (policy state coverage) | $6.75 \pm 2.22$ | $6.75 \pm 2.06$ | $\mathbf{186.75} \pm 16.21$ |
| HalfCheetah (total state coverage) | $19.50 \pm 3.87$ | $22.25 \pm 5.85$ | $\mathbf{177.75} \pm 17.10$ |
| Ant (policy state coverage) | $11.25 \pm 5.44$ | $7.00 \pm 3.83$ | $\mathbf{1387.75} \pm 77.38$ |
| Ant (total state coverage) | $107.75 \pm 17.00$ | $121.50 \pm 4.36$ | $\mathbf{6313.25} \pm 747.92$ |

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).

Reviewer comment (Theory):
How does METRA handle environments with non-stationary dynamics or environments where the temporal distance between states can change over time?

Author response:
* **“How does METRA handle environments with non-stationary dynamics or environments where the temporal distance between states can change over time?”**

METRA in its current form assumes a fixed MDP (i.e., stationary dynamics), as commonly assumed in many works in RL. We leave devising an unsupervised RL algorithm for non-stationary environments and continual learning settings for future work. We believe one potential way to deal with this issue is to combine METRA with recurrent latent state spaces [14]. We have clarified this limitation in the updated draft (Appendix A).

Reviewer comment (Experiments):
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?

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

| Environment | \| | 1-D skills | 2-D skills | 4-D skills | \| | 2 skills | 4 skills | 8 skills | 16 skills | 24 skills |
|---|---|---|---|---|---|---|---|---|---|---|
| HalfCheetah | \| | $-17.84 \pm 4.09$ | $-12.86 \pm 1.51$ | $\mathbf{-11.16} \pm 0.89$ | \| | $-26.76 \pm 12.36$ | $-10.85 \pm 1.80$ | $-11.43 \pm 0.63$ | $\mathbf{-9.97} \pm 1.52$ | $-12.10 \pm 1.75$ |
| Ant | \| | $-29.11 \pm 6.35$ | $-7.58 \pm 3.90$ | $\mathbf{-7.54} \pm 0.90$ | \| | $-23.42 \pm 4.00$ | $-13.81 \pm 3.69$ | $-10.64 \pm 4.25$ | $-8.96 \pm 1.78$ | $\mathbf{-7.28} \pm 1.71$ |

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).

Reviewer comment (Evaluation):
In Figure 8, the LEXA and DIAYN totally failed to handle most of the tasks (especially DIAYN). Is this because the skill discovery process has already failed or the learned skills is useless on downstream tasks?

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

Reviewer comment (Experiments):
DIAYN, DADS, and other previous works only consider state space tasks, so the objectives of them are not able to train the visual encoder. But the METRA objective is much more suitable to learn a visual representation. Can the authors analyze more on the advantage of the proposed method on the visual representation learning? If the vision encoders of all baselines are the same (a pretrained network), will the experiment results change?

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

Reviewer comment (Evaluation):
Similar to the last question, will the proposed method still outperform if we only consider state-based tasks? Does the advantage come from visual learning?

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

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

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