PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Sample Efficiency in Model-Free Reinforcement Learning from Images

Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, Rob Fergus

arXiv:1910.01741v3cs.LGcs.AIcs.ROstat.ML

TL;DR

Model-free RL from high-dimensional images struggles to learn useful representations from sparse rewards, while adding reconstruction losses off-policy can destabilize training. The paper studies these failures, recommends a stable autoencoder-based off-policy method, and reports competitive benchmark performance plus robustness to observational noise.

  • Problem

    Learning a compact representation from high-dimensional pixel observations using only sparse rewards is sample inefficient, while jointly adding reconstruction objectives in off-policy RL can be unstable.

  • Method

    The paper systematically studies autoencoder-based representation learning in off-policy model-free RL and recommends an end-to-end method using deterministic reconstruction with stability-oriented gradient handling.

  • Results

    The recommended approach is competitive with state-of-the-art model-based methods on traditional control benchmarks and outperforms prior methods under observational noise.

  • Takeaways & Limitations

    Reconstruction-only auxiliary learning can provide a simple, robust alternative to model-based methods without requiring a learned dynamics model.

  • Takeaways & Limitations

    The paper’s comparison defines model-based methods as those that train a dynamics model, which includes SLAC.

Abstract

from arXiv · show

Training an agent to solve control tasks directly from high-dimensional images with model-free reinforcement learning (RL) has proven difficult. A promising approach is to learn a latent representation together with the control policy. However, fitting a high-capacity encoder using a scarce reward signal is sample inefficient and leads to poor performance. Prior work has shown that auxiliary losses, such as image reconstruction, can aid efficient representation learning. However, incorporating reconstruction loss into an off-policy learning algorithm often leads to training instability. We explore the underlying reasons and identify variational autoencoders, used by previous investigations, as the cause of the divergence. Following these findings, we propose effective techniques to improve training stability. This results in a simple approach capable of matching state-of-the-art model-free and model-based algorithms on MuJoCo control tasks. Furthermore, our approach demonstrates robustness to observational noise, surpassing existing approaches in this setting. Code, results, and videos are anonymously available at https://sites.google.com/view/sac-ae/home.

1. Introduction

Pixel-based model-free RL is limited by inefficient representation learning from sparse rewards. The paper studies joint autoencoder-policy training and proposes a stable off-policy approach that performs competitively across control tasks and under observational noise.

  • Motivation: Learning representations from pixels using sparse rewards requires substantial training data, making direct pixel-based RL prohibitively slow.Existing model-free methods on Atari and DeepMind Control can require tens of millions of steps.
  • Motivation: Off-policy learning and unsupervised auxiliary objectives are identified as complementary strategies for improving sample efficiency.Off-policy methods reuse samples, while pixel reconstruction provides a simple auxiliary objective.
  • Prior Work: Earlier autoencoder approaches either separate representation and policy training or use on-policy learning, limiting joint off-policy optimization.Two-step procedures improve optimization stability but can produce suboptimal policies, while prior joint off-policy attempts performed poorly or were not reported.
  • Contribution: The paper finds that pixel reconstruction is important for learning a good representation when trained jointly, but requires careful design choices for stability.The authors revisit autoencoders specifically within off-policy model-free RL.
  • Contribution: The recommended autoencoder-based off-policy method trains end-to-end and matches state-of-the-art model-free and model-based methods on many challenging control tasks.The approach is presented as the first model-free off-policy method to jointly train latent state representations and policies in this setting.
  • Contribution: The approach is robust to observational noise and outperforms prior methods in this more practical setting.The paper also provides an open-source PyTorch implementation for researchers and practitioners.

2. Related Work

Prior visual RL work uses auxiliary losses, autoencoder pretraining, or learned dynamics, but these approaches face scalability, stability, or complexity limitations. The paper positions off-policy reconstruction-based learning as a simpler route to representation learning from pixels.

  • Motivation: Model-free RL from high-dimensional pixels is sample inefficient because sparse rewards make representation learning harder than in supervised learning.This challenge has limited scaling to complex continuous-control environments.
  • Autoencoder Methods: Early autoencoder methods learn visual feature spaces but can require recomputing features for all experiences after each update, limiting scalability.Their demonstrations were also restricted to toy problems.
  • Auxiliary Tasks: Auxiliary-task approaches in Atari commonly rely on on-policy learning and task-specific assumptions about dynamics, rewards, or changing pixel patterns.These choices can improve stability or performance but do not directly establish scalable joint off-policy representation-policy learning.
  • Definitions: Model-based methods are defined here as methods that train a dynamics model, including SLAC under this definition.This operational definition determines how the paper categorizes compared methods.
  • Variational Autoencoders: β-VAEs were used for joint representation learning from pixels, but their training difficulty led prior work to revert to alternating training procedures.Alternating optimization avoids some joint-training difficulties but separates representation and policy updates.
  • Model-Based Methods: Model-based image-based methods improve sample efficiency through world models but require balancing dynamics, reward, decoder, policy, and value objectives.This creates additional complexity compared with reconstruction-only auxiliary learning.
  • Benchmark Context: Table 1 compares SAC from pixels, PlaNet, SLAC, and SAC from proprioceptive states, with the state-based agent serving as an upper bound.The large SAC:pixel versus SAC:state gap motivates addressing representation learning in model-free off-policy RL.

3. Background

This section introduces the MDP and maximum-entropy SAC framework, then describes image-based autoencoders for learning latent states from observations. The learned latent vector replaces the unavailable true state for RL.

  • 3.1. Markov Decision Process: An MDP comprises states, actions, transitions, rewards, and a discount factor, with the policy maximizing expected cumulative reward.
  • 3.1. Markov Decision Process: The maximum-entropy objective augments expected reward with policy entropy to encourage exploration and robustness to noise.The temperature α balances reward optimization against policy stochasticity.
  • 3.2. Soft Actor-Critic: SAC is an off-policy actor-critic method that performs soft policy evaluation and improvement using replay-buffer transitions.Its policy evaluation fits a Q-function by minimizing the soft Bellman residual.
  • 3.2. Soft Actor-Critic: SAC estimates the target value from the target Q-function and improves the policy by minimizing KL divergence to a Q-induced Boltzmann distribution.The target Q-function uses exponentially moving-averaged weights to stabilize training.
  • 3.3. Image-based Observations and Autoencoders: An autoencoder maps each image observation through a convolutional encoder to a low-dimensional latent vector and reconstructs the image with a deconvolutional decoder.The encoder and decoder are trained simultaneously by maximizing expected log-likelihood.
  • 3.3. Image-based Observations and Autoencoders: In β-VAE, the variational distribution is parameterized as a Gaussian, and the resulting latent vector is used by SAC instead of the unavailable true state.

4. Representation Learning with Image Reconstruction

Image reconstruction can improve representation learning for pixel-based off-policy RL, but alternating or jointly training a β-VAE with the policy has important limitations. The proposed SAC+AE stabilizes joint learning by using a deterministic, regularized autoencoder and carefully controlling encoder updates.

  • Motivation: Pixel-based SAC is significantly outperformed by state-based SAC across the evaluated control tasks, highlighting the need for compact state representations.Model-based methods using auxiliary tasks also achieve superior performance, supporting auxiliary supervision as a route to better representations.
  • Alternating Representation Learning with a β-VAE: Alternating β-VAE and policy training improves representations as update frequency increases, but does not fully close the performance gap.The encoder is pretrained, then periodically updated using observations collected during policy interaction.
  • Joint Representation Learning with a β-VAE: Joint β-VAE representation learning is unstable in the off-policy setting and performs worse than the baseline without actor-critic gradients.The attempted encoder updates combine gradients from actor, critic, and β-VAE objectives.
  • Stabilizing Joint Representation Learning: Two factors explain the instability: β-VAE stochasticity damages performance, while actor gradients updating the shared encoder hurt performance.Smaller β values improve stability and task performance, and blocking actor gradients from reaching the encoder improves results considerably.
  • Our Approach SAC+AE: SAC+AE replaces the β-VAE with a deterministic regularized autoencoder and updates the encoder using reconstruction and soft Q-learning gradients.It also blocks actor gradients and uses a faster target-Q encoder averaging rate, τenc > τQ, to compensate for slower signal propagation.

5. Evaluation of SAC+AE

SAC+AE is evaluated across challenging image-based control tasks, standard benchmarks, noisy observations, unseen tasks, and representation recovery. It matches state-of-the-art performance on standard DMC tasks, outperforms baselines under background noise, and transfers representations to related tasks.

  • Learning Control from Pixels: The evaluation covers six domains and twelve challenging continuous-control tasks from the DeepMind Control suite.The domains include finger, cartpole, cheetah, walker, reacher, and ball in cup.
  • Learning Control from Pixels: SAC+AE matches state-of-the-art performance on DMC tasks and outperforms baselines when observations contain noise.These are the two principal results highlighted in Figure 6.
  • Performance on Noisy Observations: Forward-modeling methods perform drastically worse than SAC+AE when dynamic background distractors alter observations.The distractors consist of colored balls bouncing off each other and the frame.
  • Generalization to Unseen Tasks: An encoder pretrained on walker walk generalizes to walker stand and walker run without additional fine-tuning and significantly outperforms training from scratch.The tasks share similar image observations but have different reward functions.
  • Representation Power of the Encoder: SAC+AE representations encode proprioceptive state information more successfully than SAC:pixel representations on cheetah run.Linear projections from the learned image embeddings are compared with ground-truth proprioceptive coordinates.

6. Discussion

The discussion identifies deterministic autoencoders and restricted encoder-gradient flow as key design choices for stable off-policy pixel RL. SAC+AE is presented as a simple, competitive, robust alternative that avoids learning a dynamics model and supports further research through open-source code.

  • Discussion: Deterministic autoencoders outperform β-VAEs in off-policy RL, while actor gradients through the convolutional encoder hurt performance.The analysis attributes β-VAE difficulty to instabilities involving bootstrapping, off-policy data, and joint auxiliary-loss training.
  • Discussion: SAC+AE uses only reconstruction loss as an auxiliary task and remains competitive with state-of-the-art model-based methods without learning a dynamics model.The method is described as simpler and robust on traditional benchmarks.
  • Discussion: Ablations support joint learning, pixel reconstruction, and avoiding world-model auxiliary losses in distractor settings.The discussion contrasts joint learning with alternating training and pixel reconstruction with lower-dimensional target representations.
  • Discussion: The implementation is released as open-source PyTorch code for researchers and practitioners.The paper states that the codebase is intended to support future research in image-based reinforcement learning.
  • Experimental Scope: The experiments use DMC benchmarks spanning six domains and twelve control tasks with varied challenges and reward structures.Examples include sparse reward in ball in cup catch, high-dimensional state and action spaces in cheetah run, and exploration in reacher hard.

B.3. Training and Evaluation Setup

Training uses a fixed evaluation protocol and standardized image preprocessing across tasks, with action repeat as the main task-specific setting. The autoencoder is regularized and evaluated within a consistent SAC training setup.

  • Training and Evaluation: The agent performs one training update per new observation after collecting 1000 seed observations with a random policy.With action repeat, training observations are fewer than environment steps.
  • Training and Evaluation: Agents are evaluated every 10000 environment steps using average return over 10 episodes, taking the Gaussian policy mean.This evaluation setup is preserved throughout the experiments.
  • Autoencoder Regularization: The deterministic autoencoder adds L2 representation regularization and decoder weight decay, with λz = 10^-6 and λθ = 10^-7.These penalties extend the standard reconstruction loss.
  • Training and Evaluation: The experimental setup keeps hyperparameters fixed across tasks except for action repeat.The paper provides the remaining hyperparameters in a dedicated table.

C. Alternating Representation Learning with a β-VAE

Alternating β-VAE pretraining and policy learning improves representation learning as update frequency increases, but does not close the performance gap. The results motivate additional design changes beyond separated training.

  • Alternating Representation Learning: Iterative pretraining enables faster representation learning and boosts final performance, but joint training is additionally needed to close the gap.The paper contrasts alternating representation learning with the benefits of joint optimization.
  • Alternating Representation Learning: More frequent β-VAE updates after pretraining improve representation quality but cannot fully address the performance gap.Figure 11 varies update frequency N in environment steps while showing SAC:state as an upper bound.
  • Alternating Representation Learning: The experiment separates β-VAE and policy training without shared gradients while varying the β-VAE update frequency.The policy uses latent states from a fixed encoder during each policy-training interval.

D. Joint Representation Learning with a β-VAE

Jointly propagating actor-critic gradients into a β-VAE encoder causes instability and poorer performance than avoiding that gradient sharing.

  • D. Joint Representation Learning with a β-VAE: Joint gradient propagation from the actor-critic into the β-VAE encoder produces unstable learning.The unsuccessful setup attempts end-to-end off-policy training by sharing gradients with the encoder.
  • D. Joint Representation Learning with a β-VAE: The jointly trained SAC+VAE:pixel performs worse than the baseline SAC+VAE:pixel (iter, 1), which does not share actor-critic gradients.The comparison links the performance gap to the gradient-sharing design.

E. Stabilizing Joint Representation Learning

The paper attributes subpar joint representation learning to stochastic β-VAE behavior and non-stationary actor gradients, then examines architectural capacity as a stability-related design factor.

  • E. Stabilizing Joint Representation Learning: Figure 13 provides additional results for the experiments in Section 4.4.The cited figure summarizes the investigation of the two identified causes.
  • E. Stabilizing Joint Representation Learning: Two identified causes of subpar joint representation learning are the stochastic nature of a β-VAE and non-stationary actor gradients.These issues are presented as distinct contributors to the observed difficulty.
  • E. Stabilizing Joint Representation Learning: SAC+AE remains robust across variations in convolutional and deconvolutional autoencoder capacity.The study changes convolutional depth and filters per layer while preserving shared actor-critic convolutional weights.

G. Representation Power of the Encoder

The encoder is evaluated by whether its latent representation preserves task-relevant state information from images, including proprioceptive information, while low-dimensional reconstruction targets provide insufficient learning signal in one task.

  • G. Representation Power of the Encoder: The encoder’s latent representation successfully encodes proprioceptive state information from image observations in cheetah run.The evaluation reconstructs 8 position and 9 velocity coordinates from the learned latent space.
  • G. Representation Power of the Encoder: The representation is defined as optimal when it preserves all pixel-observation information sufficient to learn a good policy.This criterion frames encoder quality in terms of control-relevant information retention.
  • G. Representation Power of the Encoder: The state decoder reconstructs the corresponding state from the latent vector produced by the observational encoder.The encoder maps image o_t to latent z_t, and the decoder is trained jointly with the actor-critic.
  • G. Representation Power of the Encoder: Low-dimensional state supervision helps less than expected and hurts performance in ball in cup catch.The authors infer that it lacks the rich reconstruction error needed to fit the high-capacity convolutional encoder.

I. Optimality of Learned Latent Representation

The paper tests latent-representation optimality by comparing pixel- and proprioceptive-based policies trained on the same replay buffer, while showing that action repeat materially affects learning and final reward.

  • I. Optimality of Learned Latent Representation: Latent-representation optimality is assessed by training pixel and proprioceptive policies from the same fixed replay buffer.This design separates observation-space effects from differences in data collected by the policies.
  • I. Optimality of Learned Latent Representation: Fixed-buffer evaluation can prevent a policy from reaching comparable performance because it operates in an off-policy regime.The authors explicitly note this as a qualification of the comparison.
  • I. Optimality of Learned Latent Representation: The learned pixel-based latent representation outperforms proprioceptive state on the fixed buffer.The result may reflect that the buffer was collected by a pixel-trained policy, making it different from data suited to a proprioceptive policy.
  • I. Optimality of Learned Latent Representation: Action repeat has a significant effect on learning dynamics and final reward.Repeating actions can shorten the control horizon and stabilize dynamics, but may also introduce harmful bias.
  • I. Optimality of Learned Latent Representation: The action-repeat ablation averages performance over 10 random seeds while reducing training observations inversely with action repeat.It compares several action-repeat settings against PlaNet using its original per-environment setting.
Loading 1910.01741v3…