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: The paper focused on the issue of learning a policy for a given task using the learned representations a pre-trained VAE. The authors visualize that using a learned latent space of a pre-trained VAE is not good enough for learning policies and propose a solution for this problem: back-propagate gradient policies through the VAE encoder. The authors proposed two versions on this method, one with pre-training and one fully online.