Technical Summary: S3 – Stable Subgoal Selection by Constraining Uncertainty of Coarse Dynamics in Hierarchical Reinforcement Learning

1. Problem Statement

Hierarchical Reinforcement Learning (HRL) addresses long-horizon and complex tasks by separating strategic planning (high-level "Manager") from primitive execution (low-level "Worker"). While effective, HRL faces a critical challenge: non-stationarity and sparse feedback.

• Non-stationarity: The high-level agent must learn to assign subgoals while the low-level agent's capabilities evolve, and vice versa. The low-level agent faces changing subgoal assignments, while the high-level agent faces a changing execution landscape.
• Sparse Feedback: The high-level agent receives delayed, sparse extrinsic rewards only upon task completion, making exploration inefficient.
• Current Limitations: Existing methods like Hindsight Experience Replay (HER) or reachability constraints (e.g., HRAC) often treat the low-level policy as stationary or focus on reachability without considering the predictability of the outcome. They address symptoms of non-stationarity rather than the fundamental need for the Manager to select subgoals that are not just reachable, but predictable given the Worker's current competence.

The core problem is how to enable the high-level agent to perform strategic subgoal selection that accounts for the stochasticity and uncertainty of the low-level execution, thereby stabilizing the learning process.

2. Methodology: Stable Subgoal Selection (S3)

The proposed approach, S3, introduces a dynamics-aware intrinsic reward mechanism for the high-level agent. The core philosophy is that the Manager should be motivated to select subgoals that the Worker can execute with high certainty, effectively modeling the "coarse dynamics" of the environment.

2.1 Coarse Dynamics and Predictive Uncertainty

Instead of modeling primitive transitions, S3 models coarse dynamics: the transition from state $s_t$ to a landing state $s_{t+c}$ after the low-level agent executes a subgoal $g_t$ for a fixed temporal horizon $c$.

• Goal: Minimize the predictive uncertainty of these coarse transitions.
• Rationale: If the low-level agent is uncertain about reaching a specific subgoal (e.g., squeezing through a narrow door), the high-level agent should avoid assigning it. Conversely, if the outcome is predictable (e.g., reaching a hallway junction), the subgoal is preferred. This promotes "risk-averse" subgoal selection.

2.2 Intrinsic Reward via Potential-Based Shaping

To provide dense feedback without altering the optimal policy, S3 employs Potential-Based Reward Shaping (PBRS).

• Potential Function ($\Phi$): Defined over the dispersion (variance) of the coarse dynamics distribution.
• Intrinsic Reward ($r_{int}$): Calculated as the difference in potential between time steps:
  $$r_{int}^{hi} = \gamma \Phi(s_{t+c}, g_{t+c}) - \Phi(s_t, g_t)$$
• Dispersion Metric: The potential function penalizes the total variance (trace of the covariance matrix) of the predicted terminal state distribution. Lower variance (higher predictability) yields a higher intrinsic reward.

2.3 Implementation: Mixture Density Network (MDN)

Since the environment dynamics are unknown and the low-level policy is stochastic, the coarse transition distribution $p(s_{t+c} | s_t, g_t)$ is approximated using a Mixture Density Network (MDN).

• Architecture: The MDN models the distribution as a mixture of $K$ Gaussians:
  $$\hat{p}_{\theta_{MDN}}(s_{t+c} | s_t, g_t) = \sum_{k=1}^{K} \alpha_k(s_t, g_t) \mathcal{N}(s_{t+c}; \mu_k, \Sigma_k)$$
• Training: The MDN is trained in parallel with the HRL agents using tuples $(s_t, g_t, s_{t+c})$ collected from the replay buffer.
• Uncertainty Estimation: The total predictive covariance ($\Sigma_{mix}$) is decomposed into:
  1. Within-mode variance: Variability within a specific outcome scenario.
  2. Between-mode variance: Ambiguity regarding which outcome scenario will occur.
     The intrinsic reward is derived from the trace of this total covariance matrix.

2.4 Algorithm Flow

1. Execution: The Manager selects a subgoal $g_t$. The Worker attempts to reach it over $c$ steps.
2. Data Collection: The tuple $(s_t, g_t, s_{t+c})$ is stored in a replay buffer for the MDN.
3. Model Update: The MDN updates its parameters to better predict the distribution of $s_{t+c}$.
4. Reward Calculation: The MDN predicts the distribution for the current $(s_t, g_t)$, calculates the covariance trace, and computes the intrinsic reward.
5. Policy Update: The Manager is updated using the extrinsic reward plus the intrinsic shaping term.

3. Key Contributions

1. Stable Subgoal Selection (S3): A novel framework where the high-level agent learns to select subgoals based on the predictability of the low-level execution, rather than just reachability.
2. Coarse Dynamics Modeling: The introduction of a temporal abstraction scale where the environment's dynamics are modeled as a distribution over $c$-step transitions, abstracting away fine-grained primitives.
3. Dynamics-Aware Intrinsic Reward: A potential-based intrinsic reward that penalizes the dispersion of the coarse dynamics distribution. This provides dense feedback to the Manager, guiding it toward reliable, repeatable outcomes.
4. Integration with MDN: The use of Mixture Density Networks to explicitly model multi-modal outcomes and quantify uncertainty (both within and between modes) for reward shaping.
5. Generalizability: The method is designed as a modular intrinsic reward layer that can be applied to existing HRL algorithms (demonstrated on HRAC).

4. Empirical Results

The authors evaluated S3 on three MuJoCo continuous control benchmarks: Ant Maze, Ant Push, and Ant Fall.

• Ant Fall: S3 significantly outperformed baselines (HIRO, HRAC), achieving a success rate of 0.374 compared to 0.059 (HIRO) and 0.000 (HRAC). This task involves irreversible transitions (falling into a gap), where S3's ability to avoid high-uncertainty subgoals was critical.
• Ant Push: S3 achieved a success rate of 0.186, outperforming HIRO (0.134) and matching HRAC (0.180). The task requires precise block manipulation to avoid blocking the path.
• Ant Maze: S3 achieved 0.827, comparable to HRAC (0.832) and HIRO (0.790). The authors note that in purely navigational tasks with predictable locomotion, the marginal benefit of coarse dynamics modeling is lower.
• Qualitative Analysis:
  • Risk Avoidance: In Ant Push, S3 learned to avoid the "hazardous zone" where pushing a block from the south would irreversibly block the goal. The terminal state distribution became concentrated away from this danger zone over training.
  • Trajectory Stability: Visualizations showed that S3's manager trajectories were more concentrated and aligned with the task direction compared to the scattered strategies of HRAC, particularly in tasks with bottlenecks.
  • Temporal Horizon: Sensitivity analysis indicated that a manager horizon of $c=10$ provided the optimal balance between abstraction and predictability.

5. Significance and Claims

The paper claims that S3 addresses the fundamental relationship between the Manager and Worker in HRL. By explicitly modeling the predictive uncertainty of the low-level agent, S3 stabilizes the high-level policy.

• Stability: The approach mitigates the non-stationarity caused by the evolving low-level policy by rewarding the high-level agent for selecting subgoals that are currently "safe" and executable.
• Efficiency: The dense, dynamics-aware intrinsic reward accelerates learning in sparse-reward environments, particularly those with long-horizon dependencies and structural bottlenecks.
• Applicability: The method is most beneficial in environments involving physical interaction, irreversible transitions, and multi-stage dependencies where small execution errors can lead to unrecoverable states.

The authors conclude that S3 bridges the gap between model-based HRL, potential-based shaping, and information-theoretic motivation, offering a generalizable mechanism to improve hierarchical learning in complex, non-stationary environments.