The act of remembering: A study in partially observable reinforcement learning

1 INTRODUCTION . Reinforcement Learning ( RL ) agents learn policies ( i.e. , mappings from observations to actions ) by interacting with an environment . RL agents usually learn memoryless policies , which are policies that only consider the last observation when selecting the next action . In fully observable environments , learning memoryless policies is an effective strategy . However , RL methods often struggle when the environment is partially observable . Indeed , partial observability is one of the main challenges to applying RL in real-world settings ( Dulac-Arnold et al. , 2019 ) . When faced with partially observable environments , RL agents require some form of memory to learn optimal behaviours . This is usually accomplished using k-order memories ( Mnih et al. , 2015 ) , recurrent networks ( Hausknecht & Stone , 2015 ) , or memory-augmented networks ( Oh et al. , 2016 ) . In this paper , we study a lightweight alternative approach to tackle partially observability in RL . The approach consists of providing the agent with an external memory and extra actions to control it ( as shown in Figure 1 ) . The resulting RL problem is still partially observable , but if the external memory is sufficiently expressive , then optimal memoryless policies will also yield globally optimal solutions . Previous works that explored this idea using external binary or continuous memories produced poor results with standard RL methods ( Peshkin et al. , 1999 ; Zhang et al. , 2016 ) . Our work shows that the main issue is with the type of memory they were using , and that RL agents are capable of learning effective strategies to utilize external memories when structured appropriately . In what follows , we • formalize the RL problem in the context of memory-augmented environments and study the theory behind memoryless policies that jointly decide what to do and what to remember ; • propose two novel forms of external memory called Ok and OAk . These k-size buffer memories generalize k-order memories by letting the agent ( learn to ) decide whether to push the current observation into the memory buffer or not ; • empirically evaluate Ok and OAk relative to previously proposed binary ( Bk ) , k-order ( Kk ) , and LSTM memories ( the most widely used approach for partially observable RL ) . Results show that Ok and OAk memories are usually more sample efficient than LSTM memories , while being faster to train and trivial to integrate with off-the-shelf RL methods . We therefore advocate for the adoption of Ok and OAk memories for partially observable RL problems . We end with a discussion of limitations of Ok and OAk and interesting avenues for future work . 2 PRELIMINARIES . RL agents learn how to act by interacting with an environment . Often these environments are modelled as a Markov Decision Process ( MDP ) . An MDP is a tupleM = 〈S , A , R , p , γ , µ〉 , where S is a finite set of states , A is a finite set of actions , R is the finite set of possible rewards , p ( s′ , r|s , a ) defines the dynamics of the MDP , γ is the discount factor , and µ is the initial state distribution . The interaction is usually divided into episodes . At the beginning of an episode , the environment is set to an initial state s0 , sampled from µ . Then , at time step t , the agent observes the current state st ∈ S and executes an action at ∈ A . In response , the environment returns the next state st+1 and immediate reward rt sampled from p ( st+1 , rt|st , at ) . The process then repeats until the end of the episode ( when a new episode will begin ) or potentially keep going for ever in non-episodic MDPs . Agents select actions according to a policy π ( a|s ) —which is a probability distribution from states to actions . The prediction task is to estimate how “ good ” a policy is , where the policy is evaluated according to the expected discounted return in any state . This can be done by estimating the action-value function qπ of policy π , where qπ ( s , a ) represents the expected discounted return when executing action a in state s and following π thereafter . Formally , qπ ( s , a ) = Eπ [ ∞∑ k=0 γkrt+k ∣∣∣∣∣St = s , At = a ] , where Eπ [ · ] denotes the expected value of a random variable given that the agent follows policy π , and t is any time step . qπ is usually estimated using Monte Carlo samples ( Barto & Duff , 1994 ) or TD methods ( Sutton , 1988 ) . The control task involves finding the optimal policy π∗ . This is the policy that maximizes the expected discounted return in every state . To do so , most RL methods rely on the policy improvement theorem , which we discuss in Section 5 . We use a Partially Observable Markov Decision Process ( POMDP ) formulation to model partial observability . A POMDP is a tuple P = 〈S , O , A , R , p , ω , γ , µ〉 , where S , A , R , p , γ , and µ are as in the MDP above , O is a finite set of observations , and ω ( o|s ) is the observation probability distribution . Interacting with a POMDP is similar to an MDP . The environment starts from a sampled initial state s0 ∼ µ . At time step t , the agent is in state st ∈ S , executes an action at ∈ A , receives an immediate reward rt , and moves to st+1 according to p ( st+1 , rt|st , at ) . However , the agent does not observe st directly . Instead , the agent observes ot ∈ O , which is linked to st via ω ( ot|st ) . 3 RELATED WORK . Early attempts to perform RL in partially observable domains focused on learning memoryless policies . Jaakkola et al . ( 1995 ) identified an RL algorithm that was guaranteed to converge to locally optimal memoryless policies , and similar guarantees have been given in the POMDP literature ( Li et al. , 2011 ) . Unfortunately , Singh et al . ( 1994 ) showed that an optimal memoryless policy π∗ ( at|ot ) can be arbitrarily worse than the optimal history-based policy π∗ ( at|o0 , a0 , . . . , ot ) for POMDPs . Different approaches have been proposed to learn history-based policies using some form of stateapproximation technique . For example , model-based RL methods learn a state representation of histories that enables Markovian prediction of future observations , rewards , or expected returns , and then learns policies over that representation ( McCallum , 1996 ; Littman et al. , 2002 ; Poupart & Vlassis , 2008 ; Doshi-Velez et al. , 2013 ; Ghavamzadeh et al. , 2015 ; Zhang et al. , 2019 ; Toro Icarte et al. , 2019 ) . The focus of our work is on model-free methods , which are the state of the art for solving partially observable problems from low-level inputs ( such as images ) . In model-free RL , history-based policies are approximated using recurrent neural networks ( Hausknecht & Stone , 2015 ; Mnih et al. , 2016 ; Wang et al. , 2016 ; Jaderberg et al. , 2016 ) , or some form of memory-augmented neural network ( Oh et al. , 2016 ; Parisotto & Salakhutdinov , 2017 ; Khan et al. , 2017 ; Hung et al. , 2018 ) . They are usually trained using policy gradient methods . These approaches are computationally expensive because they require the backpropagation of gradients through the history of observations and actions for learning history-based policies . In comparison , our approach is much more lightweight – being faster to train than LSTMs and generally having better sample complexity . We note that it is possible to learn memoryless policies that optimally solve POMDPs . The trick is to give the agent a ( large enough ) memory and extra actions to write to it . From the agent ’ s perspective , it learns a standard memoryless policy from observations to actions , but the observations now include the state of the memory , and the actions include options for how to alter the memory . The main purpose of our work is to resurrect this simple idea by understanding why previous work were unable to make it work . We also proposed a unified framework to study agents with external memories and two novel memories that outperform existing forms of external memories . Concretely , the idea of providing some form of external memory to an agent and actions to modify it goes back to Littman ( 1993 ) , who discussed a hypothetical agent that could learn to control an external binary memory in support of solving partially observable tasks . Peshkin et al . ( 1999 ) reported the first empirical results using tabular RL to learn memoryless policies over such binary memories . While the results were promising in some small environments that required only one bit of external memory to be solved , they did not scale to more complex domains . After Peshkin et al . ( 1999 ) , there was not much work trying to push this line of research forward . We believe that the reason is that RL agents can not reliably learn to control binary memories ( as shown in our results ) . That said , there is one recent work that has further explored the idea of modifying external memories using actions . Zhang et al . ( 2016 ) proposed to use continuous memories , where each element in the array was a floating point number , instead of binary memories . However , they learned the memoryless policies using imitation learning and pointed out that standard RL methods did not work because the reward signal was insufficient supervision for the agent to understand how to appropriately control the memory . One contribution of our work is to advance our understanding of methods that provide external memory to standard RL agents , and to show that they can work well in practice . Our work is also related to neural Turing machines ( NTM ) . The idea behind NTM is to provide an external memory to neural networks which they can write to and read from ( Graves et al. , 2014 ) . All their components are differentiable and , hence , they can be trained end-to-end using gradient descent and a training set of input and output examples ( i.e. , they solve a supervised learning problem ) . Zaremba & Sutskever ( 2015 ) proposed a variation of NTMs where they used the Reinforce algorithm to control how to move the head that reads and write over a memory tape – which can be seen as a case where an RL agent learns to ( partially ) control an external memory . That said , their overall system still solves a supervised learning problem as it requires the supervision coming from input and output examples to train the rest of the components in the NTM . 4 AGENTS WITH EXTERNAL MEMORY . In this section , we formally define what it means to provide external memory to an agent , and describe several forms of external memory . We will use the following problem to aid explanation : Example 4.1 ( the gravity domain ( Toro Icarte et al. , 2019 ) ) . The gravity domain , shown in Figure 2 , consists of an agent ( purple triangle ) , a cookie , and a button . The agent can move in the four cardinal directions and receives a reward of 1 when it eats the cookie . Doing so ends the episode . There is an external force pulling the agent down—i.e. , the outcome of the “ move-up ” action is a downward movement with probability 0.9—which can be turned off ( or back on ) by pressing the button . Every episode begins with the agent in the bottom left corner and the external force on . The optimal policy for this problem is to first press the button and then to go to the cookie . Since the agent can not observe the force , optimal behaviour requires memory of the state of the button , meaning that no memoryless policies can solve this problem optimally . However , suppose that the agent was given a single bit that they could write to on every step using the special actions write-1 and write-0 . This memory can then be used to record the state of the button , and so an optimal memoryless policy for this augmented problem will optimally solve the gravity domain . Figure 1 shows a generalization of this idea . From the agent ’ s perspective , they are , as usual , performing actions in an environment and receiving observations and rewards in return . However , they are now interacting with a memory-augmented environment—which consists of a sub-environment ( i.e. , the original POMDP environment ) and a memory . The memory receives an action w ( selected by the agent ) and local information coming from the sub-environment ( o , a , r , o′ ) to update its internal state to m′ . We formalize these external memory modules as follows : Definition 4.1 ( external memories ) . Let P = 〈S , O , A , R , p , ω , γ , µ〉 be a POMDP . An external memory for P is a tuple MP = 〈M , W , Γ , η〉 , where M is a finite set of memory-states , W is a finite set of memory-writing actions , Γ ( m′|m , w , o , a , r , o′ ) is the memory-writing distribution , and η is the initial memory-state distribution . An external memory module defines the set of possible memory configurations ( M ) and how the agent can manipulate that memory ( W and Γ ) . In the one-bit example for the gravity domain , M consists of the two possible values of the bit ( 0 or 1 ) , W consists of the two possible write options of the bit ( 0 or 1 ) , and the memory-writing distribution Γ updates the bit of memory to 0 or 1 depending on which action was selected . We now define a memory-augmented environment as follows : Definition 4.2 ( memory-augmented environments ) . A memory-augmented environment is a tuple E = 〈P , MP〉 where P is a POMDP andMP is an external memory for P . The interaction between an agent and a memory-augmented environment E = 〈P , MP〉 is the same as with the original environment , just with an augmented observation and action space . At the beginning of each episode , an initial state s0 , observation o0 , and memory state m0 , are sampled according to s0 ∼ µ , o0 ∼ ω ( o0|s0 ) , and m0 ∼ η , respectively . At time step t , the agent observes ōt = 〈ot , mt〉 and executes an action āt = 〈at , wt〉 ∈ A × W in E . Consequently , the subenvironment samples an immediate reward rt and the next state st+1 according to p ( st+1 , rt|st , at ) . The sub-environment also samples the next observation ot+1 ∼ ω ( ot+1|st+1 ) . The memory state is then updated to mt+1 according to Γ ( mt+1|mt , wt , ot , at , rt , ot+1 ) . Finally , the agent receives the immediate reward rt and the next observation ōt+1 = 〈ot+1 , mt+1〉 , and the process repeats . Any standard RL algorithm can be used to find a memoryless policy for a given memory-augmented environment E = 〈P , MP〉 . We note that the optimal memoryless policy for E must be at least as good as the optimal memoryless policy for the original POMDP P . This is because E and P share a reward function , and the agent can always choose to ignore the memory . That said , if the external memory moduleMP is “ expressive enough , ” then optimal memoryless policies for 〈P , MP〉 will be just as good as the optimal policy for P . This is shown formally in Appendix A.2 .

Summary: This paper focuses on reinforcement learning in partially-observable environments, and revisits the approach that consists of extending the agent with an external memory. The main contribution of the paper is the proposal (and evaluation) of adding an action to the agent, that allows it to push its current observation (and previous action is some cases) in a k-sized queue. The observation of the agent is extended with the contents of the queue. The main argument of the authors is that learning "when to push" is easier for the agent than learning "what to push" (as done with Neural Turing Machines and memory-bit external memories), and being able not to push the current observation allows the agent to remember past observations for longer durations, as opposed to k-step observation windows approaches.