AMRL: Aggregated Memory For Reinforcement Learning

1 INTRODUCTION . We address the problem of reinforcement learning ( RL ) in tasks that require long-term memory . While many successes of Deep RL were achieved in settings that are ( near ) fully observable , such as Atari games ( Mnih et al. , 2015 ) , partial observability requires memory to recall prior observations that indicate the current state . Relying on full observability severely limits the applicability of such approaches . For example , many tasks in virtual and physical environments are naturally observed from a first-person perspective ( Oh et al. , 2016 ) , which means that an agent may need to seek out and remember task-relevant information that is not immediately observable without directly observing the entire environment . Recent research has started to address this issue , but effective learning in RL settings with long sequential dependencies remains a key challenge in Deep RL ( Oh et al. , 2016 ; Stepleton et al. , 2018 ; Parisotto & Salakhutdinov , 2018 ) . The currently most common approach to RL in partially observable settings relies on models that use memory components that were originally developed for tasks like those that occur in natural language processing ( NLP ) , e.g. , LSTMs ( Hochreiter & Schmidhuber , 1997 ) and GRUs ( Cho et al. , 2014 ) . Hausknecht & Stone ( 2015 ) first demonstrated benefits of LSTMs in RL tasks designed to test memory , and these and similar approaches have become common in Deep RL ( Wang et al. , 2016 ) , including multi-agent RL ( Rashid et al. , 2018 ; Foerster et al. , 2017 ) . In this work , we demonstrate that the characteristics of RL can severely impede learning in memory models that are not specifically designed for RL , and propose new models designed to tackle these challenges . For example , LSTMs excel in NLP tasks where the order of observations ( characters or words ) is crucial , and where influence between observations decays quickly with distance . Contrast this with a hypothetical RL example where an agent must discover a hidden passcode to escape a locked dungeon . The order of observations is highly dependent on the agent ’ s path through the dungeon , yet when it reaches the door , only its ability to recall the passcode is relevant to escaping the dungeon , irrespective of when the agent observed it and how many observations it has seen since . Figure 1 illustrates the problem . When stochasticity is introduced to a memory task , even simply as observation noise , the sample efficiency of LSTMs decreases drastically . We show that this problem occurs not just for LSTMs , but also for stacked LSTMs and DNCs ( Graves et al. , 2016 ; Wayne et al. , 2018 ) , which have been widely applied in RL , and we propose solutions that address this problem . We make the following three contributions . First , in Section 3 , we introduce our approach , AMRL . AMRL augments memory models like LSTMs with aggregators that are substantially more robust to noise than previous approaches . Our models combine several innovations which jointly allow the model to ignore noise while maintaining order-variant information as needed . Further , AMRL models maintain informative gradients over very long horizons , which is crucial for sample-efficient learning in long-term memory tasks ( Pascanu et al. , 2012 ; Bakker , 2001 ; Wierstra et al. , 2009 ) . Second , in Section 4 , we systematically evaluate how the sources of noise that affect RL agents affect the sample efficiency of AMRL and baseline approaches . We devise a series of experiments in two domains , ( 1 ) a symbolic maze domain and ( 2 ) 3D mazes in the game Minecraft . Our results show that AMRL can solve long-term memory tasks significantly faster than existing methods . Across tasks our best model achieves an increase in final average return of 9 % over baselines with far more parameters and 19 % over LSTMs with the same number of parameters . Third , in Section 6 we analytically and empirically analyze the characteristics of our proposed and baseline models with the aim to identify factors that affect performance . We empirically confirm that AMRL models are substantially less susceptible to vanishing gradients than previous models . We propose to additionally analyze memory models in terms of the signal-to-noise ratio achieved at increasing distances from a given signal , and show that AMRL models can maintain signals over many timesteps . Jointly , the results of our detailed analysis validate our modeling choices and show why AMRL models are able to effectively solve long-term memory tasks . 2 RELATED WORK . External Memory in RL . In the RL setting , work on external memory models is most relevant to our own . Oh et al . ( 2016 ) introduce a memory network to store a fixed number of prior memories after encoding them in a latent space and validate their approach in Minecraft , however the models are limited to fixed-length memories ( i.e. , the past 30 frames ) . The Neural Map of Parisotto & Salakhutdinov ( 2018 ) is similar to our work in that it provides a method in which past events are not significantly harder to learn than recent events . However , it is special-cased specifically for agents on a 2D grid , which is more restrictive than our scope of assumptions . Finally , the Neural Turing Machine ( NTM ) ( Graves et al. , 2014 ) and its successor the Differentiable Neural Computer ( DNC ) ( Graves et al. , 2016 ) have been applied in RL settings . They use an LSTM controller and attention mechanisms to explicitly write chosen memories into external memory . Unlike the DNC which is designed for algorithmic tasks , intentionally stores the order of writes , and induces sparsity in memory to avoid collisions , we write memories into order-invariant aggregation functions that provide benefits in noisy environments . We select the DNC , the most recent and competitive prior approach , for baseline comparisons . Other Memory in RL . A second and orthogonal approach to memory in Deep RL is to learn a separate policy network to act as a memory unit and decide which observations to keep . These approaches are generally trained via policy gradient instead of back-propagation through time ( BPTT ) ( Peshkin et al. , 1999 ; Zaremba & Sutskever , 2015 ; Young et al. , 2018 ; Zhang et al. , 2015 ; Han et al. , 2019 ) . These approaches are often difficult to train and are orthogonal to our work which uses BPTT . The Low-Pass RNN ( Stepleton et al. , 2018 ) uses a running average similar to our models . However , they only propagate gradients through short BPTT truncation window lengths . In fact they show that LSTMs outperform their method when the window size is the whole episode . Since we are propagating gradients through the whole episode , we use LSTMs as a baseline instead . Li et al . ( 2015 ) ; Mirowski et al . ( 2017 ) ; Wayne et al . ( 2018 ) propose the use of self-supervised auxiliary losses , often relating to prediction of state transitions , to force historical data to be recorded in memory . Along this line , model-based RL has also made use of memory modules to learn useful transition dynamics instead of learning a policy gradient or value function ( Ke et al. , 2019 ; Ha & Schmidhuber , 2018 ) . These are orthogonal and could be used in conjunction with our approach , which focuses on model architecture . Finally , several previous works focus on how to deal with storing initial states for truncated trajectories in a replay buffer ( Kapturowski et al. , 2019 ; Hausknecht & Stone , 2015 ) . Memory in Supervised Learning . In the supervised setting , there has also been significant work in memory beyond LSTM and GRUs . Similar to our work , Mikolov et al . ( 2014 ) , Oliva et al . ( 2017 ) , and Ostmeyer & Cowell ( 2019 ) use a running average over inputs of some variety . Mikolov et al . ( 2014 ) use an RNN in conjunction with running averages . However , Mikolov et al . ( 2014 ) use the average to provide context to the RNN ( we do the inverse ) , and all use an exponential decay instead of a non-decaying average . Additionally , there have been myriad approaches attempting to extend the range of RNNs , that are orthogonal to our work , given that any could be used in conjunction with our method as a drop-in replacement for the LSTM component ( Le et al. , 2015 ; Arjovsky et al. , 2016 ; Krueger & Memisevic , 2016 ; Belletti et al. , 2018 ; Trinh et al. , 2018 ) . Other approaches for de-noising and attention are proposed in ( Kolbaek et al. , 2017 ; Wöllmer et al. , 2013 ; Vaswani et al. , 2017 ; Lee et al. , 2018 ) but have runtime requirements that would be prohibitive in RL settings with long horizons . Here , we limit ourselves to methods with O ( 1 ) runtime per step . 3 METHODS . 3.1 PROBLEM SETTING . We consider a learning agent situated in a partially observable environment denoted as a Partially Observable Markov Decision Process ( POMDP ) ( Kaelbling et al. , 1998 ) . We specify this process as a tuple of ( S , A , R , P , O , Ω , γ ) . At time-step t , the agent inhabits some state , st ∈ S , not observable by the agent , and receives some observation as a function of the state ot ∈ Ω ∼ O ( ot|st ) : Ω × S → R≥0 . O is known as the observation function and is one source of stochasticity . The agent takes some action at ∈ A . The POMDP then transitions to state st+1 ∼ P ( st+1|st , at ) : S × A × S → R≥0 , and the agent receives reward rt = R ( st , at ) : S × A → R and receives a next observation ot+1 ∼ O upon entering st+1 . The transition function P also introduces stochasticity . The sequence of prior observations forms an observation trajectory τt ∈ Ωt ≡ T . To maximize ∑t=∞ t=0 γ trt , the agent chooses each discrete at from a stochastic , learned policy conditioned on trajectories π ( at|τt ) : T × A → [ 0 , 1 ] . Given P , O , and π itself are stochastic , τt can be highly stochastic , which we show can prevent learning . 3.2 MODEL . In this section we introduce our model AMRL , and detail our design choices . At a high level , our model uses an existing memory module to summarize context , and extends it to achieve desirable properties : robustness to noise and informative gradients . LSTM base model . We start from a standard memory model as shown in Figure 2a . We use the base model to produce a contextual encoding of the observation ot that depends on the sequence of prior observations . Here , we use several feed-forward layers , FF1 ( defined in A.5 ) , followed by an LSTM layer ( defined in A.6 ) : et = FF1 ( ot ) // The encoded observation ht = LSTMt ( et ) // The output of the LSTM Previous work proposed a stacked approach that combines two ( or more ) LSTM layers ( Figure 2b ) with the goal of learning higher-level abstractions ( Pascanu et al. , 2013 ; Mirowski et al. , 2017 ) . A key limitation of both LSTM and LSTM STACK approaches is susceptibility to noise . Noise can be introduced in several ways , as laid out in Section 3.1 . First , observation noise introduces variance in the input ot . Second , as motivated by our introductory example of an agent exploring a dungeon , variance on the level of the trajectory τt is introduced by the transition function and the agent ’ s behavior policy . Recall that in our dungeon example , the agent encounters many irrelevant observations between finding the crucial passcode and arriving at the door where the passcode allows escape . This variance in τt generally produces variance in the output of the function conditioning on τt . Thus , although the first part of our model makes use of an LSTM to encode previous inputs , we expect the output , ht to be sensitive to noise . Aggregators . To address the issue of noise highlighted above , we introduce components designed to decrease noise by allowing it to cancel . We call these components aggregators , labeled M in Figures 2c and 2d . An aggregator is a commutative function that combines all previous encodings ht in a time-independent manner . Aggregators are computed dynamically from their previous value : mt = g ( mt−1 , ht [ : 1 2 ] ) // The aggregated memory where ht [ : 12 ] denotes the first half of ht , and g ( ) denotes the aggregator function , the choices of which we detail below . All proposed aggregators can be computed in constant time , which results in an overall memory model that matches the computational complexity of LSTMs . This is crucial for RL tasks with long horizons . In this work we consider the SUM , AVG , and MAX aggregators defined in Table 1 . All three are easy to implement in standard deep learning frameworks . They also have desirable properties in terms of gradient flow ( Jacobian in Table 1 and signal-to-noise ratio ( SNR ) ) which we detail next . Aggregator signal-to-noise ratio ( SNR ) . Our primary design goal is to design aggregators that are robust to noise from the POMDP , i.e. , variation due to observation noise , behavior policy or environment dynamics . For example , consider the outputs from previous timesteps , ht , to be i.i.d . vectors . If we use the average of all ht as our aggregators , then the variance will decrease linearly with t. To formally and empirically assess the behavior of memory models in noisy settings , we propose the use of the signal-to-noise ratio or SNR ( Johnson , 2006 ) . The SNR is a standard tool for assessing how well a system maintains a given signal , expressed as the ratio between the signal ( the information stored regarding the relevant information ) and noise . In this section , to maintain flow , we simply state the SNR that we have analytically derived for the proposed aggregators in Table 1 . We note that the SNR decays only linearly in time t for SUM and AVG aggregators , which is empirically slower than the baselines , and has a bound independent of time for the MAX aggregator . We come back to this topic in Section 6 , where we describe the derivation in detail , and provide further empirical results that allow comparison of all proposed and baseline methods . Aggregator gradients . In addition to making our model robust to noise , our proposed aggregators can be used to tackle vanishing gradients . For example , the sum aggregator can be viewed as a residual skip connection across time . We find that several aggregators have this property : given that the aggregator does not depend on the order of inputs , the gradient does not decay into the past for a fixed-length trajectory . We can show that for a given input xi and a given output o , the gradient dot dxi ( or expected gradient ) of our proposed aggregators does not decay as i moves away from t , for a given t. We manually derived the Jacobian column of Table 1 to show that the gradient does not depend on the index i of a given past input . We see that the gradient decays only linearly in t , the current time-step , for the AVG and MAX aggregators . Given that the gradients do not vanish when used in conjunction with ht as input , they provide an immediate path back to each ht through which the gradient can flow . SET model . Using an aggregator to aggregate all previous ot yields a novel memory model that we term SET ( Figure 2c ) . This model has good properties in terms of SNR and gradient signal as shown above . However , it lacks the ability to maintain order-variant context . We address this limitation next , and include the SET model as an ablation baseline in our experiments . Combining LSTM and aggregator . In our AMRL models , we combine our proposed aggregators with an LSTM model that maintains order-dependent memories , with the goal to obtain a model that learns order-dependent and order-independent information in a manner that is data efficient and robust to noise . In order to achieve this , we reserve certain neurons from ht , as indicated by ’ / ’ in Figure 2d . We only apply our aggregators to one half of ht . Given that our aggregators are commutative , they lose all information indicating the context for the current time-step . To remedy this , we concatenate the other half of ht onto the aggregated memory . The final action is then produced by several feed-forward layers , FF2 ( defined in A.5 ) : ct = [ ht [ 1 2 : ] |ft ( mt ) ] // The final context1 at = FF2 ( ct ) // The action output by the network Straight-through connections . We showed above that our proposed aggregators provide advantages in terms of maintaining gradients . Here , we introduce a further modification that we term straight-through ( ST ) connections , designed to further improve gradient flow . Given that our proposed aggregators are non-parametric and fairly simple functions , we find that we can deliberately modify their Jacobian as follows . We pass the gradients straight through the model without any decay . Thus , in our ST models , we modify the Jacobian of g and set it to be equal to the identity matrix ( as is typical for non-differentiable functions ) . This prevents gradients from decaying at all , as seen in the ST Jacobian column of Table 1 . Our proposed models combine the components introduced above : AMRL-Avg combines the AVG aggregator with an LSTM , and uses a straight-through connection . Similarly , AMRL-Max uses the MAX aggregator instead . Below we also report ablation results for SET which uses the AVG aggregator without an LSTM or straight-through connection . In Section 5 we will see that all components of our model contribute to dramatically improved robustness to noise compared to previous models . We further validate our design choices by analyzing gradients and SNR of all models in Section 6 .

Summary: The paper proposes to add one more memory layer on top of LSTM that offers more direct integration of information over time. The goal is to demonstrate that it is very useful to aggregate previous information in a RNN-based reinforcement learning settings. The set-based aggregation improves the gradient over time and maintain good signal-to-noise ratio. The model is evaluated on Tmazes and Minecraft.