Source-linked AI summary
Visual Reinforcement Learning with Imagined Goals
Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, Sergey Levine
TL;DR
Goal-conditioned reinforcement learning struggles to learn broadly useful skills directly from high-dimensional images. RIG combines unsupervised visual representation learning with imagined goals and retroactive relabeling, enabling sample-efficient vision-based control that substantially outperforms prior methods.
Problem
Goal-conditioned policies remain difficult to learn from high-dimensional image observations, limiting general-purpose agents that must handle diverse user-specified tasks.
Method
RIG combines latent representation learning, imagined-goal sampling, latent-space rewards, and retroactive goal relabeling for off-policy goal-conditioned reinforcement learning.
Results
RIG efficiently learns simulated reaching and pushing policies from raw images without object state, substantially outperforming prior goal-conditioned methods and approaching state-based oracle performance.
Takeaways & Limitations
The resulting algorithm solves diverse goal-conditioned vision tasks without ground-truth state or reward functions, including variable-object tasks difficult to represent with fixed-length vectors.
Takeaways & Limitations
RIG uses goal images, and extending it to demonstrations or abstract representations such as language remains future work.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper introduces RIG, a framework that combines unsupervised representation learning with off-policy goal-conditioned reinforcement learning to learn general-purpose policies from raw sensory inputs. During self-supervised practice, the agent samples imagined goals and learns to achieve target observations in a latent representation.
- 1 Introduction: Standard reinforcement learning typically learns task-specific policies from hand-specified reward functions, limiting preparation for diverse tasks.Learning directly from image observations also has high sample complexity, making real-world use difficult.
- 1 Introduction: RIG jointly learns representations of raw sensory inputs and policies that achieve arbitrary goals by practicing to reach self-specified random goals.The framework is designed for automated and flexible goal-setting in complex, variable environments.
- 1 Introduction: During training, the VAE samples goals, embeds observations into a latent space, and computes latent distances; at test time, it embeds a specified goal observation for policy input.These multiple uses connect representation learning, goal selection, and goal-reaching rewards.
- 1 Introduction: A variational autoencoder provides a structured sensory representation that makes learning from images more feasible for reinforcement learning.The representation is learned with a latent variable model based on a VAE.
- 1 Introduction: RIG combines sample-efficient off-policy goal-conditioned reinforcement learning with unsupervised representation learning for general-purpose policies targeting observations.The learned latent distribution supports imagined-goal sampling for unsupervised practice and data augmentation.
2 Related Work
Prior vision-based robotic RL methods learn diverse behaviors but rely on assumptions that limit general-purpose, goal-conditioned learning from images. This work distinguishes itself through model-free goal-conditioned value learning with latent-goal relabeling and without ground-truth state or reward information during training.
- Vision-based robotic RL: Prior vision-based RL methods learn grasping, pushing, navigation, and manipulation behaviors, but assumptions limit their applicability to training general-purpose robots.The passage notes that extending time-varying-model approaches to the goal-conditioned setting with image inputs is very difficult.
- Model-based methods: Image-predictive and inverse-model methods accommodate various goals but tend to limit horizon length because of model drift.The passage states that no prior method, to its knowledge, uses model-free RL efficiently enough to train directly on real-world robots with a single goal image and without ground-truth state or reward information.
- Goal relabeling: The proposed method uses a goal-conditioned value function and retroactively relabels replay-buffer samples with goals sampled from a latent representation.Prior goal-relabeling methods addressed sparse-reward problems with known goal spaces, restricting resampled goals.
- Unsupervised representations: Prior unsupervised representation-learning methods for RL require extra information, including ground-truth rewards, expert trajectories, human demonstrations, or pretrained object-detection features.These methods use learned representations as substitutes for policy state.
3 Background
The background introduces goal-conditioned reinforcement learning for policies spanning multiple tasks and variational autoencoders for structured latent representations of high-dimensional observations. It also notes that goal-conditioned value learning can use off-policy data with standard actor-critic methods.
- Goal-conditioned reinforcement learning: Goal-conditioned reinforcement learning optimizes policies over a goal distribution to accomplish a variety of tasks rather than a single task.The formulation augments the policy and reward with a goal and optimizes expected return across goals.
- Goal-conditioned reinforcement learning: Goal-conditioned value functions use a temporal-difference loss comparing Q_w(s, a, g) with a reward plus discounted target value.The target uses a target-network parameterization Q̄_w treated as constant.
- Goal-conditioned reinforcement learning: The value-learning loss can be optimized with off-policy transitions using a standard actor-critic algorithm.The cited transition data are tuples of states, actions, next states, goals, and rewards.
- Variational Autoencoders: Variational autoencoders learn structured latent representations of high-dimensional data through an encoder that maps states to latent distributions and a decoder that reconstructs state distributions.The encoder and decoder parameters are trained jointly, with a unit Gaussian prior and a β-weighted objective.
- Variational Autoencoders: The VAE encoder represents each state with Gaussian mean and diagonal variance, while the decoder models each pixel with a Bernoulli distribution.The Bernoulli pixel parameterization corresponds to cross-entropy loss on normalized pixel values.
4 Goal-Conditioned Policies with Unsupervised Representation Learning
RIG combines β-VAE representation learning with latent-space goal-conditioned reinforcement learning to address visual RL’s representation, reward, goal-setting, and sample-efficiency challenges. It samples imagined goals, relabels experience, and trains policies from raw visual observations through learned embeddings.
- Representation Learning: RIG uses a β-VAE to embed visual states and goals into a latent space, where goal-conditioned policies and value functions operate.The encoder maps observations and goals to latent state z_t and latent goal z_g; TD3 trains the latent-space policy and Q-function.
- Representation Learning: The VAE is fine-tuned on both initial random-policy data and observations collected during exploration to address states outside its original training distribution.The paper reports that this additional training helps algorithm performance.
- Latent-Space Reward: The latent representation supplies a goal-reaching reward based on distance between the encoded state and encoded goal, avoiding pixel-space similarity.The reward is r(s, g) = −||e(s) −e(g)||_A = −||z −z_g||_A; in practice, A = I performed better than Mahalanobis distance.
- Goal Relabeling: RIG improves sample efficiency by sampling latent goals from the VAE prior and retroactively relabeling transitions with recomputed rewards.A fitted diagonal-Gaussian prior is used to mitigate mismatch between sampled latents and the VAE training distribution.
- Goal Relabeling: The final method mixes prior-sampled goals with future trajectory goals, extending beyond HER’s restriction to goals already seen along a trajectory.HER is described as limiting the number and diversity of goals available for relabeling, motivating the combined strategy.
- Imagined-Goal Practice: During self-supervised practice, the agent samples plausible latent goals from the learned distribution and trains to reach them without knowing test-time goals.The complete algorithm is called reinforcement learning with imagined goals (RIG).
5 Experiments
Experiments show that RIG efficiently learns goal-conditioned policies from images in simulation, including tasks with varying numbers of objects, and transfers to real-world robotic reaching and pushing. It substantially outperforms prior visual goal-conditioned methods, approaches the state-based oracle, and solves real-world reaching with 10,000 samples.
- Simulated Tasks: The simulated evaluation measures distance to goal from image-conditioned policies, while ground-truth positions are reserved for evaluation rather than policy training.Each policy receives an encoded goal image and raw visual observations; simulation plots report 95% confidence intervals across 5 seeds.
- Simulated Tasks: RIG efficiently learns visual reaching and pushing policies without object-state access, substantially outperforming prior methods and approaching the state-based oracle.On multi-object tasks, RIG can outperform the oracle, whose state-based reward may contain local minima.
- Ablations: A reward ablation finds that RIG’s latent Euclidean distance reward outperforms alternative reward specifications on Visual Pusher.The ablation changes only the reward function used to train the goal-conditioned value function.
- Ablations: The experiments also isolate contributions from RIG’s generative-model-based goal relabeling by comparing alternative resampling strategies.The relabeling comparison fixes the other algorithmic components while varying how transitions are assigned goals.
- Variable Numbers of Objects: Pixel-based policies handle environments with varying numbers and combinations of objects, supporting tasks whose structure is difficult to encode in a fixed-length state vector.The paper presents this flexibility as a major advantage of learning directly from images.
- Real-World Robotics: 10,000 samples solve real-world visual reaching in about an hour, while real-world visual pushing requires about 25,000 samples.RIG nearly matches the state-based oracle and far exceeds the baseline on reaching; pushing is more visually complicated and requires contact reasoning.
6 Discussion and Future Work
The paper presents an RL algorithm for efficiently solving goal-conditioned, vision-based tasks without ground-truth state or reward functions. It identifies extensions involving exploration, intrinsic motivation, multitask and meta-learning, and more flexible goal specifications.
- Contributions: The algorithm uses a generative model to embed states and goals, generate exploration goals, relabel goals and rewards, and compute latent-space distance rewards.These functions support goal-conditioned learning without access to ground-truth state or reward functions.
- Future Work: Combining prior-sampled goal generation with exploration, intrinsic motivation, multitask learning, and meta-learning could enable agents to acquire skills continuously and efficiently.The proposed future direction extends the method’s existing mechanism for autonomously generating goals.
- Future Work: Extending RIG beyond goal images to demonstrations or abstract representations such as language could make human interaction more flexible and practical.The paper specifically identifies demonstrations and language as alternative goal specifications.
Supplementary Material · 9 Complete Ablative Results
The complete ablations examine goal relabeling, reward design, online VAE training, and direct goal-space resampling against HER. Across simulated and Fetch-task evaluations, the mixed relabeling strategy and latent-distance reward perform strongly, while online representation learning and dense-reward resampling remain effective.
- 9.1 Relabeling strategy ablation: The relabeling ablation compares Future, VAE-only, mixed RIG, and no relabeling strategies for Q-function training.RIG samples goals equally from the VAE and future trajectory states.
- 9.1 Relabeling strategy ablation: RIG consistently matches or outperforms the alternative relabeling strategies in simulated final-distance-to-goal results.The comparison uses final distance to goal versus environment steps.
- 9.2 Reward type ablation: The reward ablation compares latent distance, log probability using encoder precision, and pixel MSE as goal-conditioned value-function rewards.Latent distance is the reward used in RIG.
- 9.2 Reward type ablation: RIG’s latent-distance reward consistently matches or outperforms the other reward types in simulated final-distance-to-goal results.The figure compares final distance to goal versus environment steps.
- 9.3 Online training ablation: Online VAE training updates the representation every 3000 environment steps using all images observed by the policy, producing a good policy and substantially outperforming no VAE training.The Visual Pusher ablation uses a slightly easier task with a goal space reduced to three quarters of its original lateral dimension.
- 9.4 Comparison to Hindsight Experience Replay: The HER comparison isolates direct goal-space sampling without a VAE and evaluates the same OpenAI Fetch robotics tasks using dense and sparse rewards.“Ours” mixes future goals with k = 4 and uniform environment-goal sampling, each at probability 0.5.
- 9.4 Comparison to Hindsight Experience Replay: Our resampling strategy performs best with sparse rewards across three tasks, remains reasonable with dense rewards, and improves as gradient updates per cycle increase while HER does not.Increasing updates costs more compute but reduces the hardware samples required for learning.
10 Hyperparameters
Table 1 reports the hyperparameters used in the experiments. It serves as the paper’s consolidated reference for experimental settings.
- Table 1 lists the hyperparameters used for the experiments.
- Experimental hyperparameter settings are presented in a single reference table.
- The table is titled “Hyper-parameters used for all experiments.”
11 Environment Details
The paper evaluates visual goal-conditioned control in three simulated MuJoCo environments built around a 7-DoF Sawyer arm. Tasks progress from reaching with the end-effector to pushing one or two pucks to specified positions.
- Visual Reacher: Visual Reacher uses a 7-DoF Sawyer arm whose end-effector reaches desired positions within a constrained two-dimensional workspace.Actions control end-effector velocity, with underlying state e and goal ge.
- Visual Pusher: Visual Pusher extends reaching by requiring the Sawyer arm to push a puck to a target position.The underlying state includes end-effector and puck positions, while the goal specifies desired positions for both.
- Visual Multi-Object Pusher: Visual Multi-Object Pusher adds a second puck, requiring the end-effector and both pucks to reach desired positions in separate workspace halves.Each puck and its respective goal is initialized in its corresponding half of the workspace.