Improving Sample Efficiency in Model-Free Reinforcement Learning from Images

1 INTRODUCTION . Cameras are a convenient and inexpensive way to acquire state information , especially in complex , unstructured environments , where effective control requires access to the proprioceptive state of the underlying dynamics . Thus , having effective RL approaches that can utilize pixels as input would potentially enable solutions for a wide range of real world problems . The challenge is to efficiently learn a mapping from pixels to an appropriate representation for control using only a sparse reward signal . Although deep convolutional encoders can learn good representations ( upon which a policy can be trained ) , they require large amounts of training data . As existing reinforcement learning approaches already have poor sample complexity , this makes direct use of pixel-based inputs prohibitively slow . For example , model-free methods on Atari ( Bellemare et al. , 2013 ) and DeepMind Control ( DMC ) ( Tassa et al. , 2018 ) take tens of millions of steps ( Mnih et al. , 2013 ; Barth-Maron et al. , 2018 ) , which is impractical in many applications , especially robotics . A natural solution is to add an auxiliary task with an unsupervised objective to improve sample efficiency . The simplest option is an autoencoder with a pixel reconstruction objective . Prior work has attempted to learn state representations from pixels with autoencoders , utilizing a two-step training procedure , where the representation is first trained via the autoencoder , and then either with a policy learned on top of the fixed representation ( Lange & Riedmiller , 2010 ; Munk et al. , 2016 ; Higgins et al. , 2017b ; Zhang et al. , 2018 ; Nair et al. , 2018 ) , or with planning ( Mattner et al. , 2012 ; Finn et al. , 2015 ) . This allows for additional stability in optimization by circumventing dueling training objectives but leads to suboptimal policies . Other work utilizes end-to-end model-free learning with an auxiliary reconstruction signal in an on-policy manner ( Jaderberg et al. , 2017 ) . We revisit the concept of adding an autoencoder to model-free RL approaches , but with a focus on off-policy algorithms . We perform a sequence of careful experiments to understand why previous approaches did not work well . We found that a pixel reconstruction loss is vital for learning a good representation , specifically when trained end-to-end . Based on these findings , we propose a simple autoencoder-based off-policy method that can be trained end-to-end . Our method is the first modelfree off-policy algorithm to successfully train simultaneously both the latent state representation and policy in a stable and sample-efficient manner . 1An anonymous website with code , results , and videos : https : //sites.google.com/view/sac-ae/home Of course , some recent state-of-the-art model-based RL methods ( Hafner et al. , 2018 ; Lee et al. , 2019 ) have demonstrated superior sample efficiency to leading model-free approaches on pixel tasks from ( Tassa et al. , 2018 ) . But we find that our model-free , off-policy , autoencoder-based approach is able to match their performance , closing the gap between model-based and model-free approaches in image-based RL , despite being a far simpler method that does not require a world model . This paper makes three main contributions : ( i ) a demonstration that adding a simple auxiliary reconstruction loss to a model-free off-policy RL algorithm achieves comparable results to state-of-the-art model-based methods on the suite of continuous control tasks from Tassa et al . ( 2018 ) ; ( ii ) an understanding of the issues involved with combining autoencoders with model-free RL in the off-policy setting that guides our algorithm ; and ( iii ) an open-source PyTorch implementation of our simple method for researchers and practitioners to use as a strong baseline that may easily be built upon . 2 RELATED WORK . Efficient learning from high-dimensional pixel observations has been a problem of paramount importance for model-free RL . While some impressive progress has been made applying model-free RL to domains with simple dynamics and discrete action spaces ( Mnih et al. , 2013 ) , attempts to scale these approaches to complex continuous control environments have largely been unsuccessful , both in simulation and the real world . A glaring issue is that the RL signal is much sparser than in supervised learning , which leads to sample inefficiency , and higher dimensional observation spaces such as pixels worsens this problem . One approach to alleviate this problem is by training with auxiliary losses . Early work ( Lange & Riedmiller , 2010 ) explores using deep autoencoders to learn feature spaces in visual reinforcement learning , crucially Lange & Riedmiller ( 2010 ) propose to recompute features for all collected experiences after each update of the autoencoder , rendering this approach impractical to scale to more complicated domains . Moreover , this method has been only demonstrated on toy problems . Alternatively , Finn et al . ( 2015 ) apply deep autoencoder pretraining to real world robots that does not require iterative re-training , improving upon computational complexity of earlier methods . However , in this work the linear policy is trained separately from the autoencoder , which we find to not perform as well as end-to-end methods . Shelhamer et al . ( 2016 ) use auxiliary losses in Atari that incorporate forward and inverse dynamics with A3C , an on-policy algorithm . They recommend a multi-task setting and learning dynamics and reward to find a good representation , which relies on the assumption that the dynamics in the task are easy to learn and useful for learning a good policy . Jaderberg et al . ( 2017 ) propose to use unsupervised auxiliary tasks , both observation-based and reward-based based off of real world inductive priors , and show improvements in Atari , again in the on-policy regime , which is much more stable for learning . Unfortunately , this work also relies on inductive biases by designing internal rewards to learn a good representation which is hard to scale to the real world problems . Higgins et al . ( 2017b ) ; Nair et al . ( 2018 ) use a beta variational autoencoder ( β-VAE ) ( Kingma & Welling , 2013 ; Higgins et al. , 2017a ) and attempt to extend unsupervised representation pretraining to the off-policy setting , but find it hard to perform end-to-end training , thus receding to the iterative retraining procedure ( Lange & Riedmiller , 2010 ; Finn et al. , 2015 ) . There has been more success in using model-based methods on images , such as Hafner et al . ( 2018 ) ; Lee et al . ( 2019 ) . These methods use a world model ( Ha & Schmidhuber , 2018 ) approach , learning a representation space using a latent dynamics loss and pixel decoder loss to ground on the original observation space . These model-based reinforcement learning methods often show improved sample efficiency , but with the additional complexity of balancing various auxiliary losses , such as a dynamics loss , reward loss , and decoder loss in addition to the original policy and value optimiza- tions . These proposed methods are correspondingly brittle to hyperparameter settings , and difficult to reproduce , as they balance multiple training objectives . To close the gap between model-based and model-free image-based RL in terms of sample efficiency and sidestep the issues of model learning , our goal is to train a model-free off-policy algorithm with auxiliary reconstruction loss in a stable manner . 3 BACKGROUND . A fully observable Markov decision process ( MDP ) is described by tuple 〈S , A , P , R , γ〉 , where S is the state space , A is the action space , P ( st+1|st , at ) is the probability distribution over transitions , R ( st , at , st+1 ) is the reward function , and γ is the discount factor ( Bellman , 1957 ) . An agent starts in a initial state s1 sampled from a fixed distribution p ( s1 ) , then at each timestep t it takes an action at ∈ A from a state st ∈ S and moves to a next state st+1 ∼ P ( ·|st , at ) . After each action the agent receives a reward rt = R ( st , at , st+1 ) . We consider episodic environments with the length fixed to T . The goal of standard RL is to learn a policy π ( at|st ) that can maximize the agent ’ s expected cumulative reward ∑T t=1 E ( st , at ) ∼ρπ [ rt ] , where ρπ is a state-action marginal distribution induced by the policy π ( at|st ) and transition distribution P ( st+1|st , at ) . An important modification ( Ziebart et al. , 2008 ) auguments this objective with an entropy termH ( π ( ·|st ) ) to encourage exploration and robustness to noise . The resulting maximum entropy objective is then defined as : π∗ = arg max π T∑ t=1 E ( st , at ) ∼ρπ [ rt + αH ( π ( ·|st ) ) ] , where α is a temperature parameter that balances between optimizing for the reward and for the stochasticity of the policy . We build on Soft Actor-Critic ( SAC ) ( Haarnoja et al. , 2018 ) , an off-policy actor-critic method that uses the maximum entropy framework to derive soft policy iteration . At each iteration SAC performs a soft policy evaluation step and a soft policy improvement step . The soft policy evaluation step fits a parametric soft Q-function Q ( st , at ) ( critic ) by minimizing the soft Bellman residual : J ( Q ) = E ( st , at , rt , st+1 ) ∼D [ ( Q ( st , at ) − rt − γEat+1∼π [ Q̄ ( st+1 , at+1 ) − α log π ( at+1|st+1 ) ] ) 2 ] , ( 1 ) where D is the replay buffer , and Q̄ is the target soft Q-function parametrized by a weight vector obtained using the exponentially moving average of the soft Q-function weights to stabilize training . The soft policy improvement step then attempts to learn a parametric policy π ( at|st ) ( actor ) by directly minimizing the KL divergence between the policy and a Boltzmann distribution induced by the current soft Q-function , producing the following objective : J ( π ) = Est∼D [ Eat∼π [ α log ( π ( at|st ) ) −Q ( st , at ) ] ] . ( 2 ) The policy π ( at|st ) is parametrized as a diagonal Gaussian to handle continuous action spaces . When learning from raw images , we deal with the problem of partial observability , which is formalized by a partially observable MDP ( POMDP ) . In this setting , instead of getting a low-dimensional state st ∈ S at time t , the agent receives a high-dimensional observation ot ∈ O , which is a rendering of potentially incomplete view of the corresponding state st of the environment ( Kaelbling et al. , 1998 ) . This complicates applying RL as the agent now needs to also learn a compact latent representation to infer the state . Fitting a high-capacity encoder using only a scarce reward signal is sample inefficient and prone to suboptimal convergence . Following prior work ( Lange & Riedmiller , 2010 ; Finn et al. , 2015 ) we explore unsupervised pretraining via an image-based autoencoder . In practice , the autoencoder is represented as a convolutional encoder fenc that maps an image observation ot to a low-dimensional latent vector zt , and a deconvolutional decoder fdec that reconstructs zt back to the original image ot . The optimization is done by minimizing the standard reconstruction objective : J ( AE ) = Eot∼D [ 1 2 ||fdec ( zt ) − ot||22 ] where zt = fenc ( ot ) . ( 3 ) Or in the case of β-VAE ( Kingma & Welling , 2013 ; Higgins et al. , 2017a ) , where the variational distribution is parametrized as diagonal Gaussian , the objective is defined as : J ( VAE ) = Eot∼D [ 1 2 ||fdec ( zt ) − ot||22 − β 2 ∑ i ( 1 + logσ2t , i − z2t , i − σ2t , i ) ] , ( 4 ) where zt = fenc ( ot ) and σ2t = fenc std ( ot ) . The latent vector zt is then used by an RL algorithm , such as SAC , instead of the unavailable true state st. To infer temporal statistics , such as velocity and acceleration , it is common practice to stack three consecutive frames to form a single observation ( Mnih et al. , 2013 ) . We emphasize that in contrast to model-based methods ( Ha & Schmidhuber , 2018 ; Hafner et al. , 2018 ) , we do not predict future states and solely focus on learning representations from the current observation to stay model-free . 4 A DISSECTION OF LEARNING STATE REPRESENTATIONS WITH β-VAE In this section we explore in a systematic fashion how model-free off-policy RL can be made to train directly from pixel observations . We start by noting a dramatic performance drop when SAC is trained on pixels instead of proprioceptive state ( Section 4.2 ) in the off-policy regime . This result motivates us to explore different ways of employing auxiliary supervision to speed up representation learning . While a wide range of auxiliary objectives could be added to aid effective representation learning , for simplicity we focus our attention on autoencoders . We follow Lange & Riedmiller ( 2010 ) ; Finn et al . ( 2015 ) and in Section 4.3 try an iterative unsupervised pretraining of an autoencoder that reconstructs pixels and is parameterized by β-VAE as per Nair et al . ( 2018 ) ; Higgins et al . ( 2017a ) . Exploring the training procedure used in previous work shows it to be sub-optimal and points towards the need for end-to-end training of the β-VAE with the policy network . Our investigation in Section 4.4 renders this approach useless due to severe instability in training , especially with larger β values . We resolve this by using deterministic forms of the variational autoencoder ( Ghosh et al. , 2019 ) and a careful learning procedure . This leads to our algorithm , which is described and evaluated in Section 5 .

Summary: This paper proposes an approach to make the model-free state-of-the-art soft actor-critic (SAC) algorithm for proprioceptive state spaces sample-efficient in higher-dimensional visual state spaces. To this end, an encoder-decoder structure to minimize image reconstruction loss is added to SAC's learning objectives. Importantly, the encoder is shared between the encoder-decoder architecture, the critic and the policy. Furthermore, Q-critic updates backpropagate through the encoder such that encoder weights need to trade off image reconstruction and critic learning. The approach is evaluated on six tasks from the DeepMind control suite and compared against proprioceptive SAC, pixel-based SAC, D4PG as well as to the model-based baselines PlaNet and SLAC. The proposed method seems to achieve results competitive with the model-based baselines and significantly improves over raw pixel-based SAC. Further ablation studies are presented to investigate the information capacity of the learned latent representation and generalization to unseen tasks.