PaperlayerView exact revision on arXiv

Source-linked AI summary

METRA: Scalable Unsupervised RL with Metric-Aware Abstraction

Seohong Park, Oleh Rybkin, Sergey Levine

arXiv:2310.08887v2cs.LGcs.AIcs.RO

TL;DR

Unsupervised RL seeks useful behaviors for downstream tasks, but the scalability of prior exploration and skill-discovery approaches in complex, high-dimensional environments remains unclear. METRA maximizes coverage of a compact latent space linked to states by temporal distance; experiments report diverse useful behaviors, outperforming prior methods across state-based and pixel-based settings, including pixel-based Quadruped and Humanoid locomotion.

  • Problem

    Prior unsupervised RL methods leave scalability to complex, high-dimensional environments unclear, while common mutual-information skill discovery can yield limited state coverage.

    Source: §1 · p. 1, §2 · p. 3

  • Method

    METRA learns diverse behaviors by maximizing coverage of a compact latent space connected to states through temporal distance.

    Source: §1 · p. 2, §4 · p. 4, §4.2 · pp. 5–6

  • Results

    Across five state-based and pixel-based continuous-control environments, METRA learned diverse useful behaviors and outperformed previous unsupervised RL methods, including locomotion behaviors in pixel-based Quadruped and Humanoid.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The learned compact latent space can be used to solve various downstream tasks in a zero-shot manner.

    Source: §1 · p. 2

    In practical implementation, METRA uses a relatively small update-to-data ratio, and the authors see room to improve sample efficiency.

    Source: §6 · p. 10

Abstract

from arXiv · show

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 $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-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/

1 INTRODUCTION

METRA aims to make unsupervised RL scalable by learning diverse behaviors through coverage of a compact latent space rather than exhaustive state-space coverage. Experiments report useful behaviors across state-based and pixel-based environments, including locomotion in pixel-based Quadruped and Humanoid.

1 INTRODUCTION

  • The paper seeks unsupervised pre-training that discovers potentially useful behaviors and could support more efficient learning across downstream tasks.Source: §1 · p. 1
  • Prior approaches leave scalability to complex, high-dimensional environments unclear: exhaustive exploration may be infeasible, while skill discovery may lack exploration incentives or pixel-based scalability.The stated concerns include limited state coverage for some skill-discovery methods and lack of direct scalability to pixel-based control for others.Source: §1 · pp. 1–2
  • METRA maximizes coverage of a compact latent space connected to states by temporal distances, aiming for approximate state-space coverage without representing every pose.The Figure 1 illustration shows behaviors running in all directions in a two-dimensional latent space, without covering every possible leg pose.Source: §1 · p. 2, Figure 1 · p. 2
  • Across five state-based and pixel-based continuous-control environments, METRA learned diverse useful behaviors and a compact latent space for zero-shot downstream tasks, outperforming previous unsupervised RL methods.The paper also reports diverse locomotion behaviors in pixel-based Quadruped and Humanoid.Source: §1 · p. 2

2 WHY MIGHT PREVIOUS UNSUPERVISED RL METHODS FAIL TO SCALE?

The section contrasts two common unsupervised RL strategies and explains why their objectives may not scale to complex environments. METRA instead emphasizes behavior differences measured by temporal distance.

2 WHY MIGHT PREVIOUS UNSUPERVISED RL METHODS FAIL TO SCALE?

  • The Figure 2 comparison contrasts exhaustive state coverage and MI-based discriminability with Wasserstein dependency maximization of a distance metric.The proposed comparison uses temporal distance to encourage skills spanning longer state-space subspaces.Source: Figure 2 · p. 3
  • Pure exploration aims to cover every state or capture full dynamics, goals that may be computationally infeasible in large state spaces.The paper reports that these approaches even fail to cover the state space in the state-based 29-dimensional Ant environment.Source: §2 · p. 3
  • Mutual-information skill discovery can produce distinguishable behaviors yet limited state coverage, because MI is metric-agnostic and does not prioritize how different behaviors are.The section says prior methods may also fail to scale when they require capturing the entire MDP or assume compact, structured states.Source: §2 · p. 3

3 PRELIMINARIES AND PROBLEM SETTING

The problem setting defines unsupervised skill learning in a reward-free controlled Markov process. The goal is to learn useful behaviors without supervision, data, or prior knowledge.

3 PRELIMINARIES AND PROBLEM SETTING

  • The setting is a reward-free controlled Markov process with state space S, action space A, initial-state distribution μ, and transition kernel p.Source: §3 · p. 3
  • A skill is a latent vector z paired with a latent-conditioned policy, sampled from a prior and held fixed throughout an episode.Source: §3 · pp. 3–4
  • The objective is to learn diverse, useful behaviors without supervision, data, or prior knowledge.Source: §3 · p. 4

4 A SCALABLE OBJECTIVE FOR UNSUPERVISED RL

METRA seeks scalable unsupervised RL by maximizing coverage in a compact latent space linked to states through a metric. It instantiates this idea with Wasserstein dependence and temporal distance, yielding a tractable objective for learning diverse behaviors.

4 A SCALABLE OBJECTIVE FOR UNSUPERVISED RL

  • The objective targets broad state coverage through a compact latent space and a latent-conditioned policy, rather than covering every state directly.Wasserstein dependence is metric-aware: it encourages distinct skill trajectories to be far apart under the chosen metric.Source: §4 · p. 4

4.1 TRACTABLE OPTIMIZATION

  • A tractable optimization uses a factorized score, then a last-state objective with ψ(z) = z to yield reward from latent-state displacement projected onto z.The resulting reward is r(s, z, s′) = (ϕ(s′) − ϕ(s))⊤z, avoiding extra samples of z for each data point.Source: §4.1 · pp. 4–5

4.2 FULL OBJECTIVE: METRIC-AWARE ABSTRACTION (METRA)

  • METRA uses temporal distance—the minimum environment steps from one state to another—as a representation-invariant metric suited to pixel observations.Under this metric, the constraint is imposed on adjacent state pairs: ∥ϕ(s) − ϕ(s′)∥2 ≤ 1.Source: §4.2 · p. 5
  • The learned abstraction is interpreted as nonlinear principal components that preserve temporal distances; under simplifying assumptions, linear squared METRA is equivalent to PCA under that metric.The latent dimensions emphasize state-space manifolds with long shortest paths.Source: §4.2 · pp. 5–6
  • WDM connects to LSD and CSD objectives using Euclidean distance and can be simplified into Wasserstein variants of DIAYN, DADS, and CIC.The connection also explains the zero-mean prior used in discrete LSD.Source: §4.2 · p. 6
  • For zero-shot goal reaching, the difference ϕ(g) − ϕ(s) specifies a skill direction, avoiding a separate goal-conditioned policy.For continuous skills, METRA normalizes the difference; for discrete skills, it selects the dimension with the largest difference.Source: §4.2 · p. 6
  • The symmetric latent distance may be overly restrictive in highly asymmetric environments, and setting ψ(z) = z may restrict behavior diversity to some degree.The authors suggest an asymmetric quasimetric for the former issue and the full WDM objective for the latter.Source: §4.2 · pp. 6–7
  • Figure 3 reports METRA as the only of 11 methods to discover diverse locomotion skills in pixel-based Quadruped and Humanoid.The figure shows learned-policy trajectories for locomotion and coincidental success rates for six predefined Kitchen tasks.Source: Figure 3 · p. 7

5 EXPERIMENTS

Across five state-based and pixel-based locomotion and manipulation environments, the experiments assess METRA’s learned behaviors, coverage, and downstream usefulness. METRA achieves strong coverage and downstream performance, including diverse pixel-based locomotion skills.

5.1 EXPERIMENTAL SETUP

  • The evaluation uses five robotic locomotion and manipulation environments, including state-based Ant and HalfCheetah and pixel-based Quadruped, Humanoid, and Kitchen.The experiments report comparisons against 11 prior methods across skill discovery, exploration, and goal-reaching groups.Source: §5.1 · p. 7, §5 · p. 7

5.2 QUALITATIVE COMPARISON

  • METRA discovers diverse behaviors across state-based and pixel-based domains, and is the only compared method to learn locomotion skills in pixel-based Quadruped and Humanoid.Pure exploration methods mostly exhibit chaotic, random behaviors; LSD discovers state-based locomotion but does not scale to pixel-based settings.Source: §5.2 · p. 8, Figure 7 · p. 8

5.3 QUANTITATIVE COMPARISON

  • METRA achieves the best coverage in most domains and the best downstream performance on most tasks, while uniquely learning meaningful skills in pixel-based settings.Coverage measures x, x-y, or task coverage; downstream tasks use a high-level controller over frozen skills.Source: Figure 5 · p. 8, Figure 6 · p. 8, §5.3 · p. 9
  • METRA achieves the best coverage in most environments, while pure exploration methods fail to fully explore state-based Ant and pixel-based Humanoid.Kitchen is an exception where pure exploration methods also work decently.Source: §5.3 · p. 9, Figure 8 · p. 9
  • METRA achieves the best performance on all five goal-reaching tasks, and its skills solve them in a zero-shot manner.LEXA also performs non-trivially on three tasks but struggles with state-based Ant and pixel-based Humanoid.Source: Figure 9 · p. 9, §5.3 · p. 10

6 CONCLUSION

METRA covers a compact latent skill space connected to states by temporal distance and learns diverse useful behaviors. The conclusion also identifies practical sample-efficiency and evaluation-scope limitations.

6 CONCLUSION

  • METRA covers a compact latent skill space connected to states by temporal distance and learns diverse useful behaviors in locomotion and manipulation, including pixel-based Quadruped and Humanoid.The authors describe the work as a step toward broadly applicable unsupervised RL and large-scale pre-training with minimal supervision.Source: §6 · p. 10
  • Vanilla SAC and relatively small update-to-data ratios leave room to improve sample efficiency.The paper reports UTD ratios of 1/4 for Kitchen and 1/16 for Quadruped and Humanoid.Source: §6 · p. 10
  • Evaluation covers locomotion and manipulation but not other environment types such as Atari, and the fixed-MDP assumption leaves non-stationary or non-Markovian dynamics unaddressed.The authors leave evaluation on more diverse environments and extensions to those dynamics for future work.Source: §6 · p. 10

A EXTENDED RELATED WORK

  • Related work notes a suspected connection to graph-Laplacian option discovery and shared temporal-distance constraints with goal-conditioned RL.The authors say they have not found a formal connection to graph-Laplacian methods.Source: A EXTENDED RELATED WORK · pp. 17–18

B.1 UNIVERSALITY OF INNER PRODUCT DECOMPOSITION

  • For continuous f on compact Hausdorff spaces, finite sums of products ϕ_i(x)ψ_i(y) can approximate f arbitrarily closely.With dense function classes such as universal neural-network approximators, the inner product ϕ(x)⊤ψ(y) can approximate f as D grows.Source: §B.1 · p. 18

B.2 LIPSCHITZ CONSTRAINT UNDER THE TEMPORAL DISTANCE METRIC

The section proves an equivalence between a bound on adjacent-state embedding distances and a bound using temporal distance.

B.2 LIPSCHITZ CONSTRAINT UNDER THE TEMPORAL DISTANCE METRIC

  • Conversely, the temporal-distance bound implies the adjacent-state bound because distinct adjacent states have temporal distance 1.For identical states, the embedding distance is zero.Source: §B.2 · p. 19
  • Bounding embedding distances for adjacent states by 1 implies the temporal-distance bound for any reachable state pair, by the triangle inequality along a shortest path.Unreachable pairs satisfy the temporal-distance condition trivially.Source: §B.2 · p. 19

C A CONNECTION BETWEEN METRA AND PCA

The appendix connects squared linear METRA to PCA under specified temporal-embedding conditions, then relates WDM objectives to DIAYN, DADS, and CIC.

C A CONNECTION BETWEEN METRA AND PCA

  • A temporally consistent embedding represents temporal distance in a potentially higher-dimensional Euclidean space, and PointMass is given as an example.The PCA derivation also assumes Gaussian latent skills, deterministic initial state and transitions, and reachability within T steps.Source: C A CONNECTION BETWEEN METRA AND PCA · p. 19
  • With a temporally consistent embedding and ellipsoidal embedding space, linear squared METRA selects the top-d eigenvectors of A, yielding PCA in the temporal embedding space.The analysis assumes a squared objective for mathematical convenience; the authors caution that practical METRA may be nonlinear and need not have a temporally consistent embedding.Source: C A CONNECTION BETWEEN METRA AND PCA · pp. 19–20

D CONNECTIONS BETWEEN WDM AND DIAYN, DADS, AND CIC

D.1 DIAYN

  • The DIAYN analogy treats the reward ϕL(st)⊤z as a Lipschitz, inner-product variant of DIAYN.The comparison uses a Gaussian variational distribution q(z|s).Source: §D.1 · p. 21

D.2 DADS

  • The sample-based objective is presented as a Lipschitz, inner-product variant of DADS.The comparison uses the variational model q(s′−s|s,z)=N(ψ(s,z),I).Source: §D.2 · p. 21

D.3 CIC

  • The CIC connection yields a Lipschitz variant of CIC, while the paper uses ψL(z)=z as a simpler working objective.The authors leave combining the more general contrastive-like objective with existing contrastive-learning techniques for future work.Source: §D.3 · p. 22

E.1 FULL QUALITATIVE RESULTS

Qualitative results show seed-robust diverse locomotion and latent spaces that reflect temporally important state dimensions. Greater latent capacity is associated with more diverse skills.

E.1 FULL QUALITATIVE RESULTS

  • As latent-space capacity grows, METRA's skills become more diverse under its objective of maximizing state coverage within that capacity.Figure 14 compares skills learned with different latent-space sizes.Source: Figure 14 · p. 23
  • Across eight seeds, METRA discovers diverse locomotion behaviors in state-based Ant and HalfCheetah and pixel-based Quadruped and Humanoid.The experiments use 2-D skills for Ant and Humanoid, 4-D skills for Quadruped, and 16 discrete skills for HalfCheetah.Source: §E.1 · pp. 22–23

E.2 LATENT SPACE VISUALIZATION

  • In state-based Ant and pixel-based Humanoid, the learned 2-D latent space captures x-y coordinates, the most temporally spread-out dimensions in those state spaces.With higher-dimensional latent spaces, especially discrete ones, METRA also captures more diverse, nonlinear behaviors.Source: §E.2 · p. 23, Figure 13 · p. 23

E.3 ABLATION STUDY OF LATENT SPACE SIZES

The study varies the size and type of METRA’s skill space on Ant and HalfCheetah; skill diversity generally increases with latent-space capacity.

E.3 ABLATION STUDY OF LATENT SPACE SIZES

  • Skill diversity generally increases as the latent-space capacity grows.Figure 14 compares learned skills across the tested latent-space sizes.Source: §E.3 · p. 23
  • The experiments compare 1-D, 2-D, and 4-D continuous skills on Ant and HalfCheetah.The study also tests 2, 4, 8, 16, and 24 discrete skills.Source: §E.3 · p. 23
  • The skill-capacity ablation includes both continuous and discrete skill configurations.Figure 14 presents skills learned at the tested latent-space sizes.Source: §E.3 · p. 23

F EXPERIMENTAL DETAILS

The experimental details describe the implementations and benchmark settings, including environments, coverage measures, and downstream tasks used for evaluation.

F EXPERIMENTAL DETAILS

  • METRA, LSD, CIC, DIAYN, and DADS are implemented on the same codebase, while exploration baselines use original implementations or pixel-based Dreamer variants.LEXA uses its original implementation; the reported runs use an internal cluster of A5000 GPUs.Source: F EXPERIMENTAL DETAILS · pp. 23–24

F.1 ENVIRONMENTS

  • Benchmarks include state-based Ant and HalfCheetah, pixel-based Quadruped and Humanoid, and pixel-based Kitchen.Pixel observations are 64 × 64 × 3 with no proprioceptive state input; the experiments use no prior knowledge or supervision.Source: §F.1 · p. 24
  • Coverage counts occupied locomotion position bins or achieved Kitchen tasks, using distinct trajectory sets for policy, queue, and total coverage.Policy coverage uses 48 deterministic trajectories; queue coverage uses the most recent 100000 training trajectories, while total coverage uses all training trajectories.Source: §F.1 · p. 24
  • Skill-discovery methods are compared on five downstream tasks: AntMultiGoals, HalfCheetahGoal, HalfCheetahHurdle, QuadrupedGoal, and HumanoidGoal.The tasks evaluate reaching goals or jumping hurdles, with task rewards and goal-sampling rules specified in the passage.Source: §F.1 · p. 24
  • The LEXA comparison uses five goal-conditioned tasks, evaluating final-state distance in locomotion and successful tasks in Kitchen.Locomotion goals are sampled from environment-specific ranges; Kitchen uses six single-task goal images and tasks.Source: §F.1 · p. 24

F.2 IMPLEMENTATION DETAILS

Implementation settings cover skill choices, baseline configurations, downstream controllers, and zero-shot goal-conditioned evaluation procedures.

F.2 IMPLEMENTATION DETAILS

  • Skill-discovery comparisons use environment-specific skill counts and dimensions, with METRA and LSD continuous skills normalized to unit length.CIC uses 64-D continuous skills, which the authors found gave better state coverage than 2-D or 4-D skills.Source: §F.2 · pp. 24–25
  • Unsupervised exploration methods and LEXA use their original implementations and hyperparameters.For LEXA’s achiever, the temporal-distance variant is used because it performs better than the cosine-distance variant.Source: §F.2 · p. 25
  • Downstream high-level controllers select a skill every 25 steps in Ant and HalfCheetah or every 50 steps in Quadruped and Humanoid.The high-level policy uses PPO for discrete skills and SAC for continuous skills; skill-discovery and PPO hyperparameters are listed in Tables 1 and 2.Source: §F.2 · p. 25, Table 1 · p. 25, Table 2 · p. 25
  • For zero-shot goal-conditioned evaluation, METRA and LSD select skills using the described procedure, while DIAYN selects from its discriminator and LEXA uses its achiever policy.Skills are recomputed each step in locomotion but held fixed through Kitchen episodes.Source: §F.2 · pp. 25–26

Related papers

Loading 2310.08887v2…