[Paper Review] Visual Reinforcement Learning with Imagined Goals

Ashvin Nair, Vitchyr H. Pong+4 more|arXiv (Cornell University)|Jul 12, 2018

Reinforcement Learning in RoboticsComputer Science30 references181 citations

TL;DR

The paper presents RIG, a framework that learns goal-conditioned policies from raw images by coupling off-policy RL with unsupervised latent representations, imagined goals, and retroactive goal relabeling to achieve user-specified image goals.

ABSTRACT

For an autonomous agent to fulfill a wide range of user-specified goals at test time, it must be able to learn broadly applicable and general-purpose skill repertoires. Furthermore, to provide the requisite level of generality, these skills must handle raw sensory input such as images. In this paper, we propose an algorithm that acquires such general-purpose skills by combining unsupervised representation learning and reinforcement learning of goal-conditioned policies. Since the particular goals that might be required at test-time are not known in advance, the agent performs a self-supervised "practice" phase where it imagines goals and attempts to achieve them. We learn a visual representation with three distinct purposes: sampling goals for self-supervised practice, providing a structured transformation of raw sensory inputs, and computing a reward signal for goal reaching. We also propose a retroactive goal relabeling scheme to further improve the sample-efficiency of our method. Our off-policy algorithm is efficient enough to learn policies that operate on raw image observations and goals for a real-world robotic system, and substantially outperforms prior techniques.

Motivation & Objective

• Develop a general-purpose, goal-conditioned RL framework that operates on raw image observations without ground-truth rewards or states.
• Learn a structured latent representation to facilitate goal sampling, state encoding, and reward shaping.
• Enable self-supervised goal imagination and relabeling to improve sample efficiency during training.
• Demonstrate scalability to variable object numbers and real-world robotic manipulation tasks.

Proposed method

• Train a beta-VAE to learn latent representations of observations and goals.
• Embed states and goals into a latent space and train a goal-conditioned Q-function Q(z, a, z_g) with a latent-space policy pi(z, z_g).
• Use latent distance as a shaped reward r = -||z - z_g||, enabling vision-based control without true state rewards.
• Sample goals from the VAE prior to drive self-supervised practice and to relabel experiences (latent goal relabeling) to improve data efficiency.
• Apply off-policy TD3-style learning; optionally fine-tune the VAE during training for better representation alignment.
• Employ a mixture of prior-sampled and future-state relabeling to maximize data efficiency.

Experimental results

Research questions

• RQ1 Can vision-based, goal-conditioned policies be learned directly from images without ground-truth state or reward signals?
• RQ2 How do latent representations influence goal sampling, reward shaping, and sample efficiency in vision-based RL?
• RQ3 Does the approach scale to tasks with variable numbers of objects and transfer to real-world robotic manipulation?
• RQ4 What is the impact of different goal relabeling strategies on learning efficiency?

Key findings

• RIG substantially outperforms prior vision-based goal-conditioned RL methods on simulated image-based tasks.
• Latent-space rewards (distance in latent space) provide a better-shaped signal than pixel MSE or Mahalanobis distance for learning.
• Relabeling with a mixture of latent-goal samples and future-state strategies yields the best sample efficiency.
• The method handles variable-object scenarios without changing the input representation.
• RIG achieves competitive performance with a real-world Sawyer robot on visual reaching and pushing tasks with modest data.