How memory architecture affects learning in a simple POMDP: the two-hypothesis testing problem

Reinforcement learning is generally difficult for partially observable Markov decision processes ( POMDPs ) , which occurs when the agent ’ s observation is partial or noisy . To seek good performance in POMDPs , one strategy is to endow the agent with a finite memory , whose update is governed by the policy . However , policy optimization is non-convex in that case and can lead to poor training performance for random initialization . The performance can be empirically improved by constraining the memory architecture , then sacrificing optimality to facilitate training . Here we study this trade-off in a two-hypothesis testing problem , akin to the two-arm bandit problem . We compare two extreme cases : ( i ) the random access memory where any transitions between M memory states are allowed and ( ii ) a fixed memory where the agent can access its last m actions and rewards . For ( i ) , the probability q to play the worst arm is known to be exponentially small in M for the optimal policy . Our main result is to show that similar performance can be reached for ( ii ) as well , despite the simplicity of the memory architecture : using a conjecture on Gray-ordered binary necklaces , we find policies for which q is exponentially small in 2m , i.e . q ∼ α2m with α < 1 . In addition , we observe empirically that training from random initialization leads to very poor results for ( i ) , and significantly better results for ( ii ) thanks to the constraints on the memory architecture . 1 INTRODUCTION . Reinforcement learning is aimed at finding the sequence of actions that should take an agent to maximise a long-term reward ( Sutton & Barto ( 2018 ) ) . This sequential decision-making is usually modeled as a Markov decision process ( MDP ) : at each time step , the agent chooses an action based on a policy ( a function that relates the agent ’ s state to its action ) , with the aim of maximizing its value ( the expected discounted sum of rewards ) . Deterministic optimal policies can be found through dynamic programming ( Bellman ( 1966 ) ) when MDPs are discrete ( both states and actions belong to discrete sets ) and the agent fully knows its environment ( Watkins & Dayan ( 1992 ) ) . A practical difficulty arises when the agent only have a partial observation of its environment or when this observation is imperfect or stochastic . The mathematical framework is then known as a partially observable Markov decision process ( POMDP ) ( Smallwood & Sondik ( 1973 ) ) . In this framework , the agent ’ s state is replaced by the agent ’ s belief , which is the probability distribution over all possible states . At each time step , the agent ’ s belief can be updated through Bayesian inference to account for observations . In the belief space , the problem becomes fully observable again and the POMDP can thus be solved as a “ belief MDP ” . However , the dimension of the belief space is much larger than the state space and solving the belief MDP can be challenging in practical problems . Some approaches seek to resolve this difficulty by approximating of the belief and the value functions ( Hauskrecht ( 2000 ) ; Roy et al . ( 2005 ) ; Silver & Veness ( 2010 ) ; Somani et al . ( 2013 ) ) , or use deep model-free reinforcement learning where the neural network is complemented with a memory ( Oh et al . ( 2016 ) ; Khan et al . ( 2017 ) ) or a recurrency ( Hausknecht & Stone ( 2015 ) ; Li et al . ( 2015 ) ) to better approximate history-based policies . Here we focus on the idea of Littman ( 1993 ) , who proposed to give the agent a limited number of bits of memory , an idea that has been developed independently in the robotics community where it is known as a finite-state controller ( Meuleau et al . ( 1999 ; 2013 ) ) . These works show that adding a memory usually increases the performance in POMDPs . But to this day , attempts to find optimal memory allocation have been essentially empirical ( Peshkin et al . ( 2001 ) ; Zhang et al . ( 2016 ) ; Toro Icarte et al . ( 2020 ) ) . One central difficulty is that the value is a non-convex function of policy for POMDPs ( Jaakkola et al . ( 1995 ) ) : learning will thus generally get stuck in poor local maxima for random policy initialization . This problem is even more acute when memory is large or when all transitions between memory states are allowed . To improve learning , restricting the policy space to specific memory architectures where most transitions are forbidden is key ( Peshkin et al . ( 2001 ) ; Zhang et al . ( 2016 ) ; Toro Icarte et al . ( 2020 ) ) . However , there is no general principles to optimize the memory architectures or the policy initialization . In fact , this question is not understood satisfyingly even in the simplest tasks- arguably a necessary step to later achieve a broad understanding . Here , we work out how the memory architecture affects optimal solutions in perhaps the simplest POMDP , and find that these solutions are intriguingly complex . Specifically , we consider the twohypothesis testing problem . At each time step , the agent chooses to pull one of two arms that yield random rewards with different means . We compare two memory structures : ( i ) a random access memory ( RAM ) in which all possible transitions between M distinct memory states are allowed ; ( ii ) a Memento memory in which the agent can access its last m actions and rewards . When the agent is provided with a RAM memory , we study the performance of a “ column of confidence ” policy ( CCP ) : the agent keeps repeating the same action and updates its confidence in it by moving up and down the memory sites until it reaches the bottom of the column and the alternative action is tried . The performance of this policy is assessed through the calculation of the expected frequency q to play the worst arm ( thus the smaller q , the better ) . For the CCP , q can be shown to be exponentially small in M . This result is closely related to the work of Hellman & Cover ( 1970 ) on hypothesis testing and its extension to finite horizon ( Wilson ( 2014 ) ) . In practice , we find that learning a policy with a RAM memory and random initialization leads to poor results , far from the performance of the column of confidence policy . Restricting memory transitions to chain-like transitions leads to much better results , although still sub-optimal . Our main findings concerns the Memento memory architecture . Surprisingly , despite the lack of flexibility of the memory structure , excellent policies exist . Specifically , using a conjecture on Grayordered binary necklaces ( Degni & Drisko ( 2007 ) ) , we find a policy for which q is exponentially small in 2m —which is considerably better than q ∼ ln ( m ) /m , optimal for an agent that only plays m times . For Memento memory , we also observe empirically that learning is faster and perform better than in the RAM case . The code to reproduce the experiments is available at https : //anonymous.4open . science/r/two-hypothesis-BAB3 , and uses a function defined here https : // anonymous.4open.science/r/gradientflow/gradientflow . The experiments where executed on CPUs for about 10 thousand CPU hours . 2 POMDPS AND THE TWO-HYPOTHESIS TESTING PROBLEM . 2.1 GENERAL FORMULATION . Definition 2.1 ( POMDP ) . A discrete-time POMDP model is defined as the 8-tuple ( S , A , T , R , Ω , O , p0 , γ ) : S is a set of states , A is a set of actions , T is a conditional transition probability function T ( s′|s , a ) where s′ , s ∈ S and a ∈ A , R : S → R is the reward function1 , Ω is a set of observations , O ( o|s ) is a conditional observation probability with o ∈ Ω and s ∈ S , p0 ( s ) : S → R is the probability to start in a given state s , and γ ∈ [ 0 , 1 ) is the discount factor . A state s ∈ S specifies everything about the world at a given time ( the agent , its memory and all the rest ) . The agent starts its journey in a state s ∈ S with probability p0 ( s ) . Based on an observation o ∈ Ω obtained with probability O ( o|s ) the agent takes an action a ∈ A . This action causes a transition to the state s′ with probability T ( s′|s , a ) and the agent gains the reward R ( s ) . And so on . 1In the literature , R also depends on the action : R : S × A → R. Our notation is not a loss of generality . The set of state can be made bigger S → S ×A in order to contain the last action . Definition 2.2 ( Policy ) . A policy π ( a|o ) is a conditional probability of executing an action a ∈ A given an observation o ∈ Ω . Definition 2.3 ( Policy State Transition ) . Given a policy π , the state transition Tπ is given by Tπ ( s ′|s ) = ∑ o , a∈Ω×A T ( s′|s , a ) π ( a|o ) O ( o|s ) . ( 1 ) Definition 2.4 ( Expected sum of discounted rewards ) . The expected sum of future discounted rewards of a policy π is Gπ = E s0 ∼ p0 s1 ∼ Tπ ( ·|s0 ) s2 ∼ Tπ ( ·|s1 ) . . . [ ∞∑ t=0 γtR ( st ) ] . ( 2 ) Note that a POMDP with an expected sum of future discounted rewards with discount factor γ can be reduced to an undiscounted POMDP ( Altman ( 1999 ) ) , as we now recall ( see proof in Appendix A ) : Lemma 2.1 . The discounted POMDP defined in 2.1 with a discount γ is equivalent to an undiscounted POMDP with a probability r = 1− γ to be reset from any state toward an initial state . In the undiscounted POMDP , the agent reaches a steady state p ( s ) which can be used to calculate the expected sum of discounted rewards Gπ = 1rEs∼p [ R ( s ) ] . 2.2 OPTIMIZATION ALGORITHM . To optimize a policy algorithmically , we apply gradient descent on the expected sum of discounted rewards . First , we parametrize a policy with parameters w ∈ R|A|×|Ω| , normalized to get a probability using the softmax function πw ( a|o ) = exp ( wao ) ∑ b exp ( wbo ) . Then , we compute the transition matrix T̃π , from which we obtain the steady state p using the power method ( See Appendix B ) . Finally , we calculate Gπ by the Lemma 2.1 . Using an algorithm that keeps track of the operations ( we use pytorch Paszke et al . ( 2017 ) ) , we can compute the gradient of Gπ with respect to the parameters w and perform gradient descent with adaptive time steps ( i.e . a gradient flow dynamics ) : d dt w = d dw Gπ ( w ) . ( 3 )

Summary: This paper tackles a two-armed bandit problem (of means $Ber(1/2+\mu)$ and $Ber(1/2-\mu)$, respectively) when memory is limited. Specifically, they model this problem as a POMDP (or, alternatively, hypothesis testing), where the hidden state determines the mean of each arm. The author claim that this simple model might provide some insights on memory architecture in general POMDPs (namely, how the history representation affects the learning process). The authors describe two memory architectures:  (i) A general finite-state machine ('RAM'), which transitions based on the played arm and its reward. Specifically, they focus on the column of confidence policy (CCP). The memory structure - a chain where each end of the chain represents an arm. $+1$ reward moves the state towards an arm and $-1$ moves it away, except for the end states that have low escape probability. The policy - playing the arm that its end-state closer to the current state.   (ii) History is presented as a memory buffer with the last arm plays and their rewards('Memento'). In this case, the authors present a policy that plays the same action sequence cyclically ('necklaces'). These cycles are chained in a Gray-ordering and the policy moves between cycles (change one action) only if a full cycle indicates that this switch is beneficial.  The authors analyze the (stationary) probability of playing the worse actions and show that the RAM model enjoys better (lower) probability. Finally, the authors try to learn a policy when relying on either RAM or Memento architectures. They showed that a random initialization, the Memento memory performs better, while imposing structural constraints that correspond with the suggested policies in (i) and (ii) make RAM memory perform better.