Self-Attentional Credit Assignment for Transfer in Reinforcement Learning

1 INTRODUCTION . To some , intelligence is measured as the capability of transferring knowledge to unprecedented situations . While the notion of intellect itself is hard to define , the ability to reuse learned information is a desirable trait for learning agents . The coffee test ( Goertzel et al. , 2012 ) , presented as a way to assess general intelligence , suggests the task of making coffee in a completely unfamiliar kitchen . It requires a combination of advanced features ( planning , control and exploration ) that would make the task very difficult if not out of scope for the current state-of-the-art Reinforcement Learning ( RL ) agents to learn . On the other hand , it is solved trivially by humans , who exploit the universally invariant structure of coffee-making : one needs to fetch a mug , find coffee , power the coffee machine , add water and launch the brewing process by pushing the adequate buttons . Thus , to solve the coffee test , transfer learning appears necessary . Were we to possess a random kitchen simulator and a lot of compute , current transfer methods would still fall short of consistently reusing structural information about the task , hence also falling short of efficient adaptation . Credit assignment , which in RL refers to measuring the individual contribution of actions to future rewards , is by definition about understanding the structure of the task . By structure , we mean the relations between elements of the states , actions and environment rewards . In this work , we investigate what credit assignment can bring to transfer . Encouraged by recent successes in transfer based on supervised methods , we propose to learn to assign credit through a separate supervised problem and transfer credit assignment capabilities to new environments . By doing so , we aim at recycling structural information about the underlying task . To this end , we introduce SECRET ( SElf-attentional CREdit assignment for Transfer ) , a transferable credit assignment mechanism consisting of a self-attentive sequence-to-sequence model whose role is to reconstruct the sequence of rewards from a trajectory of agent-environment interactions . It assigns credit for future reward proportionally to the magnitude of attention paid to past state-action pairs . SECRET can be used to incorporate structural knowledge in the reward function without modifying optimal behavior , as we show in various generalization and transfer scenarios that preserve the structure of the task . Existing backward-view credit assignment methods ( Arjona-Medina et al. , 2019 ; Hung et al. , 2018 ) require to add auxiliary terms to the loss function used to train agents , which can have detrimental effects to the learning process ( de Bruin et al. , 2018 ) , and rely on an external memory , which hinder the generality of their approach . SECRET does neither . Also , as we show in Sec . 3.1 , the architecture we consider for SECRET has interesting properties for credit assignment . We elaborate about our novelty with respect to prior work in Sec . 4 . We insist on the fact that the focus of our work is on transfer and that it is not our point to compete on credit assignment capabilities . We would like to emphasize several aspects about the generality of SECRET : 1 ) our method does not require any modification to the RL algorithm used to solve the tasks considered , 2 ) it does not require any modification to the agent architecture either and 3 ) it does not alter the set of optimal policies we wish to attain . Moreover , our method for credit assignment is offline , and as a result , it can use interaction data collected by any mean ( expert demonstrations , replay memories ( Lin , 1992 ) , backup agent trajectories . . . ) . We believe that this feature is of importance for real-world use cases where a high number of online interactions is unrealistic but datasets of interactions exist as a byproduct of experimentation . Background We place ourselves in the classical Markov Decision Process ( MDP ) formalism ( Puterman , 1994 ) . An MDP is a tuple ( S , A , γ , R , P ) where S is a state space , A is an action space , γ is a discount factor ( γ ∈ [ 0 , 1 ) ) , R : S ×A×S → R is a bounded reward function that maps state-action pairs to the expected reward for taking such an action in such a state . Note that we choose a form that includes the resulting state in the definition of the reward function over the typical R : S ×A→ R. This is for consistency with objects defined later on . Finally , P : S × A → ∆S is a Markovian transition kernel that maps state-action pairs to a probability distribution over resulting states , ∆S denoting the simplex over S. An RL agent interacts with an MDP at a given timestep t by choosing an action at ∈ A and receiving a resulting state st+1 ∼ P ( ·|st , at ) and a reward rt = R ( st , at , st+1 ) from the environment . A trajectory τ = ( si , ai , ri ) i=1 , ... , T is a set of state-action pairs and resulting rewards accumulated in an episode . A subtrajectory is a portion of trajectory that starts at the beginning of the episode . The performance of an agent is evaluated by its expected discounted cumulative reward E [ ∑∞ t=0 γ trt ] . In a partially observable MDP ( POMDP ) , the agent receives at each timestep t an observation ot ∼ O ( ·|st ) that contains partial information about the underlying state of the environment . 2 SECRET : SELF-ATTENTIONAL CREDIT ASSIGNMENT FOR TRANSFER . SECRET uses previously collected trajectories from environments in a source distribution . A selfattentive sequence model is trained to predict the final reward in subtrajectories from the sequence of observation-action pairs . The distribution of attention weights from correctly predicted nonzero rewards is viewed as credit assignment . In target environments , the model gets applied to a small set of trajectories . We use the credit assigned to build a denser and more informative reward function that reflects the structure of the ( PO ) MDP . The case where the target distribution is identical to the source distribution ( in which we use held-out environments to assess transfer ) will be referred to as generalization or in-domain transfer , as opposed to out-of-domain transfer where the source and the target distributions differ . 2.1 SELF-ATTENTIONAL CREDIT ASSIGNMENT . Credit assignment as offline reward prediction We learn to assign credit through an offline reward prediction task , based on saved trajectories of agent-environment interactions . We create a sequence-to-sequence ( seq2seq ) model ( Sutskever et al. , 2014 ) that takes as input the sequence of observation-action pairs and has to reconstruct the corresponding sequence of environment rewards . Being offline , the reward prediction task is learned separately from the RL task , and the reward prediction model does not share representations with the agent . This way , the representations learned for credit assignment do not affect or get mixed with the representations learned for control . Operating offline brings several advantages : one can directly interact with the replay memory of agents and even use expert demonstrations or arbitrary saved transitions as a source of supervision , which could be useful in settings where on-policy interactions are costly , such as robotics . We equip our seq2seq model with an attention mechanism ( Bahdanau et al. , 2015 ) and view the attention weights of the reward reconstruction task as our primary source of assigned credit . The motivation to do so is that the seq2seq model looks into the past to find predictive signal in order to reconstruct the reward , so observation-action pairs it attends to should be those which reduce its uncertainty about the future , in other words those that explain future reward and should be credited . On the use of observations In MDPs , environment states follow the Markov property : they summarize the history of previous interactions and are sufficient to predict the future . As such , predictive models are highly biased towards focusing on the current sequence element , which hinders credit assignment . Under that consideration , when dealing with MDPs , we turn states into observations by applying transformations that hide a certain amount of information from states and break the Markov assumption . For instance , in gridworlds with visual states , we crop the image and get a player centered image with a given window size . Doing so encourages the model to look into the past to find predictive signal , and allow us to track the relative importance given to each element to reconstruct the credit assigned . In POMDPs , this might be unnecessary depending on the amount of information shared between observations and true states . Self-attention for credit assignment Unlike other seq2seq architectures , self-attentive models like Transformers ( Vaswani et al. , 2017 ) have direct computational paths between pairs of sequence elements , due to their representations that depend on projections of all sequence elements . This feature is key to long-term credit assignment . As an illustration , consider an RL task where the terminal reward depends only on the first observation , which is drawn randomly . Predicting the reward correctly requires to remember the first observation , which would be very challenging for a recurrent architecture whose memory goes through O ( n ) transformations , n being the size of the sequence . On the other hand , a self-attentive model directly accesses the value of the initial observation , which makes credit assignment easier . Reward prediction architecture We use a Transformer decoder with a single self-attention layer ( Lin et al. , 2017 ) and a single attention head . The model input is a sequence of observationaction couples ( ot , at ) t=0 , ... , T . Each observation goes through a series of convolutional layers ( for visual inputs ) followed by a series of feed-forward layers . Each action representation , a one-hot vector in the discrete action case , is concatenated to the learned observation embedding . Those representations of dimensionality di are combined with positional encoding ( PE ) , fed to a self-attention layer and then to a position-wise feedforward layer that outputs logits for reward prediction classes . PE encodes the relative positions of sequence elements , see Appendix A for details . Self-attention is an attention mechanism with parameterization ( Wk , Wq , Wv ) , each matrix belonging to Rdi×dk , that puts sequence elements in relation by computing non-linear similarity scores for all pairs of elements in the sequence . To do so , each sequence element is mapped to a query vector that is matched against keys and values obtained from the previous elements . To be consistent with the goal of assigning credit , the model should not be able to peek into the future . Thus , we restrict the computational window of each sequence element to the information stored in representations of the previous elements in the sequence and its own by applying a causal mask Mc to the result of the pairwise similarity computations , assigning a value of 0 to masked elements after the softmax . Let X = ( xt ) t=0 , ... , T ∈ RT×di denote the input sequence in a matrix form , xt being the result of internal computations of the model on its tth input . In the same fashion , we note Z = ( zt ) t=0 , ... , T ∈ RT×dk the sequence resulting from the application of self-attention . We then have Z = softmax ( Mc ( QKT ) − C ( 1−Mc ) √ dk ) V , where Q = XWq ∈ RT×dk stores queries , K = XWk ∈ RT×dk keys , and V = XWv ∈ RT×dk values as linear projections of the input ; dk stands for the dimension of the key vectors , Mc ∈ { 0 , 1 } T×T is a binary matrix that acts as a causal mask ( a lower triangular matrix ) , is the Hadamard product and C is a large constant ( 109 in practice ) . Notably , the resulting observation-action representation can be viewed as a linear combination of the values of previous elements : zt = ∑t i=0 αi←tvi where α·←t = ( αi←t ) i=1 , ... , t ∝ exp ( 〈qt , ki〉/ √ dk ) . The vector αt contains the normalized attention weights for the prediction at timestep t and sums to 1 . Since observations contain only a portion of their initial information , the fact that the model succeeds in the prediction task indicates that it reconstructed the missing information from its past . Therefore , attention weights themselves can be viewed as a form of credit assignment , and will be used as such in what follows . While performing regression on the rewards could also be an option , our experiments found that regression tends to converge to poor local optima . Consequently , we predict the sign of the experi- enced rewards : q ( r ) = sign ( r ) with sign ( 0 ) = 0 . We chose the sign as the classification target for its invariance to the scale of the rewards . We use a weighted sequential cross-entropy as the loss function over the class-wise model predictions fθ , c , writing τ ( o , a ) the subtrajectory of τ ending with the observation-action couple ( o , a ) to translate the effect of the binary mask : Lθ ( τ ) = − ∑ c∈ { −1,0,1 } w ( c ) |τ | ∑ ( o , a , r ) ∈τ 1 { q ( r ) = c } log ( fθ , c ( τ ( o , a ) ) ) . We have found class weighting to be very important to keep the variance of prediction metrics across a variety of datasets of sampled trajectories low . Generating trajectories To train SECRET , we generate a dataset of trajectories which contains a certain proportion of successful trajectories . If source environments are simple enough so that the task has sufficient chance to be solved by acting randomly , we use a random policy to generate trajectories . For more complex distributions of environments , we use a RL agent ( either trained or in the learning phase ) to generate trajectories . We think purely exploratory methods like Ecoffet et al . ( 2019 ) could have advantages over using an RL agent and leave the study of their use for future work .

Summary: This work focuses on credit assignment using a self-attention module for transfer RL problems. Specifically, the reward signals are assigned backward to previous states according to the attention weights. This can be helpful especially when the reward signal is sparse. Experiments on the newly proposed Triggers environment and the DMLlab keys & doors environment show that the proposed algorithm, SECRET, can speed up training in the transferred environment.