Source-linked AI summary
Goal-Conditioned Reinforcement Learning with Imagined Subgoals
Elliot Chane-Sane, Cordelia Schmid, Ivan Laptev
TL;DR
Goal-conditioned reinforcement learning struggles with tasks requiring temporally extended reasoning. RIS predicts imagined subgoals during training, uses them to form a KL-regularized prior, and reports faster learning and stronger performance on challenging robotic tasks without using subgoals at test time.
Problem
Goal-conditioned reinforcement learning can learn a wide range of skills but struggles with temporally extended reasoning.
Method
RIS jointly trains a goal-conditioned policy, critic, and high-level policy, using value-guided imagined subgoals to define a prior for KL-regularized policy learning.
Results
RIS greatly accelerates learning and outperforms competing approaches on challenging simulated navigation and vision-based manipulation environments.
Takeaways & Limitations
The resulting flat policy can solve temporally extended goal-reaching tasks without requiring subgoals during test-time execution.
Takeaways & Limitations
The high-level policy can predict subgoals outside the valid state distribution, including unreachable robot poses or unrealistic images.
Abstract
from arXiv · showhide
Goal-conditioned reinforcement learning endows an agent with a large variety of skills, but it often struggles to solve tasks that require more temporally extended reasoning. In this work, we propose to incorporate imagined subgoals into policy learning to facilitate learning of complex tasks. Imagined subgoals are predicted by a separate high-level policy, which is trained simultaneously with the policy and its critic. This high-level policy predicts intermediate states halfway to the goal using the value function as a reachability metric. We don't require the policy to reach these subgoals explicitly. Instead, we use them to define a prior policy, and incorporate this prior into a KL-constrained policy iteration scheme to speed up and regularize learning. Imagined subgoals are used during policy learning, but not during test time, where we only apply the learned policy. We evaluate our approach on complex robotic navigation and manipulation tasks and show that it outperforms existing methods by a large margin.
1. Introduction
The paper addresses the difficulty of temporally extended goal-conditioned reinforcement learning by using imagined intermediate subgoals to guide policy learning. RIS predicts subgoals during training, incorporates them through a KL-regularized prior, and produces a flat policy that operates without subgoals at test time.
- Goal-conditioned reinforcement learning can learn many skills but struggles with temporally extended reasoning.
- RIS automatically predicts intermediate subgoals and uses them to direct policy search for distant goals.The approach progressively extends the agent’s effective horizon by applying subgoals across state-goal pairs.
- The learned flat policy reaches distant goals at test time without access to subgoals.Imagined subgoals are used only during training, while the resulting policy directly generates actions during evaluation.
- A high-level policy predicts subgoals while using the goal-conditioned value function as a relative reachability measure.The high-level policy is trained simultaneously with the goal-conditioned policy and critic.
- RIS defines a prior from actions for reaching intermediate subgoals and constrains policy iteration toward that prior with KL regularization.The subgoals guide adaptation of simpler behaviors into complex goal-reaching behavior without being explicitly pursued in the environment.
- RIS accelerates policy learning on simulated robotics environments involving motor control and temporally extended reasoning.
2. Related Work
Prior work addresses long-horizon goal-conditioned tasks with hierarchical subgoals, value-based planning, policy priors, and guided policy search. RIS differs by using high-level subgoals to guide training of a non-hierarchical policy rather than executing subgoals at test time.
- Goal-conditioned methods learn policies that predict actions toward desired goals, often using hindsight experience replay.
- Hierarchical reinforcement learning uses high-level policies and intermediate subgoals to decompose long-horizon tasks into simpler problems.
- Jointly learning high-level and low-level policies can introduce training instability in hierarchical methods.
- Value functions have been used to measure state distances and plan subgoal sequences, including methods that optimize subgoals at test time.RIS also uses value-based distances but avoids relying on subgoals during test-time execution.
- Policy-prior methods constrain target policies toward action distributions, while behavior priors can reduce value overestimation for out-of-distribution actions.
- Guided policy search alternates between expert trajectory generation and policy improvement, whereas RIS guides policy search with high-level-policy subgoals.
3. Method
RIS uses imagined midpoint subgoals to guide goal-conditioned policy learning without requiring explicit subgoal reaching. A value-based high-level policy predicts valid intermediate states, while KL regularization incorporates their induced actions into policy improvement.
- High-Level Policy: RIS uses a high-level policy to predict intermediate subgoals halfway along paths from current states to goals.The subgoals are defined using the current policy’s value function as a state-distance measure.
- Algorithm Summary: Imagined subgoals guide training but are unnecessary at test time, when the learned goal-conditioned policy directly generates actions.The policy, critic, and high-level policy are trained within an off-policy actor-critic framework, with an exponential moving average used for the prior to support convergence.
- High-Level Policy: The value norm |V π(si, sj)| estimates expected discounted steps between valid states and changes as the goal-reaching policy improves.This policy-dependent distance guides midpoint subgoal selection.
- High-Level Policy: A KL-regularized high-level objective keeps predicted subgoals near the valid state distribution, avoiding infeasible or adversarial states.Replay-buffer states provide samples for estimating the valid-state distribution, and implicit KL enforcement yields advantage-weighted maximum likelihood.
- Policy Improvement with Imagined Subgoals: RIS defines a prior policy as the distribution of actions for reaching high-level-policy subgoals, then constrains target-policy improvement toward that prior with KL regularization.The prior supplies an initial guess based on simpler subgoal-reaching behaviors while the target policy still maximizes goal-conditioned value.
4. Experiments
The experiments evaluate RIS on challenging ant navigation and vision-based manipulation tasks, including ablations and comparisons with prior methods. RIS learns imagined subgoals that guide policy learning and achieves strong performance on temporally extended configurations.
- Experimental Setup: The study evaluates RIS on four increasingly difficult ant mazes and a vision-based 2 DoF robotic manipulation task.Ant experiments use unknown obstacles and challenging initial-state/goal configurations; manipulation requires controlling an arm to move a puck.
- Experimental Setup: Success is measured by proximity to the desired goal, with agents trained for 1 million environment steps and results averaged over four random seeds.For ant navigation, success uses x-y Euclidean distance; manipulation evaluates whether the arm and puck reach their respective desired positions.
- Ablative Analysis: RIS learns subgoals coherent with oracle trajectories during training, while implicit regularization prevents predicted subgoals from diverging and helps guide policy learning.The high-level policy’s probability mass approaches oracle subgoals, whereas removing implicit regularization causes substantial x-y divergence and failure to guide the task.
- Ablative Analysis: Ablations show that uniform and moving-average priors can learn encountered configurations but fail on the most difficult Ant U-maze test configurations.The moving-average action prior also fails to learn during training, highlighting the role of subgoal-informed priors.
- Comparison to the State of the Art: Across four ant mazes, RIS significantly improves sample efficiency, often solving the most challenging configurations with fewer than 500 thousand environment interactions.LEAP requires significantly more interactions, while SAC, HIRO, and TDM fail or perform worse in especially challenging settings.
- Comparison to the State of the Art: As planning horizons increase in S-shaped and ω-shaped mazes, RIS consistently solves increasingly complex configurations, unlike SAC and LEAP.After 1 million training steps, RIS maintains performance as goal complexity increases, whereas the other methods degrade.
- Comparison to the State of the Art: In vision-based manipulation, RIS achieves a higher success rate than LEAP, while SAC fails to solve the temporally extended task consistently.RIS uses one image-encoder and actor forward pass per interaction, whereas LEAP additionally performs expensive image-subgoal planning.
5. Conclusion
RIS imagines subgoals self-supervisedly and uses them to facilitate training. Experiments show faster learning on temporally extended tasks and stronger performance than competing approaches, while better Q-value estimation remains future work.
- RIS imagines possible subgoals self-supervisedly and uses them to facilitate goal-conditioned policy training.
- The value function trains a high-level state-space policy, whose imagined subgoals define a prior incorporated into policy learning.
- Experiments on simulated navigation and vision-based manipulation show greatly accelerated learning of temporally extended tasks and outperforming competing approaches.
- Future work could use subgoals to obtain better Q-value estimates and improve exploration across potential states.
A. High-level Policy Derivation
The high-level policy derivation optimizes a subgoal advantage objective regularized against a valid-state distribution, then projects its closed-form solution into a parametric policy.
- The derivation considers a subgoal advantage function, a valid-state distribution, and an experience-induced state-goal distribution.
- The high-level policy improvement objective combines subgoal advantage with logarithmic regularization terms and normalization constraints.
- The closed-form solution includes a normalizing partition function to ensure the state distribution constraint sums to one.
- The closed-form solution is projected into parametric policies by minimizing reverse KL divergence.
- Choosing the valid-state distribution from replay-buffer experience turns high-level improvement into weighted maximum likelihood over sampled visited states.
B. Implementation Details
RIS uses actor-critic learning with KL regularization, neural high-level subgoal policies, smoothed prior estimation, Monte Carlo KL estimates, shared visual encoders, and hindsight relabelling.
- The actor-critic implementation removes entropy during evaluation and replaces it with KL divergence to the prior during policy improvement.
- The high-level policy parameterizes a Laplace distribution and uses batch advantage softmax weights instead of estimating its normalizing factor.
- The prior policy uses exponentially moving-averaged policy weights, while its density and KL term are estimated by Monte Carlo.
- The KL estimate uses I = 10 samples, and the implementation uses N = 1 with reparameterization for backpropagation.
- Hindsight Experience Replay relabels minibatch goals using 20% original goals, 40% replay-buffer states, and 40% future trajectory states.
- Vision-based tasks share an image encoder across the policy, high-level policy, and Q-function, with fixed goal, next-state, and candidate-subgoal representations during updates.
C. Environments
The evaluation uses ant navigation in multiple mazes and vision-based robotic manipulation with a puck, defined by explicit spatial goal criteria.
- Ant Navigation tasks use a four-legged ant that must learn to navigate U-shaped, S-shaped, Π-shaped, and ω-shaped mazes.
- Navigation goals are reached when the ant’s x-y position lies within 0.5 units of the desired location, with reward −1 until success.
- The manipulation task controls an arm robot in 2D to move a puck using 84 × 84 RGB top-down observations.
- Manipulation succeeds when both the arm and puck are within 5cm of their respective target positions.
D. Hyperparameters
The experiments specify optimizer and evaluation settings for RIS and SAC, alongside environment-specific network and policy-horizon configurations.
- RIS uses the Adam optimizer and reports results after one million environment interactions.
- SAC automatically tunes policy entropy to match the target entropy of −dim(A).
- RIS uses a convolutional image encoder with four layers and a 16-dimensional fully connected output in the vision-based environment.
- The implementation details include hidden-layer and VAE architectures for LEAP, TDM, and ant-navigation experiments.The supplied passages identify these architectures but do not provide a complete consolidated hyperparameter table.
E. Additional Results
Additional results visualize imagined-subgoal reconstructions for random and additional configurations in the vision-based robotic manipulation environment.
- Additional configurations of the vision-based robotic manipulation environment are used to provide image reconstructions of imagined subgoals.
- Figure 10 shows reconstructed imagined subgoals alongside current states and desired goals.
- The reconstructions cover different random configurations of the vision-based robotic manipulation environment.