[Paper Review] Goal-Conditioned Reinforcement Learning with Imagined Subgoals

Elliot Chane-Sane, Cordelia Schmid+1 more|arXiv (Cornell University)|Jul 1, 2021

Reinforcement Learning in Robotics26 citations

TL;DR

The paper introduces RIS, a method that trains a goal-conditioned policy with imagined subgoals generated by a high-level policy and regularized via KL-divergence to a prior, accelerating learning for temporally extended tasks; at test time only the flat policy is used.

ABSTRACT

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.

Motivation & Objective

• Motivate goal-conditioned RL to handle temporally extended reasoning
• Propose imagined subgoals to guide and accelerate policy learning
• Learn a high-level policy predicting subgoals and a prior for KL-regularized policy updates
• Enable a single flat policy at test time that can reach distant goals
• Demonstrate sample-efficient learning on simulated robotics navigation and manipulation tasks

Proposed method

• Introduce imagined subgoals predicted by a high-level policy trained jointly with the target policy and its critic
• Define a prior policy as the distribution over actions needed to reach imagined subgoals, and regularize the goal-conditioned policy with a KL constraint toward this prior
• Use the current value function to define a distance between states, guiding subgoal selection
• Predict subgoals as midpoints along the path from current state to goal to reduce horizon
• Train a non-parametric high-level policy and project it into the parametric space via a KL-based update (Eq. 6–7)
• Employ an off-policy actor-critic framework with HER for sparse rewards (Algorithm 1)

Experimental results

Research questions

• RQ1 Can imagined subgoals accelerate learning for goal-conditioned RL on long-horizon tasks?
• RQ2 How can a high-level policy predict subgoals that effectively regularize and guide a flat policy?
• RQ3 Does KL-regularization toward a subgoal-informed prior improve stability and sample efficiency?
• RQ4 Does RIS generalize from navigation tasks to vision-based manipulation with pixel inputs?
• RQ5 Is test-time performance achieved with a single policy despite training with subgoals?

Key findings

• RIS substantially improves sample efficiency over SAC+HER, LEAP, and TDMs across ant navigation mazes and vision-based manipulation
• High-level policy learns subgoals coherent with oracle (midpoints along trajectories) during training
• Implicit KL-regularization keeps imagined subgoals within valid state distributions, preventing divergence
• Imaged subgoals in vision tasks enable RIS to outperform SAC and LEAP with fewer computations (single forward pass vs planning)
• Averaged subgoal priors (via moving average) stabilize training and support convergence
• RIS solves increasingly difficult horizon tasks where traditional methods struggle