Near-Optimal Representation Learning for Hierarchical Reinforcement Learning

• Goal-conditioned hierarchical designs let higher-level policies communicate goals while lower-level policies learn to reach corresponding states.Source: §1 · p. 1
• Existing approaches either learn representations end-to-end or use the state space directly as the goal space.Source: §1 · p. 1
• End-to-end representations can be under-defined because not all possible sub-tasks are expressible as goals.Source: §1 · p. 1
• Using the full state avoids information loss but becomes difficult to scale, especially when observations are images and goals must be target images.Source: §1 · p. 1
• The paper proposes measuring representation loss by the ability to support near-optimal policies rather than by reconstruction quality.Source: §1 · pp. 1–2
• 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.Source: §1 · p. 2

• The framework models a two-level hierarchical policy on an MDP, with the higher-level policy choosing goals for the lower-level policy.Source: §2 · p. 2
• A state representation f : S → R^d defines the goal space G = R^d used by the higher-level controller.Source: §2 · p. 2
• 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.Source: §2 · p. 2
• The lower-level policy is trained with a goal-conditioned reward based on the distance between the next state representation and the desired goal.Source: §2 · p. 2
• The mapping Ψ identifies the low-level behavior that a communicated goal induces, rather than merely representing a goal to be reached.Source: §2 · p. 2
• The framework restricts environment rewards to state-only functions and defines Rmax as the maximal absolute reward.Source: §2 · p. 3

• The central question is how much performance is lost when the higher-level policy can act only through Ψ.Source: §3 · p. 3
• Because Ψ restricts the types and number of lower-level behaviors, the optimal policy may not be expressible by the constrained higher-level policy.Source: §3 · p. 3
• The paper compares the optimal hierarchical policy using Ψ with an optimal hierarchical policy that is agnostic to Ψ.Source: §3 · p. 3
• The unrestricted higher-level policy can select any low-level behavior, allowing it to express all possible low-level behaviors.Source: §3 · p. 3
• Sub-optimality is defined through the difference between the state values of the optimal unrestricted and Ψ-constrained hierarchical policies.Source: §3 · p. 3
• 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.Source: §3 · p. 3

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.Source: §4 · p. 3, §4.2 · p. 5

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.Source: §4.1 · p. 3

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.Source: §4.1 · p. 4

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.Source: §4.1 · p. 4, §4.2 · p. 4

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.Source: §4.2 · p. 4
• 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.Source: §4.2 · p. 4

5.1 LEARNING GOOD REPRESENTATIONS

• Representations fθ and auxiliary functions ϕθ are implemented as separate neural networks parameterized jointly by θ.The parameterization uses θ = [θ1, θ2].Source: §5.1 · p. 5
• 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.Source: §5.1 · p. 5
• 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θ.Source: §5.1 · p. 5
• 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.Source: §5.1 · pp. 5–6
• 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.Source: §5.1 · p. 6

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.Source: §5.2 · p. 6

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.Source: §6 · p. 6

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.Source: §6 · pp. 6–7

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.Source: §6 · p. 7
• 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.Source: §6 · p. 7

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.Source: §7 · p. 7

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.Source: Figure 2 · p. 7, §7 · p. 9

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.Source: Figure 3 · p. 8, §7 · p. 9

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.Source: §7 · p. 8
• 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.Source: §7 · p. 8

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.Source: Figure 4 · p. 8, §7 · p. 9

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.Source: §8 · p. 9
• 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.Source: §8 · p. 9
• 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.Source: A PROOF OF THEOREM 3 (GENERALIZATION… · pp. 12–13
• 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.Source: B PROOF OF CLAIM 4 (GENERALIZATION OF… · pp. 13–14

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.Source: §C.1 · pp. 14–15
• 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.Source: §C.1 · p. 14
• 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.Source: Figure 5 · p. 14, §C.1 · pp. 14–15

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.Source: §C.2 · p. 15
• 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.Source: §C.2 · p. 15

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.Source: Figure 6 · p. 15
• 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.Source: Figure 6 · p. 15

β-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.Source: Figure 7 · p. 16

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.Source: Figure 9 · p. 16
• 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.Source: Figure 10 · p. 17

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 πφ.Source: H COMPARISON TO HIRO · p. 18
• 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.Source: H COMPARISON TO HIRO · p. 18