Source-linked AI summary

Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model

Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, Sergey Levine

arXiv:1907.00953v4cs.LGcs.AIstat.ML

TL;DR

High-dimensional image-based RL must learn both useful representations and task policies, a difficulty that standard end-to-end training addresses jointly. SLAC separates these roles with a stochastic sequential latent model and performs maximum-entropy actor-critic learning in the learned latent space. Across image-based continuous-control tasks, the paper reports improved sample efficiency and final performance over prior model-free and model-based methods, while noting that safety concerns remain unaddressed.

  • Problem

    High-dimensional image observations make it difficult to learn directly from inputs and to extract compact task-relevant representations for reinforcement learning.

  • Method

    SLAC learns a stochastic sequential latent representation through predictive modeling and performs maximum-entropy actor-critic reinforcement learning in that learned latent space.

  • Results

    SLAC achieves improved sample efficiency and final task performance compared with prior model-free and model-based reinforcement learning methods on image-based continuous-control tasks.

  • Takeaways & Limitations

    Explicit representation learning is used to combine predictive latent modeling with task learning for complex image-based continuous-control policies.

  • Takeaways & Limitations

    The work does not explicitly address safety concerns in automated decision-making systems.

Abstract

from arXiv · show

Deep reinforcement learning (RL) algorithms can use high-capacity deep networks to learn directly from image observations. However, these high-dimensional observation spaces present a number of challenges in practice, since the policy must now solve two problems: representation learning and task learning. In this work, we tackle these two problems separately, by explicitly learning latent representations that can accelerate reinforcement learning from images. We propose the stochastic latent actor-critic (SLAC) algorithm: a sample-efficient and high-performing RL algorithm for learning policies for complex continuous control tasks directly from high-dimensional image inputs. SLAC provides a novel and principled approach for unifying stochastic sequential models and RL into a single method, by learning a compact latent representation and then performing RL in the model's learned latent space. Our experimental evaluation demonstrates that our method outperforms both model-free and model-based alternatives in terms of final performance and sample efficiency, on a range of difficult image-based control tasks. Our code and videos of our results are available at our website.

1 Introduction

SLAC separates representation learning from task learning by learning a predictive latent model and performing reinforcement learning in its latent space. The approach targets more sample-efficient image-based control while integrating stochastic sequential modeling with RL.

  • High-dimensional image observations make both direct learning and extracting compact task-relevant representations difficult.
  • SLAC separates representation learning from task learning by explicitly learning a latent representation through predictive model learning.Predictive learning provides informative supervision before the agent has made task progress.
  • The stochastic sequential model represents high-dimensional observations as consequences of a latent process with Gaussian priors and latent dynamics.Its stochastic latent state represents uncertainty in a partially observed Markov decision process.
  • SLAC integrates stochastic sequential model learning and reinforcement learning, performing RL in the model’s learned latent space.The method derives its objective through variational inference in a control-as-inference formulation for POMDPs.
  • SLAC substantially outperforms prior model-free and model-based RL algorithms on image-based continuous-control benchmark tasks.The reported benefits include improved sample efficiency and asymptotic performance.

2 Related Work

Prior work addresses representation learning, partial observability, and sequential latent modeling through distinct auxiliary, recurrent, model-based, or latent-variable approaches. SLAC combines predictive latent representations with model-free RL while using a different latent-state treatment from several related methods.

  • Prior work identifies a representation-learning bottleneck in end-to-end deep RL and explores auxiliary supervision to learn representations before task success.
  • Research on partial observability includes exact or approximate POMDP solutions, recurrent end-to-end RL, and latent dynamical models used with model-based RL.
  • Related latent-variable approaches provide model-free RL with representations from forward-model samples, particle-filter beliefs, or learned belief-space models.
  • Sequential latent-variable models differ in their factorizations, architectures, and training objectives.SLAC is compatible with such models when they provide a mechanism for sampling latent states from the belief.

3 Preliminaries

SLAC addresses policy learning from high-dimensional observations in POMDPs by combining variational latent-state learning with maximum-entropy reinforcement learning. Its foundations include stochastic sequential models and the actor-critic framework underlying SAC.

  • The paper frames its problem as jointly learning a latent representation of the underlying MDP state and a maximum-entropy policy in a POMDP.
  • Maximum-entropy RL maximizes expected rewards together with policy entropy, with temperature α controlling the trade-off.
  • SAC uses soft policy iteration and parameterized Q-functions and policies to handle continuous action spaces.
  • Latent-variable models learn compact representations by maximizing observation likelihood through variational inference with encoders, decoders, priors, and approximate posteriors.
  • Sequential latent models incorporate actions and temporal structure because a single observation may not contain enough information to infer the latent state.
  • SLAC combines a fully stochastic sequential latent-variable model with maximum-entropy actor-critic learning for image-based continuous control.

4 Joint Modeling and Control as Inference

The paper extends control as inference from fully observed MDPs to POMDPs by jointly modeling latent states, observations, and optimality. Maximizing the resulting marginal likelihood combines representation learning with maximum-entropy control.

  • In the fully observed setting, binary optimality variables connect maximum-entropy policy optimization to approximate inference.
  • For POMDPs, the graphical model includes latent variables and assigns optimality likelihoods to latent states and actions.
  • SLAC maximizes the marginal likelihood of past observations and future optimality variables rather than optimality likelihood alone.This objective models observed data from the past and future action optimality together.
  • The variational distribution uses future dynamics to prevent controlling transitions or selecting optimistic actions, while the action posterior represents the policy.
  • The resulting ELBO contains separate model and policy objective terms, with reward represented through the log-likelihood of optimality.

5 Stochastic Latent Actor Critic

SLAC combines a stochastic latent-variable model with maximum-entropy actor-critic learning, training the model from observations and the critic on sampled latent states. Its policy remains conditioned on observation-action histories while latent states support the critic and Bellman backup.

  • Algorithm: SLAC maximizes an ELBO whose terms split into a latent-variable model objective and a maximum-entropy reinforcement-learning objective.The model objective is optimized directly, while the RL objective is optimized through approximate message passing using Q-function messages.
  • Algorithm: The latent-variable model uses Gaussian prior and transition distributions with neural-network parameterized means and variances, plus an observation likelihood and inference model.The filtering distribution is sampled using the reparameterization trick.
  • Algorithm: The policy is optimized with a soft actor-critic-style loss but is conditioned on past observations and actions rather than latent states.This avoids policies with perfect latent-state access that could produce over-optimistic behavior.
  • Algorithm: Algorithm 1 alternates environment interaction and gradient updates for the model, two critics, and actor using replayed observation-action-reward sequences.The procedure initializes from an environment reset, stores transitions, samples sequences and latent states, then updates ψ, θ_i, and φ.
  • Inputs: The method receives 64 × 64 image observations from continuous-control tasks spanning DeepMind Control and OpenAI Gym benchmarks.The images are inputs to both the latent-variable model and the policy.
  • Algorithm: The actor-critic component trains soft Q-functions on tuples of sampled latent states, actions, rewards, and successor latent states.Two Q-functions and automatic temperature tuning follow the actor-critic design used by prior work.

6 Experimental Evaluation

SLAC is evaluated on image-based continuous-control tasks from DeepMind Control Suite and OpenAI Gym, alongside model-based, model-free, and latent-variable alternatives. It also examines latent-model design, actor–critic inputs, pretraining, and update frequency.

  • Comparative Evaluation: SLAC is evaluated on four DeepMind Control Suite and four OpenAI Gym tasks using raw image observations.The comparisons include model-based and model-free algorithms, including SAC, D4PG, MPO, DVRL, PlaNet, and DrQ.
  • Comparative Evaluation: SLAC converges to higher performance than PlaNet and SAC on all four OpenAI Gym image-based tasks.These episodes have variable lengths because they terminate based on task conditions such as failure.
  • Comparative Evaluation: SLAC achieves comparable or better sample efficiency than model-based and model-free alternatives on the DeepMind Control Suite.It also substantially outperforms DVRL and achieves comparable or slightly better performance than DrQ.
  • Comparative Evaluation: SLAC exceeds PlaNet on all four DeepMind Control Suite tasks and outperforms PlaNet by a large margin on harder image-based OpenAI Gym tasks.On the Gym tasks, prior methods generally performed poorly despite considerable hyperparameter tuning.
  • Ablation Experiments: The ablations compare latent-variable models, actor and critic inputs, model pretraining, and training updates per iteration within SLAC.Additional results extend these comparisons to five further tasks.
  • Ablation Experiments: The fully stochastic latent model outperforms alternatives, while temporal dependencies provide the largest improvement, followed by latent factorization and full stochasticity.Critics perform significantly worse with history inputs than latent samples, pretraining accelerates learning, and two updates per iteration provide the best tradeoff.

7 Conclusion

SLAC combines representation learning with task learning in a stochastic latent space, achieving improved sample efficiency and final performance over prior model-based and model-free methods.

  • SLAC combines off-policy model-free RL with representation learning through a sequential stochastic state-space model.
  • Representation learning and task learning jointly improve sample efficiency and final task performance compared with prior model-based and model-free RL methods.
  • A natural extension would use model predictions to generate synthetic samples in a mixed model-based and model-free method.
  • Explicit representation learning may support faster training, more complex tasks, and reuse or transfer of learned representations across tasks.

Broader Impact

The paper connects image-based RL to broader robotics and societal applications while acknowledging unresolved safety and explainability concerns.

  • Learning directly from vision could reduce reliance on precise and expensive sensor setups in robotic systems.
  • Standard end-to-end RL often produces brittle, hyperparameter-sensitive, slow, and inefficient solutions when learning representations and tasks together.
  • More efficient and easier-to-train deep RL could support wider use in real-world robotics.
  • The work does not explicitly address safety concerns or explainability, although safety controllers could be used alongside it.

A Derivation of the Evidence Lower Bound and SLAC Objectives

The appendix derives SLAC by applying variational inference to control as inference in a POMDP, decomposing the ELBO into model and maximum-entropy RL objectives.

  • Control as inference casts optimal-policy learning as inference conditioned on optimal behavior in a POMDP.
  • SLAC maximizes the marginal likelihood of past observations and future optimality variables, adding observation likelihood to supervise latent representations.
  • For future steps, the variational distribution uses the stochastic dynamics to prevent optimistic action selection.
  • The maximum-entropy RL objective is expressed through Q-function messages and optimized using Bellman-residual learning.
  • The resulting optimal policy is optimal under the belief-averaged Q value and is equivalent to the Q-MDP heuristic.
  • The ELBO decomposes into a model objective and a maximum-entropy RL objective, yielding model, policy, and critic losses.

B Latent Variable Factorization and Network Architectures

SLAC uses a fully stochastic sequential latent-variable model whose factorization and network design balance expressivity, trainability, and useful representations for RL.

  • The latent variable is factorized into two stochastic variables to obtain more expressive latent distributions and share parts of the prior and posterior.
  • The model uses a generative process and an inference model, represented respectively by solid and dashed arrows in Figure 7.
  • The generative model comprises an initial standard-normal latent distribution and neural-network-parameterized conditional distributions.
  • The image-model networks use convolutional and fully connected layers with shared convolutional parameters across distributions.
  • The inference model is factorized so that one transition distribution exactly matches the generative model, removing that term from the ELBO KL divergence.
  • The actor uses a Gaussian layer and tanh bijector to constrain actions to the bounded space [−1, 1].

C Training and Evaluation Details

The experiments use task-specific training and evaluation settings, including action repeats, stochastic policies, and multiple trials with averaged returns.

  • Optimization: The control algorithm follows SAC hyperparameters but uses a 100000-step replay buffer because image observations require more memory.
  • Optimization: DeepMind Control Suite uses 1 gradient step per environment step, while OpenAI Gym uses 3.The Q-function and policy use learning rate 0.0003 and batch size 256; model training uses 0.0001 and batch size 32.
  • Environment interaction: Action repeats reduce the actual samples used relative to plotted benchmark environment steps.For example, 1 million plotted cheetah steps correspond to 250000 samples with an action repeat of 4.
  • Evaluation: The policy uses the same stochastic policy for behavior and evaluation because deterministic greedy evaluation was comparable or worse.
  • Evaluation: Results average returns over 10 evaluation episodes across 10 DeepMind Control Suite trials and 5 OpenAI Gym trials.
  • Evaluation: The DeepMind Control Suite experiments select the best per-task pixel-decoder variance from σ2 ∈ {0.04, 0.1, 0.4}.

D Ablation Experiments

The ablations show that temporal and stochastic structure in the latent model is important, while several training and decoder choices are task-sensitive. Pretraining helps learning, but excessive updates and generic augmentation can hurt or fail to help.

  • Pixel decoder: Decoder variance σ2 = 0.1 performs well overall, while walker walk prefers σ2 = 0.4 and ball-in-cup catch prefers σ2 = 0.04.The variance controls the relative weighting of reconstruction loss and KL-divergence, and the best setting depends on task complexity.
  • Data augmentation: Random cropping improves learning only for reacher easy; it does not improve performance on the other evaluated tasks.
  • Latent variable model: Temporal dependencies provide the largest latent-model improvement, followed by autoregressive factorization and a fully stochastic model.
  • Actor and critic inputs: Using the latent sample as critic input generally outperforms using observation-action history, while actor-input choice is usually indifferent.The exceptions are cartpole swingup and reacher easy.
  • Model pretraining: Pretraining with random data before task learning accelerates learning and can improve asymptotic performance.Little or no pretraining produces slower learning and, in some cases, worse asymptotic performance.
  • Training updates: More training updates per iteration slightly speed learning, but too many increase trial variance and can slightly worsen asymptotic performance.
  • Latent-model predictions: Generated prior samples use no ground-truth image frames after initialization, whereas conditional-prior samples use the true first frame.Posterior samples reconstruct sequences from latents sampled using encoded observations.
Loading 1907.00953v4…