Source-linked AI summary

Near-Optimal Representation Learning for Hierarchical Reinforcement Learning

Ofir Nachum, Shixiang Gu, Honglak Lee, Sergey Levine

arXiv:1810.01257v2cs.AI

TL;DR

The paper studies how to learn compact goal representations in hierarchical reinforcement learning without making sub-tasks inexpressible or goals impractically high-dimensional. It defines representation sub-optimality, derives bounds linking that quantity to representation-learning objectives, and reports better representations and hierarchical policies on difficult continuous-control tasks.

  • Problem

    Choosing a goal representation is difficult because end-to-end representations can be under-defined, while full-state goals are hard to scale.

  • Method

    The paper defines representation sub-optimality through the expected-return gap between optimal hierarchical policies and derives practical representation-learning objectives that bound this gap.

  • Results

    The proposed representation-learning approach yields qualitatively better representations and quantitatively better hierarchical policies across difficult continuous-control tasks, with only the proposed method approaching the XY oracle in one challenging experiment.

  • Takeaways & Limitations

    A principled representation objective can produce compact representations whose hierarchical-policy return approaches the optimal policy within bounded error.

  • Takeaways & Limitations

    The framework restricts the environment reward function to be defined only on states.

Abstract

from arXiv · show

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

1 INTRODUCTION

Goal-conditioned hierarchical reinforcement learning uses representations to communicate goals between higher- and lower-level policies, but choosing that representation involves a trade-off between expressiveness and scalability.

  • Goal-conditioned hierarchical designs let higher-level policies communicate goals while lower-level policies learn to reach corresponding states.
  • Existing approaches either learn representations end-to-end or use the state space directly as the goal space.
  • End-to-end representations can be under-defined because not all possible sub-tasks are expressible as goals.
  • Using the full state avoids information loss but becomes difficult to scale, especially when observations are images and goals must be target images.
  • The paper proposes measuring representation loss by the ability to support near-optimal policies rather than by reconstruction quality.
  • A theoretical objective yields representations whose hierarchical-policy return approaches the optimal policy within bounded error, and the resulting practical method performs well on difficult continuous-control tasks.

2 FRAMEWORK

The framework uses a two-level policy in which a higher-level controller selects goals and a lower-level policy executes behaviors conditioned on those goals. A representation maps states to the goal space and determines how goal signals shape lower-level behavior.

  • The framework models a two-level hierarchical policy on an MDP, with the higher-level policy choosing goals for the lower-level policy.
  • A state representation f : S → R^d defines the goal space G = R^d used by the higher-level controller.
  • The higher-level policy samples a goal every c steps, while the lower-level policy converts that goal into actions over the following c-step interval.
  • The lower-level policy is trained with a goal-conditioned reward based on the distance between the next state representation and the desired goal.
  • The mapping Ψ identifies the low-level behavior that a communicated goal induces, rather than merely representing a goal to be reached.
  • The framework restricts environment rewards to state-only functions and defines Rmax as the maximal absolute reward.

3 HIERARCHICAL POLICY SUB-OPTIMALITY

The paper defines representation sub-optimality by comparing the best hierarchical policy constrained by the goal-to-behavior mapping with an optimal policy unrestricted by that mapping. It then seeks bounds connecting this loss to the representation itself.

  • The central question is how much performance is lost when the higher-level policy can act only through Ψ.
  • Because Ψ restricts the types and number of lower-level behaviors, the optimal policy may not be expressible by the constrained higher-level policy.
  • The paper compares the optimal hierarchical policy using Ψ with an optimal hierarchical policy that is agnostic to Ψ.
  • The unrestricted higher-level policy can select any low-level behavior, allowing it to express all possible low-level behaviors.
  • Sub-optimality is defined through the difference between the state values of the optimal unrestricted and Ψ-constrained hierarchical policies.
  • Although Ψ depends on the representation f, the relationship is indirect, so the paper derives bounds to connect sub-optimality with a practical learning objective for f.

4 GOOD REPRESENTATIONS LEAD TO BOUNDED SUB-OPTIMALITY

The paper bounds the sub-optimality caused by a goal representation by relating it to recoverability of lower-level behavior and representation-learning objectives. The results extend from single-step deterministic policies to temporally abstracted settings and general policy classes.

  • Main theoretical contribution: Theoretical bounds connect representation-induced sub-optimality to goal-conditioned policy objectives and representation learning objectives.Claim 4 is identified as the main theoretical contribution and applies to any class of low-level policies, including closed-loop and open-loop policies.
  • Single-step deterministic policies: SubOpt(Ψ) ≤ Cϵ, with C = 2γ (1−γ)^2 Rmax, in the single-step deterministic setting.The bound is stated for the restricted case c = 1, where the higher-level policy chooses a goal at every step and Ψ translates it into a low-level action.
  • Recoverability and sub-optimality: An invertible mapping Ψ yields zero sub-optimality because every action can be recovered through a goal.More generally, the theory requires recoverability of action effects rather than full invertibility.
  • Representation-learning objective: The representation objective matches true transition dynamics to dynamics implied by f and ϕ, rather than optimizing reconstruction loss.The distribution K is interpreted as a dynamics model determined by the representation and auxiliary inverse goal model.
  • Temporal abstraction and general policies: For temporal abstraction, the higher-level goal is translated into a c-step lower-level behavior, with divergences weighted across future steps.The general formulation uses a discounted weighted sum in place of single-step divergences and rewards.
  • Temporal abstraction and general policies: SubOpt(Ψ) ≤ Cϵ, with C = 2γ/(1−γ^c) Rmaxw, in the temporally abstracted setting.The bound covers higher-level actions selected every c steps and translated into lower-level policies.

5 LEARNING

The paper turns its theoretical bound into a practical representation-learning objective using replay-buffer trajectories and auxiliary neural networks. The same representation supports a low-level policy objective and connects to mutual-information estimation.

  • 5.1 LEARNING GOOD REPRESENTATIONS: Representations fθ and auxiliary functions ϕθ are implemented as separate neural networks parameterized jointly by θ.The parameterization uses θ = [θ1, θ2].
  • 5.1 LEARNING GOOD REPRESENTATIONS: Replay-buffer states and subsequent c actions approximate the theoretical supremum over state-policy pairs.This restricts candidate policies to c-step deterministic open-loop action sequences and simplifies ϕθ to a standard neural network.
  • 5.1 LEARNING GOOD REPRESENTATIONS: The inner objective compares expected future-state energy under experienced dynamics with energy under replay-buffer states sampled from ρ.The objective uses Eθ(s′, s, π) = exp(−D(fθ(s′), ϕθ(s, π))) and the induced distribution Kθ.
  • 5.1 LEARNING GOOD REPRESENTATIONS: Gradients are estimated from experienced future states and mini-batches of independently sampled replay-buffer states.The replay buffer supplies both the numerator estimate and the state samples used to approximate the denominator.
  • 5.1 LEARNING GOOD REPRESENTATIONS: The objective resembles CPC-style mutual-information estimation, but conditions on actions or policy rather than only on state pairs.The paper interprets the objective as maximizing MI(st+k; st, π) through an energy function.
  • 5.2 LEARNING A LOWER-LEVEL POLICY: The low-level policy is trained with a reward derived from the representation objective, using sampled action sequences to represent π.This differs from standard low-level rewards, which use only the first term and are unweighted.

6 RELATED WORK

The paper situates its approach within MDP abstraction, state-abstraction bounds, and mutual-information representation learning. It distinguishes its contribution by connecting mutual-information estimators to hierarchical-RL representation learning with sub-optimality guarantees.

  • MDP abstraction: Goal-conditioned hierarchical representations can be viewed as action abstractions because goals identify behaviors for the lower-level policy.The hierarchy also relates to abstractions of state, action, and temporal spaces.
  • State-abstraction bounds: Exact state abstractions incur no performance loss, whereas approximate abstractions generally provide bounded sub-optimality.This places the paper’s bounded-error analysis within a broader abstraction literature.
  • Mutual-information objectives: The proposed objective resembles CPC and MINE because it uses mutual-information-style energy modeling.The paper relates its bounds to model-based predictive-error bounds in representation learning for MDPs.
  • Mutual-information objectives: The paper claims to be the first to connect these mutual-information estimators to hierarchical-RL representation learning with theoretical sub-optimality guarantees.Earlier information-theoretic measures had been used for exploration rather than this representation-learning framework.

7 EXPERIMENTS

Experiments evaluate the learned representation against oracle, reconstruction-based, dynamics-based, end-to-end, and raw-observation baselines on continuous-control tasks. The method recovers near-optimal coordinates, approaches oracle performance during online hierarchical learning, and emphasizes task-relevant block coordinates.

  • Experimental setup: The evaluation compares the proposed objective with XY, VAE, E2C, E2E, and Whole obs baselines.XY uses the agent’s x, y position; Whole obs uses the raw observation without representation learning.
  • Qualitative representations: In Figure 2, the method recovers a near-ideal 2D x, y representation without supervision, including from top-down images, whereas other approaches do not properly recover one.The embeddings are colored by episode time-step from black at the beginning to yellow at the end.
  • Quantitative performance: In Figure 3, only the proposed method approaches the XY oracle outside simple point environments during concurrent online representation and hierarchical-policy learning.Results use 10M training steps, medians over 10 trials, and 30th and 70th percentiles.
  • Experimental setup: Experiments cover MuJoCo navigation and pushing tasks, including Ant Maze, Ant Push, Ant Fall, Ant Block, and Ant Block Maze.The tasks include both point and ant environments and a ⊃-shaped corridor.
  • Experimental setup: Raw observations include agent coordinates, orientation, limbs, and—when applicable—block coordinates; image variants replace agent coordinates with a 5 × 5 × 3 top-down image.The image experiments test representation learning from more difficult raw inputs.
  • Representation sensitivity: On the block-moving task, learned representations emphasize block x, y coordinates more than agent coordinates.Sensitivity is assessed by perturbing ant versus block coordinates, and block position primarily determines external reward.

8 CONCLUSION

The paper concludes that its principled sub-optimality analysis yields a practical representation-learning objective for hierarchical RL. The resulting method is evaluated on high-dimensional continuous-control tasks and is supported by a bounded-error theoretical framework.

  • 8 CONCLUSION: The approach measures representation sub-optimality through the loss in expected return relative to optimal state values.The direct sub-optimality measure is intractable, so the paper derives a mathematical relationship to a practical learning objective.
  • 8 CONCLUSION: The paper’s practical objective is intended to learn representations whose hierarchical-policy return approaches optimal-policy return within bounded error.The conclusion describes the resulting objective as practical and effective across difficult continuous-control tasks.
  • 8 CONCLUSION: The theory bounds the hierarchical policy’s value difference by comparing state visitation frequencies and policy values.The proof constructs a corresponding hierarchical policy and bounds total-variation divergence before bounding value differences.
  • 8 CONCLUSION: The analysis relates representation-induced transition modeling to divergences between the true dynamics and the dynamics implied by f and ϕ.The derivation uses the energy-based distribution K and inequalities including Jensen’s and Pinsker’s.

C.1 ENVIRONMENTS

The paper evaluates navigation tasks involving agents or movable blocks, including image-based variants, and uses consistent hierarchical-RL training settings across environments.

  • Environments: Training targets are sampled randomly from possible environment locations, while evaluation uses a single difficult target point.For Ant Fall, the target additionally includes a z coordinate.
  • Environments: Image variants remove the agent’s x, y coordinates and append a low-resolution 5 × 5 × 3 top-down environmental view.The channels encode immovable blocks, movable blocks, and chasms.
  • Environments: The considered tasks are navigation problems in which an agent or movable block must reach a target location.Tasks include Ant Maze, Ant Push, Ant Fall, Point Maze, Ant Block, and Ant Block Maze environments.
  • Training details: The implementation inputs the whole observation to the lower-level policy and uses a Huber distance function for low-level rewards.These choices differ from the referenced prior setup, which zeroed out x, y coordinates for the lower-level policy.
  • Training details: The higher-level policy uses a two-dimensional goal with actions in [−10, 10]2, Gaussian exploration with standard deviation 5, and neural representations trained by Adam.The representation network has two 100-unit hidden layers, while ϕθ uses 400- and 300-unit hidden layers; the learning rate is 0.0001.

D OBJECTIVE FUNCTION EVALUATION

The evaluation compares the proposed distance-based objective with CPC-style and β-VAE variants, examining whether objective design affects learned representations and performance.

  • Objective comparisons: The proposed objective is compared with variants implemented in a style more closely resembling CPC mutual-information objectives.The comparison concerns alternative formulations of the representation-learning objective.
  • Objective comparisons: Using a dot product instead of the distance function D is detrimental, while several distance-based variants may perform similarly.The caption distinguishes the dot-product substitution from alternative distance-based formulations.
  • β-VAE comparisons: β-VAE variants do not perform well outside the simple point-mass environment, even with the additional hyperparameter.The stated drawback is reconstructing the entire observation, including unimportant or high-variance features such as ant joint velocities.

G ADDITIONAL QUALITATIVE RESULTS

Additional results test robustness to weaker training data and compare the proposed method with original HIRO, while the accompanying algorithm specifies replay-buffer training and objective terms.

  • Additional results: With data collected by a random higher-level policy, the method recovers near-ideal representations despite weaker alignment between representation and task objectives.This replicates the Figure 2 results under a setting with less task-directed data.
  • Additional results: The original HIRO representation behaves similarly to the XY oracle in non-image tasks but performs much worse when observations are images.The original formulation uses position-based, rather than velocity-based, observation components as sub-goals.
  • Training algorithm: Algorithm 1 trains representation and policy parameterizations from replay-buffer experience using batches of c-step transitions.The listed inputs include replay buffer D, training steps N, batch size B, and parameterizations fθ, ϕθ, and πφ.
  • Training algorithm: The objective combines an attractive distance term with a repulsive exponential term and an additional L term.The listed implementation also computes low-level rewards and includes higher- and lower-level reinforcement-learning updates.
Loading 1810.01257v2…