Guiding Representation Learning in Deep Generative Models with Policy Gradients

1 INTRODCTION . Reinforcement Learning ( RL ) gained much popularity in recent years by outperforming humans in games such as Atari ( Mnih et al . ( 2015 ) ) , Go ( Silver et al . ( 2016 ) ) and Starcraft 2 ( Vinyals et al . ( 2017 ) ) . These results were facilitated by combining novel machine learning techniques such as deep neural networks ( LeCun et al . ( 2015 ) ) with classical RL methods . The RL framework has shown to be quite flexible and has been applied successfully in many further domains , for example , robotics ( Andrychowicz et al . ( 2020 ) ) , resource management ( Mao et al . ( 2016 ) ) or physiologically accurate locomotion ( Kidziński et al . ( 2018 ) ) . The goal of representation learning is to learn a suitable representation for a given application domain . Such a representation should contain useful information for a particular downstream task and capture the distribution of explanatory factors ( Bengio et al . ( 2013 ) ) . Typically , the choice of a downstream task influences the choice of method for representation learning . While Generative Adversarial Network ( GAN ) s are frequently used for tasks that require high-fidelity reconstructions or generation of realistic new data , auto-encoder based methods have been more common in RL . Recently , many such approaches employed the Variational Auto Encoder ( VAE ) ( Kingma & Welling ( 2013 ) ) framework which aims to learn a smooth representation of its domain . Most of these approaches follow the same pattern : First , they build a dataset of states from the RL environment . Second , they train the VAE on this static dataset and lastly train the RL mode using the VAE ’ s representation . While this procedure generates sufficiently good results for certain scenarios , there are some fundamental issues with this method . Such an approach assumes that it is possible to collect enough data and observe all task-relevant states in the environment without knowing how to act in it . As a consequence , when learning to act the agent will only have access to a representation that is optimized for the known and visited states . As soon as the agent becomes more competent , it might experience novel states that have not been visited before and for which there is no good representation ( in the sense that the experienced states are out of the original learned distribution and the mapping is not appropriate ) . Another issue arises from the manner the representation is learned . Usually , the VAE is trained in isolation , so it decides what features are learned based on its own objective function and not on what is helpful for the downstream task . Mostly , such a model is tuned for good reconstruction . Without the information from the RL model , such a representation does not reflect what is important for the downstream task . As a consequence , the VAE might omit learning features that are crucial for good performance on the task because they appear negligible with respect to reconstruction ( Goodfellow et al . ( 2016 ) , Chapter 15 , Figure 15.5 ) . For example , small objects in pixel-space are ignored as they affect a reconstruction based loss only marginally . Thus , any downstream task using such a representation will have no access to information about such objects . A good example for such a task is Atari Breakout , a common RL benchmark . Figures 1a and 1b show an original Breakout frame and its reconstruction . While the original frame contains the ball in the lower right hand corner , this crucial feature is missing completely in the reconstruction . We approach this issue through simultaneously learning representation and RL task , that is by combining the training of both models . As an advantage , this abolishes the need of collecting data before knowing the environment as it combines VAE and RL objectives . In consequence the VAE has an incentive to represent features that are relevant to the RL model . The main contributions of this paper are as follows : First we show that combined learning is possible and that it yields good performing policies . Second , we show that jointly trained representations incorporate additional , task-specific information which allows a RL agent to achieve higher rewards then if it was trained on a static representation . This will be shown indirectly by comparing achieved rewards as well as directly through an analysis of the trained model and its representation . 2 RELATED WORK . Lange & Riedmiller ( 2010 ) explored Auto Encoder ( AE ) ( Lecun ( 1987 ) ; Bourlard & Kamp ( 1988 ) ; Hinton & Zemel ( 1994 ) ) as a possible pre-processor for RL algorithms . The main focus in their work was finding good representations for high dimensional state spaces that enables policy learning . As input , rendered images from the commonly used grid world environment were used . The agent had to manoeuvre through a discretized map using one of four discrete movement actions per timestep . It received a positive reward once reaching the goal tile and negative rewards elsewhere . The AE bottleneck consisted only of two neurons , which corresponds to the dimensionality of the environemnt ’ s state . Fitted Q-Iteration ( FQI ) ( Ernst et al . ( 2005 ) ) was used to estimate the Q-function , which the agent then acted -greedy upon . Besides RL , they also used the learned representation to classify the agents position given an encoding using a Multi-Layer Perceptron ( MLP ) ( Rumelhart et al . ( 1985 ) ) . For these experiments , they found that adapting the encoder using MLP gradients lead to an accuracy of 99.46 % . However , they did not apply this approach to their RL task . A compelling example for separate training of meaningful representation is provided by Higgins et al . ( 2017b ) who proposed a framework called DARLA . They trained RL agents on the encoding of a β-VAE ( Higgins et al . ( 2016 ) ; Higgins et al . ( 2017a ) ) with the goal of zero-shot domain transfer . In their approach , β-VAE and agent were trained separately on a source domain and then evaluated in a target domain . Importantly , source and target domain are similar to a certain extent and only differ in some features , e.g . a blue object in the source domain might be red in the target domain . During training of the β-VAE , the pixel-based reconstruction loss was replaced with a loss calculated in the latent space of a Denoising Auto Encoder ( DAE ) ( Vincent et al . ( 2008 ) ) . Thereby their approach avoids missing task relevant feature encodings at the cost of training another model . For one of their evaluation models , they allowed the RL gradients to adapt the encoder . Their results show that subsequent encoder learning improves performance of Deep Q-Learning ( DQN ) but decreases performance of Asynchronous Advantage Actor-Critic ( A3C ) ( Mnih et al . ( 2016 ) ) . Ha & Schmidhuber ( 2018 ) proposed a combination of VAE , Recurrent Neural Networks ( RNN ) ( Hochreiter & Schmidhuber ( 1997 ) ) and a simple policy as a controller . They hypothesized that by learning a good representation of the environment and having the ability to predict future states , learning the policy itself becomes a trivial task . Like in most other models , the VAE was pre-trained on data collected by a random policy . Only the RNN and the controller were trained online . The compressed representation from the VAE was passed into a RNN in order to estimate a probability density for the subsequent state . The controller was deliberately chosen as a single linear layer and could thus be optimized with Covariance Matrix Adaptation - Evolution Strategy ( CMA-ES ) ( Hansen ( 2006 ) ) . This work demonstrated how a VAE can provide a versatile representation that can be utilized in reinforcement learning . In addition , such an approach allows to predict the subsequent encoded state . While these findings encourage the usage of VAE in conjunction with RL , this is only possible in environments where the state space can be explored sufficiently by a random policy . However , if the policy can only discover important features after acquiring a minimal level of skill , sampling the state space using a random policy will not yield high-performing agents . Learning such features would only be possible if the VAE is continuously improved during policy training . Another interesting combination of VAEs and RL was recently proposed by Yang et al . ( 2019 ) , with their so called Action-Conditional Variational Auto-Encoder ( AC-VAE ) . Their motivation for creating this model was to train a transparent , interpretable policy network . Usually , the β-VAEs decoder is trained to reconstruct the input based on the representation the encoder produced . In this work though , the decoders objective was to predict the subsequent state st+1 . As input it got the latent space vector z combined with an action-mapping-vector , which is the action vector at with a zero-padding to match the latent spaces dimensionality . Inspecting the decoder estimates for st+1 when varying one dimension of the latent space showed , that each dimension encoded a possible subsequent state that is likely to be encountered if the corresponding action from this dimension was taken . Unfortunately , the authors did not report any rewards they achieved on Breakout , hence it was not possible for us to compare model performances . 3 COMBINATION OF REINFORCEMENT AND REPRESENTATION LEARNING OBJECTIVES . In this section , we will first revisit the fundamentals of RL and VAEs and discuss their different objective functions . Then , we propose a joint objective function that allows for joint training of both models using gradient descent based learning methods . 3.1 REINFORCEMENT LEARNING WITH POLICY OPTIMIZATION . RL tries to optimize a Markov Decision Process ( MDP ) ( Bellman ( 1957 ) ) that is given by the tuple 〈S , A , r , p , γ〉 . S denotes the state space , A the action space and p : S × R × S × A → [ 0 , 1 ] the environment ’ s dynamics function that , provided a state-action pair , gives the state distribution for the successor state . r : S × A → R is the reward and γ ∈ [ 0 , 1 ) the scalar discount factor . The policy πθ ( a|s ) is a stochastic function that gives a probability distribution over actions for state s. θ denotes the policy ’ s parameter vector which is typically subject to optimization . A trajectory τ = ( s0 , a0 , ... , sT , aT ) consisting of an alternating sequence of states and actions can be sampled in the environment , where T stands for the final timestep of the trajectory and ai ∼ πθ ( ai|si ) . The overarching goal of RL is to find a policy that maximizes the average collected reward over all trajectories . This can be expressed as the optimization problem maxEτ∼p ( τ ) [ ∑ t r ( s , a ) ] , which can also be written in terms of an optimal policy parameter vector θ∗ = argmaxθ Eτ∼p ( τ ) [ ∑ t r ( s , a ) ] . When trying to optimize the policy directly be searching for θ∗ , policy optimization algorithms like A3C , Actor-Critic with Experience Replay ( ACER ) ( Wang et al . ( 2016a ) ) , Trust Region Policy Optimization ( TRPO ) ( Schulman et al . ( 2015a ) ) or Proximal Policy Optimization ( PPO ) ( Schulman et al . ( 2017 ) ) are commonly used . The fundamental idea behind policy optimization techniques is to calculate gradients of the RL objective with respect to the policy parameters : ∇θJ ( θ ) = E τ∼p ( τ ) [ ∇θ logπθ ( τ ) r ( τ ) ] ( 1 ) where we defined ∑T t=0 r ( s , a ) = r ( τ ) for brevity . However , most policy optimization methods introduce heavy modifications to this vanilla gradient in order to achieve more stable policy updates . Throughout our work , we have used PPO as RL algorithm because it is quite sample efficient and usually produces stable policy updates . For an in-depth description of PPO , we refer to our A.1 or the original work Schulman et al . ( 2017 ) .

Summary: This paper proposes a method for reinforcement learning with representations learned from a VAE. The VAE is used to encode the states (images) and the mean of the posterior is used as input to the policy. The VAE is trained using the variational lower bound and the policy is optimized using PPO. The model can be jointly trained with VAE and RL losses from scratch, or the VAE can be pre-trained and fixed, or pre-trained and finetuned.