[Paper Review] Near-Optimal Representation Learning for Hierarchical Reinforcement Learning

Ofir Nachum, Shixiang Gu+2 more|arXiv (Cornell University)|Oct 2, 2018

Reinforcement Learning in RoboticsComputer Science36 references63 citations

TL;DR

This paper introduces a sub-optimality measure for representations in goal-conditioned hierarchical RL, derives bounds linking representation quality to policy performance, and proposes practical learning objectives that yield near-optimal hierarchical policies.

ABSTRACT

We study the problem of representation learning in goal-conditioned hierarchical reinforcement learning. In such hierarchical structures, a higher-level controller solves tasks by iteratively communicating goals which a lower-level policy is trained to reach. Accordingly, the choice of representation -- the mapping of observation space to goal space -- is crucial. To study this problem, we develop a notion of sub-optimality of a representation, defined in terms of expected reward of the optimal hierarchical policy using this representation. We derive expressions which bound the sub-optimality and show how these expressions can be translated to representation learning objectives which may be optimized in practice. Results on a number of difficult continuous-control tasks show that our approach to representation learning yields qualitatively better representations as well as quantitatively better hierarchical policies, compared to existing methods (see videos at https://sites.google.com/view/representation-hrl).

Motivation & Objective

• Motivate representation learning in goal-conditioned hierarchical RL to balance state compression with expressiveness for sub-tasks.
• Define a formal sub-optimality measure based on the reward of the optimal hierarchical policy using a given representation.
• Derive bounds showing how representation choices affect hierarchical policy performance and connect these to practical learning objectives.
• Propose a representation learning objective and auxiliary inverse model that bound sub-optimality and can be optimized in practice.

Proposed method

• Model a two-level hierarchical policy where the higher level chooses a goal in a learned representation space and a lower level acts to reach that goal.
• Introduce a mapping Psi that translates state-goal pairs into lower-level behaviors and analyze sub-optimality with respect to the gap to the fully expressive baseline.
• Derive theoretical bounds (Theorems 1 and 3, Claims 2 and 4) linking SubOpt(Psi) to divergences between true dynamics and representation-based dynamics.
• Define a representation learning objective based on KL/Divergence between true next-state distributions and a learned energy-based model K via a distance function D.
• Develop a practical training objective (Equation 11–15) that optimizes f (the representation) and an auxiliary model varphi to bound sub-optimality.
• Extend the framework to temporal abstraction (c>1) and generalized policies (Equation 9) with corresponding bounds.

Experimental results

Research questions

• RQ1 How does the choice of representation f affect the sub-optimality of a hierarchical policy when using a fixed lower-level mapping Psi?
• RQ2 Can we bound the loss in return due to representation-induced constraints and translate this bound into a practical representation-learning objective?
• RQ3 How does temporal abstraction (c > 1) influence sub-optimality bounds and learnable representations?
• RQ4 What is the relationship between mutual information-based objectives and the proposed representation learning bound in hierarchical RL?
• RQ5 Can learned representations achieve near-optimal performance on challenging continuous-control tasks compared to using full state or oracle representations?

Key findings

• A representation learning objective with a principled bound yields hierarchical policies whose returns approach the optimal policy within a bounded error.
• The framework connects representation learning directly to sub-optimality, enabling practical training objectives that reduce information loss in state-to-goal mappings.
• Empirical results on continuous-control tasks show qualitative and quantitative benefits over baselines, including scenarios with image observations.
• Learned representations tend to emphasize task-relevant coordinates (e.g., block positions) over others, aligning with reward structure.