The act of remembering: A study in partially observable reinforcement learning
Authors: Rodrigo Toro Icarte
Richard Valenzano
Toryn Q. Klassen
Phillip Christoffersen
Amir-massoud Farahmand
Sheila A. McIlraith
Venue:  (2021)
Keywords: Reinforcement Learning
Partial Observability
Memory Representations
External Memories
POMDPs.

Abstract: Partial observability remains a major challenge for reinforcement learning (RL). In fully observable environments it is sufficient for RL agents to learn memoryless policies. However, some form of memory is necessary when RL agents are faced with partial observability. In this paper we study a lightweight approach: we augment the environment with an external memory and additional actions to control what, if anything, is written to the memory. At every step, the current memory state is part of the agent’s observation, and the agent selects a tuple of actions: one action that modifies the environment and another that modifies the memory. When the external memory is sufficiently expressive, optimal memoryless policies yield globally optimal solutions. We develop the theory for memory-augmented environments and formalize the RL problem. Previous attempts to use external memory in the form of binary memory have produced poor results in practice. We propose and experimentally evaluate alternative forms of k-size buffer memory where the agent can decide to remember observations by pushing (or not) them into the buffer. Our memories are simple to implement and outperform binary and LSTM-based memories in well-established partially observable domains.

Under review as a conference paper at ICLR 2021
THE ACT OF REMEMBERING: A STUDY IN PARTIALLY OBSERVABLE REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Partial observability remains a major challenge for reinforcement learning (RL). In fully observable environments it is sufficient for RL agents to learn memoryless policies. However, some form of memory is necessary when RL agents are faced with partial observability. In this paper we study a lightweight approach: we aug- ment the environment with external memory and additional actions control what, if anything, is written to the memory. At every step, the current memory state is part of the agent's observation, and the agent selects tuple of actions: one action that modifies the environment and another that modifies the memory. When the external memory sufficiently expressive, optimal memoryless policies yield globally optimal solutions. We develop the theory for memory-augment environments and formalize the RL problem. Previous attempls to use external memory in the form of binary memory have produced poor results in practice. We propose and experimentally evaluate alternative forms of k-size buffer memory where the agent can decide to remember observations by pushing (or not) them into the buffer. Our memories are simple to implement and outperform binary and LSTTM-base memories in well-established partially observable domains.
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 environ- ments, 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 chal- lenges to applying RL in real- settings Dbulac-Aroole 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 al. 2015). recurrent networks (Hausknecht & Stone 2015), or memory-angmennted 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 (as shown in Figure [ The resulting RL problem still partially observable, but 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 LST memories (the most widely used approach for partially observable RL.).

---

Under review as a conference paper at ICLR 2021
Memory-Augmented Environment
á (a, w)
RL agent
Environment
Memory
a
(0,,,r,o))
0 = "o',m')
m'
Figure 1: A diagram of a Memory-Augmented Environment.
Results show that Ok and OAk memories 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 mod- elled as Markov Decision Process (MDP). An MDP is tuple M (S, A, R. p, Y, 1b), where S is a finite set of states, A is finite set of actions, R is the finite set of possible rewards, p(s',r|s, defines the dynamics of the MDP, The interaction is usually divided into episodes. At the beginning of an episode, the environment is set to an initial state So, sampled from μ.. Then, at time step t. the agent observes the current state 8, E S and executes an action a, E A. In response, the environment returns the next state $1+1 and immediate reward rt sampled from p(s++1, rel8s,, az). The process then repeats until the end of the episode (when new episode will begin) or potentially keep going for ever in non-episodic MDPs.
the discount factor, and μr is the initial state distribution.
Agents select actions according to a policy T(a whicc which is probability distribution from states to actions. The prediction task is to estimate how "good" a policy is, where the policy is evalu- ated according to the expected discounted return in any state. This can be done by estimating the function 9x of policy π, where qu (s, represents the expected discounted return when executing action a in state s and following æ thereafter. Formally,
In (n a) = EET
Σ
Tetk S A,
where E, [-] denotes the expected value of random variable given that the agent follows policy m, and L is any time step. 4a is usually estimated using Monte Carlo samples (Barto & Duff 1994) or TD methods (Sutton 1988). The control task involves finding the optimal policy क. This 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, A, R, p. 7/ p), where 5, A, R. p, ?. and l are in the MDP above, O is finite set of observations, and w(ols) is the observation probability distribution. Interacting with a POMDP is similar to an MDP. The environment starts from a sampled initial state 80 J.. At time step the agent is in state 8, E S, executes an action a, E A, receives an immediate reward re, and moves to 8t+1 according to p($(+1, relSe, at). However, the agent does not observe s directly. Instead, the agent observes o E O, which is linked to s, via w(orss:).
3 RELATED WORK
Early attempts to perform RL in partially observable domains focused learning memoryless poli- cics. 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 ( (1994) showed that an optimal memoryless policy x (a, ox) can be arbitrarily worse than the optimal history- based policy n² (ar/on dg...., o.) for POMDPs.
Different approaches have been proposed to learn history-based policies using some form of state- approximation technique. For example, model-based RI. methods learn state representation of
2

---

Under review as a conference paper at ICLR 2021
histories that enables Markovian prediction of future observations, rewards, or expected returns, and
then learns policies over that representation (McCallum 1996; Littman et al. [202 Poupart & Vlas-
sis, 2008 Doshi- Dosei-Velez et al. 2013] Ghavamzadeh al. 2015 Zhang al.] 2019 Toro Icarte al.)
2019). y. The focus of our work is on model-free methods, which are the state of the art for solving par-
tially observable problems from low-level inputs (such as images). In model-free RI., history-basee
policies are approximated using recurrent neural networks (Hausknecht & Stone 2015] Miih et al.)
2016 Wang et al. 2016 Jaderberg al., 2016). or some form of memory-augmented neural net- P Parisotto & Salakhutdinov 2017] Khan et al., 2017 HHung et al. 2018). They are usually trained using policy gradient methods. These approaches are computationaly expensive because they require the backpropagation of gradients through the history of observations and ac- tions for learning history-based policies. comparison, our approach much more lightweight - being faster to train than LSTMs and generally having better sample complexity.
work (Oh et al. 2016
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 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 make it work. We also proposed 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 an agent and actions to modify it goes back to Littman (1993), who discussed a hypothetical agent that could learn to control an ex- ternal binary memory in support of solving partially observable tasks, Peshkin 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 RI. agents cannot 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 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 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 supervised learning problem). Zaremba & Sutskever (205))) proposed a variation of NTMs where they used the Reinforce algorithm to control how to move the head that reads and write over 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 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 consists of an agent (purple triangle), cookie, and button. The agent can move in the four cardinal directions and receives a reward of I 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 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 cannot observe the force, optimal behaviour requires memory of the state of the button,
3

---

Under review as a conference paper at ICLR 2021
gravity domain
g-learning
5-step actor-critic
Memories:
K1
B1
O1
Ave
2
OA1
None
0.2
0.4
0.6
0.8
2
10 14114115 18:20
Training steps (in millions)
Training steps (in millions)
Figure 2: Experiments in the gravity domain. We reported the avg. reward per 100 steps.
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 every step using the special actions wr te-1 and write-0. 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 D shows a generalization of this idea. From the agent's perspective, they are, as usual, per- forming actions in environment and receiving observations and rewards in return. However, they enviroument-we -which consists sub- environmrsn
are now interacting with a ceenory-aggment (i.e.. the original POMDP environment) and memory. The memory receives action (selected by the agent) and local information coming from the -ub-enviromment (o, a, o') to update its in- ternal state m'. We formalize these external memory modules follows:
Definition 4.1 (external memories). Let P = (5, A, R. w, ",), be a POMDP. An external memory for P tuple MM (M, W, W,r, n),, where M is finite set of memory-states, W is finite set of memory-writing actions, ['mm [(mm|mm, w, o') is the memory-writing distribution, and 7 is the initial memory-ste state distribution.
external memory module defines the set of possible memory configurations (M) and how the agent can manipulate that memory (W and I'). In the one-bit example for the gravity domain, M consists of the two possible values the bit (0 or 1), W consists of the two possible write options of the bit (0 or 1), and the memory- wwittigg distribution I updates the bit of memory to 0 or depending on which action was selected. We now define a memory- auggented environment as follows:
Definition 4.2 (memory-augmented environments). A memory augmeeted environment is tuple = (P, MP) where P is a POMDP and Mp an external memory for P.
The interaction between an agent and a nemorr-aagmeent environment È = (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 sp, observation on, and memory state mag are sampled according to 8g μl, Og ~ (S)). and ma ~ 1) respectively. At time step i, the agent observes o (or, m²) and executes an action ā (a,, wx) A x W in E. Consequently, the sub- environment samples an immediate reward r, and the next state 8:41 according to p(s++1, I' Fr/St, at) The sub-environment also samples the next observation 0111 w(Or411 (8t+1). The memory state is then updated to m++1 according to [(me+1 |m, Wt. 0₁, a, Tt, 0++1). Finally, the agent receives the immediate reward F1 and the next observation 0 (0141, mc+1). and the process repeats.
Any standard RL algorithm can be used to find a memoryless policy for a given weemory-aggmentt environment (P, Mp). We note that the optimal memoryless policy for & must be at least as good as the optimal memoryless policy for the original POMDP P. This is because & and P share a reward function, and the agent can always choose ignore the memory. That said, if the external memory module Mp is "expressive enough," then optimal memoryless policies for (P,Mp) will be just as good as the optimal policy for P. This shown formally in Appendix A.2
4.1 EXTERNAL MEMORY MODULES
Let consider several examples of external memory modules. We begin by showing how binary memories (Littman 1993 1994 Peshkin et al.) 1999) can be expressed using this formalism. We use the notation Bk refer binary memory of bits:
4

---

Under review as a conference paper at ICLR 2021
Definition 4.3 (Bk memories). Let P 0,
7. R.
be POMDP. A Bk memory for P
is a k-bit external memory Mp = (M, W.I. m), where M = {0, 1 W = {0, 1J" nn(0)) = (zero otherwise), and I'mm |m, w. o, a, d') 1 if and only if m' = (zero otherwise).
Bk memories are especially attractive given how flexible and expressive they are. Unfortunately, learning to control Bk memories is difficult.
2/oowss the performance of tabular q-learning (Watkins& & Dayan 1992) and 5-step actor-critic (Grondman et al. 2012) in the gravity domain using different types of external memories. In the figure, None represents not using any external memory, and K1, O1, and OA1 are explained below. Notice that the agent using q-learning or 5-step actor-critic was unable to learn how to use the B1 memory to consistently solve the gravity domain.
There are two main problems with Bk memories. First, the action space grows exponentially with k. Second, Bk memories can be too flexible in that the agent can modify the memory arbitrarily and irrespective of what has actually happened, and thereby completely alter what the agent believes about its current situation. For example, recall that in the gravity domain, the agent should use the memory to record whether gravity is on (0) or off (1). However, if the agent incorrectly decides to record that the gravity is off prematurely (i.e., before touching the button), it will believe it has transitioned from a state with low expected reward (where it first has to go to the button) to a state with high expected reward (where the agent wrongly believes that it can go directly the cookie without any opposition from gravity). This can lead to an unstable learning process, as shown below.
The main motivation behind our proposed Ok memories is to alleviate these issues. Ok memories are a generalization of k-order memories, which are buffers of fixed size that contain the last k observations. We refer to k-order memories as Kk memories, where the second 'k" indicates the size of the buffer. We formally describe them as external memories in Appendix A.3 Note that K1 represents a 2-order memory since actions are taken over (o, m). The main disadvantage of k-order memories is that they do not allow the agent to remember events that occurred more than k steps in the past. Ok memories solve this issue by letting the agent decide whether to push the current observation into the k-order buffer or not. Since the agent can only push into the buffer observations that did occur, Ok memories are unable to imagine events that have not yet happened.
Definition 4.4 (Ok memories). Let P (8, O, A, R. p, w, Y, n) be POMDP. An Ok memory for P is a memory buffer (of size k) .Mp = (M, W.T, m) where M = (O U {0) )a W = [T,L).
(70)) 1 (zero otherwise), and ['m' [m, w. a, d') 1 if w 1 and m' m. m = (o
, and m = (0²,... , , o) (zero otherwise).
Ok memories have strong empirical performance in the gravity domain (see Figure 2). outperform- ing B1 and KI. That said, Ok memories are insufficient in domains where the history of actions matters. For such domains, we propose OAk memories. An OAk memory is similar to an Ol mem- ory but when the agent chooses to push to its buffer, the information includes the current observation and the action that is executed in the sub-environmenn. OAk memories are defined in Appendix A.3
We note that optimal memoryless policies over Bk, OAK, Ok, and Kk will optimally solve the original POMDP for some value of k. under some assumptions. This is shown in Appendix A.4
5 LEARNING POLICIES IN
ENVIRONMENTS
The objective of this section is to understand the theory behind learning memoryless policies over environments and to provide insights into why Ok and OAk memories tend to perform better than Bk memories. We begin with the following example.
Example 5.1 (a recall task). The recall task is a partially observable environment with only one possible observation, o (i.e, all states appear the same), and three actions, a1, 02, and 0g. The episode ends after performing three actions. If the agent executes actions 01, ag, and ag (in that order), it gets reward of I; otherwise it gets reward of 0.
The purpose of the recall task is to show that even if a memoryless policy for a memory-augmente environment itself might not be an MDP.
environment is globally optimal, the memory-augment Figure 3 shows a transition diagram for the recall task using an OAI memory. Since the observa- tion is always the same, the different states that the agent encounters only differ by the state in the memory. In the diagram, nodes represent the memory states and the transitions show how the mem- ory is updated by the agent's 's actions. Note that node i represents that the memory buffer contains (o, and that the buffer starts empty (0). For the action labels, the first number indicates the action
5

---

Under review as a conference paper at ICLR 2021
OAI env.
g-learning (OAI)
q-learning (B2)
B2 env.
Reward per op.
Reward per ep.
11.21331
IT
2T
3T
2T
12.
IT
3T
22.
32
2T
0.5
1.5
0.5
1
1.5
3T
Training steps (in millions)
Training steps (in millions)
13,23,33
Figure 3: Performance of greedy policy every 10, 000 training steps in the recall task.
number in the sub-environment (1, or 3). The second character represents the memory action. For instance, 2T represents that the agent executed 02 in the sub-environment and saved that action into the memory buffer. The label o/w stands for otherwise. The blue arrows show a deterministic memoryless policy that optimally solves this problem. That is, execute 1T, then 2T, and finally 3T.
Notice that this memory-augmentter environment has a memoryless policy that optimal for the original POMDP, but it is not an MDP. The reason is that the reward given by the bottom transition depending on the history. If the agent follows the blue path, gets reward of one.
3T will If the agent follows the red arrow, it gets a reward of zero. Something similar occurs when using B2 memory, which is the smallest Bk memory that can encode an optimal policy for this task. This 'non-Markovianee" impacts the performance of RL agents that explicitly exploit the Markovian assumption. For example, we run q-learning and evaluate the performance the greedy policy (i.e., without exploration) every 10, 000 steps, see that q-learning does not converge. Instead, q-learning jumps between an optimal policy and a zero reward policy, as shown in Figure 3
Now that we know that memory-augmennt augmented environments not MDPs, we focus on proving that they are POMDPs. Such proof can be found in Appendix A.1 ...1 and has important repercussions. In particular, all the theory for learning memoryless policies for POMDPs [Littman] 1994 [Singh al.] 1994 Jaakkola et al. 1995] Li et al.] 1 Azizzadenesheli et al.) 2018) also applies to memory- augmented environments. We explore this further two parts: the prediction problem and the control problem. We then discuss the practical implications that follow from the theory.
5.1 How TO EVALUATE POLICIES IN MEMORY-AUGMERN
ENVIRONMENTS
For a given POMDP P and a memoryless policy (all)), the policy prediction problem consists of estimating (o, a). Here, 4- (o, a) is defined over observations and represents the expected dis- counted return when executing action a given observation o (at any time step f) and following thereafter: qu (0. a) E >>= 7 T O2 = o, A, al It is known that Monte-Carlo estimates guaranteed to converge the real values of q. (o, a), though they do have high variance. In contrast, TD estimates have lower variance but might not converge to gx (o a) (Singh et al. 1994).
Failing to correctly estimate (= a) is the reason behind q-learning's 's instability in the recall task (Figure 3) For instance, let и be the optimal policy represented by blue arrows in OA1, 1, then the real value for the red arrow grr (0, 2 2T is zero (the agent gets reward it executes 02 the first action). However, a one-step TD estimate would converge to 4a (0, 2T) O + 74a 3T) This is a problem since now 4t (0, 2T) qx (0, IT) = 7² (for Y (0, and so q-learning will move from the current optimal policy to the zero reward policy that executes 2T in O. We refer to these types of transitions as Markkanoo shorteuts. Note that, as Figure shows, the B2 memory has more non-Markovian shortcuts than OAI. This is why q-learning over B2 is more unstable than q-learning over OAI in this domain. More generally, we would expect that Bk memories introduce more non-Markovian shortcuts than OAΣ memories since they more flexible, which could partially explain the better empirical performance OAk and Ok memories.
T,
There are two approaches that can mitigate this problem. The first is use n-step TD estimates, with a large enough value of n. As Figure 4 shows, the performance of 20-step actor-critic in the gravity domain is far superior to 5-step actor-critic. The second is to increase the size of the memory, since doing so tends to remove Mon-Markovian shortcuts. This is also shown in Figure as 5-step actor-critic performs better when using 02 OA2, than when using O1 OA1.

---

Under review as a conference paper at ICLR 2021
5 5-step actor-critic
5-step actor-critic
20-step actor-critic
Avg. reward
6
+
84y
2
2
0
2
4 6 8
10 18
4
6 8 10121 14 16 18
2
10121 14
18 20
Training steps (in millions)
Training steps (in millions)
Training steps (in millions)
None
K2
B2
02
OA2
K1
B1
01 OA1
Figure 4: Experiments in the gravity domain. We reported the avg. reward per 100 steps.
5.2 How TO IMPROVE POLICIES IN MEMORY-AUGMENTE ENVIRONMENTS
We now focus our attention on the second part of the problem: how to use q-value estimates to find better policies. To do so, most RL algorithms exploit the policy improvement theorem. For MDPs, this theorem guarantees that updating the current policy π by any amount towards the greedy policy
7 (s)
arg ma (=(8,0 a) will lead to better policies (Watkins] 1989 Sutton Barto 2018).
When learning memoryless policies for POMDPs, it is known that the policy improvement theorem only works locally (Jaakkola al.) 1995). To see why, note that the q- q-values over observations can be written in terms of q- q-values over states: 4n a) Exes (s (s, a). for all E O, a E A o) is the probability of being in state given that the observation o (at any time step), when following policy л. Intuitively, the policy improvement theorem does not work generally here because moving =(0) towards r(o) = arg max 9- (o, increases the expectation over 4- (s, a) without considering how P=(s)o) might change. Conversely, the policy improvement theorem works locally because updating by small amount will also only have small effect on Pa (s)o), making such difference insignificant. Therefore, a policy learning method that takes small update steps is guaranteed to converge locally optimal memoryless policies explaining why actor-critic converges smoothly in the gravity domain (Figure 2). Unfortunately, convergence optimal memoryless policies is not guaranteed for general POMDPs.
(Sigh et al. 1994). Here, (s)
Since memory-augmented environments are a form of POMDP, this local convergence guarantee also applies to them. As such, if the memory can represent the optimal policy, then that solution will be stable given accurate q-value estimations. This raises the question of what conditions for memory would guarantee convergence topic, we considered an idealized version of Jaakkola et al. (1995)ss approach. This agent starts from a random policy π and uses an oracle to compute q- (o, a) for all
globally optimal memoryless policy. To investigate this
and a E A. Then, moves T (() towards arg m A (o, a) a small step 8 for all o E O, and repeats.
Figure shows the behaviour of this algorithm on a variant of the recall task. The rewards were se- lected to encourage convergence to suboptimal solutions (more details in Appendix B). In this envi- ronment, OAI and B1 are enough to encode memoryless policy that is globally optimal. However, environments
note that OA1 converges to a suboptimal solution. Therefore, memory-augmentt might converge to suboptimal solutions even the memory is expressive enough to encode globally optimal policies. We do note that this problem vanishes as we increase the size of the memory in this domain. Unfortunately, convergence to an optimal memoryless policy cannot be guaranteed, even for memories that can model the belief states, as we prove for Bk in Appendix C
5.3 SUMMARY: FROM THEORY TO PRACTICE
The theory suggests that the best approaches for learning effective memoryless policies in memory- augmented environments are methods that exploit the policy improvement theorem locally and eval- uate policies using Monte-Carlo estimates (or n-step TD methods), such as n-step actor-critic, A3C (Mnih et al. 2016), or PPO (Schulman et al. (017). Our empirical evidence also suggests the use of Ok or OAk memories over Bk memories. While the core of our experimental analysis uses PPO, also tested pure TD methods, including Sarsa(A) (A) (Seijen & Sutton 2014) and DDQN (Van Hasselt] et al. 2016). Those results, which shown in Appendix [D6 also favor Ok memories. Finally,

---

Under review as a conference paper at ICLR 2021
Memories:
Memories:
Reward
OA2 OAI
Reward
B5
-10
-10
B2
B1
Other:
Optimal
0 2
5
6
0
1
2
5
6
Policy updates (in thousands)
Policy updates (in thousands)
Figure 5: Tabular experiments in a variation of the recall task (details in Appendix B
note that integrating external memories into existing RL toolkits is trivial. For instance, it takes less than 40 lines of code to integrate each external memory into OpenAl gym (Brockman et al. 2016).
EXPERIMENTAL EVALUATION
We ran experiments on variety of environments with different types of external memory, including our new Ok and OAk memories, as well as the existing k-order memories (Kk) (Mnih et al. 2015) and binary memories (Bk) (Littman 1993). Below, we present results when using PPO (Schulman ee al. 2017) and these memories. We also experimented with using memory (None) and when using an LSTM. TigureBJshows the results. Each line the average reward per episode over runs and the shadow area represents half standard deviation. Details of the domains, hyperparameters, and network architectures can be found in Appendix D We will release our code upon publication.
The left column shows results in the Hallway environments. These environments have been shown to be difficult for PPO with LSTM-ba -bsedd memory in previous work (Toro Icarte et al. 2019) and, indeed, were also unable to get PPO with LSTMs perform any better than a random policy. In contrast, PPO with OA6 and 06 memories is able to solve these tasks.
The middle column shows results in the MiniGrid environment (Chevalier-ooisvert & Willems 2018). We experimented with the RedBlueDoors and MemoryS7 environments because they were specifically designed test the agent's 's memory capabilities. We also decreased the agent's field of view from 8x8 to 3x3 cells to make these problems more challenging. In both cases, O3 and OA3 perform best, as they consistently converge good solutions on all runs. In contrast, the LSTM performance was unreliable: we note that around half of the LSTM runs converged to poor policies.
The previous results used feed-forward networks for function approximation in grid-like domains. To test our approach in visually complex domains using convolutional networks, we also experi- mented with two Atari games: Pong and Seaquest. For these domains, we only gave the agent one frame of the game at time (aside from the current memory state) and followed Machado et al.) 2018)'s recommendations for making the environment stochastic. These domains are almost fully- observable, it is unreasonable to expect Ok, BA, or OAK outperform k-order memory. Still, 03 has comparable performance to K3 in Pong and outperformed LSTMs in Seaquest. This shows that Ok memories Atari when trained by DDQN (see Appendix D.6) but does not when using PPO.
work well in visually complex domains. Note that OA3 performs well in
Finally, note that learning memoryless policies is usually faster than learning histery-based poli- cies. In fact, training PPO with an Ok memory was between 1.06 to 9.85 times faster than training PPO with an LSTM when using CPUs and between 1.71 2.94 times faster when using GPUs. The complete list of speedups can be found in Appendix D.5
7 DISCUSSION, LIMITATIONS, AND FUTURE WORKS
Partial observability is a major challenge when applying RL in the real world. The most popular approaches for tackling partial observability are k-order memories and recurrent neural networks. These approaches are simple to implement, have strong empirical performance, and can be used off sff-th without extra coding effort. The idea of allowing agent to modify external mem- ory using primitive actions is similarly attractive since it can be easily implemented and combined with any RI. method. Unfortunately, historically, such techniques have not worked well in practice
8

---

Under review as a conference paper at ICLR 2021
Hollway-Cookie
MeniiriddReBBBueDoor
Atari-Pong
Avg. reward per ep.
200
pouue pou 22
Ava rown pop ep.
20
150
0.8
10
0.6
100
Avg. poyaal
0
0.4
50
0.2
- 10
20
100 150 200 250
2
8
Training steps (in millions)
Training steps (in millions)
Training steps (in millions)
Hallway- Keys
MiniGridd-emmory57
Atari-Seaquest
Avg. reward per ep.
per dP
Avg. reward per ep:
120
100
0.9
1,500
0.8
1,000
(ene Ava
0.7
0.5
200
50
100
150
200
250
300
Training steps (in millions)
Training steps (in millions)
Training steps (in millions)
K6 B6
06
OA6
None
B3
03
OA3
LSTM
Figure 6: Results over partially observable benchmarks using PPO and different memories.
1999: Zhang al. 2016). We revisited this idea and contributed a theoretical frame- work and new types of external memories. Notably, these memories outperformed k-order memories 2015). binary memories (Peshkin et al., 1999), and LSTM memories (Hausknecht & Stone 2015 in most of our experiments and were faster to train than LSTMs.
Peshhin al.
Mnih al.
We view Ok and OAk memories as the first step towards studying more expressive forms of mem- ories. Ok and OAk are limited by the size of their buffers and that the agent can only push obser- vations into the buffer. This restricts the problems that can be solved. For instance, Ok memories cannot keep track of whether the current time step is odd or even (which a B1 memory could do). We believe that many of the limitations of Ok memories can be overcome by letting the agent determine position in the buffer in which to save (or remove) LSTM policy to control an Ok memory. It will also be interesting to study existing forms of memo- ries in the context of eeemory-augmentee environments. For instance, McCallum instance.McCallummm (1996) showed the effectiveness of tree-based memories in model- RL and tape-like memories and stacks have worked well in supervised learning (Joulin & Mikolov 2015 Zaremba & Sutskever 2015).
observation or alternatively by training an
Another avenue for future work is to further study the theory behind mneoory-augmentar environ- ments. We focused our analysis on the POMDP literature, but it is known that all theoretical guar- antees given for function approximation in RL also apply to partial observability (Sutton & Barto 2018 Chapter 17.3). Although the main conclusions that can be drawn from that body of work similar to those described in Section 5.3 some recent works provide some interesting guarantees. For instance, the con-deluioonal q learrnng algorithm, while impractical, guaranteed converge to globally optimal memoryless policies in memory-augment environments (Lu et al. 2018
CONCLUDING REMARKS
This work presented a lightweight approach tackling partially observable RL. We provided the agent with an external memory and extra actions to write to it, and then used RL to learn memo- ryless policy that jointly decides what to do and what to remember. This idea has been around since the 90s, but this is the first work to show how to make it work well in practice. The key step was to study the theory behind forms of memories that support learning. Using the same RL agent, formed LSTM memories while being faster to train and trivial to implement. Our results suggests a broad array of topics for future exploration in the theory and practice of partially observable RI.
environments and to use that theory to propose novel
external memories outper-

---

Under review as a conference paper at ICLR 2021
REFERENCES
Kamyar Azizzadenesheli, Manish Kumar Bera, and Animashree Anandkumar. Trust region policy optimization for POMDPs. arXiv preprint arXiv:1810.07900 2018.
Andrew Barto and Michael Duff. Monte Carlo matrix inversion and reinforcement learning. In Pro- ceedings of the 7th Conference on Advances in Neural Information Processing Systems (NIPS), pp. 687-694, 1994.
Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning envi- ronment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl gym. CoRR, abs/1606.01544 2016. /ttp:///aaax org/
abs/1606.01544
Maxime Chevalier-Boisvert and Lucas Willems. Minimalistic gridworld environment for OpenAI gym.
2018.
Finale Doshi- Velez. David Pfau, Frank Wood, and Nicholas Roy. Bayesian nonparametric meth- ods for partially-bbseraable reinforcement learning. IEEE transactions on pattern analysis and machine intelligence,
2013.
Gabriel Dulac-Amold, Daniel Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. arXiv preprint arXiv: /2901, 2019.
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, Aviv Tamar, al. Bayesian reinforcement learning: A survey. Foundations and Trends in Machine Learning, 8(5-6):359 483, 2015.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv: 1410.5001, 2014.
arXiv preprint
Ivo Grondman, Lucian Busoniu, Gabriel A. D. Lopes, and Robert Babuska. survey of actor- critic reinforcement learning: Standard and natural policy gradients. IEEE Transactions Systems, Man, and Cybernetics, Part C (Applications and Reviews), 1216:1291 2012.
Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable MDPs, In AAAI Fall Symposium on Sequential Decision Making for Intelligent Agents (AAAI-SDMIAJ5), 2015.
Christopher Hesse, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu. OpenAl baselines.
com/ http//opnaa///// wwwaaaiaaaaeeenneeeeee 2017.
Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne, Optimizing agent behavior over long time scales by transporting value. CoRR, abs/1810.06721, 2018. URL http///xivvvvo// org abs 11101111110 06721
Tommi Jaakkola, Satinder P. Singh, and Michael I. Jordan. Reinforcement learning algorithm for partially observable Markov decision problems. In Proceedings of the 8th Conference on Ad- vances in Neural Information Processing Systems (NIPS), pp. 345-352. 1995.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czamecki, Tom Schaul, Joel Z. Leibo. David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint axXiv:1611.05397, 2016.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. In Proceedings of the 28th Conference on Advances in Neural Information Processing Systems (NIPS), pp. 190-198, 2015.
Arbaaz Khan, Clark Zhang, Nikolay Atanasov, Konstantinos Karydis, Vijay Kumar, and Daniel D. Lee. Memory augmented control networks. arXiv preprint 17Xiv:1709.05706, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1:Xiv:/412.600 2014.
10

---

Under review as a conference paper at ICLR 2021
Yanjie Li, Baoqun Yin, and Hongsheng Xi. Finding optimal memoryless policies of POMDPs under the expected average reward criterion. European Journal of Operational Research, 211(3):556- 567, 2011.
Michael L. Littman. An optimization baaed categorization of reinforcement learning environments. In From Animals to Animuls P Proceedings of the Second International Conference on Simulation of Adaptive Behavior, pp. 262-270, 1993.
Michael L. Littman. Memoryless policies: Theoretical limitations and practical results. In Proceed- ings the 3rd International Conference on Simulation of Adaptive Behavior (SAB), pp. 238-245, 1994.
Michael L. Littman, Richard S. Sutton, and Satinder Singh. Predictive representations of state. In Proceedings of the 15th Conference on Advances in Neural Information Processing Systems (NIPS), pp. 1555-1561, 2002.
Tyler Lu, Dale Schuurmans, and Craig Boutilier. Non-delusional q-learning and value-iteration. In Proceedings of the 31st Conference on Advances in Neural Information Processing Systems (NeurIPS), pp. 9949-9959, 2018.
Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Andrew Kachites McCallum. Reinforcement Learning with Selective Perception and Hidden State. PhD thesis, University of Rochester, 1996.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, Nature, (1807540):529-5333 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of the 33rd International Conference on Machine Learning (ICML), pp. 1928-1937, 2016.
Junhyuk Oh, Valliappa Chockalingam, Satinder Singh, and Honglak Lee. Control of memory, active perception, and action minecraft. In Proceedings of the 33rd International Conference on Machine Learming (ICML), pp. 2790-2799) 2016.
Emilio Parisotto and Ruslan Salakhutdinov. Neural map: Structured memory for deep reinforcement learning. arXiv preprint 1702.0/000 2017.
Leonid Peshkin, Nicolas Meuleau, and Leslie Pack Kaelbling, Learning policies with external mem- ory. In Proceedings of the 16th International Conference on Machine Learning (ICML), pp. 307- 314, 1999.
Pascal Poupart and Nikos Vlassis. Model-based Bayesian reinforcement learning in partially ob- servable domains. In Proceedings of the 10th International Symposium on Artificial Intelligence and Mathematics (ISAIM), pp. 1-2, 2008
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06:0 06447, 2017.
Harm Seijen and Richard S. Sutton. True online TD (lambda). In Proceedings of the 31st Interna- tional Conference on Machine Learning (ICML), pp. 692-700, 2014.
Satinder Singh, Tommi Jaakkola, and Michael I. Jordan. Learning without state-estimation in partially observable Markovian decision processes. In Proceedings of the 11th International Con- ference on Machine Learning (ICML), pp. 284-292. 1994.
Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3(1):9 1988.
11

---

Under review as a conference paper at ICLR 2021
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
Rodrigo Toro Icarte, Ethan Waldie, Toryn Q. Klassen, Rick Valenzano, Margarita P. Castro, and Sheila A. Mcllraith. Learning reward machines for partially observable reinforcement learning. In Proceedings of the 32nd Conference Advances Neural Information Processing Systems (NeurIPS), pp. 15523-1553,, 2019.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. In Proceedings of the 30th AAAI Conference on Artificial Intelligence (AAAI), pp. 2094- 2100, 2016.
Ziyu Wang, Victor Bapst, Nicolas Heess, Volodymyr Mnih, Remi Munos, Koray Kavukcuoglu, and Nando de Freitas. Sample efficient actor-critic with experience replay. arXiv preprint
1rXiv:1611.01224 2016.
Christopher H. Watkins and Peter Dayan. Q-learning. 1992.
Machine learning, 8(3-4):27--292,
Christopher John Cornish Hellaby Watkins. Learning from delayed rewards. PhD thesis, King's College, Cambridge, 1989.
Lucas
Willems.
RL starter files.
https::/gittubb.com//eee 1emms/
r starttr-rileeee 2019.
Wojeiech Zaremba and Ilya Sutskever. Reinforcement learning neural turing machines-revised. arXiv preprint
2015.
Amy Zhang, Zachary C. Lipton, Luis Pineda, Kamyar Azizzadenesheli, Anima Anandkumar, Lau- rent Itti, Joelle Pineau, and Tommaso Furlanello. Learning causal state representations of partially observable environments. arXiv preprint ar 1906::9966.0444 2019.
Marvin Zhang, Zoe McCarthy, Chelsea Finn, Sergey Levine, and Pieter Abbeel. Learning deep neural network policies with continuous memory states. In Proceedings of the 2016 IEEE Inter- national Conference on Robotics and Automation (ICRA), pp. 520-527, 527, 2016.
A FORMAL ANALYSIS
A.1 MEGORY-AUGMENTED ENVIRONMENTS AS POMDPS
In this section, we show how to define memory-augmenter environments as POMDPs. Given a environment & (P, Mpp,, where P A, 7, n) is POMDP and Mp (M, W,T.M .)) is an external memory module for P, we define the POMDP p' = (S', O', A', R' w/, ,','',,'' ,4,) that corresponds to memory-augmentt environment & (P,Mp)
as follows:
= MX0
xx=
Ax
v'((s', ml'i,,,'','),, rs, w)) = p(8', rls,, ".arr
((mm',o')(ss o)) = { m' = m and =
otherwise
= ml m. u(s)n(m)w(ols)
Note that has been included in S' to ensure the consistency of p' and w/ in the case that the next observation is stored in memory. Note also that, by construction, P' is equivalent to & because both environments generate rewards and observations following the same probability distributions from the agent's perspective.
12

---

Under review as a conference paper at ICLR 2021
A.2 THE EXPRESSIVENESS OF AN EXTERNAL MEMORY MODULE
In this section, show that if an external memory module is expressive enough, then the optimal memoryless policy for the ceemor--augmentee environment corresponds to an optimal history-based solution to the original POMDP. We begin with the following definition:
Definition A.1 eeeemor--uddatt function). Given memory-augmented environment, & (P.Mp), %, p) and external memory module Mp = (M, W. m).
where POMDP P = (5, O, A, R,p,W, %/ memory-update function for is defined as f:M x
A + W.
The memory-update function is akin to a deterministic policy for manipulating the memory (and only the memory). It does not determine how we select environment actions, it only dictates what memory-writing action to take once an environment action has been selected, given the current memory state, observation, and the action to be executed. Below, define criteria for external memory module to be able to represent the optimal history-based policy for a POMDP, through the existence of a memory-updating function (i.e., way for manipulating the memory) that can sufficiently summarize the interaction's 's history.
Let 7 og. ao...... Gr-1. Ot be the observation -nttiiationaatiinnn history for the original POMDP P. We can use the memory update function f to construct a corresponding observation-action history Ht' in the memory-augmented environment, where each o replaced by tuple (m;; and each a; replaced by a tuple w.). In particular, we can sample an initial memory state ma according to y, set the initial observation in the nemory-augmented environment as (mo, 00), set the first action as (ao: (mo, oo, do)), and set the next observation as mi) where m, is sampled from I', This process can then be continued until the end of H to create a history H' for the memory-augmented environment.
In the analysis below, we will need to refer to the last memory state m, of such a history H." However, since 17 and I' can be non-deterministic, there may be multiple valid histories H'' that can be generated in this way for any H. Thus, this generation process may result in any one of set of last memory states. For any given history H. we will refer to this set as Oy(H) E 24.
O, A, A, 7, is function of the history Recall that the optimal policy for POMDP (,,
of an interaction: n ="aeeoo lon, ao...., at-1, 0t). When the POMDP model is available, it can also be expressed in terms of belief states. A belief state at step is probability distribution over (being in) each of the states in S. defined as b : :S + [0, 1]. The initial belief state is computed
using the initial observation Op: bs(s) x w(ools) for all s E S. The belief state by=1 then determined from the previous belief state by, the executed action dg and the resulting observation ot+1 as (8') α w(or+1|s) Is') E.es P(s, ae s' (s) for all s' S. In this way, the belief state correctly summarizes the history of an interaction, meaning that the optimal policy for P can then be written as policy of the belief states T (br(b.). For convenience, we will let 6(H) be the belief'
state be for history H Ou, ag. Qt-1,
We can now use the notion of a belief state to define the following:
Definition A.2 (sufficiently expressive). Given a weemory-augmentt environment & = (P,Mp), an external memory module Mp is sufficiently expressive for P if there exists a memory- update function f for & such that for any two histories H ob. db. ... a of and H22 ,, where of = of, the following holds:
og. .f.....a....
if (((*) # b(H(?), then 22(H() ( (H)
Intuitively, given a memory augmented environment, & = (P,Mp) the external memory module Mp sufficiently expressive, if there is a way use it to distinguish between belief states in the POMDP P. If this holds, then there is a memoryless policy for the memory-augmented environ- ment that can act differently (as needed) for each individual belief state. Thus, the following holds immediately:
Proposition A.1. If an external memory module Mp is sufficiently expressive for P, then the opti- mal memoryless policy for the optimal history-based policy for P.
environment È = (P,Mp ) is equivalent to the
13

---

Under review as a conference paper at ICLR 2021
A.3 FORMAL DEFINITIONS FOR KA AND OAK MEMORIES
In this section, we formally define k-order memories and OAk memories as external memory mod- ules. We begin with k-order memories:
Definition A.3 (Kk memories). For POMDP P = (S, A, p, w. n, u), Kk external memory
module for P of size k is defined as Mp (M, W,T.w n), where M (0 [0])* W {T}. 7(00)) = (zero otherwise), and [(m'|m, m, w. 0, 0, = 1f m = (o',0" ... 0²) and m = (0" o .) (zero otherwise).
We now define OAk memories as follows:
Definition A.4 (OAk memories), For POMDP (5, 0, p. 7),,), OAk memory for P defined as MM (M, W, T n),, where M ((O x (0))* W [T, (0h) (zero otherwise), and ['mm m. O, a.T, if = 1 and m' m, or w = T,m (e" = e²,),
and m' (e - (o, a)) where e (o',a') e (zero otherwise).
A.4 SUFFICIRNCY OF KA, BK, OK, AND OAK
In this section, we formally show that the external memory modules described in this work are sufficiently expressive if the buffer size is large enough:
Proposition A.2. For any POMDP P, the following holds:
for 1. Bk is sufficiently expressive for P if k. 2 [log2 1011+ [log, |A[l [log2 u] where u is the maximum number of possible belief states in the set of all histories that end any given observation o. That is, u is defined as
max ((((L) 7 is history that ends in o}
(1)
oEO
2. OAk is sufficiently expressive for P if the belief state of any history is dependent on only the last i Ibssevvation-aciinn pairs.
3. Ok and Kk are sufficiently expressive the belief state of any history only depends on the last i observations.
Proof. For the case of BA; recall that to make belief state update, we use the last belief state the last observation o, the last action a, and the current observation o'. Thus, we define the memory- update function to use the binary memory to store the first three of these components. Specifically, assign a unique integer to each of the possible observations, and use the first [log2 |O] bits to store the integer corresponding to o. We will do the same for the actions, and store the integer uniquely assign integer
corresponding to the last action in the next [log2 |A]] bits. Finally, to each of the belief states that we could have been in when at and we will store that integer in the next [log2 ui] bits. Given the last belief state, last observation, and last action, the agent will therefore be able to distinguish between the belief state that they are in currently with observation o.. Thus, the memory- update function that correctly uses the bits in this way satisfies the conditions in definition A.2 and so Bk is sufficiently expressive for P. Therefore, the statement holds for Bk.
For the case of OAK, if the belief state only depends on the last & observation-action pairs, then the memory-saving function that always saves will be able to distinguish between belief states. Thus OAK is sufficiently expressive for in this case. An analogous argument then applies for Ok and Kk in the conditions outlined above.
B VARIATION OF THE RECALL TASK
We proposed a variation of the recall task to study whether OAk memories can converge to subopti- mal memoryless policies by locally improving its policy, as discussed in Section 5.2 This problem has only one observation o, two actions = {0, 1}. and the discount factor is 1. The episode ends after executing actions. The agent always receives reward except when executing the last tion. The final reward depends on the three actions executed during the episode according to the following table (where a, represents the i-th action):
14

---

Under review as a conference paper at ICLR 2021
(11, 02)
R
(11, (2)
R
(0, 0)
-5
(2, 0)
0
Memories:
(0, 1)
0.5
(2, 1)
0.5
Reward
0.5
OA1
(0, 2)
(2, 2)
-5
B5
(0, 3)
0.5
(2, 0.5
(I, 0)
0
(3, 0)
0
-0.5
Other:
(1, 1)
0.5
(3, 1)
0.5
Optimal
(1,
-0.5
(3, 2)
-5
2
8
10
Policy updates (in thousands)
(1. 0.75
(3, 0.5
Figure 7: Variation of the recall task where Bk always converges to a suboptimal policy.
01, 02, a3)
R
(11, (1,02 03)
(0, 0, 0)
0
(1, 0, 0)
-100
(0, 0, 1)
2
(1, 0, 1)
-100
(0, 1,0) 0)
3
11,
-10
(0, 1)
1
1,
-10
While this problem can be optimally solved using an OA1 memory, the ideal agent from Section 5.2 converges suboptimal policy. The reason is that an optimal policy would execute action (0. T} when the memory buffer is 0, action (1, TT when the buffer contains (o, 0), and action (0, T) when
the buffer contains (o, 1). However, if the agent, while exploring, executes action (1, T) in the first action, then executing (0, Т) from (o. 1) will give a expected return of -100. As this large penalty
will be considered when estimating qu ((,, 1), (0, T)), the agent will prefer action (1, T) over the optimal action (0, T) in (o, 1) - causing the convergence to suboptimal policy.
C CONVERGENCE TO SUBOPTIMAL SOLUTIONS
In this section, we show that having sufficiently expressive external memory module is not enough guarantee that a poliey-improvement based algorithm will find a memoryless policy that is optimal for the original POMDP. We do so by showing that the Jaakkola et al. (1995) algorithm will get stuck local minimum when using BY memory, for any value of k. on variant of the recall task described in Section 5 We refer to this problem as the 4-action recall task. In this case, an episode consists of 2 steps, and there are four possible actions that can be taken. We denote these actions as A = 1, 3). The rewards received at the end of the episode is shown in the table in Figure 7 That figure also shows (empirically) that a B5 memory converges to suboptimal policy in this problem. We also note that the optimal policy is given by selecting action 0 and then action 2. If any other action is taken as the first action, then the optimal action is to take action 3.
We now show that B2 is sufficiently expressive for the 4-action recall task, by identifying policy for the task. The memory state will start with a value 00. When the memory is 00, action 0 should be taken along with the write action 01. Then action 2 should be taken for maximal reward. The memory write action can be arbitrary at this point since the episode over. If, however, any action other than 0 is taken in the initial state when the memory 00, then the write action should be 11. In the resulting state, action 3 will be taken along with any arbitrary memory writing action.
We will see that the q-values for the policy that always takes (1, 3) a local optimum this problem. This shown for B1 in Lemma C.2 and generalized later. As a result, algorithms that follow a policy improvement scheme can get stuck this optimum. We show that this can happen when using an idealized version of the Jaakkola et al. (1995) algorithm, even when starting from uniform policy over all actions. This algorithm alternates between policy evaluation and policy improvement. During policy improvement, the current policy æ is moved step closer to the policy Tº that is greedy on the current v-values of π. That is, for any observation o and action a, the new policy n' is given by r' o) (1 - è) TT(a | + o). where is an input parameter where 0 < < 1. For our analysis, we assume that before every policy improvement update, we computed the actual q-values of the current policy. We refer to this algorithm as the Jaakkola et al. (1993) algorithm with perfect q q-value estimates. For simplicity, we assume the algorithm breaks ties between write actions in favour of the higher valued binary number (i.e. I over 0 or 11 over 10). Under this scenario, we will
15

---

Under review as a conference paper at ICLR 2021
that a sufficiently expressive external memory module is not enough to guarantee convergence to policy that as good as the optimal history- based policy. This is shown Theorem C.6
We first prove that this algorithm fails to converge the optimal policy from the starting state in the case of I,, and then show how the argument can be extended to any Bk. Let us now introduce some notation for the case of l: 1. We will Paw short for (a,, w) (o, m)) and Q short for E, Ro + yR, +... m, a, o, w/. Note that we omit the environment observation o since there is only one observation in this problem and
discount factor of We also define Pa the probability given memory state m of action a marginalized over memory-write actions (i.e. Peb + Pai).
Lemma C.1. Letb, =
I+ Pio So
and t = Σ,
ph m
The q-values for a given policy a for the 4-action
recall task with BI memory are given by:
Pau
Σ 0gig3
riab; +
Σ by USiS3
ra OSi53
Σ
Q
liFia
05i53
Proof. We consider the cases of m = 0 and m separately. Let s denote the initial state of the POMDP, and So, 81, 82, as the states reached after performing actions 0, 1, 2, respectively, in the initial state.
Case m 0:
La
Σ
s = 0, + P = s*IM = O|E, + + + ...M. = w, S ss
05i53
P S = M = 0, + + M = 0] Σ T.
Tia
0Si53
05153
Now, let us turn to [x[S |S = M = 0]:
1 s" M = =1 Σ
|
P s;IM = 0]
05:53
For each Pr[S st|
0]. we now get the following:
P S
8;, M 0
Piu
P. [5 8; M = 0]
b,
P
M
0
+
Ej
P/D
We note that the 1/2 factors come from the fact that if we an infinite number of episodes of this task, half the encountered states would be s* (since every episode starts there), and half would be the result of taking an action in s*
By substituting this into the expression above, we see the statement holds in this case.
Case m 1:
ES [RO + RR ...IM 1 a. W w)
Σ
riaPa s, s
05i53
Again, when we compute P=[S = si| MM 1). we get the following:
P M 1,5 1 8i]
S s, M
t2
Pr M =
Eos
E
The 1/2 factors emerge for the same reason as discussed above. Again, with substitution, this recovers the desired q-value expression.
16

---

Under review as a conference paper at ICLR 2021
Let show that the suboptimal policy that selects action when the memory state is 0 and action 3 when the memory state is 1 is a local optimum:
Lemma C.2. Let иi be the policy for the 4-action recall task with BI memory with pid and
Ps1 = L. Then the policy that is greedy on the q q-values of is Tı itself.
Proof. Consider the q-values of मæ given in Lemma C.1 Since b, = 0. Qa simplifies to
Losic3 Tais Since P3 = and p 0 for i * 3. this means that Qi 1 r13 3/4,
and Qa Taa 1/2. Since pf and p. 0 for o 1, we also have that Qao 1/2 for all
Now notice that 1 1, and 1, = for # such Qu simplifies to Tia-
Let us now consider the greedy policy T on these q-values. In the case that the memory state is O, the greedy action over Q _h
3/ 4 and @o... 1/2 for any a and w where a o or w * 1 is clearly to set pfi = 1. In the case that the memory state is 1, the greedy action is the one that maximizes ria which is a = 3 as that gives 3/4. Therefore, T9 will set p1d = I since it tiebreaks favour of higher memory write actions. Since ~ is clearly the same policy nı, the statement is true. D
We will now show that the Jaakkola et al. (1995) algorithm with perfect q-value estimation will con- verge this local optimum. We begin with the following lemma, which will used for the inductive step in the full result.
Lemma C.3. Let E be a constant such that 0 S S 1. Suppose that a is a policy for the 4-action recall task with B1 memory, such that the policy that is greedy on the q-values of is mi Let ă be defined such that for any ã = - नt + ETi Then the policy that is greedy on the q q-values of ñ is also Th.
Proof. In this proof, we will use Qm for the q-values for म. We will similarly define Paw. b, and as Pm b;, and f; were defined above. Notice that p9 (1 E)p + E, P3i and p... = (1 =)p" otherwise. We will also express i = (1 - E)pli + EII where Il the
E)P + E,
indicator function. We now consider the q-values in the two cases of 0 and
Case m = 0:
We begin by expressing 6; in terms of b;-
P
(1 3)p%
1 1+ P/O (1+))
b, 1 Σ, P0
-E)ab;
1+ (1 (1- =) E,
E, P/O 1 E, P/0
I+ <) E,p,
(2)
where a -11 [1+ E1 (1 E, P/)) In addition, we have the following expression:
1 Σ b
1+ Σ E0s Pi0
Σ
P
1
(3)
1+ Σ Eosis3
1+
Eosiss P/D
1+ Eosisa
0SIS3
05i53
Clearly, similar expression exists for 1 E2 b. We will now these expressions in the following derivation, which begins with the v-vale expression from Lemma C.1
2o =
Σ
Tiabs +
Σ
6
P
rai
05i53
05i53
0SiS3
=
Σ
riabs +
1
Σ Pa
1+ Σ, 50
03153
05153
Σ
Diab;
1+ - -E)E, E, P..
1
Σ
P rai
05153
05153
17

---

Under review as a conference paper at ICLR 2021
Σ 05153
riali + 1 2 P/o
Σ Pi Tai
05153
=
Σ
ria - E)ab; + +a(1
Σ b.)
Σ
Pr Tai
05153
05153
US153
Recall that Di
(Pi P..)
Thus, we can continue the derivation follows:
Qa aw
Σ
(1 s)ab; + a(1 Σ
b,)
Σ
Po + PaIrac
USi53
05153
0S153
(1 - E)a
riabs
USi53
+ a(1
Σ bi) clw-orat ellu=1ras
+
(1 Σ
(Po P Pri
05i53
05153
- =)a
Σ
riab +
Σ b,) Σ ras
05i53
05i53
05i53
+ Ea(11 Σ
bi) (lw=0Tal + Iw=3)
USiS3
(1 e)Quu + +e(1
Σ 05153
b.) (Tw-oTal 11
Notice that the above is merely linear combination (with non-negative coefficients) of Qa and al= 1733 The remaining scalars do not depend on w. Now Q has its maximum when a = and w = 1 because this was the greedy action over the q-values of æ. We also see that Iw==Taa1 + I=== has its maximum when up. Thus, the greedy policy over the q-values of in memory state 0, takes a = 1, which is the same as n.
and w by the way the reward function is set
Case m 1:
By Lemma C.I we have the following:
Ph E, P31
Σ 05153
Tio
(1-5) Eossss Pia + er1o
5) E, p +
Recall that Q1 had the maximal value for = 3 and 1, since these correspond to the greedy action in ny. Since the denominator in t, is constant, this means that the numerator, Σosis3 Ph Tias takes on its maximum value when numerator in the expression above is maximized when 3 and w 1. Since the denominator is constant, Q1 is also maximized with these values. So the greedy policy on म, when m = 1, also agrees with म..
and w Further, Fia maximized for a = 3. Thus, the
Since these are all the possible memory states in B1, the greedy policy for ã must also be mi
this lemma as the inductive step to prove that the Jaakkola al. (1995) algorithm with perfect valuu estimates will converge to this suboptimal policy on the 4-action recall task when using B memory.
We
18

---

Under review as a conference paper at ICLR 2021
Lemma C.A. The Jaakkola al. (1995) algorithm with perfect q-value estimates will converge a suboptimal policy on the 4-action recall task with Bl external memory, when initialized with a uniform policy.
Proof. Let e be the fixed value that the policy is moved towards the greedy policy on every step, where 0 S 1. The proof is by induction. We will show that every step of the algorithm, स will be the policy that is greedy on the current policy. Thus, on every step, the current policy will get closer to ...
Base Case: We begin at the uniform policy Tu. Putting the values in to the expressions in Lemma
we see that the q-values for this policy are
a,
m
(0. 0
0.916
-1.25
(0,
-0.916
-1.25
(1.0)
0.2916
0.5
(1,1)
0.2916
0.5
(2.0)
-1.4583
2.375
(2.1)
-1.4883
-2375
(3.0)
-0.47916
0.5625
(3,1)
0.47991
0.5625
Given the tiebreaking rules defined, we see that the greedy actions the same as those taken by Ti.
Inductive Step: Assume that the policy that is greedy the q-values of the current policy æ is equivalent to m. Thus, the next policy is क (1 2) + em. By Lemma C.3 the policy that is greedy on the q-values of f is also mi
Thus, from initialization of Tus we see that the on every policy improvement step, the policy will get closer to ni, which it much reach in the limit. Since this policy is not as good as the best history- policy, the statement holds.
One may be tempted to think that adding more memory would solve this problem, i.e., it will allow the agent to learn policy that as good as the optimal history-based policy for mePOMDP P. since Bk is sufficiently expressive for k > 2. However, the policy improvement mechanism cannot guarantee that this is learned. We see this below.
Lemma C.S. The Jaakkola et al. (1993) algorithm with perfect q-value estimates will converge to suboptimal policy on the 4-action recall task with Bl external memory, when initialized with uniform policy, for any k 2 1.
Proof. We generalize the argument in Lemmas C.3 and C.4 to Bk: memory. This gives 2 possible memory states O...., 2 1. We note that in this case, due to tiebreak rules, the locally optimal, globally suboptimal policy क, is given by p P1(2=-1) = P = (where m 2 1), an Paw. = oth- erwise. Notice that analogous formulas as given in Lemma C.T for the q-values can be constructed for an arbitrary k. In the case of Q the expression stays as is, there are just more options for w. For convenience, we will refer to b, from Lemma C.I as b,c.
In the case of Q for n 1 (the analogue of Q in Lemma C.I the only change is that t; is po replaced by th Σ, only two points of doccision-making in this task, and the memory state can only be non-zere zerzerr in the second one since it is always initialized to 0.
- We note that the reason this has stayed so similar is because there
Now we will prove a similar statement as Lemma C.4 when the Jaakkola et al. (1995) algorithm is the policy closer to
initialized to the uniform policy, the policy improvement step will always Ti, since ni is the policy that is greedy on the current q- q-values. The proof by induction.
We begin with the uniform policy. We will proceed by identifying the relation between various quantities in the case of = 1 and for an arbitrary >> This will allow us to re-use the computation performed in the base case of Lemma C.4
Base Case:
19

---

Under review as a conference paper at ICLR 2021
Notice that because there always 4 environment actions applicable and 2 memory write actions, there are 2 possible actions in every state. Thus Paw = 1/2 Let Law = 11223 be this value in the case that A = 1. Clearly, Paw 4
Now, let us compare the value of by relative by (which was b; in Lemma C.1 for the case of k 1):
P4o
1+ Losisa Go
3
6. = + E, 1%
2 x 1+ - Losisa 950
by = 26-1 b2
1+ 2 E,
where 8h = (1 + EEosss (035s3P00//0 P100 20-1
Eosies
We will now use this expression, and the fact that the marginalized probabilities p. qa 1/4, to do the following derivation:
Paw = Σ reabf +
1
Σ by
Σ
Peroi
P rai
(4)
05153
05153
05i53
g6
riaby + Bgg
1-
Σ
by
Σ
rai
Tai
(5)
2
USiS3
USiS3
05i53
18
riaby +
Σ
by
Σ
ras ras
(6)
2k-1
05153
05153
05i53
2
gh
Σ
6
Σ 05ig3
q
Tai
2k-1
05153
We note that the second line follows by a similar identity to one used in expression 3
The first term in this last expression (ie. on line 6 is just the same the q-value in the case that k = 1 is multiplied by constant. It is therefore maximized by the values a 1 and w 1 as shown in the base case of Lemma C.2] The second term (ie. line 7 is maximized for the average rewards seen after any first action (since the uniform policy is in use). By inspection, we can see this happens for a = 1 and it is equal for the different memory writing actions. Thus, by our tie-breaking definition, the greedy action will be a 1 and = 2 1 and so the base case holds when the memory state is O.
Let us now consider Qth for m 2 1. Here, the 1.t are exactly the same as the value of I in the case that h 1, when the uniform policy is in use. So applying the same tie-breaking rules means that the greedy action is a = 3, 2* 1, as predicted.
Inductive step: We now consider how Bk for k > is updated. For m = 0, we can use the same
argument as used in Lemma C.3 simply replacing cIu=1Tas with Σ 2*
Iw=iTas. For m 2 1, we have that Q is updated in the exact same way for m 2 as in the proof in Lemma C.3 for m = 1, and further that 1 m 2$ 1 have Q-values which remain unchanged, as the conditional probability 1 of having taken action on having written m does not change since pa.. are all simply scaled by - e for these values m. As such, the greedy actions will again be those given by ny.
Thus, every step of policy improvement during the Jaakkola et al. (1995) algorithm with perfect q- value estimates with move closer to ग, when starting from the uniform policy. Since is suboptimal, the statement is true.
We can now prove our main result: that the Jaakkola et al. (1995) algorithm is not guaranteed to converge to a policy that is as good as best history-basse based policy, even when using sufficiently expressive memory and perfect q-value estimations. This follows from Lemma C.5)and and the fact that B2 is sufficiently expressive for the recall task used in this section.
C.5
Theorem C.6. There exists a POMDP P. such that even with a sufficiently expressive external memory module M, the Jaakkola al. (1995) algorithm with perfect g-value estimates will not
20

---

Under review as a conference paper at ICLR 2021
converge to the optimal memoryless policy from an initially uniform policy Tu. when used on the environment (P.M)
Further, another result can seen from this very same domain, by modifying P.
Proposition C.I. There exists a POMDP P, such that even with a sufficiently expressive external memory module M, the Jaakkola el al. ( (1995) algorithm with perfect q-value estimates will converge to policy that arbitrarily far from the optimal memoryless policy from an initially uniform policy Tu when used the memory auggnettdd environment (P.M)
Proof. This proof uses the same technique given by Singh et al. (1944). To do so, we merely have to scale the rewards used in the 4-action recall task by any constant 2 1. Since the ordering of the action values remains the same, the same arguments will apply. Thus the algorithm will converge to policy that receives & less reward than the optimal history-basee policy, which can be made arbitrarily large by choice of a.
The above counterexample P may further give some indication why OA-type memory may be empirically outperforming B-type memory in our experiments. large part of the reason the above counterexample fails hold is because the states after the first action become confounded, as the agent does not learn to discern, even with arbitrarily large memory, between the different actions to trigger the storage of a given memory state.
For empirical confirmation of this, see Figure ( B5 converging to and OA1 converging to n in POMDP P. Intuitively, more pre-packaged knowledge memory, and more structure given to the agent via the memory write/read mechanism, means the megory-augmented agent has fewer ways that it can confound itself learning how to use memory, which is why choosing memory that comes with more built-in structure helps avoid local minima which caused by this confusion.
D EXPERIMENTAL EVALUATION
This section provides further details about our experimental section.
D.I TABULAR EXPERIMENTS IN THE GRAVITY DOMAIN AND RECALL TASK
In the experiments with tabular q-learning, all the memories were tested using the same hyper- parameters: They explored using e-greedy with 0.01, the discount factor was 0.95, and the learning rate was 0.1. We also used an optimistic initialization for the q-values.
In the experiment with n-step actor-critic, all the memories were also tested using the same hyper- parameters. The discount factor was 0.95. We used soft-max exploratory policy initialized to a uniform distribution. The learning rate for the policy 0.1 and the learning rate for the value function was 0.001.
D.2 THE HALLWAY DOMAINS
Figure 8 shows an overview of the Hallway environments (Toro Icarte et al. 2019). These environ- ments consist of three rooms connected by hallway. The agent (shown as a purple triangle) can move in the four cardinal directions but its actions fail with 5% probability. These are partially observable environments since the agent can only see what it is in the room that it currently oceu- pies, as shown in Figure 8a What makes these tasks difficult is the hallway. The hallway forces the agent to observe long sequences of identical observations (multiple times) to solve task. However, depending previous observations, the optimal actions (and expected rewards) will be completely different when the agent in the hallway.
The cookie domain is shown in Figure 8b In this domain, there is a button in the yellow room that, when pressed, causes a cookie to randomly appear in the red or blue room. The agent receives a reward of +1 for cating the cookie and may then go and press the button again. Pressing the button before eating the cookie removes the existing cookie and delivers a new cookie. Each episode is 10, 000 steps long, during which the agent should attempt eat many cookies possible.
21

---

Under review as a conference paper at ICLR 2021
(a) Agent's view.
(b) Cookie domain.
(c) Keys domain.
Figure 8: The hallway environments.
Figure [c shows another hallway environment, called the keys domain. Here, the agent receives reward +1 when reaches the coffee. To do so, must open the two doors (shown in brown). Each door requires a different key open it, and the agent can only carry one key at time. Initially, the two keys are randomly located either the blue room, the red room, or split between them. When the coffee is reached, the agent is relocated in the middle of the hallway, the doors are locked, and the two keys are randomly placed in the blue and red rooms. Each episode 10, 000 steps long, during which the agent should attempt to reach the coffee many times possible.
While interacting with the environment, the agent is given a "top-down" view of the world repre- sented as a
binary matrices. The first matrix has I in the current location of the agent, the second has a 1 in only those locations that are currently observable, and the remaining matrices each correspond to object in the environment. These remaining matrices only have at those loca- tions that were both currently observable and contained that object (i.e., locations in other rooms are "blacked out").
We ran experiments using the OpenAl baseline implementation of PPO (Hesse et al. 2017 For O6, OA6, K6, and B6, the neural network used had 4 fully connected layers with 512 neurons per layer with tanh activation functions. The LSTM baseline used the same network with an extra LSTM layer of 512 neurons. All approaches were trained using the Adam optimizer (Kingma & Ba 2014).
The approaches 06, OA6, K6, and B6 used learning rate of 1e-5, clipping range of 0.1, 16 training minibatches per update, and one training epoch per update. The value used for the n-step TD updates was 2048. The rest of the hyperparameters were set their default values (Hesse (et al. 2017). The LSTM baseline used the same hyperparameters, but with a learning rate of le-3 and value of of 128 for the n-step TD updates. We found these values worked marginally better than other configurations for the LSTM.
D.3 THE MINTGRID DOMAINS
The MiniGrid domains (Chevalier-Boisver & Willems 2018) consist of agent (red triangle) in grid environment. The agent can only see what is near to it, as shown in Figure The environment contains many objects that the agent can interact with. The agent has 7 actions: turn left, turn right, move forward, pick up an object, drop the object being carried, toggle (open doors, interact with objects), and done (to complete a task). We ran experiments on the two environments that were designed to test the agent's 's memory capabilities: RedBlueDoors and MemoryS7.
Figure 9a shows the RedBlueDoors domain. this environment, the agent randomly located in room with red and blue door. The agent is rewarded by opening the red door and then the blue door (in that order). Concretely, the agent receives a reward of 1 minus a small discount, which proportional to the length of the episode, for solving this problem (and zero otherwise). The episode ends after opening the blue door or after reaching a time limit of 1280 steps.
Figure shows the MemoryS7 domain. In this environment, the agent starts in small room where sees an object and then it has to navigate to the matching object at the other side of the room. Going the right object results in reward of 1. Going to the wrong object results in zero reward. The episode ends when the agent reaches any of the candidate objects or after reaching a time limit of 245 steps.
22

---

Under review as a conference paper at ICLR 2021
9
(a) RedBlucDoors domain.
(h) MemoryS7 domain.
Figure 9: The minigrid environments.
(a) Pong domain. Figure 10: Atari environments.
(b) Seaquest domain.
The observations encoded using one-hot representations. This encoding is binary tensor that contains the objects in each tile that the agent can see. The objects are represented by three one-hot vectors that describe the object's type, color, and status (e.g., a closed red door).
We ran experiments using the OpenAl baseline implementation PPO (Hesse et al. 2017). For O3, OA3, K3, and B3, the neural network used had 5 fully connected layers with 128 neurons per layer with tanh activation functions. The LSTM baseline used the same network with an extra LSTM layer of 128 neurons. All approaches were trained using the Adam optimizer (Kingma & Ba 2014).
The approaches O3, OA3, K3, and B3 used a learning rate of le-5, a value of 128 for n in the n-step TD updates, 8 training minibatches per update, and training epochs per update. The rest of the hyperparameters were set to their default values (Hesse et al.) [2017). The LSTM baseline used the same hyperparameters, but learning rate of 1e-3 in the MemoryS7 and le-5 in the Red- BlueDoors. Those learning rates improved the L.STM performance considerably. Note that these are the hyperparameters recommended by Willems (2019) to train PPO with LSTMs over MiniGrid environments. The only exception is the learning rate, which we fine-tuned.
D.4 THE ATARI DOMAINS
Figure 10 shows snapshots of two atari games: Pong and Seaquest. The objective of Pong is to control the green paddle hit the white ball and beat the opponent (red paddle) by getting the ball passed the opponent's paddle. The objective of Seaquest is to control a submarine to shoot enemies and rescue divers. We ran experiments using the OpenAl's interface (Brockman al., 2016) of the arcade learning environment (Bellemare et al. 2013). To incorporate some partial observability, we gave the agent only one frame as input. Hence, the agent cannot know the direction where objects are moving. Following Machado et al. (2018)'s recommendations, we made these environments stochastic by including sticky actions with 0.25 probability and forcing the agent to randomly execute up to 30 no-op actions at the beginning of each episode. We also used three standard practices introduced by DeepMind (Mnih et al. 2015) for training agents in atari domains. These are to reduce the input image grayscale image of 84x84 pixels, to clip the rewards, and to end the episode when the agent loses life. In these domains, we used the same (convolutional)
neural network and hyperparameters proposed by Hesse al. (017)
23

---

Under review as a conference paper at ICLR 2021
Table 1: Speedups using PPO, Each value shows the speedup relative training an LSTM policy.
Domain Cookie
Memory
CPU
GPU
Domain Keys
Memory
CPU
GPU
None
10.87
2.32
None
7.40
3.45
K6
8.63
2.43
K6
5.13
2.97
B6
10.13
2.56
B6
5.86
3.09
06
9.85
2.03
06
5.37
2.93
OA6
8.24
1.92
OA6
4.32
2.75
LSTM
1.00
1.00
LSTM
1.00
1.00
RBDoors
None
2.99
3.15
MemoryS7
None
2.91
2.24
K3
2.26
2.89
K3
2.73
2.78
B3
2.90
2.92
B3
3.01
2.13
03
2.18
2.94
03
2.80
2.07
OA3
2.48
2.34
OA3
2.40
2.76
LSTM
1.00
1.00
LSTM
1.00
1.00
Pong
None
1.51
2.18
Seaquest
None
1,40
2.72
K3
1.31
1.21
K3
1.26
2.11
B3
1.33
1.76
B3
1.19
1.63
1.51
1.71
1.06
2.11
OA3
0.57
0.67
OA3
0.21
0.40
LSTM
1.00
1.00
LSTM
1.00
1.00
D.5 SPEEDUP COMPARISON
Learning memoryless policies for memory-auee -ugmmettdd environments is usually faster than learning history-based policies. To show this, Table I reports the relative speedups of PPO using different forms of external memories with respect to PPO using an LSTM. To compute these values, we ran each agent for 200, 000 steps. In the CPU experiments, we used 16 logical cores from AMD Ryzen Threadripper 2990WX processor. In the GPU experiments, we used one NVIDIA Tesla P100 GPU and 8 CPUs. In all these cases, learning memoryless policy for memory -ugmmetee environment was faster than learning a history-based policy using an LSTM (sometimes over 10 times faster). The only exception OA3 in Atari games which was surprisingly slow.
D.6 EXPERIMENTS USING TEMPORAL DIFFERENCE METHODS
As discussed in the paper, the POMDP theory suggests that policy learning methods should be pre- ferred when facing difference methods still take advantage of Ok and OAk memories. Here show few results using Sarsa(A) (Seijen & Sutton 2014) and DDQN (Van Hasselt et al. 2016).
environments. However, we have found that pure temporal
D.6.1 EXPERIMENTS USING SARSA(A)
This section discusses results using Sarsa(A) the gravity domain. Sarsa(A) a strong candidate learn policies over memory-augmented environments because the value of e [0, controls the degree in which this algorithm behaves as pure TD method (A 0) or a pure Monte Carlo environments because pure TD methods might be unable to accurately evaluate policies (see Section In fact, Deet,Peeshkin et al. (999) recommended to use Sarsa(A) with a À close to to learn policies over memories.
method 1). This is an important feature for memory-augment
In these experiments, we used e-greedy exploration (with 0.01), discount factor of 0.95, and learning rate of 0. We used an optimistic initialization for the q-values. Tigure[[sshow the results for À equal 0, 0.5, and 1. Note that Sarsa(X) learns to control O1 and OA1 memories regardless of the value of X. However, the most stable performance was obtained when using Sarsa(0.5) with OI. Note that Sarsa(1.0) was able to properly control the B1 memory, but its performance unstable.
D.6.2 EXPERIMENTS USING DDQN
We also experiments using the OpenAl baseline implementation (Hesse al. 2017) of DDQN Van Hasselt et al., 2016). Figure [2uure[2]soore shows the results. Each line the average reward per episode over 3 runs and the shadowed area represents half standard deviation. In general, Ok memories
24

---

Under review as a conference paper at ICLR 2021
Sarsa(O.0)
Sarsa(O.5)
Sarsa(1.0
Avg. 1
Rewwaa avy
4
84V
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Training steps (in millions)
Training steps (in millions)
Training steps (in millions)
None
K1
B1
01
OAI
Figure 11: Tabular experiments in the gravity domain. We reported the avg. reward per 100 steps.
Hallway-Cooke
MiniGrid-RedBlueDDonn RedBlueDoors
Atari-Pong
0.6
Avg. reward per ep.
100
Avg. tevaal per 4p
Aewal p po d0
20
80
0.4
10
60
10
0.2
20
0.1
BAV
0
2
7
2
4
6
2
Training steps (in millions)
Training steps (in millions)
Training steps (in millions)
None
K3
B3
03
OA3
Figure 12: Results over partially observable benchmarks using DDQN and different memories.
outperform the other forms of memories, except for Pong where K3 performs the best (as expected). Note that OA3 performs well in Pong, while it did not when trained using PPO.
In the Cookie domain, the neural network used had fully connected layers with 128 neurons per layer with tanh activation functions. On every step, we trained the network using 64 sampled experiences from a replay buffer of size 100, 000 using the Adam optimizer (Kingma & Ba 2014), a discount factor of 0.99, and learning rate of 1e-5. The exploratory policy was c-greedy with = 0.1. The target network was updated every 100 steps.
In the RedBlueDoors domain, the neural network used had 4 fully connected layers with 128 neurons layer with tanh activation functions. On every step, we trained the network using 32 sampled experiences from a replay buffer of size 100, 000 using the Adam optimizer (Kingma & Ba 2014). discount factor of 0.9, and a learning rate of le-5. The exploratory policy was c-greedy with (..1 The target network was updated every 100 steps.
In Pong, we used the same (convolutional) neural network and hyperparameters proposed by Hesse al. (207) for DDQN.
25


Reviews:
[{"review_id": "LSBm1UX0ISh", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper extends the agent actions with an ability to write to an external memory. The paper does a nice survey of the previous approaches. The paper explains the difficulties with bootstrapping and policy improvement in POMDPs. The paper proposes simple memories for storing a buffer of k observations. The agent has the ability to push or not to push the current observation to the buffer. The whole content of the external memory is visible to the agent at each time step.\n\nThe paper is nicely written and clear. I enjoyed reading the survey of the related work and the explanation of the problems.\nI welcome that the paper tries to use RL to control the external memory. This topic deserves more research.\n\nSuggestions to increase the paper impact:\n1) The survey of the existing literature can be made more valuable by mentioning RL methods suitable for function approximation. As explained in \"Reinforcement Learning: An Introduction\", Section 17.3 \"Observations and State\" by Sutton and Barto, partial observability is a special case of function approximation.\n2) It would be nice to mention the limitations of the different memories. Maybe that would help you to design better memories.\nFor example, if the task requires to count to N, you would need a memory with log2(N) bits. You can also discuss the need to use long n-step returns, e.g., n=2048.\n3) The proposed Ok, OAk memories are not better at all tasks than LSTM. Consider using these memories *together* with LSTM.\nIt would be nice to get the benefits of both approaches.", "extra_scores": {}, "extra_text": {}}, {"review_id": "L1J86AFsebT", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work proposes a lightweight approach to control memory in POMDPs. It is an alternative to heavier overhead approaches such as recurrent networks or memory-augmented networks (Oh et al. 2016).\n\nThe main contribution is to resurrect the old idea of simple lightweight memories and address what made them fail in complex domains by developing novel memory systems. The novel type of memory developed is the $Ok$ memory, wherein the agent is given a choice whether to push the most recent frame into the memory buffer or not. The commonly used $Kk$ memories in contrast push the most recent state into memory by default. \n\n+ The paper develops the theory of memory augmented POMDPs from basic principles and formalizes the learning problem.\n+ A simple idea is methodically and thoroughly explored through experimentation. \n+ The paper is mostly well written and easy to follow with helpful toy problems and useful illustrations.\n\nAn obvious flaw with limited capacity buffer memory systems that store original observations is that longer term dependencies are harder to capture. This work offers a simple fix by including the decision of whether to store an observation or not into the agent's action space. But the system remains limited by having to store a fixed amount of full observations, as compared to more complex memory systems that can chose *what* to write along with when to write (eg. Neural map by Parisotto and Salakhutdinov 2017). Therefore, while the theory and evaluation are extensive, the memory system itself is limited and inefficient for environments where certain features of the state may need to be extracted and tracked for a long duration (not just a few frames).  \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "8_WqDURMhH-", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors propose a series of methods to tackle policy learning in POMDPs. \n\nAt the core of the proposed method sits an augmented memory that the agent is allowed to write on and read from at the time of decision making. \n\nThis allows the agent to store some of its past to be used for future decision making. \n\nThe authors raise a few training, stability, and limitation issues in prior work. And argue that their proposed method improves over the prior works.\n\nThe idea of using augmented memory is exciting and fundamental and definitely worth expanding.\n\nHowever, I found a few issues with the presentation and contribution of this paper that I would be happy to share. \n\n1) Second line of abstract:\n\"Learn- ing memoryless policies is efficient and optimal in fully observable environments.\" It is not clear what the authors mean here. Is the learning of memoryless policy efficient and optimal? If yes, it would great if the authors could parse it.\n\nIf they mean memoryless policies are optimal, then I recommend the authors to restate this statement since it is incorrect. \n\n2) The authors state \"can solve problems that were unsolvable using LSTMs\" since it is an impossibility statement, I would recommend either proving a reference for such a statement or provide proof.\n\n3) \"Notice that neither q-learning nor 5-step actor-critic were able to understand how to use the B1 memory to consistently solve the gravity domain.\" I guess the authors mean the agent using q-learning or 5-step actor-critic were not able to learn how to use ... .\n\n4) I strongly recommend the authors to use more concrete notation. It seems that they study episodic POMDP or maybe a fixed horizon. They mentioned episodic but did not define it. Also, it is not clear what are state, action, and observation spaces. Are they finite? In the analysis I found in the appendix, it seems the authors approach finite ones. But it would be useful to mention it in the main body since there are high-dimensional exps in the paper. Also, q is not defined. I checked the referenced paper,  Jaakkola et al. 1995, there it was also not clear what is q. They first define it for time step zero. Then later use it for any time step. They seem to not define P_\\pi(s|m). It would be great if the authors could define these quantities. The authors state that \"Pπ (s|o) is the probability of being in state s given that the observation is o, when following policy π\" well, at what time step? Please define these terms. \n\n5) I did not understand the point of a recall task example on page 5. \nAs the authors know, there are examples of POMDPs that constant actions are optimal. I am not sure what would a significant conclusion one can draw from such examples.\n\n\n----------------\nGeneral evaluation:\n6) Almost all the theoretical statements are straightforward results of definitions and provided for justification. They are useful in understanding the paper. I appreciate the authors for including them. \n\n7) I again encourage the authors to make their notation a bit more concrete. If they study fixed horizon POMDPs, where the stages (time step in the episode) are encoded?\n\n8) While I appreciate the augmented memory type policy, it seems the authors' proposed method is quite fragile. For example, for the OAk setting, if the agent needs to store k pairs of (o_t,a_{t-1}) to achieve a good solution, then the method breaks? \n\n9) Such fragileness mentioned in 8 seems not to be an issue in methods that learn latent states. How would the authors handle that?\n\n10) Since the contribution is empirical, I would be happy if the authors provide a study against existing baselines, e.g., those referred to in the related works. \n\n11) Regarding memoryless policies in pomdps, I recommend the authors to take a look at Policy Gradient in Partially Observable Environments: Approximation and Convergence. They seem to have some convergence analysis that might be useful.\n\n\n\n---------\nPost rebuttal.\nI would like to thank the authors for their clarifying reply. I will discuss with other reviewers and AC, and update my evaluation further. \n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "5WGYZuVKmvc", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe ideas proposed in the paper are simple and elegant. It seems surprising that this idea has not yet been proposed, but the paper shows great effort from the authors to look for related work, and discuss it. It just seems that a discussion of the work of McCallum in the late 90's is absent from this paper. McCallum, in his PhD thesis (and several papers before that), presents an algorithm that allows the agent to learn what observations to keep in a memory, and for how long. The implementation is, contrary to this paper, based on discrete data structures (trees for instance) and is more complicated than adding a \"push current observation\" action:\n\nMcCallum, R. \"Reinforcement learning with selective perception and hidden state.\" (1997).\n\nWork on Reinforcement Learning Neural Turing Machines is also related to what this paper present, but considers cases where the agent is trained (with backpropagation) on \"what\" to push on a memory tape, instead of simply being able to write the current observation on the memory tape:\n\nZaremba, Wojciech, and Ilya Sutskever. \"Reinforcement learning neural turing machines-revised.\" arXiv preprint arXiv:1505.00521 (2015).\n\nClarity: my main complain with this paper is that it takes a very long time for the reader to see what the main contribution of the paper will be. The abstract and introduction both say that the paper will present a method that works in POMDPs, but does not state the method (the Ok memory). I would strongly advise the authors to explicitly mention, very early in the paper, that they propose to extend the agent with an action that allows the current observation to be enqueued in a fixed-length queue. A second (more minor) remark is that an intuition of what the Ok memory is (in addition to Definition 4.4) would help the reader understand that the agent observes the contents of the memory, and has an extra action to push the current observation onto it (is it really the case, by the way?)\n\nOriginality: the paper seems original. It is quite interesting that so much related work proposed ideas around what this paper proposes, but never exactly the formalism proposed in this paper.\n\nSignificance: the formalism presented in this paper could be applicable to a variety of POMDPs for which a finite set of past (possibly-continuous) observations help with executing good actions. Even though the paper does not point at real-world tasks that fit this definition, I believe that many \"not that Markovian\" real-world tasks could be solved using small Ok memories. As such, the significance of the work presented in this paper could be quite high.\n\nTo summarize my review:\n\n- pros: simple idea, good empirical performance, applicable to many POMDPs\n- cons: many pages of the paper must be read before understanding what its main contribution will be, and a discussion of industrial or real-world tasks that fit the class of POMDPs solvable with Ok memories would have been nice.\n\nI therefore recommend (borderline) acceptance, but strongly suggest that the authors state their contribution in the abstract and introduction.\n\nAuthor response: the updated paper is much clearer, and its abstract and introduction allow to clearly see what will be the contributions of this paper. I thank the authors for having followed my advice about this point. With this problem addressed, I recommend accepting this paper.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a refreshening insight into the classical idea of using external memory for reinforcement learning agents that learn and act in partially observable environments. The authors investigate a number of different memory architectures (Ok, OAk, Kk) and provide an insightful discussion on why we want to restrict the structure of the memory. 

Reviewers generally appreciated the technical contribution of the paper, although not very convinced that this work will have a significant impact on future work. AC is also not sure about the conclusion drawn from the paper, where policies with external memory could have better sample complexity compared to rnn-based policies. BTTT is computationally expensive, but it shall give better direction of which state to jump to, compared to the authors approach where the gradients are stopped at every timestep. So there should be pros and cons about this approach, and AC suspects that the sample complexity improvement actually comes from the fact that authors are explicitly limiting what can be stored in the memory, e.g. O or OA. This advantage can be broken in some other domains. AC admits that this is only a speculation at this point, but the motivation to use the external memory framework proposed in the paper needs to be more carefully investigated.

Decision: Reject