[Paper Review] Universal Planning Networks

Aravind Srinivas, Allan Jabri+3 more|arXiv (Cornell University)|Apr 2, 2018

Reinforcement Learning in RoboticsComputer Science51 references92 citations

TL;DR

UPNs embed a differentiable gradient-based planner inside a goal-conditioned policy, learning plannable latent representations that improve planning, transfer, and reward design for visuomotor tasks. The approach enables image-based goal specification and transfer across morphologies, with latent spaces usable for RL rewards.

ABSTRACT

A key challenge in complex visuomotor control is learning abstract representations that are effective for specifying goals, planning, and generalization. To this end, we introduce universal planning networks (UPN). UPNs embed differentiable planning within a goal-directed policy. This planning computation unrolls a forward model in a latent space and infers an optimal action plan through gradient descent trajectory optimization. The plan-by-gradient-descent process and its underlying representations are learned end-to-end to directly optimize a supervised imitation learning objective. We find that the representations learned are not only effective for goal-directed visual imitation via gradient-based trajectory optimization, but can also provide a metric for specifying goals using images. The learned representations can be leveraged to specify distance-based rewards to reach new target states for model-free reinforcement learning, resulting in substantially more effective learning when solving new tasks described via image-based goals. We were able to achieve successful transfer of visuomotor planning strategies across robots with significantly different morphologies and actuation capabilities.

Motivation & Objective

• Learn representations that are effective for goal-directed planning and control from visual inputs.
• Embed a differentiable gradient-descent planner within a neural policy and train end-to-end.
• Demonstrate that learned latent representations support image-based goal specification and transfer across morphologies.
• Show that latent spaces can provide distance-based rewards for model-free RL on new tasks.

Proposed method

• Encode current and goal observations into a latent space via f_phi and simulate transitions with a forward model g_theta.
• Use a gradient descent planner (GDP) in the latent space to produce action sequences that minimize the distance to the encoded goal.
• Back-propagate planning gradients through the full planning computation graph to update phi and theta.
• Train with an outer imitation learning objective to match expert demonstrations (Algorithm 2).
• Optionally re-plan at test time using MPC-style replanning to handle longer horizons.
• Demonstrate latent-space rewards r(o_t,o_g) = -||f_phi(o_t)-f_phi(o_g)||^2 with Huber loss and RL fine-tuning.

Experimental results

Research questions

• RQ1 Does embedding a gradient descent planner inside a policy yield better visuomotor imitation from pixels than reactive or autoregressive baselines?
• RQ2 Can the learned latent space support transfer to different robot morphologies and more complex tasks via image-based goals?
• RQ3 Can planning updates at test time improve performance and approach expert success?
• RQ4 Are the learned representations useful for defining distance-based rewards in model-free RL on new tasks?

Key findings

• UPN learns effective visual goal-directed policies more efficiently (data-limited) than traditional imitation learners.
• The latent representations provide a meaningful, obstacle-aware distance metric useful for transfer and reward design.
• Allowing more test-time GDP updates improves planning performance and can reach expert levels with sufficient demonstrations.
• Latent-space rewards derived from UPN representations enable RL to outperform other feature spaces (VAE, RIL, AIL) on transfer tasks.
• UPN representations trained on one morphology can improve RL performance on unseen morphologies and even more complex tasks.
• Reinforcement learning using UPN-derived rewards often surpasses expert performance in some transfer settings.