METRA: Scalable Unsupervised RL with Metric-Aware Abstraction
Authors: Seohong Park
Oleh Rybkin
Sergey Levine
Venue: ICLR 2024 oral (2024)
Keywords: reinforcement learning
Primary area: reinforcement learning

Abstract: Unsupervised pre-training strategies have proven to be highly effective in natural language processing and computer vision. Likewise, unsupervised reinforcement learning (RL) holds the promise of discovering a variety of potentially useful behaviors that can accelerate the learning of a wide array of downstream tasks. Previous unsupervised RL approaches have mainly focused on pure exploration and mutual information skill learning. However, despite the previous attempts, making unsupervised RL truly scalable still remains a major open challenge: pure exploration approaches might struggle in complex environments with large state spaces, where covering every possible transition is infeasible, and mutual information skill learning approaches might completely fail to explore the environment due to the lack of incentives. To make unsupervised RL scalable to complex, high-dimensional environments, we propose a novel unsupervised RL objective, which we call Metric-Aware Abstraction (METRA). Our main idea is, instead of directly covering the entire state space, to only cover a compact latent space $\mathcal{Z}$ that is metrically connected to the state space $\mathcal{S}$ by temporal distances. By learning to move in every direction in the latent space, METRA obtains a tractable set of diverse behaviors that approximately cover the state space, being scalable to high-dimensional environments. Through our experiments in five locomotion and manipulation environments, we demonstrate that METRA can discover a variety of useful behaviors even in complex, pixel-based environments, being the first unsupervised RL method that discovers diverse locomotion behaviors in pixel-based Quadruped and Humanoid. Our code and videos are available at https://seohong.me/projects/metra/

Published as a conference paper at ICLR 2024
METRA: SCALABLE UNSUPERVISED RL WITH METRIC-AWARE ABSTRACTION
Seohong Park' Oleh Rybkin'
Sergey Levine'
Uniiversity of California, Berkeley
seohong@berkeley.co
ABSTRACT
Unsupervised pre-training strategies have proven to be highly effective in natural language processing and computer vision. Likewise, unsupervised reinforcement learning (RL) holds the promise of discovering variety of potentially useful be- haviors that can accelerate the learning of wide array of downstream tasks. Pre- vious unsupervised RL approaches have mainly focused on pure exploration and mutual information skill learning. However, despite the previous attempts, mak- ing unsupervised RL truly scalable still remains a major open challenge: pure exploration approaches might struggle in complex environments with large state spaces, where covering every possible transition is infeasible, and mutual infor- mation skill learning approaches might completely fail explore the environ- ment due to the lack of incentives. To make unsupervised RL scalable to com- plex, nigh-dimensional environments, we propose novel unsupervised RL ob- jective, which we call Metric-Aware Abstraction (METRA). Our main idea is, instead of directly covering the entire state space, to only cover a compact latent space Z that is metrically connected to the state space S by temporal distances. By learning to move in every direction in the latent space, METRA obtains a tractable set of diverse behaviors that approximately cover the state space, being scalable to high digh-dimensionll environments. Through our experiments in five lo- comotion and manipulation environments, cover a variety of useful behaviors even in complex, pixel-based environments, being the first unsupervised RL method that discovers diverse locomotion behav- iors in pixel-based Quadruped and Humanoid. Our code and videos are available at https: seohong wwe projecttstttteeepo wttpeeetttttteeeeeeeeeeeeeeeeeee /metra/
demonstrate that METRA can dis-
INTRODUCTION
Unsupervised pre-training has proven transformative in domains from natural language processing to computer vision: contrastive representation learning (Chen et al., 2020) can acquire effective features from unlabeled images, and generative autoregressive pre-training (Brown et al., 2020) can enable language models that can be adapted to plethora of downstream applications. could derive an equally scalable framework for unsupervised reinforcement learning (RL) that au- tonomously explores the space of possible behaviors, then we could enable general-purpose unsu- pervised pre-trained agents serve an effective foundation for efficiently learning broad range of downstream tasks. Hence, our goal in this work is to propose a scalable unsupervised RL ob- jective that encourages an agent explore its environment and learn breadth of potentially useful behaviors without any supervision.
we
While this formulation of unsupervised RL has been explored in a number of prior works, making fully unsupervised RL truly scalable still remains a major open challenge. Prior approaches to un- supervised RL can be categorized into two main groups: pure exploration methods (Burda et al., 2019; Pathak et al., 2019; Liu & Abbeel, 202 Mendonca et al., 2021: Rajeswar et al., 2023) and unsupervised skill discovery methods (Eysenbach et al., 2019a; Sharma et al., 2020; Laskin et al., 2022; Park et al., 2022). While these approaches have been shown to be effective in several unsupervised RL benchmarks (Mendonca et al., 2021; I; Laskin et al., 2021), it is not entirely clear whether such methods can indeed be scalable to complex environments with high intrinsic dimen- sionality. Pure exploration based unsupervised RL approaches aim to either completely cover the entire state space (Burda et al., 2019; Liu & Abbeel, 202 fully capture the transition dynamics of the Markov decision process (MDP) (Pathak et al., 2019; Sekar et al., 2020; Mendonca et al., 2021; Rajeswar al., 2023). However, in complex environments with large state space, it may be

---

Published as a conference paper at ICLR 2024
s = R 6444 64x3 (Pixels)
==
Temporally close
Temporally distaul
METRA covers the most "Lemporally spread-out" manifold
Figure 1: Illustration of METRA. Our main idea for scalable unsupervised RL to cover only the most "important" low-dimensional subset of the state space, analogously to PCA. Specifically, METRA covers the most "temporally spread-out" (non-linear) manifold, which would lead to appreximate coverage of the state space S. In the example above, the two-dimensional 2 space captures behaviors running in all directions, not necessarily covering every possible leg pose.
infeasible to attain either of these aims. In fact, we will show that these methods fail to cover the state space even in the mate- 29-dimensional MuJoCo Ant environment. On the other hand, unsupervised skill discovery methods aim to discover diverse, distinguishable behaviors, e.g., by maximizing the mutual information between skills and states (Gregor et al., 2016; Eysenbach et al., 2019a). While these methods do learn behaviors that are mutually different, they either do not nec- essarily encourage exploration and thus often have limited state coverage in the complete absence of supervision (Eysenbach et al., 2019a; Sharma et al., 2020), or are not directly scalable to pixel-based control environments (Park al., 2022; 2023b).
In this work, we aim to address these challenges and develop an unsupervised RI. objective, which call Metric-Aware Abstraction (METRA), that scales to complex, image-based environments with high intrinsic dimensionality. Our first main idea is to learn diverse behaviors that maximally cover not the original state space but a compact latent metric space defined by a mapping function 6 : S + Z with a metric d. Here, the latent state is connected by the state space by the metric d, which ensures that covering latent space leads to coverage of the state space. Now, the question becomes which metric to use. Previous metric-basee skill learning methods mostly used the Eu- clidean distance (or its scaled variant) between two states (He et al., 2022; Park et al., 2022; 2023b). However, such state-based metrics are not directly applicable to complex, high-dimensional state space (e.g., images). Our second main idea is therefore to use temporal distances (i.e., the num- ber of minimum environment steps between two states) as a metric for the latent space. Temporal distances are invariant to state representations and thus applicable to pixel-based environments as well. As result, by maximizing coverage in the compact latent space, we can acquire diverse be- haviors that approximately cover the entire state space, being scalable to high-dimensional, complex environments (Figure 1).
Through our experiments on five state-based and pixel-based continuous control environments, we demonstrate that method learns diverse, useful behaviors, as well compact latent space that can be used to solve various downstream tasks in a zero-shot manner, outperforming previous unsupervised RL methods. To the best of our knowledge, METRA is the first unsupervised RL method that demonstrates the discovery of diverse locomotion behaviors in baxel Quadruped and Humanoid environments.
2 WHY MIGHT PREVIOUS UNSUPERVISED RL METHODS FAIL TO SCALE?
The goal of unsupervised RL is to acquire useful knowledge, such as policies, world models, or exploratory data, by interacting with the environment in an unsupervised manner (i.e., without tasks or reward functions). Typically, this knowledge is then leveraged to solve downstream tasks more efficiently. Prior work in unsupervised RL can be categorized into two main groups: pure explo- ration methods and unsupervised skill discovery methods. Pure exploration methods aim to cover the entire state space or fully capture the environment dynamics. They encourage exploration by maximizing uncertainty (Pathak et al., 2017; Shyam et al., 2019; Burda et al., 2019; Pathak et al., 2019: Sekar et al., 2020; Mazzaglia et al., 2022) or state entropy (Lee et al., 2019; Pong et al., 2020; Liu & Abbeel, 2021b; Yarats al., 2021). Based the data collected by the exploration
2

---

Published as a conference paper at ICLR 2024
I(S:Z)
Iw(S: 2)
Pure Exploration RND, Disag, APT, DIAYN, DADS,
[(S: 2) [w 2) METRA (ours)
MI vs. WDM
Sketch comparing different unsupervised RL objectives. Pure exploration approaches to cover every possible state, which is infeasible in complex environments (e.g., such methods might be "stuck" forever finding novel joint angle configurations of a robot, without fully exploring the environment; see Figure 3). The mutual information (S; 2) has no underlying distance metrics, and thus does not prioritize coverage enough, only focusing on skills that are discriminable. In contrast, our proposed Wasserstein de- pendeney measure (S; Z) maximizes the distance metrie d, which choose to be the temporal distance, forcing the learned skills to span the "longest" subspaces of the state space, analogously to (temporal, non- linear) PCA.
Figure 2:
policy, these methods learn a world model (Rajeswar et al., 2023), train a goal-conditioned pol- icy (Pong al., 2020; Pitis et al., 2020; Mendonca al., 2021; Hu al., 2023), learn skills via trajectory autoencoders (Campos Camúñezz et al., 2020; Mazzaglia et al., 2023), or directly fine- tune the learned exploration policy (Laskin et al., 2021) to accelerate downstream task learning. While these pure exploration-basee approaches are currently the leading methods in unsupervised RL benchmarks (Mendonca et al., 2021; Laskin et al., 2021; Mazzaglia et al., 2023; Rajeswar al., 2023), their scalability may be limited in complex environments with large state spaces because it often computationally infeasible to completely cover every possible state or fully capture the dy- namics. In Section 5, we empirically demonstrate that these approaches even fail to cover the state space of the state-based 29-dimensional Ant environment.
Another line of research in unsupervised RL aims to learn diverse behaviors (or skills) that are distinguishable from one another, and our method also falls into this category. The most common approach to unsupervised skill discovery is to maximize the mutual information (MI) between states and skills (Gregor et al., 2016; Eysenbach et al., 2019a; Sharma et al., 2020; Hansen et al., 2020):
(S: (pl 2) [p(s)p(z)).
By associating different skill latent vectors a with different states these methods learn diverse skills that are mutually distinct. However, they share the limitation that they often end up discover- ing simple, static behaviors with limited state coverage (Campos Camúñez et al., 2020; Park et al., 2022). This because MI is defined by a KI. divergence (Equation (1)), which is a metric-agnostic quantity (e.g., MI is invariant to scaling: see Figure 2). As a result, the MI objective only focuses on the distinguishability of behaviors, regardless of "how different" they are, resulting in limited state coverage (Campos Camúñez et al., 2020; Park et al., 2022). To address this limitation, prior works combine the MI objective with exploration bonuses (Campos Camúñez et al., 2020; Strouse et al., 2022; Park & Levine, 2023) or propose different objectives that encourage maximizing dis- tances in the state space (He et al., 2022; Park et al., 2022; 2023b). Yet, it remains unclear whether these methods can scale complex, high-dimensional environments, because they either attempt to completely capture the entire MDP (Campos Camúñez et al., 2020; Strouse et al., 2022; Park & Levine, 2023) or assume a compact, structured state space (He et al., 2022; Park et al., 2022: 2023b). Indeed, to the best of our knowledge, no previous unsupervised skill discovery methods have suc- ceeded in discovering locomotion behaviors on pixel-based locomotion environments. Unlike these approaches, our method learns compact set of diverse behaviors that are maximally different in terms of the temporal distance. As a result, they can approximately cover the state space, even in a complex, high-dimensional environment. We discuss further related work in Appendix B.
3 PRELIMINARIES AND PROBLEM SETTING
We consider controlled Markov process, an MDP without reward function, defined as M (S, A. μ., p). S denotes the state space, AA denotes the action space, /4 A(S) denotes the initial state distribution, and p S x A + A ((A) denotes the transition dynamics kernel. We consider set of latent vectors = E Z, which can be either discrete continuous, and latent-conditioned policy T [a]s, 2). Following the terminology in unsupervised skill discovery, we refer to latent vectors z (and their corresponding policies T[als, 2)) skills. When sampling a trajectory, a skill from the prior distribution, 2 ~ p(2). and then roll out a trajectory with π [[als, 2)),, where is fixed for the entire episode. Hence, the joint skill-trajectory distribution is given p(T, 2)
first sample
3

---

Published as a conference paper at ICLR 2024
p(z)p(su) Io T 2))]sss)))(((++)))) at). where denotes (su. dg, 81, .i....... sy). Our goal in this work is to learn a set of diverse, useful behaviors x [[ass, 2), without using any supervision, data, or prior knowledge.
4 A SCALABLE OBJECTIVE FOR UNSUPERVISED RL
Desiderata. We first state our two desiderata for scalable unsupervised RL objective. First, instead of covering every possible state in given MDP, which is infeasible in complex environments, we want to have compact latent space of a tractable size and a latent-conditioned policy a that translates latent vectors into actual behaviors. Second, we want the behaviors from different latent vectors to be different, collectively covering as much of the state space as possible. words, we want to maximize state coverage under the given capacity of Z. An algorithm that satisfies these two desiderata would be scalable to complex environments, because we only need to learn compact set of behaviors that approximately cover the MDP.
other
Objective. Based on the above, we propose the following novel objective for unsupervised RL:
Iw(S: 2) = W(p(s, 2), ((s))(s))))
where (w; 2) is the Wasserstein dependency measure (WDM) (Ozair et al., 2019) between states and skills, and W is the 1-Wasserstein distance the metric space (8 x 2, with distance metric d. Intuitively, the WDM objective in Equation (2) can be viewed as a "Wasserstein variant" of the previous MI objective (Equation (I)), where the KL divergence in MI is replaced with the Wasserstein distance. However, despite the apparent similarity, there exists significant difference between the two objectives: MI is completely agnostic the underlying distance metric, while WDM is a metric-aware quantity. As result, the WDM objective (Equation (2)) not only discovers diverse skills that are different from one another, as in the MI objective, but also actively maximizes distances d between different skill trajectories (Figure 2). This makes them collectively cover the state space as much as possible (in terms of the given metric d). The choice of metric for d is the state space
critical for effective skill discovery, and simple choices like Euclidean metrics would generally not be effective for non-metric state representations, such as images. Therefore, instantiating this approach with the right metric is an important part of our contribution, we will discuss in Section 4.2. Until then, we assume that we have a given metric d.
4.1 TRACTABLE OPTIMIZATION
While our objective (S; 2) has several desirable properties, is not immediately straightforward to maximize this quantity in practice. In this section, we describe a simple, tractable objective that can be used maximize (w; Z) in practice. We begin with the Suntorovich-Rubenstee duality (Villani et al., 2009; Ozair et al., 2019), which provides tractable way to maximize the Wasserstein dependency measure:
(w 2) sup Ep(,,s)[[(s, =)] ((s,))(s) I1fl
where Il denotes the Lipschitz constant for the function f 2 under the given distance metric d, i.e., llf = (1 [f(s1, 21) - f(22, =2)l/d(()ss)))))))))))))))/))) /d(s/s(((s)))))))))))))))(())))) (82, 22)). Intuitively, f is score function that assigns larger values to (s, tuples sampled from the joint distribution and smaller values to (s, 2) tuples sampled independently from their marginal distributions. We note can jointly train a L--ipschitz-constriee score function f(s, 2) using gradient descent and a skill policy π [aass, (, 2) using RL, with the reward function being an empirical estimate of Equation (3), r(s,2) f(s, - N Σ i= (s) where 21, 22, , Nn are N independent random samples from the prior distribution p(=).
that Equation (3) is already a tractable objective,
However, since sampling N additional 2s for each data point is computationally demanding, we will further simplify the objective to enable more efficient learning. First, we consider the parameteri-
zation f(,,2) = Φ(s)² (2() with Φ S + IRD and v Z + R with independent 1-Lipschitz constraints which yields the following objective:
Iw(
sup
(o(s))) T +(2)] - [[s)] (x))]
(4)
'While || < 1, اRاا < 1 is not technically equivalent to || < 1, we use the former as it is more tractable. Also, we note that llfll can be upper-hounded in terms of [íőÍl,, ||дğÍд,, sup, (1x(s)ll,, and sup, [[a(2)] under d((s1, (82, 22)) (sup, [10())]2)] 22) + (sup, [[((=)]])][[)]]dd
4

---

Published as a conference paper at ICLR 2024
Here, note that the decomposition 2s) = d(s))(((2) is universal; i.e., the expressiveness of f(s, 2) is equivalent to that of ((s) (2) when D cc. The proof can be found in Appendix C.
Next, consider a variant of the Wasserstein dependency measure that only depends the last state: (w(r 2), similarly to VIC (Gregor et al., 2016). This allows us to further decompose the objective with telescoping sum follows:
Iw(Sr: Z) ~0
sup
(p(r)))(((T)) [x)] (p(>))((sT)] (ppnsll [x(=)((()]
(5)
| I.lii...
sup P.1
Σ
(E (Eptre)l(((se1)) ((x))) (x)) (p(r)(((411)) -(s)) p (x)))((=)))
(6)
where we also use the fact that p(0)) and p(z) are independent. Finally, we set (=)) ) to z. While this makes v less expressive, it allows derive the following concise objective:
Iw(ST: se sup Ep(r.)) 10
Σ (0($(+1) - o(s))) 8) 1=
where 2 = Ep(z)[=] = |z]. Here, since we can always shift the prior distribution p(z) to have a zero mean, we can assume Σ 0 without loss of generality. This objective can now be easily maximized by jointly training (ss) and rTaas,)) 2) with r(s, z, 8') = (Ø(s') x(s)) under the constraint ||0l| S 1. Note that we do not need any additional random samples of z, unlike Equation (3).
4.2 FULLI ORJECTIVE: METRIC-AWARD ABSTRACTION (METRA)
So far, we have not specified the distance metric d for the Wasserstein distance in WDM (or equiv- alently for the Lipschitz constraint | S 1). Choosing am appropriate distance metric crucial for learning compact set of useful behaviors, because it determines the priority by which the be- haviors are learned within the capacity of 2. Previous metric-based skill discovery methods mostly employed the Euclidean distance (or its scaled variant) as a metric (He al., 2022: Park al., 2022; 2023b). However, they are not directly scalable to high-dimensional environments with pixel-based observations, in which the Euclidean distance is not necessarily meaningful.
In this work, we propose to use the temporal distance (Kaelbling, 1993; Hartikainen et al., 2020; Du- rugkar et al., 2021) between two states as a distance metric diemp ($1, $2), the minimum number of environment steps to reach 82 from 81. This provides a natural way to measure the distance between two states, as it only depends on the inherent transition dynamics of the MDP, being invariant to the state representation and thus scalable to pixel-based environments. Using the temporal distance metric, we can rewrite Equation (7) as follows:
T-11 (d($))1))) - ((s:)) -((s)))))((s))
sup Ep(raa) n,
s.t. (1((s) o(s')12 1, V(s,s') E Sadj,
(8)
where Sudj denotes the set of adjacent state pairs in the MDP. Note that || 1 is equivalently converted into ([p(s) - (ss')l2 I 1 under the temporal distance metric (see Theorem C.3).
Intuition and interpretation. We next describe how the constrained objective in Equation (8) may be interpreted. Intuitively, a policy x[ass, 2) that maximizes our objective should learn to move as far possible along various directions in the latent space (specified by =). Since distances in the latent space, [(()))) - ((*2) |2. are always upper-bounded by the corresponding temporal distances in the MDP. given by dremp (31, (31, s2). the learned latent space should assign its (limited) dimensions
to the manifolds in the original state space that are maximally "spread out",, in the sense that shortest paths within the set of represented states should be as long as possible. This conceptually resembles "principal components" of the state space, but with respect to shortest paths rather than Euclidean distances, and with non-linear é rather than linear σ. Thus, we would expect é to learn to abstract the state space in lossy manner, preserving temporal distances (Figure 9), and emphasizing those degrees of freedom of the state that span the largest possible "temporal" (non-linear) manifolds call our method Metric-Aware Abstraction (METRA). In Appendix D, we derive a formal connection between METRA and principal component analysis (PCA) under the temporal distance metric under several simplifying assumptions.
(Figure 1). Based on this intuition,
Theorem 4.1 (Informal statement of Theorem D.2), Under some simplifying assumptions, linear squared METRA is equivalent to PCA under the temporal distance metric.
5

---

Published as a conference paper at ICLR 2024
Algorithm Metric-Aware Abstraction (METRA)
I: Initialize skill policy s(als, 2), representation function ((8), Lagrange multiplier X, replay buffer D
2: for i +- 1 to (# epochs) do
3: for x 1 to (#f episodes per epoch) do
4: Sample skill p(=)
5: Sample trajectory 7 with I(aa,, z) and add to replay buffer 'D
6: end for
7: Update (s)) to maximize B(o,,."")> ((((s)) - ((s)) 2 À min(e, 1 - (10(8) - ((5) - 9: Update =a]ls,)) using SAC (Haarnoja et al., 2018a) with reward r(s,,, (o(s') ((s))
8: Update À to minimize B(+,,,")>>/ (A- min(s, 1 -
10: end for
Connections to previous skill discovery methods. There exist several intriguing connections between our WDM objective (Equation (2)) and previous skill discovery methods, including DI- AYN (Eysenbach et al., 2019a), DADS (Sharma al., 2020), CIC (Laskin al., 2022), LSD (Park et al., 2022), and CSD (Park et al., 2023b). Perhaps the most apparent connections are with LSD and CSD, which also use similar constrained objectives to Equation (7). In fact, although not shown by the original authors, the constrained inner product objectives of LSD and CSD are also equivalent to (Sr: 2), but with the Euclidean distance (or its normalized variant), instead of the temporal distance. Also, the connection between WDM and Equation (7) provides further theoretical insight into the rather "ad-hoc" choice of zero-centered one-hot vectors used in discrete LSD (Park et al., 2022); we must use zero-mean prior distribution due to the Σ term in Equation (7). There exist several connections between our WDM objective and previous MI-based skill discovery meth- ods as well. Specifically, by simplifying WDM (Equation (2)) in three different ways, "Wasserstein variants" of DIAYN, DADS, and CIC. We refer to Appendix for detailed derivations.
obtain
Zero-shot goal-reaching with METRA. Thanks to the state abstraction function ((s), METRA provides a simple way to command the skill policy to reach goal state in a zero-shot manner, as in LSD (Park et al., 2022). Since é abstracts the state space preserving temporal distances, the difference vector <(g) 6(s) tells us the skill we need to select to reach the goal state from the current state s. As such, by simply setting = (o(g) ((s))/ll((()
Φ(s)]12 (for continuous skills) or = arg max (ø(g) o(s)) (for discrete skills), we can find the skill that leads to the goal. With this technique, METRA can solve goal- conditioned tasks without learning a separate goal-conditioned policy, as we will show in Section 5.3.
Implementation. We optimize the constrained objective in Equation (8) using dual gradient descent with Lagrange multiplier À and small relaxation constant & > O, similarly to Park et al. (2023b); Wang et al. (2023). We provide pseudocode for METRA in Algorithm
Limitations. One potential issue with Equation (8) is that we embed the temporal distance into the symmetric Euclidean distance the latent space, where the temporal distance can be asymmetric. This makes our temporal distance abstraction more "conservative" in the sense that it considers the minimum of both temporal distances, i.e., [Ø(st) - o(s2)ll2 S min(dremp (31, 82), deemp ($2, $1)). While this conservatism is less problematic in our benchmark environments, in which transitions are mostly "symmetric", it might be overly restrictive in highly asymmetric environments. To resolve this, we can replace the Euclidean distance (lø(s1) of o(52)ll2 in Equation (8) with an asymmetric quasimetric, as in Wang et al. (2023). We leave this extension for future work. Another limitation is that the simplified WDM objective in Equation (7) only considers behaviors that move linearly in the latent space. While this does not necessarily imply that the behaviors are also linear in the original state space (because : S + Z nonlinear mapping). this simplification, which stems from the fact that we set (z) = z. might restrict the diversity of behaviors to some degree. We believe this can be addressed by using the full WDM objective in Equation (4). Notably, the full objective (Equation (4)) resembles contrastive learning, and we believe combining it with scalable contrastive learning techniques is an exciting future research direction (see Appendix E.3). We refer to Appendix A for practical limitations regarding our implementation of METRA.
EXPERIMENTS
Through our experiments in benchmark environments, we aim to answer the following questions: (1) Can METRA scale to complex, nigh-dimensional environments. including domains with image observations? (2) Does METRA discover meaningful behaviors in complex environments with no supervision? (3) Are the behaviors discovered by METRA useful for downstream tasks?

---

Published as a conference paper at ICLR 2024
METRA
DIAYN DADS
CIC
ICM
APS P2E/ Disag LBS'
Ant (States)
HalEChectah (States)
Quadroped (Piseels)
Humanoid
(Pixels)
Kitchen (Pixels)
Figure 3: Examples of behaviors learned by 11 unsupervised RL methods. For locomotion environments, we plot the x-y 2)) trajectories sampled from learned policies. For Kitchen, measure the coincidental success rates for six predefined tasks. Different colors represent different skills z. METRA is only method that discovers diverse locomotion skills in pixel-based Quadruped and Humanoid. We refer to Figure 11 for the complete qualitative results (8 seeds) of METRA and our project page for videos.
5.1 EXPERIMENTAL SETUP
We evaluate our method on five robotic lo- comotion and manipulation environments (Fig- 4): state-based Ant and HalfCheetah from Gym (Todorov et al., 2012; Brockman et al., 2016), pixel-based Quadruped and Humanoid from the DeepMind Control (DMC) Suite (Tassa al., 2018), and pixel-based version of Kitchen from Gupta et al. (2019); Mendonca et al. (2021). For pixel-based DMC locomotion environments, colored floors allow the agent to infer its location from pixels, similarly to Hafner et al. (2022); Park et al. (2023a) (Figure 10). Throughout the experiments, we do not use any prior knowl- edge, data, or supervision (e.g., observation restriction, early termination, etc.). As such, in pixel- based environments, the agent must learn diverse behaviors solely from 64 x 64 x 3 camera images.
HallChoctah (Ststex)
Ant (Statie)
Qundeuped (Plonle)
Humazood (Pivela)
Kitchen (Piook)
Figure 4: Benchmark environments.
We compare METRA against previous methods in three groups: (1) unsupervised skill discov- ery. (2) unsupervised exploration, and (3) unsupervised goal-reaching methods. For unsupervised skill discovery methods, we compare against two MI-based approaches, DIAYN (Eysenbach et al., 2019a) and DADS (Sharma al., 2020), one hybrid method that combines MI and an exploration bonus, CIC (Laskin et al., 2022), and one metric- basid approach that maximizes Euclidean dis- tances. LSD (Park et al., 2022). For unsupervised exploration methods, we consider five pure explo- ration approaches, ICM (Pathak et al., 2017), RND (Burda et al., 2019), Plan2Explore (Sekar et al., 2020) (or Disagreement (Pathak al., 2019)), APT (Liu & Abbeel, 2021b), and LBS (Mazzaglia et al., 2022), and one hybrid approach that combines exploration and successor features, APS (Liu & Abbeel, 2021a). We note that the Dreamer (Hafner et al., 2020) variants of these methods (especially LBS (Mazzaglia et al., 2022)) are currently the state-of-the-ar methods in the pixel-based unsuper- vised RL benchmark (Laskin al., 2021; Rajeswar al., 2023). For unsupervised goal-reaching methods, we mainly compare with state-of tet-thh-ar art unsupervised RL approach, LEXA (Mendonca al., 2021), as well as two previous skill discovery methods that enable zero-shot goal-reaching. DIAYN and LSD. We use 2-D skills for Ant and Humanoid, 4-D skills for Quadruped, 16 discrete skills for HalfCheetah, and 24 discrete skills for Kitchen. For CIC, we use 64-D skill latent vectors for all environments, following the original suggestion (Laskin et al., 2022).
5.2 QUALITATIVE COMPARISON
We first demonstrate examples of behaviors (or skills) learned by our method and the 10 prior unsu- pervised RL methods on each of the five benchmark environments in Figure 3. The figure illustrates that METRA discovers diverse behaviors in both state-based and pixel-based domains. Notably, METRA the only method that successfully discovers locomotion skills in pixel-based Quadruped and Humanoid, and shows qualitatively very different behaviors from previous unsupervised RL methods across the environments. Pure exploration methods mostly exhibit chaotic, random behav- iors (videos), and fail fully explore the state space (in terms of æ-y coordinates) even in state-based

---

Published as a conference paper at ICLR 2024
HallChrotah (States)
Ant (Statew)
Quadraped (Pieels)
Humamoid (Pioele)
Kitchem (Pixnls)
Pal St Oragr
[S5
DIAYS
DADS
Skgn
Figure 5: Quantitative comparison with unsupervised skill discovery methods (8 seeds), We measure the state/task coverage of the policies learned by five skill discovery methods, METRA exhibits the best coverage across all environments, while previous methods completely fail to explore the state spaces of pixel- based locomotion environments. Notably, METRA is the only method that discovers locomotion skills in pixel-basse Quadruped and Humanoid.
AusMulsCools
HallChectahhooll
HattChectabturrde
QuadrupedCioal
HumaneldGoal
METRA
15D
CIC
DLAYN
DADS
Figure 6: Downstream task performance comparison of unsupervised skill discovery methods (4 seeds). To verify whether learned skills are useful for downstream tasks, we train hicrarchical high-level controller on top of the frozen skill policy to maximize task rewards. METRA exhibits the best or near-best performance across the five tasks, which suggests that the behaviors learned by METRA are indeed useful for the tasks.
Ant and HalfCheetah. This is because it is practically infeasible to completely cover the infinitely many combinations of joint angles and positions in these domains. MI-based skill discovery meth- ods also fail to explore large portions of the state space due to the metric-agnosticity of the KI. divergence (Section 2), even when combined with exploration bonus (i.e., CIC). LSD, a previous metric-based skill discovery method that maximizes Euclidean distances, does discover locomotion skills in state-based environments, but fails to scale to the pixel-based environments, where the Eu- clidean distance on image pixels does not necessarily provide a meaningful metric. In contrast to these methods, METRA learns various task behaviors by maximizing temporal distances in diverse ways. On our project page, we show additional qualitative results of METRA with dif- ferent skill spaces. We note that, when combined with a discrete latent space, METRA discovers even more diverse behaviors, such as doing a backflip and taking a static posture, in addition to locomotion skills. We refer to Appendix F for visualization of learned latent spaces of METRA.
5.3 QUANTITATIVE COMPARISON
Next, we quantitatively compare METRA against three groups of 11 previous unsupervised RL approaches, using different metrics that are tailored to each group's primary focus. For quantitative results, we use 8 seeds and report 95% confidence intervals, unless otherwise stated.
Comparison with unsupervised skill discovery methods. We first compare METRA with other methods that also aim to solve the skill discovery problem (i.e., learning a latent-conditioned policy [alss, 2) that performs different skills for different values of 2). These include LSD, CIC, DIAYN, and DADS?. We implement these methods on the same codebase as METRA. A. For comparison, employ two metrics: policy coverage and downstream task performance. Figure 5 presents the policy coverage results, where we evaluate the skill policy's 's x coverage (HalfCheetah), z-y coverage (Ant, Quadruped, and Humanoid), or task (Kitchen) coverage at each evaluation epoch. The results show that METRA achieves the best performance in most of the domains, and is the only method that successfully learns meaningful skills in the pixel-based settings, where previous skill discovery methods generally fail. In Figure 6, we evaluate the applicability of the skills discovered by each method downstream tasks, where the downstream task is learned by hierarchical controller T² (2]s) that selects (frozen) learned skills to maximize the task reward (see Appendix G for details). METRA again achieves the best performance on most of these tasks, suggesting that the behaviors learned by METRA not only provide greater coverage, but also are more suitable for downstream tasks in these domains.
Comparison with pure exploration methods. Next, we quantitatively compare METRA to five unsupervised exploration methods, which do not aim to learn skills but only attempt to cover the state space, ICM, LBS', RND, APT, and Plan2Explore (or Disagreement), and one hybrid method
We do not compare against DADS in pixel-based environments due the computational cost of its skill dynamics model p(s' [s, s), which requires predicting the full next image.
SSince LBS requires world model, we only evaluate it pixel-based environments, where use the Dreamer variants of pure exploration methods (Rajeswar et al., 2023).
8

---

Published as a conference paper at ICLR 2024
HaltChheetah (btaaee)
Ast (Statea)
Quadioped [Pixiia]
Hunonold (Misaaa)
Kitcben (PMicala)
Kitchem (Ptsola)
METRA
RND
10M
APS
PSK/DDeG
Figure 7: Quantitative comparison with pure exploration methods (8 seeds). We compare METRA with six unsupervised exploration methods in terms of state coverage. Since it is practically infeasible to completely cover every possible state transition, pure exploration methods struggle to explore the state space of complex environments, such as pixel-base based Humanoid or state-based Ant.
HalfChoetah (States)
Ant (States)
Quadruped (Pixels)
Humamoid (Pixels)
Kiteben (Pixols)
METRA
LSD
LEXA
DIAYN
Than
Figure 8: Downstream task performance comparison with LEXA (8 seeds). We compare METRA against LEXA, a state-of-the-art unsupervised goal-reaching method, on five goal-conditioned tasks. The skills Icarned by METRA can be employed to solve these tasks in a zcro-shot manner, achieving the best performance.
that combines exploration and successor features, APS. We use the original implementations by Laskin et al. (2021) for state-based environments and the Dreamer versions by Rajeswar et (2023) for pixel-based environments. As the underlying RL backbones are very different (e.g.. Dreamer is model-based, while METRA uses model-free SAC), we compare the methods based on wall clock time. For the metric, instead of policy coverage (as in Figure 5), we measure total state coverage (i.e., the number of bins covered by any training trajectories up to each evaluation epoch). This metric is more generous toward the exploration methods, since such methods might not cover the entire space on any single iteration, but rather visit different parts of the space on different iterations (in contrast our method, which aims produce diverse skills). In Kitchen, we found that most methods max out the total task coverage metric, and we instead use both the queue coverage and policy coverage metries (see Appendix G for details). Figure 7 presents the results, showing that METRA achieves the best coverage in most of the environments. While pure exploration methods also work decently the pixel-based Kitchen, they fail to fully explore the state spaces of state- based Ant and pixel-based Humanoid, which have complex dynamics with nearly infinite possible combinations of positions, joint angles, and velocities.
Comparison with unsupervised goal-reaching methods. Finally, we compare METRA with LEXA, ofate-of-the-art unsupervised goal-reaching method. LEXA trains an exploration policy with Plan2Explore (Sekar et al., 2020), which maximizes epistemic uncertainty in the transition dynamics model, in parallel with goal-conditioned policy π (a)ss, g) on the data collected by the exploration policy. We compare the performances of METRA, LEXA, and two previous skill dis- covery methods (DLAYN and LSD) on five goal-reaching downstream tasks. We use the procedure described in Section 4.2 to solve goal-conditioned tasks in zero-shot manner with METRA. Fig- ure presents the comparison results, where METRA achieves the best performance on all of the five downstream tasks. While LEXA also achieves non-trivial performances in three tasks, it strug- gles with state-based Ant and pixel-based Humanoid, likely because it is practically challenging to completely capture the transition dynamics of these complex environments.
CONCLUSION
In this work, we presented METRA, scalable unsupervised RL method based on the idea of cover- ing a compact latent skill space that is connected to the state space by temporal distance metric. We showed that METRA learns diverse useful behaviors in various locomotion and manipulation envi- ronments, being the first unsupervised RI. method that learns locomotion behaviors pixel-based Quadruped and Humanoid.
Final remarks. In unsupervised RL, many excellent prior works have explored pure exploration or mutual information skill learning. However, given that these methods not necessarily readily scalable to complex environments with high intrinsic state dimensionality, as discussed in Section 2, may need completely different approach to enable truly scalable unsupervised RL. We hope that this work serves as a step toward broadly applicable unsupervised RL that enables large-scale pre-training with minimal supervision.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We would like to thank Youngwoon Lee for an informative discussion, and RAIL members and anonymous reviewers for their helpful comments. This work was partly supported by the Korea Foundation for Advanced Studies (KFAS), ARO W911NF-21-1-0097 and the Office of Naval Re- search. This research used the Savio computational cluster resource provided by the Berkeley Re- search Computing program UC Berkeley.
REPRODUCIBILITY STATEMENT
We provide our code at the following repository: http://///gtthbb.oooeeeeoeee cttptuub.ccc//eeeeeeeeeeeeeeeeee METRA. We provide the full experimental details in Appendix G.
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. ArXiv, abs/1807. 10299, 2018.
Adríà Puigdomènech Badia, Pablo Sprechmann, Alex Vitvitskyi, Daniel Guo, Bilal Piot, Steven Kapturowski, Olivier Tieleman, Martín Arjovsky, Alexander Pritzel, Andew Bolt, and Charles Blundell. Never give Learning directed exploration strategies. In International Conference on Learning Representations (ICLR), 2020.
Richard F Bass. Real analysis for graduate students. Createspace Ind Pub, 2013.
Kate Baumli, David Warde-Farley, Steven Stenberg Hansen, and Volodymyr Mnih. Relative varia- tional intrinsic control. AAAI Conference on Artificial Intelligence (AAAI), 2021.
Marc G. Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Rémi Munos. Unifying count-based exploration and intrinsic motivation. In Neural Information Pro- cessing Systems (NeurIPS), 2016.
Glen Berseth, Daniel Geng, Coline Devin, Nicholas Rhinehart, Chelsea Finn, Dinesh Jayaraman, and Sergey Levine. Smirl: Surprise minimizing reinforcement learning in unstable environments. In International Conference on Learning Representations (ICLR), 2021.
Rajendra Bhatia. Matrix analysis. Springer Science Business Media, 2013.
G. Brockman, Vicki Cheung, Ludwig Pettersson, J. Schneider, John Schulman, Jie Tang, and W. Zaremba. OpenAl Gym. ArXiv, abs/1606.01540,, 2016.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voor Gretchen Krueger, T. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, ılya Sutskever, and Dario Amodei. Language models are few-shot learners. In Neural Informa- tion Processing Systems (NeurIPS), 2020.
Yuri Burda, Harrison Edwards, Amos J. Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations (ICLR), 2019.
Víctor Campos Camúñez, Alex Trott, Caiming Xiong, Richard Socher, Xavier Giró Nieto, and Jordi Torres Viñals. Explore, discover and leam: unsupervised discovery of state-covering skills. In International Conference on Machine Learning (ICML), 2020.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning (ICML), 2020.
10

---

Published as a conference paper at ICLR 2024
Wenze Chen, Shiyu Huang, Yuan Chiang, Tingling Chen, and Jun Zhu. Dgpo: Discovering multiple strategies with enersity-guides policy optimization. In AAAI Conference on Artificial Intelligence (AAAI), 2024.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith W. Ross. Randomized ensembled double q- learning: Learning fast without model. In International Conference Learning Representa- fions (ICLR), 2021.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Gu. Variational empow- erment as representation learning for goal-conditioned reinforcement learning. In International Conference Machine Learning (ICMI.), 2021.
John D. Co-Reyes, Yuxuan Liu, Abhishek Gupta, Benjamin Eysenbach, P. Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajec- tory embeddings. In International Conference on Machine Learning (ICML), 2018.
Ishan Durugkar, Mauricio Tec, Scott Niekum, and Peter Stone. Adversarial intrinsic motivation for reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2021.
Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. First return, then explore. Nature, 590:580-586, 2020.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without a reward function. In International Conference on Learning Representa- tions (ICLR), 2019a.
Benjamin Eysenbach, Ruslan Salakhutdinov. and Sergey Levine. Search on the replay buffer: Bridg- ing planning and reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2019b.
Carlos Florensa, Yan Duan, and Abbeel. Stochastic neural networks for hierarchical reinforcement learning. In International Conference Learning Representations (ICLR), 2017.
Carlos Florensa, Jonas Degrave, Nicolas Manfred Otto Heess, Jost Tobias Springenberg, and Mar- tin A. Riedmiller. Self sulf-super learning of image embedding for continuous control. ArXiv, hbbs/190.009%/ 00994, 2019.
Justin Fu, John D. Co-Reyes, and Sergey Levine. Ex2: Exploration with exemplar models for deep reinforcement learning. Neural Information Processing Systems (NeurIPS), 2017.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. ArXiv, abs/1611.07500 2016.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. Conference on Robot Learning (CoRL), 2019.
Michael U Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2010.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. International Confer- on Machine Learning (ICML), 2018a.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, G. Tucker, Sehoon Ha, Jie Tan, Vikash Ku- mar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Soft actor-critic algorithms and applications. ArXiv, abx/1812.055e 2018b.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn ing behaviors by latent imagination. International Conference on Learning Representations (ICLR), 2020.
11

---

Published as a conference paper at ICLR 2024
Danijar Hafner, Kuang-Huei Lee, Ian S. Fischer, and P. Abbeel. Deep hierarchical planning from pixels. In Neural Information Processing Systems (NeurIPS), 2022.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In International Conference on Machine Learning (ICML), 2022.
S. Hansen, Will Dabney, André Barreto, T. Wiele, David Warde-Farlcy, and V. Mnih. Fast task inference with variational intrinsic successor features. International Conference on Learning Representations (ICLR), 2020.
Kristian Hartikainen, Xinyang Geng, Tuomas Haarnoja, and Sergey Levine. Dynamical distance learning for semi-super and unsupervised skill discovery. In International Conference on Learning Representations (ICLR), 2020.
Elad Hazan, Sham M. Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In International Conference on Machine Learning (ICML), 2019.
Shuncheng He, Yuhang Jiang, Hongchang Zhang, Jianzhun Shao, and Xiangyang Ji. Wasserstein unsupervised reinforcement learning. In AAAI Conference Artificial Intelligence (AAAI), 2022.
Takuya Hiraoka, Takahisa Imagawa, Taisei Hashimoto, Takashi Onishi, and Yoshimasa Tsuruoka. Dropout q-functions for doubly efficient reinforcement learning. In International Conference on Learning Representations (ICLR), 2022.
Rein Houthooft, Xi Chen, Yan Duan, John Schulman, Filip De Turck, and P. Abbeel. Vime: Varia- tional information maximizing exploration. In Neural Information Processing Systems (NeurIPS), 2016.
Edward S. Hu, Richard Chang, Oleh Rybkin, and Dinesh Jayaraman. Planning goals for exploration. In International Conference on Learning Representations (ICLR), 2023.
Zheyuan Jiang. Jingyue Gao, and Jianyu Chen. Unsupervised skill discovery via recurrent skill training. In Neural Information Processing Systems (NeurIPS), 2022.
Leslie Pack Kaelbling. Learning to achieve goals. In International Joint Conference Artificial Intelligence (ICCAI), 1993.
Alexad-
Kamienny, Jean Tarbouriech, Alessandro Lazaric, and Ludovic Denoyer. Direct then diffuse: Incremental unsupervised skill discovery for state covering and goal reaching. In International Conference on Learning Representations (ICLR), 2022.
Jackyeom Kim, Seohong Park, and Gunhee Kim. Unsupervised skill discovery with bottleneck option learning. In International Conference Machine Learning (ICML), 2021.
Scongun Kim, Kyowoon Lee, and Jaesik Choi. Variational curriculum reinforeement learning for unsupervised discovery of skills. In International Conference on Machine Learning (ICML), 2023.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Martin Klissarov and Marlos C. Machado. Deep laplacian-based options for temporally-extendded exploration. In International Conference Machine Learning (ICML), 2023.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representa- tions (ICLR), 2021.
Saurabh Kumar, Aviral Kumar, Sergey Levine, and Chelsea Finn. One solution is not all you need: Few-shot extrapolation via structured maxent rl. In Neural Information Processing Sys- tems (NeurIPS), 2020.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Abbeel. Urlb: Unsupervised reinforcement learning benchmark. In Neural Informa- fion Processing Systems (NeurIPS) Datasets and Benchmarky Track, 2021.
12

---

Published as a conference paper at ICLR 2024
Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and P. Abbeel. Un- supervised reinforeement learning with contrastive intrinsic control. In Neural Information Pro- cessing Systems (NeurIPS), 2022.
Yann LeCun, Bernhard E. Boser, John S. Denker, Donnie Henderson, Richard E. Howard, Wayne E. Hubbard, and Lawrence D. Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1:541-5511 1989.
Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric P. Xing, Sergey Levine, and Ruslan Salakhut- dinov. Efficient exploration via state marginal matching. ArXiv, abss/1906.05274, 2019.
Mengdi Li, Xufeng Zhao, Jae Hee Lee, Comelius Weber, and Stefan Wermter. Internally rewarded reinforcement learning. In International Conference Machine Learning (ICML), 2023.
Hao Liu and Pieter Abbeel. APS: Active pretraining with successor features. In International Conference on Machine Learning (ICML), 2021a.
Hao Liu and Pieter Abbeel. Behavior from the void: Unsupervised active pre-training. In Neural Information Processing Systems (NeurIPS), 2021b.
Marlos C. Machado, Mare G. Bellemare, and Michael Bowling. A laplacian framework for option discovery in reinforcement learning. In International Conference Machine Learning (ICML), 2017.
Marlos C. Machado, Clemens Rosenbaum, Xiaoxiao Guo, Miao Liu, Gerald Tesauro, and Murray Campbell. Eigenoption discovery through the deep successor representation. In International Conference on Learning Representutions (ICLR), 2018.
Pietro Mazzaglia, Ozan Çatal, Tim Verbelen, and B. Dhoedt. Curiosity-driven exploration via latent bayesian surprise. In AAAI Conference on Artificial Intelligence (AAAI), 2022.
Pietro Mazzaglia, Tim Verbelen, B. Dhoedt, Alexandre Lacoste, and Sai Rajeswar. Choreographer: Learning and adapting skills in imagination. In International Conference on Learning Represen- tations (ICLR), 2023.
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discover- ing and achieving goals via world models. Neural Information Prucessing Systems (NeurIPS), 2021.
Shakir Mohamed and Danilo J. Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2015.
Mirco Mutti, Lorenzo Pratissoli, and Marcello Restelli. Task-agnostic exploration via policy gra- dient of non-paraeetric state entropy estimate. In AAAI Conference on Artificial Intelligence (AAAI), 2021.
OpenAl OpenAI, Matthias Plappert, Raul Sampedro, Tao Xu, 1lge Akkaya, Vineet Kosaraju, Peter Welinder, Ruben D' DaS Arthur Petron, Henrique Pondé de Oliveira Pinto, Alex Paino, Hyeonwoo Noh, Lilian Weng, Qiming Yuan, Casey Chu, and Wojeiech Zaremba. Asymmetric self-play for automatic goal discovery in robotic manipulation. ArXiv, abs/2101 .048220 2021.
Georg Ostrovski, Marc G. Bellemare, Aäron van den Oord, and Rémi Munos. Count-based explo- ration with neural density models. In International Conference on Machine Learning (ICML), 2017.
Sherjil Ozair, Corey Lynch, Yoshua Bengio, Aäron van den Oord, Sergey Levine, and Pierre Ser- manet. Wasserstein dependency measure for representation learning. In Neural Information Processing Systems (NeurIPS), 2019.
Seohong Park and Sergey Levine. Predictable mdp abstraction for unsupervised model-based rl. In International Conference Machine Learning (ICML), 2023.
Scohong Park, Jongwook Choi, Jaekyeom Kim, Honglak Lee, and Gunhee Kim. Lipschitz- constrained unsupervised skill discovery. In International Conference on Learning Represen- fations (ICI.), 2022.
13

---

Published as a conference paper at ICLR 2024
Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. Hiql: Offline goal- conditioned with latent states as actions. In Neural Information Processing Systems (NeurIPS), 2023a.
Seohong Park, Kimin Lee, Youngwoon Lee, and P. Abbeel. Controllability-aware unsupervised skill discovery. In International Conference on Machine Learning (ICML), 2023b.
Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International Conference on Machine Learning (ICML), 2017.
Deepak Pathak, Dhiraj Gandhi, and Abhinav Kumar Gupta. Sele-super exploration via dis- agreement. In International Conference on Machine Learning (ICML), 2019.
Silviu Pitis, Harris Chan, S. Zhao, Bradly C. Stadie, and Jimmy Ba. Maximum entropy gain ex- ploration for long horizon multi-goal reinforcement learning. International Conference on Machine Learning (ICML), 2020.
Vitchyr H. Pong, Murtaza Dalal, S. Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew-Fit: State-covering self- supervisdd reinforcement learning. In International Conference on Machine Learning (ICML), 2020.
Ben Poole, Sherjil Ozair, Aäron van den Oord, Alexander A. Alemi, and G. Tucker. On variational bounds of mutual information. International Conference on Machine Learning (ICML), 2019.
A. H. Qureshi, Jacob Johnson, Yuzhe Qin, Taylor Henderson, Byron Boots, and Michael C. Yip. Composing task-agnostic policies with deep reinforcement learning. In International Conference on Learning Representations (ICLR), 2020.
Sai Rajeswar, Pietro Mazzaglia, Tim Verbelen, Alexandre Pich'e, B. Dhoedt, Aaron C. Courville, and Alexandre Lacoste. Mastering the unsupervised reinforcement learning benchmark from pixels. In International Conference on Machine Learning (ICML), 2023.
Nick Rhinchart, Jenny Wang, Glen Berseth, John D. Co-Reyes, Danijar Hafner, Chelsea Finn, and Sergey Levine. Information is power: Intrinsic control via information capture. In Neural Infor- mation Processing Systems (NeurIPS), 2021.
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametric topological memory for navigation. In International Conference on Learning Representations (ICLR), 2018.
Tom Schaul, Dan Horgan, Karol Gregor, and David Silver. Universal value function approximators.
In International Conference on Machine Learning (ICML), 2015.
John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and P. Abbeel. High-dimensional continuous control using generalized advantage estimation. In International Conference on Learning Representations (ICIR), 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv, abs/1707. 063407777777...7 2017.
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, P. Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International Conference Machine Learning (ICML), 2020.
Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, P. Abbeel, and Kimin Lee. State entropy maximization with random encoders for efficient exploration. International Conference on Machine Learning (ICMI), 2021.
Nur Muhammad (Mahi) Shafiullah and Lerrel Pinto. One after another: Learning incremental skills for a changing world. In International Conference on Learning Representations (ICLR), 2022.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics- aware unsupervised discovery of skills. In International Conference on Learning Representations (ICLR), 2020.
14

---

Published as a conference paper at ICLR 2024
Pranav Shyam, Wojeiech Jaškowski, and Faustino J. Gomez. Model-based active exploration. In International Conference on Machine Learning (ICML), 2019.
DJ Strouse, Kate Baumli, David Warde-Farley, Vlad Mnih, and Steven Stenberg Hansen. Learning more skills through optimistic exploration. In International Conference on Learning Representa- tions (ICLR), 2022.
Sainbayar Sukhbaatar, Iyy Kostrikov, Arthur D. Szlam, and Rob Fergus. Intrinsic motivation and automatic curricula via asymmetric self-play. In International Conference on Learning Represen- tations (ICLR), 2018
Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, Xi Chen, Yan Duan, John Schulman, Filip De Turck, and Abbeel. #exploration: A study of count-based exploration for deep rein- forcement learning. In Neural Information Processing Systems (NeurIPS), 2017.
Yuval Tassa, Yotam Doron, Alistair Muidal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillierap, and Martin A. Riedmiller. Deepmind control suite. ArXiv, abs/1801.00690 2018.
Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2012.
Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards. In Neural Information Processing Systems (NeurIPS), 2021.
Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? In International Conference on Learning Representations (ICLR), 2023.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. ArXiv, abs/ 0tos/180.000 00744, 2018
Cédric Villani et al. Optimal transport: old and new. Springer, 2009.
Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal goal-reaching rein- forcement learning via quasimetric learning. In International Conference Machine Learning (ICML).
David Warde-Farley, Tom Van de Wiele, Tejas Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih. Unsupervised control through -onpaaraettric discriminative rewards. In In- ternational Conference on Learning Representations (ICLR), 2019.
Rushuai Yang, Chenjia Bai, Hongyi Guo, Siyuan Li, Bin Zhao, Zhen Wang, Peng Liu, and Xuelong Li. Behavior contrastive learning for unsupervised skill discovery. In International Conference on Machine Learning (ICML), 2023.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with pro- totypcal representations. In International Conference on Machine Learning (ICML). 2021.
Tom Zahavy, Yannick Schroecker, Feryal M. Behbahani, Kate Baumli, Sebastian Flennerhag, Shaobo Hou, and Satinder Singh. Discovering policies with domino: Diversity optimization maintaining near optimality. In International Conference on Learning Representations (ICLR), 2023a.
Tom Zahavy, Vivek Veeriah, Shaobo Hou, Kevin Waugh, Matthew Lai, Edouard Leurent, Nenad Tomasev, Lisa Schut, Demis Hassabis, and Satinder Singh. Diversifying ai: Towards creative chess with alphazero. ArXiv, abs/2308.09175. 2023b.
Jesse Zhang, Haonan Yu, and Wei Xu. Hierarchical reinforcement learning by discovering intrinsic options. In International Conference Learning Representations (ICLR), 2021.
Andrew Zhao, Matthicu Lin, Yangguang Li, Y. Liu, and Gao Huang. A mixture of surprises for unsupervised reinforcement learning. Neural Information Processing Systems (NeurIPS), 2022.
Zihan Zhou, Wei Fu, Bingliang Zhang, and Yi Wu. Continuously discovering novel strategies via reward- switching policy optimization. International Conference on Learning Representations (ICLR), 2022.
15

---

Published as a conference paper at ICLR 2024
preserves temporal distances denotes adjacent states)
Learn to move every direction in 2,
which approximately covers S
S = R64X64X3
= R² s(s,'') (1((s) S
S = 164X64X3
3 R=R2 max ( o(s))
Figure 9: Intuitive interpretation of METRA. Our main idea is to only cover compact latent space Z that metrically connected the space S. Specifically, METRA learns lossy, compact representation function Φ S Z, which preserves temporal distances (leff), and learns to move every direction Z, which leads to approximate coverage the state space S (right).
Observation
Global view
Observation
Global view
(a) Quadruped
(b) Humanoid
Figure 10: Visualization of pixel-based DMC Quadruped and Humanoid. We gradient-colored floors allow the agent to infer its location from pixel observations, similarly Hafner et al. (2022); Park al. (2023a).
A LIMITATIONS
Despite its state-of-the-art performance in several benchmark environments, METRA, in its current form, has limitations. We refer to Section 4.2 for the limitations and future research directions regarding the METRA objective. In terms of practical implementation, METRA, like other similar unsupervised skill discovery methods (Sharma al., 2020; Park et, 2022; 2023b), uses relatively small update-to-dae to-dat (UTD) ratio (i.e., the average number of gradient steps per environment step); e.g., use 1/4 for Kitchen and 1/16 for Quadruped and Humanoid. Although we demonstrate that METRA learns efficiently in terms of wall clock time, we believe there is room for improvement in terms of sample efficiency. This is mainly because we use vanilla SAC (Haarnoja al., 201 8a) as its RL backbone for simplicity, and we believe increasing the sample efficiency of METRA by combining it with recent techniques in model-free RL (Kostrikov et al., 2021; Chen al., 2021; Hiraoka et al., 2022) model-based RI. (Hafner et al., 2020; Hansen et al., 2022) is interesting direction for future work.
Another limitation of this work is that, while we evaluate METRA various locomotion and ma- nipulation environments, following prior work in unsupervised RL and unsuperivsed skill discov- ery (Eysenbach et al., 2019a; Sharma et al., 2020; Mendonca et al., 2021; Laskin et al., 2021; He et al., 2022; Park et al., 2022; Zhao et al., 2022: Shafiullah & Pinto, 2022; Laskin et al., 2022; Park al., 2023b; Yang al., 2023). we have not evaluated METRA on other different types of environ- ments, such as Atari games. Also, since we assume fixed MDP (i.e., stationary. fully observable dynamics, Section 3), METRA in its current form does not particularly deal with non-stationary
16

---

Published as a conference paper at ICLR 2024
non-Markovian dynamics. We leave applying METRA to more diverse environments extending the idea behind METRA to non-stationary or non-Markovian environments for future work.
B EXTENDED RELATED WORK
In addition to unsupervised skill discovery (Mohamed & Rezende, 2015; Gregor et al., 2016; Flo- rensa et al., 2017; 7; Co-Reyes et al., 2018; Achiam et al., 2018; Eysenbach et al., 2019a; Warde-Farley et al., 2019; Shyam et al., 2019; Lee et al., 2019; Sharma et al., 2020; Campos Camúñez et al., 2020;
Hansen et al., 2020; Pong et al., 2020; Baumli et al., 2021; Choi et al., 2021; Yarats et al., 2021;
Kim et al., 2021; Zhang ct al., 2021; He et al., 2022: Strouse et al., 2022: Laskin et al., 2022; Park et al., 2022: Shafiullah & Pinto, 2022: Jiang et al., 2022; Zhao et al., 2022; Kamienny et al., 2022; Park & Levine, 2023; Park et al., 2023b; et
et al., 2023; Kim et al., 2023) and pure exploration (or unsupervised gondicond RL) methods (Houthooft et al., 2016; Bellemare et al., 2016; Tang et al., 2017; Ostrovski et al.. 2017: Fu et al., 2017; Pathak et al.. 2017: Hazan et al., 2019; Shyam et al., 2019; Burda al., 2019; Pathak et al., 2019; Lee al., 2019; Ecoffet et al., 2020; Pitis al., 2020; Badia et al., 2020; Mutti et al., 2021; Liu & Abbeel, 2021b; Mendonca et al., 2021; Yarats al., 2021; Sco et al.. 2021; Mazzaglia al., 2022: 2023; Hu et al., 2023; Rajeswar et al., 2023), there have also been proposed other types of unsupervised RL approaches, such as ones based on asymmetric self-play (Sukhbaatar et al., 2018; OpenAl et al., 2021), surprise minimization (Berseth et al., 2021; Rhinehart et al., 2021), and forward- barkd-bawrrdddddr representations (Touati & Ollivier, 2021; Touati et al., 2023). One potentially closely related line of work graph Laplacian- based option discovery methods (Machado et al., 2017; 2018; Klissarov & Machado, 2023). These methods learn a set of diverse behaviors based on the eigenvectors of the graph Laplacian of the MDP's adjacency matrix. Although we have not found formal connection these methods, cxist a deep, intriguing connection between METRA and graph Laplacian- based methods, given that they both discover behaviors based the temporal dynamics of the MDP. METRA is also related to several works in goal-conditioned RL that consider temporal distances (Kaelbling, 1993; Schaul et al., 2015; Savinov et al., 2018; Eysenbach al., 2019b; Florensa et al., 2019; Hartikainen et al., 2020; Durugkar et al., 2021; Wang et al., 2023). In particular, Durugkar et al. (202 I); Wang et al. (2023) use similar temporal distance constraints to ours for goal-conditioned RL.
suspect there might
C THEORETICAL RESULTS
C.I UNIVERSALITY OF INNER PRODUCT DECOMPOSITION
Lemma C.1. Let X and y be compact Hausdorff spaces (e.g., compact subsets in R and C(A) be the set of real-valued continuous functions on A. For any function f(x, y) E C(X y) and e > O, there exist continuous functions ((x) : X - IR and p(y) y - RRD with D 2 1 such that sUy (f(x,,) <xxx)
& (b())l <.
Proof. We invoke the Stone Weierstrass theorem (Bass (2013), Theorem 20.44), which implies that the set of functions T {E { (i())d(()) D N, < i < D, oi E C(X), viy e C())} is dense in C(X' x )) if T is an algebra that separates points and vanishes at no point. The only non-trivial part is to show that T is closed under multiplication. Consider
(x, = E.r g0a) (z)es (y) E T and (x)))) = EE
Φi (2) (x)x (v) E ET. We have
(1, and ( (m) () E C(C) for all i.j. Hence, (0) (a, u) (I, ET.
EE EP Pi (x)o)) (o)x (1) (w))) (P) (v), where of (z)))) (x) E C(X)
Dj (1)
(2)
Theorem C.2 (ó(x:) (yy) is universal approximator of f(x 8)). Let X' and y be compact Haus- dorff spaces and &h C C(X) and c CCY) be dense sets in C(X) and C(Y), respectively. Then, := E (x(x))s())
D E N,VI S & D, di E $ the E w} is also dense in C(X x Y). In other words, o(x) vyy) can approximate [(x, y) to arbitrary accuracy if o and to are modeled with universal approximators (e.g. neural nefworks) and D + xc.
Proof. By Lemma C.1, for any f E C(X' x y) and E > O, there exist D E N, di E C(A), and w/ E C()) for 1 < i S D such that sup wey (sxx, Ebn (x))(())) < e/3. Define
17

---

Published as a conference paper at ICLR 2024
M, == supigisD [ba())). Since + is dense, for each S i D, there exists Φs & such that
supzex (di(z) (z(2) < c3D3MMy) Define M₂ == suPpsisD,zex (di(z)). Similarly, for each
1 ssis D, there exists dh v such that supyey lvv(y))
Now, have
D (xx.v) - - ((())))))) - [xx,0) Σ (((()))))))) + Σ (x))v())))) - (9)
Σ
Σ ((,.(0) -
- vv(v))) +
(x))
(10)
3 +
(11)
(12)
for any x E X' and y E V. Hence, T is dense in C(A's x )))
C.2 LIPSCHITZ CONSTRAINT UNDER THE TEMPORAL DISTANCE METRIC
Theorem C.3. The following two conditions are equivalent:
(a) (b(()) - <(v)ll2 S dremp for all u, v E
(b) [0(s) - 6(s))l2 2 - for all (s, 5') E Sadj-
Proof. We first show (a) implies (b). Assume (a) holds. Consider (8, 8)) E Sudj. I + s. by (a), have lØ(s) (s S dremp s') 1. Otherwise, i.e., = (0((s) o(s') 0 1. Hence, (a) implics (b).
Next, we show (b) implies (a). Assume (b) holds. Consider u, v E S. If chemp (u, v) (i.e., is
not reachable from w). (a) holds trivially. Otherwise, let be dtemp (u, v). By definition, there exists (su u, $1...., Sk-1, Sk v) such that (s,, 8141) Sadj for all 0 S i S k 1. Due to the triangle inequality and (b), we have [lő(u) - <(v)ll2 I E/E [(o(s)) ((s++)ll s k = dremp (u, v). Hence, (b) implies (a).
D A CONNECTION BETWEEN METRA AND PCA
In this section, we derive theoretical connection between METRA and principal component anal- ysis (PCA). Recall that the METRA objective can be written as follows:
sup E Ep(r.e) 9
Σ (6(s111) -(()))))(((s)))==+ = -pp(ra)) [d(ST) Tz]
(13)
s.t. (ςØ(u) - 0(v)l22 S dremp (u. v), Vu,
(14)
where dhemp denotes the temporal distance between two states. To make a formal connection be- tween METRA and PCA, we consider the following squared variant of the METRA objective in this section.
sup Ep(r.2) [(0(s)) [2)]] s.t. [lő(u) - -(v))l2 S cuemp v), Vu, U e S, F
(15)
which is almost the same as Equation (13) but the objective is now squared. The reason we consider this variant simply for mathematical convenience.
Next, we introduce the notion of a temporally consistent embedding.
Definition D.1 (Temporally consistent embedding). We call that an MDP M admits a temporally consistent embedding if there exists vi(s) S + RM such that
temp (u, v) (()) (b)))l2, Vu, S.
(16)
18

---

Published as a conference paper at ICLR 2024
Intuitively, this states that the temporal distance metric can be embedded into a (potentially very nigh-dimensional) Euclidean space. We note that s is different from ø in Equation (13), and R" can be much higher-dimensional than Z. An example of an MDP that admits temporally consistent emhedding is the PointMass environment: if an agent in R can move in any direction up to a unit satisfies d'emp (u, v) I|u 0:12 for all u, v E IR (with slightly generalized notion of temporal distances in continuous time) and thus the MDP admits the temporally consistent embedding of y. A pixel-based PointMass environment is another example of such an MDP.
speed, vb(x;)
Now, we formally derive connection between squared METRA and PCA. For simplicity, we as- sume = Rd and p(z) (0, Ia). where (0 La) denotes the d-dimensional isotropic Gaussian distribution. We also assume that .M has deterministic initial distribution and transition dynamics function, and every state reachable from the initial state within steps. We denote the set of x n positive definite matrices as + the operator norm of a matrix A as ||A||opp and the m-dimensional unit &2 ball as B.
Theorem D.2 (Linear squared METRA is PCA in the temporal embedding space). Let M be MDP that admits a temporally consistent embedding v * IRm If : S
is a linear mapping from the embedding space, i.e., o(s) W (ss) with W E Rmxxd and the embedding space t {
E S} forms an ellipse, i.e., 3. E $4 s.f. T = fx E R x A S 1}, then W (a 02 aa) maximizes the squared METRA objective in Equation (15), where a ....
aa are the top-d eigenvectors of A.
Proof. Since M admits a temporally consistent embedding, we have
llØ(u) - d(v)ll2 & dremp (u, v) Vu, v E (v(()) (()))l22 < [1b(u) v(v)l22 Vu | (u 2)](2 < ||u u112 Vu [Wlp I 1,
(17)
(18)
(19)
(20)
where we use the fact that tp is surjection from S to v and that A is positive definite. Now, have
= sup (p(r.))(((sT)) 2)²1
(21)
=
sup *.I[W[lapS1
E ((((T))(((TT) W2)
(22)
= sup
E [f(=)[((()) T W2)² (:: Every state is reachable within T steps)
(23)
sup
E [p(=)(((())) (Vo VAW 2) ( = VA-Ts(())
(24)
S1
= sup
Eptal sup (g(=) VAW2)²]
(25)
[W
9m8m-3m
= =sup
E Eptal
sup (u T <AW2)²]
(26)
[W
|ul|aSS1
sup
E Ex(e) AW
(:: Dual norm)
(27)
[W $1
= sup
Ex(=]22 W AW2]
(28)
[W $1
= sup
E [t WT AAW]
(29)
[W] $1
= sup
tr(E Tw AW)
(30)
= sup tr(WWT A). [1
(31)
Since WWT is positive semidefinite matrix with rank at most d and || S 1, there exists d eigenvalues 0 < .,,..., da < and the corresponding orthonormal eigenvectors .I,..., Va such that WWT Ed_ ABURU Hence, tr(WWT A) =Et dRVY Avg, and maximize this,
19

---

Published as a conference paper at ICLR 2024
X2 1 as A is positive definite. The remaining problem is to find d orthonor- mal vectors ....... Vg that maximize Σx= U/ Aek. By the Ky Fan's maximum principle (Bhatia, 2013), its solution is given as the d eigenvectors corresponding to the d largest eigenvalues of A. Therefore, W [a] ag ag), where d1...., ad are the top-d principal components of A, maxi- mizes the squared METRA objective in Equation (15).
must set X1
Theorem D.2 states that linear squared METRA is equivalent PCA in the temporal embedding be nonlinear, the shape of v can be arbitrary, and the MDP may
space. In practice, however, not admit any temporally consistent embeddings. Nonetheless, this theoretical connection hints at the intuition that the METRA objective encourages the agent to span the largest "temporal" mani- folds in the state space, given the limited capacity of Z.
CONNECTIONS BETWEEN WDM AND DIAYN, DADS, AND CIC
In this section, we describe connections between WDM objectives (either Iw(S; Z) or [w(Sr; Z)) and previous mutual information skill learning methods, DIAYN (Eysenbach et al., 2019a), DADS (Sharma et al., 2020), and CIC (Laskin et al., 2022). Recall that the (w Z) objective (Equation (4)) maximizes
Σ (LxT).))oz(ss)) dz(2)) Exc)) (x())(())))
(32)
and the (w(S:: 2) objective (Equation (6)) maximizes
Σ -Bxcrell00((441) - (x(ss))T((++)))) - [ Ep(s)Vr())) (xo (33)
where we use the notations Φ1 and t to denote that they are Lipschitz constrained. By simplifying Equation (32) or Equation (33) in three different ways, we will show that can obtain "Wasser- stein counterparts" of DIAYN, DADS, and CIC. For simplicity, we assume p(2) (N I), where ((0, I) denotes the standard Gaussian distribution.
E. DIAYN
If we set (2) = in Equation (32), we get
Te = 02 (bt)))
(34)
This is analogous to DIAYN (Eysenbach et al., 20192)1))) which maximizes
- -H(2|S) + H(2)
(35)
2 Ep(r,>)
log (=(ss)
(36)
Ep((,
Σ
-||2 (s
(37)
rE
DIAYN - ((((s))) =11
(38)
where '2' and '22" respectively denote '>' and "=" up constant scaling or shifting, and we assume that the variational distribution {=][x) modeled as (d(ss), I). By comparing Equation (34) and Equation (38), DIAYN, This analogy will become clearer later.
see that Equation (34) can be viewed Lipschitz, inner-product variant of
E.2 DADS
If we set (x) in Equation (33), we get
Y2
(8t+1 8t) T (2((2) - (8141 - 8t) T Ep(=)((n(z))
(39)
1
(8t+1 ST) (z)(2))
(8411 se) tb (z)),
(40)
1=1
20

---

Published as a conference paper at ICLR 2024
where we use L independent samples from (0, 21, 2g,...,, &L, to approximate the expectation.
This is analogous to DADS (Sharma et al., 2020), which maximizes:
ZIS)= - (S Z) + SS)
(41)
Z (TTr))
log ( 118, - log log ((s(+1())) M
(42)
1
~ R B)(()))
log (000(((+1//)) 18,,
Σ
log (0 23)
(43)
L
==1
L
DADS -1((8+11 88) (888,, =)1
+ T. ;=1
Σ
|1(5+1 81) - --),
(44)
where we assume that the variational distribution g(s'| s, =) is modeled as q(s' s s|8, z)
(ib(s, 2), I), as in the original implementation (Sharma et al., 2020). We also use the same sample- based approximation as Equation (40). Note that the same analogy also holds between Equation (40) and Equation (44) (i.e., Equation (40) is Lipschitz, inner-product variant of DADS).
E.3
CIC
If we do not simplify Φr or the in Equation (32), we get
= (st) (2)))) (s)))) T Exte)(b
(45)
~ (s()s) (2())) 1 Σ (s)))) (z)),,
(46)
where we use the same sample-base based approximation as Equation (40). By Jensen's 's inequality, Equa- tion (46) can be lower-bounded by
I log Σ exp (xx()))))))x Φr. (s:) (()))T(((())))))))(())))))())( ()(())
(47)
I
as in WPC (Ozair et al., 2019). This is analogous to CIC (Laskin et al., 2022), which estimates the MI via noise contrastive estimation (Gutmann Hyvärinen, 2010; van den Oord et al., 2018; Poole et al., 2019):
I
(S; 2 E Ep(r,=)
((s)) (e))))))=)>)))))))))>))) - log Σ 1=1
exp (ő(st) ((xi) (d(ss)
(48)
L
-C = o(s:) T (1)))= log Σ exp (d(st) T (zx))
(49)
Note that Equation (47) can be viewed as a Lipschitz variant of CIC (Equation (49)).
In this work, found this variant to work well while being simple, but we believe exploring these other variants is an interesting future research direction. particular, given that Equation (47) resembles the standard contrastive learning formulation, combining this (more general) objective with existing contrastive learning techniques may lead another highly scalable unsupervised RL method, which we leave for future work.
the vhy, (2) z simplification with Equation (33) (i.e., Equation (7)).
F ADDITIONAL RESULTS
F.1 FULL QUALITATIVE RESULTS
Figure 1l shows the complete qualitative results of behaviors discovered by METRA state-based Ant and HalfCheetah, and pixel-based Quadruped and Humanoid (8 seeds for each environment). We use 2-D skills for Ant and Humanoid, 4-D skills for Quadruped, and 16 discrete skills for HalfCheetah. The full qualitative results suggest that METRA discovers diverse locomotion be- haviors regardless of the random seed.
21

---

Published as a conference paper at ICLR 2024
Ant (States)
HalfCheetah (States)
Quadruped (Pixels)
Humanoid (Pixels)
Figure 11: Full qualitative results METRA (8 seeds). METRA learns diverse locomotion behaviors regardless of the random seed.
S
R 29
2 R
S R04x64x3
2 = R
Ant (States)
Humanoid (Pixels)
2-y trajectories
((s) trajectories
trajectories
((3) trajectories
Latent space visualization. METRA learns capture r-y coordinates in doo-dimensionall latent spaces in both state-based Ant and pixel-based Humanoid, they the most temporally spread-out dimen- sions the state space. We note that, with a higher-dimensional latent space (especially when Z is discrete), METRA not only learns locomotion skills but also captures more diverse behaviors, as shown in the Cheetah and Kitchen videos on our project page.
Figure 12:
F.2 LATENT SPACE VISUALIZATION
METRA simultaneously learns both the skill policy (a]ss and the representation function o(s), to find the most "temporally spread-out" manifold in the state space. We train METRA on state- based Ant and pixel-based Humanoid with 2-D continuous latent spaces Z, and visualize the learned latent space by plotting o(s) trajectories in Figure 12. Since the xzy plane corresponds to the most temporally "important" manifold in both environments, METRA learns to capture the z-y coordi- nates in two-dimensional Φ, regardless of the input representations (note that Humanoid is pixel- based). We also note that, with a higher-dimensional latent space (especially when Z is discrete), METRA not only learns locomotion skills but also captures more diverse, non-linear behaviors, as shown in the Cheetah and Kitchen videos on our project page.
F.3 ABLATION STUDY OF LATENT SPACE SIZES
To demonstrate how the size of the latent space Z affects skill learning, we train METRA with 1- D. 2-D, and 4-D continuous skills and 2, 4, 8, 16, and 24 discrete skills on Ant and HalfCheetah. Figure 13 compares skills learned with different latent space sizes, which suggests that the diversity of skill generally increases as the capacity of 2 grows.
F.4 ADDITIONAL BASELINES
In the main paper, compare METRA with 11 previous unsupervised exploration and unsuper- vised skill discovery methods. In this section, we additionally compare METRA with DGPO (Chen et al., 2024), method that aims to find diverse behaviors that maximize task reward (Kumar et al., 2020; Zhou et al., 2022: Zahavy et al., 2023a;b; Chen et al., 2024). Since we consider a con- trolled Markov process without external rewards, use only the intrinsic reward part of DGPO for
22

---

Published as a conference paper at ICLR 2024
= R!
=
==
2= 2221
2= (1,3,34,44
2 1l..,88 2= Du.. [0] 2 [l.., .4)
Ant (States)
HalfCheetah (States)
Figure 13: Skills learned with different latent space sizes. Since METRA maximizes state coverage under the capacity the latent space Z, skills become more diverse as the capacity of Z grows.
Table 1: Comparison with DGPO. We compare METRA with an additional baseline for discrete skill learn- ing, DGPO (Chen et al., 2024). METRA exhibits the best state coverage in both Ant and HalfChectah. We use 4 random seeds and 't' denotes standard deviations.
Environment (Metric)
DIAYN
DGPO
METRA (ours)
HalfCheetah (policy state coverage)
6.75 2.22
6.75
186.75
HalfCheetah (total state coverage)
19.50 3.87
22.25
= 5.85
177.75 = 17.10
Ant (policy state coverage)
11.25 5.44
7.00
1387.75
Ant (total state coverage)
107.75 + 17.00
121.50 = 4.38
6313.25
747.92
comparison:
(2218141)
rt
DGPO
min
log (141) + ((
(50)
2'EZ,"'
where g is a skill discriminator (Eysenbach et al., 2019a) and DGPO assumes that Z is discrete space. Intuitively, this objective encourages each behavior to be maximally different from the most similar other behavior.
Table presents the comparison results on HalfCheetah and Ant, where we train DIAYN, DGPO, and METRA with 16 discrete skills for 10000 epochs (16M steps). Even though DGPO maximizes "worst-case" diversity (Equation (50)), still maximizes metric-agnostic KI. divergence between different skills (Chen et al., 2024), which leads to limited state coverage, in DIAYN. In contrast, METRA maximizes metric-aware Wasserstein distance and thus shows significantly better state coverage.
G EXPERIMENTAL DETAILS
We implement METRA on top of the publicly available LSD codebase (Park al., 2022). Our implementation available at /ttps://gittee implement LSD (Park et al., 2022), CIC (Laskin et al., 2022), DI-
cwm httm///ehhnngaa///eee. /METRA. For unsuper-
vised skill discovery methods, AYN (Eysenbuch et al., 2019a), and DADS (Sharma et al., 2020) on the same codebase as METRA. For six exploration methods, ICM (Pathak et al., 2017), LBS (Mazzaglia et al., 2022), RND (Burda et al., 2019), APT (Liu & Abbeel, 2021b), APS (Liu & Abbeel, 2021a), and Plan2Explore (Sekar use the original implementations by Laskin
al., 2020) (or Disagremeent (Pathak al., 2019)). et al. (2021) for state-based environments and the Dreamer (Hafner et al., 2020) variants by Rajeswar al. (2023) for pixel-based environments. For LEXA (Mendonca al., 2021) in Section 5.3, use the original implementation by Mendonca et al. (2021). We run our experiments on an internal cluster consisting of A5000 GPUs. Each run in Section 5.3 takes no more than 24 hours.
G.1 ENVIRONMENTS
Benchmark environments. For state-based environments, use the same MuJoCo HalfCheetah and Ant environments (Todorov et al., 2012; Brockman et al., 2016) as previous work (Sharma et al., 2020; Park et al., 2022; 2023b). HalfCheetah has 18 18-dimensional state space and Ant has a 29-dimensional state space. For pixel- environments, we use pixel- Quadruped and Humanoid from the DeepMind Control Suite (Tassa et al., 2018) and basel-based version of
23

---

Published as a conference paper at ICLR 2024
Kitchen by Gupta al. (2019); Mendonca et al. (2021). In DMC locomotion environments, use gradient-colored floors to allow the agent to infer its location from pixels, similarly to Hafner et al. (2022); Park et al. (2023a). In Kitchen, we use the same camera setting as LEXA (Mendonca et al., 2021). Pixel-based environments have an observation space of 64 x 64 x 3, and we do not use any proprioceptive state information. The episode length is 200 for Ant and HalfCheetah, 400 for Quadruped and Humanoid, and 50 for Kitchen. We use an action repeat of 2 for pixel-based Quadruped and Humanoid, following Mendonca et al. (2021). In our experiments, we do not use any prior knowledge or supervision, such as the z-y prior (Eysenbach et al., 2019a; Sharma et al., 2020), or early termination (Park et al., 2022).
Metrics. For the state coverage metric in locomotion environments, we count the number of 1 x 1- sized x-y bins (Ant, Quadruped, and Humanoid) or 1-sized x bins (HalfCheetah) that are occupied by any of the target trajectories. In Kitchen, we count the number of pre-defined tasks achieved by any of the target trajectories, where we use the same six pre-defined tasks as Mendonca et al. (2021): Kettle (K), Microwave (M), Light Switch (LS), Hinge Cabinet (HC), Slide Cabinet (SC), and Bottom Burner (BB). Each of the three types of coverage metrics, policy state coverage (Figures and 7), queue state coverage (Figure 7), and total state coverage (Figure 7), uses different target trajectories. Policy state coverage, which is mainly for skill discovery methods, is computed by 48 deterministic trajectories with 48 randomly sample skills at the current epoch. Queue state coverage is computed by the most recent 100000 training trajectories up to the current epoch. Total state coverage is computed by the entire training trajectories up to the current epoch.
Downstream tasks. For quantitative comparison of skill discovery methods (Figure 6). we use five downstream tasks, AntMultiGoals, HalfCheetahGoal, HalrCheetahHurdlle,, QuadrupedGoal, and HumanoidGoal, mostly following the prior work (Park et al., 2022). In HalfCheetahGoal, QuadrupedGoal, and HumanoidGoal, the task is to reach a target goal (within a radius of 3) ran- domly sampled from [-100, 100). 1-7.5, 7.5]², and [-5, 5]2, respectively. The agent receives a re- ward of 10 when it reaches the goal. In AntMultiGoalss, the task is reach four target goals (within a radius of 3), where each goal is randomly sampled from [s> 7.5, + 7.5] x [sy 7.5, Sy 7.5). where (sz, sy) is the agent's 's current z-y position. The agent receives a reward of 2.5 whenever reaches the goal. A new goal is sampled when the agent either reaches the previous goal or fails to reach it within 50 steps. In HalfCheetaHHurdle (Qureshi et al., 2020), the task is to jump over multiple hurdles. The agent receives a reward of 1 whenever it jumps over hurdle. The episode length is 200 for state- based environments and 400 for pixel-based environments.
For quantitative comparison with LEXA (Figure 8), we use five goal-conditioned tasks. In locomo- tion environments, goals are randomly sampled from [-100, 100) (HalfCheetah), [-50, 50)2 (Ant), (-15, 15]2 (Quadruped), or [010, 10 (Humanoid). We provide the full state as goal g. whose dimensionality is 18 for HalfCheetah, 29 for Ant, and 64 x 64 x 3 for pixel-based Quadruped and Humanoid. In Kitchen, we use the same six (single-task) goal images and tasks Mendonca et al. (2021). We measure the distance between the goal and the final state in locomotion environments and the number of successful tasks in Kitchen.
G.2 IMPLEMENTATION DETAILS
Unsupervised skill discovery methods. For skill discovery methods, we use 2-D continuous skills for Ant and Humanoid, 4-D continuous skills for Quadruped, 16 discrete skills for HalfCheetah, and 24 discrete skills for Kitchen, where continuous skills are sampled from the standard Gaussian distribution, and discrete skills are uniformly sampled from the set of zero-centered one-hot vec- tors (Park et al., 2022). METRA and LSD use normalized vectors (i.e., =/|[=1[2) for continuous skills, as their objectives are invariant to the magnitude of z. For CIC, we use 64-D continuous skills for all environments, following the original suggestion (Laskin et al., 2022), and we found that using 64-D skills for CIC leads to better state coverage than using 2-D or 4-D skills. We present the full list of hyperparameters used for skill discovery methods in Table 2.
Unsupervised exploration methods. For unsupervised exploration methods and LEXA, we use the original implementations and hyperparameters (Laskin et al., 2021; Mendonca et al., 2021; Rajeswar et al., 2023). For LEXA's goal-condin goal-conditioned policy (achiever), we test both the temporal distance and cosine distance variants and use the former as it leads to better performance.
24

---

Published as a conference paper at ICLR 2024
Table Hyperparameters for unsupervised skill discovery methods.
Hyperparameter
Value
Learning rate
0.0001 Adam (Kingma & Ba, 2015)
Optimizer
# episodes per epoch
200 (Quadruped, Humanoid), 100 (Kitchen), 50 (Ant, HalfCheetah)
# gradient steps per epoch
Minibatch size Discount factor 7
256
0.99
10 (Ant, HalfCheetah), 10 (Kitchen), 3 x 10 (Quadruped. Humanoid)
Replay buffer size
Encoder
CNN (LeCun et al., 1989)
# hidden layers
# hidden units per layer
1024
Target network smoothing coefficient
0.995 0.01 (Kitchen),
Entropy coefficient
auto-adjust (Haarnoja et al., 2018b) (others)
METRA E
10
METRA initial
30
Table 3: Hyperparameters for PPO high-level controllers.
Hyperparameter
Value
episodes per epoch
# gradient steps per episode
10
Minibatch size Entropy coefficient
256
0.01
GAE À (Schulman et al., 2016)
0.95
PPO clip threshold
0.2
High-level controllers for downstream tasks. Figure 6, we evaluate learned skills on down- stream tasks by training high-level controller n (2|s,, stask) that selects skill every K = 25 (Ant and HalfCheetah) or K = 30 (Quadruped and Humanoid) environment steps, where 8 task denotes the task-specific information: the goal position ("-Goal' or "-MultiGoals' tasks) the next hurdle position and distance (HalfChhetahHuurdle) At every steps, the high-level policy selects a skill 2, and then the pre-trained low-level skill policy T 2) executes the same for K steps. We train high-level controllers with PPO (Schulman et al., 2017) for discrete skills and SAC (Haarnoja al., 2018s 8a) for continuous skills. For SAC, we use the same hyperparameters as unsupervised skill discovery methods (Table 2), and we present the full list of PPO-specific hyperparameters in Table
Zero-shot goal-conditioned RL. In Figure 8, we evaluate the zero-shot performances of METRA, LSD. DIAYN, and LEXA on goal-conditioned downstream tasks. METRA and LSD use the proce- dure described in Section 4.2 to select skills. We re-compute every step for locomotion environ- use the same 2 selected at the first step throughout the episode, as we find
ments, but in Kitchen, that this leads to better performance. DIAYN chooses hased on the output of the skill discriminator the goal state (i.e., q(=lg), where q denotes the skill discriminator of DIAYN). LEXA uses the goal-conditioned policy (achiever), TT)ssg)).
25


Reviews:
[{"review_id": "dV4eTS0xav", "reviewer": "Reviewer_CufN", "summary": "This paper introduces a novel unsupervised RL objective called Metric-Aware Abstraction (METRA). The objective is to only learn to explore on a compact latent space which is metrically connected to the state space by temporal distances. The learned skills on this latent space are scalable to a variety of downstream control tasks. METRA is the first unsupervised RL method that demonstrates the discovery of diverse locomotion behaviors in pixel-based tasks.", "questions": "- 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? \n\n- 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?\n\n- 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?\n\n- In Figure 7, some methods are very unstable on the Kitchen tasks (variance of different seeds is large). Can the authors give any reason?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The empirical study of this paper is very sound and solid. The paper evaluates the method on various control tasks, including locomotion and manipulation tasks. Besides, the paper aims to address the unsupervised RL problem on visual-based tasks, which are much more challenging in the area. The paper also compare the results to multiple previous works, showing the significant improvement on skill discovery.\n\n- The methodology part is very organized. The authors aim to maximize state converage under a specific metrics, which should be scalable to pixel-based tasks. Then temporal distance makes sense and is easy to be turned to an constrainted optimization problem.\n\n- The paper is very well-written. The background in Section 2 and 3 clearly shows the motivation of this work and the connections to the previous methods. The method and empirical study both illustrate many details and the source code is linked, which make this work easy to understand and follow.", "weaknesses": "- The paper can be more impactful and solid if the method is deployed on the real world tasks, like locomotion control on a real robot. Besides, as the authors have already listed in Appendix A, the method can be combined to more recent RL works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fBk5jjlrHd", "reviewer": "Reviewer_1oDh", "summary": "The paper presents a novel unsupervised reinforcement learning (RL) method, Metric-Aware Abstraction (METRA), which aims to make unsupervised RL scalable to complex, high-dimensional environments. The authors propose a new unsupervised RL objective that encourages an agent to explore its environment and learn a breadth of potentially useful behaviors without any supervision. The key idea is to cover a compact latent space that is metrically connected to the state space by temporal distances, instead of directly covering the state space. The authors demonstrate that METRA can discover a variety of useful behaviors in complex environments, outperforming previous unsupervised RL methods.\n\nI have read the response and the authors address my concerns. I have raised my rating to accept.", "questions": "1. How does METRA handle environments with non-stationary dynamics or environments where the temporal distance between states can change over time?\n2. 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?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper introduces a novel unsupervised RL objective, METRA, which is a significant contribution to the field. The idea of using temporal distances as a metric for the latent space is innovative and provides a new perspective on unsupervised RL.\n2. The paper is technically sound, and the proposed method is well-motivated and clearly explained. The authors provide a thorough theoretical analysis of their method, including a connection to principal component analysis (PCA).\n3. The paper is well-written and organized. The authors do a good job of explaining the motivation behind their method, the details of the method itself, and the experimental setup.", "weaknesses": "1. 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.\n2. 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.\n\n- [1] MG Bellemare, Y Naddaf, J Veness, and M Bowling. “The arcade learning environment: An evaluation platform for general agents.” Journal of Artificial Intelligence Research (2012).\n- [2] Kurach, Karol, et al. \"Google research football: A novel reinforcement learning environment.\" Proceedings of the AAAI conference on artificial intelligence. Vol. 34. No. 04. 2020.\n- [3] Zhou, Zihan, et al. \"Continuously discovering novel strategies via reward-switching policy optimization.\" arXiv preprint arXiv:2204.02246 (2022).\n- [4] Chen, Wenze, et al. \"DGPO: Discovering Multiple Strategies with Diversity-Guided Policy Optimization.\" arXiv preprint arXiv:2207.05631 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cToNTv7bh0", "reviewer": "Reviewer_N85D", "summary": "This work proposes a novel objective for learning diverse skills in unsupervised skill discovery. In particular, this objective can enforce good policy coverage and is scalable to high-dimensional environments. The authors theoretically analyze the learning process of the Wasserstein dependency measure, and the analysis is accompanied by convincing experiments. The experiments and theoretical analysis show the effectiveness of the proposed approach.", "questions": "1. Does METRA have the potential to be applied to environments with temporal dependencies (NetHack, MineDojo, ...)?\n\n2. If we apply METRA to a traditional RL setting with external rewards, how would the discovered skills help in finding task-specific policies?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well-structured. The authors first analyze the common limitations of existing unsupervised RL approaches and then provide solid theoretical and empirical evidence to show why and how the proposed method works, making this paper understandable.\n\n2. Experiments are well-described and highly reproducible. Experiments have good coverage. The selection of baselines and environments is reasonable and convincing.", "weaknesses": "There are no significant weaknesses in this paper. The theoretical explanations of why choosing WDM as the objective might be a little complicated for readers lacking corresponding background. Some explicit examples or pictures may help.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AtmkRA0uSQ", "reviewer": "Reviewer_YkQo", "summary": "This work introduces an innovative algorithm designed to discover approximate state-covering behaviors in a task-agnostic manner for reinforcement learning agents. Building upon existing skill-learning techniques that use the mutual information between skills and states to learn distinguishable and state-covering skills, this work introduces a modification to the objective. The modification involves constraining the latent space of the skills to prioritize temporally compact representations. This novel approach ensures that the skills are estimated within a manifold preserving the temporal properties of the state space, allowing for maximization of spread-out trajectories over time. The incorporation of a metric, specifically temporal distance, grounds the latent space to be compact, distinguishing it from previous metric-agnostic alternatives relying solely on mutual information. The study employs a \"Wasserstein variant\" of the mutual information objective, adeptly modified to facilitate tractable optimization and online learning. One of the key advantages of pre-training RL agents using this objective lies in its scalability, enabling the learning of a compact skill space even from high-dimensional observations, such as images, where prior methods faced challenges in scalability. Moreover, when training a hierarchical controller to employ these learned skills, METRA demonstrates superior downstream task performance in previously unsolved benchmarks. In essence, the paper makes a significant contribution by introducing an effective new algorithm and strengthens its credibility by providing essential theoretical underpinnings for its design.", "questions": "The temporal distance is a very natural choice for a distance metric for reinforcement learning. However, are there any other metrics that are promising for their applications to the reinforcement learning setting in the METRA framework? Were these considered for this work?\n\nMETRA forces the latent variables to maintain linear relationships in the latent space. Although this allows for non-linear policies in the state space, can the latter limit the diversity of the learned skills in other challenging environments? (e.g. stochastic, multi-agent). Could METRA be modified to model non-linear skills even in the latent space?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This work builds upon an extensive body of research focused on skill discovery and learning, leveraging variants of mutual information to guide the pre-training process. This work tackles a limitation inherent in the existing MI objective, which is metric-agnostic, and hence does not directly incentivize skills to maximize an explicit metric for state coverage. The proposed modification facilitates the operation of skills within a temporally compact space, wherein the maximization of diversity aligns with the adoption of state-covering behaviors. The incorporation of a well-defined theoretical formulation enhances the paper's credibility, providing a strong foundation for the proposed objective.\n\nThe key contribution of the proposed algorithm is its scalability, particularly when dealing with high-dimensional observations - an ongoing challenge in the broader field of reinforcement learning. By addressing this limitation, the paper contributes significantly to the unsupervised RL domain. This work presents the development of an algorithm that stands as a substantial and impactful contribution to the field of unsupervised RL, and that is clearly theoretically justified.", "weaknesses": "One weakness of the paper is the absence of a comparative analysis of the proposed algorithm's performance in alternative benchmarks, especially those where purely exploratory algorithms like RND have demonstrated exceptional results (e.g., in Atari). Specifically, METRA's performance remains unclear in settings such as discrete control or, more significantly, in stochastic environments. An important aspect that is missing is a demonstration of how the temporally compact space learned by METRA enables reward-free pre-training in these challenging scenarios. Addressing these points would significantly enhance the paper's impact, making it a more substantial and comprehensive contribution to the field.\n\nThe paper iteratively simplifies the proposed objective to enable tractable optimization. However, it remains unclear if these modifications impact the performance of the algorithm. An ablation study of these could also provide insightful details of the proposed objective. (e.g. if training with the formulation that requires N rollouts for each latent, would the obtained skills be more diverse?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)