Predictive Coding for Boosting Deep Reinforcement Learning with Sparse Rewards

1 INTRODUCTION . Recent progress in deep reinforcement learning ( DRL ) has enabled robots to learn and execute complex tasks , ranging from game playing ( Jaderberg et al. , 2018 ; OpenAI , 2019 ) , robotic manipulations ( Andrychowicz et al. , 2017 ; Haarnoja et al. , 2018 ) , to navigation ( Zhang et al. , 2017 ) . However , in many scenarios learning depends heavily on meaningful and frequent feedback from the environment for the agent to learn and correct behaviors . As a result , reinforcement learning ( RL ) problems with sparse rewards still remain a difficult challenge ( Riedmiller et al. , 2018 ; Agarwal et al. , 2019 ) . In a sparse reward setting , the agent typically explores without receiving any reward , until it enters a small subset of the environment space ( the ” goal ” ) . Due to lack of frequent feedback from the environment , learning in sparse reward problems is typically hard , and heavily relies on the agent entering the ” goal ” during exploration . A possible way to tackle this is through reward shaping ( Devlin & Kudenko , 2012 ; Zou et al. , 2019 ; Gao & Toni , 2015 ) , where manually designed rewards are added to the environment to guide the agent towards finding the ” goal ” ; however , this approach often requires domain knowledge of the environment , and may bias learning if the shaped rewards are not robust ( Ng et al. , 1999 ) . RL problems often benefit from representation learning ( Bengio et al. , 2013 ) , which studies the transformation of raw observations of an environment ( sensors , images , coordinates etc ) into a more meaningful form , such that the agent can more easily extract information useful for learning . Intuitively , raw states contain redundant or irrelevant information about the environment , which the agent must take time to learn to distinguish and remove ; representation learning directly tackles this problem by either eliminating redundant dimensions ( Kingma & Welling , 2013 ; van den Oord et al. , 2017 ) or emphasizing more useful elements of the state ( Nachum et al. , 2018a ) . Much of the prior work on representation learning focuses on generative approaches to model the environment , but some recent work also studies optimizations that learn important features ( Ghosh et al. , 2018 ) . In this paper , we tackle the challenge of DRL to solve sparse reward tasks : we apply representation learning to provide the agent meaningful rewards without the need for domain knowledge . In particular , we propose to use predictive coding in an unsupervised fashion to extract features that maximize the mutual information ( MI ) between consecutive states in a state trajectory . These predictive features are expected to have the potential to simplify the structure of an environment ’ s state space : they are optimized to both summarize the past and predict the future , capturing the most important elements of the environment dynamics . We show this method is useful for model-free learning from either raw states or images , and can be applied on top of any general deep reinforcement learning algorithms such as PPO ( Schulman et al. , 2017 ) . Although MI has traditionally been difficult to compute , recent advances have suggested optimizing on a tractable lower bound on the quantity ( Hjelm et al. , 2018 ; Belghazi et al. , 2018 ; Oord et al. , 2018 ) . We adopt one such method , Contrastive Predictive Coding ( Oord et al. , 2018 ) , to extract features that maximize MI between consecutive states in trajectories collected during exploration ( one thing worth noting is that this method is not restricted to specific predictive coding schemes such as CPC ) . Such features are then used for simple reward shaping in representation space to provide the agent better feedback in sparse reward problems . We demonstrate the validity of our method through extensive numerical simulations in a wide range of control environments such as maze navigation , robot locomotion , and robotic arm manipulation ( Figure 4 ) . In particular , we show that using these predictive features , we provide reward signals as effective for learning as handshaped rewards , which encode domain and task knowledge . This paper is structured as follows : We begin by providing preliminary information in Section 2 and discussing relevant work in Section 3 ; then , we explain and illustrate the proposed method in Section 4 ; lastly , we present experiment results in Section 5 , and conclude the paper by discussing the results and pointing out future work in Section 6 . 2 PRELIMINARIES . Reinforcement Learning and Reward Shaping : This paper assumes a finite-horizon Markov Decision Process ( MDP ) ( Puterman , 1994 ) , defined by a tuple ( S , A , P , r , γ , T ) . Here , S ∈ Rd denotes the state space , A ∈ Rm denotes the action space , P : S ×A×S → R+ denotes the state transition distribution , r : S × A → R denotes the reward function , γ ∈ [ 0 , 1 ] is the discount factor , and finally T is the horizon . At each step t , the action at ∈ A is sampled from a policy distribution πθ ( at|st ) where s ∈ S and θ is the policy parameter . After transiting into the next state by sampling from p ( st+1|at , st ) , where p ∈ P , the agent receives a scalar reward r ( st , at ) . The agent continues performing actions until it enters a terminal state or t reaches the horizon , by when the agent has completed one episode . We let τ denote the sequence of states that the agent enters in one episode . With such definition , the goal of RL is to learn a policy πθ∗ ( at|st ) that maximizes the expected discounted reward Eπ , P [ R ( τ0 : T−1 ) ] = Eπ , P [ ∑T−1 0 γ tr ( st , at ) ] , where expectation is taken on the possible trajectories τ and starting states x0 . In this paper , we assume model-free learning , meaning the agent does not have access to P . Reward shaping essentially replaces the original MDP with a new one , whose reward function is now r′ ( st , at ) . In this paper , reward shaping is done to train a policy πr′ that maximizes the expected discounted reward in the original MDP , i.e . Eπr′ , P [ R ( τ0 : T−1 ) ] . Mutual Information and Predictive Coding : Mutual information measures the amount of information obtained about one random variable after observing another random variable ( Cover & Thomas , 2012 ) . Formally , given two random variables X and Y with joint distribution p ( x , y ) and marginal densities p ( x ) and p ( y ) , their MI is defined as the KL-divergence between joint density and product of marginal densities : MI ( X ; Y ) = DKL ( p ( x , y ) ‖p ( x ) p ( y ) ) = Ep ( x , y ) [ log p ( x , y ) p ( x ) p ( y ) ] . ( 1 ) Predictive coding in this paper aims to maximize the MI between consecutive states in the same state trajectory . As MI is difficult to compute , we adopt the method of optimizing on a lower bound , InfoNCE ( Oord et al. , 2018 ) , which takes the current context ct to predict a future state st+k : MI ( st+k ; ct ) ≥ ES [ log f ( zt+k , ct ) f ( zt+k , ct ) + ∑ sj∈S f ( zj , ct ) ] ( 2 ) Here , f ( x , y ) is optimized through cross entropy to model a density ratio : f ( x , y ) ∝ P ( x|y ) P ( x ) . zt+k is the embedding of state xt+k by the encoder , and ct is obtained by summarizing the embeddings of previous n states in a segment of a trajectory , zt−n+1 : t , through a gated recurrent unit ( Cho et al. , 2014 ) . Intuitively , the context ct pays attention to the evolution of states in order to summarize the past and predict the future ; thus , it forces the encoder to extract only the essential dynamical elements of the environment , elements that encapsulate state evolution . 3 RELEVANT WORKS Our paper uses the method of Contrastive Predictive Coding ( CPC ) ( Oord et al. , 2018 ) , which includes experiments in the domain of RL . In the CPC paper , the InfoNCE is applied to the LSTM component ( Hochreiter & Schmidhuber , 1997 ) of an A2C architecture ( Mnih et al. , 2016 ; Espeholt et al. , 2018 ) . The LSTM maps every state observation to an embedding , which is then directly used for learning . This differs from our approach , where we train on pre-collected trajectories to obtain embeddings , and only use these embeddings to provide rewards to the agent , which still learns on the raw states . Our approach has two main advantages : 1 ) Preprocessing states allows us to collect exploration-focused trajectories , and obtain embeddings that are suitable for multi-tasking . 2 ) Using embeddings to provide rewards is more resistant to noises in embeddings than using them as training features , since in the former case we care more about the accumulation of rewards across multiple states , where the noises are diluted . Applying representation learning to RL has been studied in many prior works ( Nachum et al. , 2018a ; Ghosh et al. , 2018 ; Oord et al. , 2018 ; Caselles-Dupré et al. , 2018 ) . In a recent paper on actionable representation ( Ghosh et al. , 2018 ) , representation learning is also applied to providing the agent useful reward signals . In actionable representation paper , states are treated as goals , and embeddings are optimized in a way that the distance between two states reflects the difference between the policies required to reach them . This is fundamentally different from our approach , which aims to extract features that have predictive qualities . Furthermore , computing actionable representation requires trained goal-conditioned policies as a part of the optimization , which is a strict requirement , while this paper aims to produce useful representations without needing access to trained policies . Lastly , VQ-VAE ( van den Oord et al. , 2017 ) is a generative approach that provides a principled way of extracting low-dimensional features . In contrast to VAE ( Kingma & Welling , 2013 ) , it outputs a discrete codebook , and the prior distribution is learned rather than static . VQ-VAE could be useful for removing redundant information from raw states , which may speed up learning ; however , since the goal of VQ-VAE is reconstruction , it does not put emphasis on features that are particularly useful for learning , nor does it attempt to understand the environment dynamics across long segments of states . Our use of predictive coding is thus be a better fit for reinforcement learning , as we emphasize on features that help understand the evolution of states rather than reconstruct each individual state .

Summary:   This paper proposes using the features learned through Contrastive Predictive Coding as a means for reward shaping. Specifically, they propose to cluster the embedding using the clusters to provide feedback to the agent by applying a positive reward when the agent enters the goal cluster. In more complex domains they add another negative distance term of the embedding of the current state and goal state. Finally, they provide empirical evidence of their algorithm working in toy domains (such as four rooms and U-maze) as well as a set of control environments including AntMaze and Pendulum.