Source-linked AI summary
Near-Optimal Representation Learning for Hierarchical Reinforcement Learning
Ofir Nachum, Shixiang Gu, Honglak Lee, Sergey Levine
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 · showhide
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.