A First-Occupancy Representation for Reinforcement Learning

Both animals and artificial agents benefit from state representations that support rapid transfer of learning across tasks and which enable them to efficiently traverse their environments to reach rewarding states . The successor representation ( SR ) , which measures the expected cumulative , discounted state occupancy under a fixed policy , enables efficient transfer to different reward structures in an otherwise constant Markovian environment and has been hypothesized to underlie aspects of biological behavior and neural activity . However , in the real world , rewards may move or only be available for consumption once , may shift location , or agents may simply aim to reach goal states as rapidly as possible without the constraint of artificially imposed task horizons . In such cases , the most behaviorally-relevant representation would carry information about when the agent was likely to first reach states of interest , rather than how often it should expect to visit them over a potentially infinite time span . To reflect such demands , we introduce the firstoccupancy representation ( FR ) , which measures the expected temporal discount to the first time a state is accessed . We demonstrate that the FR facilitates exploration , the selection of efficient paths to desired states , allows the agent , under certain conditions , to plan provably optimal trajectories defined by a sequence of subgoals , and induces similar behavior to animals avoiding threatening stimuli . 1 INTRODUCTION . In order to maximize reward , both animals and machines must quickly make decisions in uncertain environments with rapidly changing reward structure . Often , the strategies these agents employ are categorized as either model-free ( MF ) or model-based ( MB ) ( Sutton & Barto , 2018 ) . In the former , the optimal action in each state is identified through trial and error , with propagation of learnt value from state to state . By contrast , the latter depends on the acquisition of a map-like representation of the environment ’ s transition structure , from which an optimal course of action may be derived . This dichotomy has motivated a search for intermediate models which cache information about environmental structure , and so enable efficient but flexible planning . One such approach , based on the successor representation ( SR ) ( Dayan , 1993 ) , has been the subject of recent interest in the context of both biological ( Stachenfeld et al. , 2017 ; Gershman , 2018 ; Momennejad et al. , 2017 ; Vértes & Sahani , 2019 ; Behrens et al. , 2018 ) and machine ( Kulkarni et al. , 2016 ; Barreto et al. , 2017b ; a ; 2018 ; Machado et al. , 2020 ; Ma et al. , 2020 ; Madarasz & Behrens , 2019 ) learning . The SR associates with each state and policy of action a measure of the expected rate of future occupancy of all states if that policy were to be followed indefinitely . This cached representation can be acquired through experience in much the same way as MF methods and provides some of the flexibility of MB behaviour at reduced computational cost . Importantly , the SR makes it possible to rapidly evaluate the expected return of each available policy in an otherwise unchanging environment , provided that the reward distribution remains consistent . However , these requirements limit the applicabilty of the SR . In the real world , rewards are frequently non-Markovian . They may be depleted by consumption , frequently only being available on the first entry to each state . Internal goals for control—say , to pick up a particular object—need to be achieved as rapidly as possible , but only one at a time . Furthermore , while a collection of SRs for different policies makes it possible to select the best amongst them , or to improve upon them all by considering the best immediate policy-dependent state values ( Barreto et al. , 2018 ) , this capacity still falls far short of the power of planning within complete models of the environment . Here , we propose a different form of representation in which the information cached is appropriate for achieving ephemeral rewards and for planning complex combinations of policies . Both features arise from considering the expected time at which other states will be first accessed by following the available policies . We refer to this as a first-occupancy representation ( FR ) . The shift from expected rate of future occupancy ( SR ) to delay to first occupancy makes it possible to handle settings where the underlying environment remains stationary , but reward availability is not Markovian . Our primary goal in this paper is to formally introduce the FR and to highlight the breadth of settings in which it offers a compelling alternative to the SR , including , but not limited to : exploration , unsupervised RL , planning , and modeling animal behavior . 2 REINFORCEMENT LEARNING PRELIMINARIES . Policy evaluation and improvement In reinforcement learning ( RL ) , the goal of the agent is to act so as to maximize the discounted cumulative reward received within a task-defined environment . Typically , task T is modelled as a finite Markov decision process ( MDP ; ( Puterman , 2010 ) ) , T = ( S , A , p , r , γ , µ ) , where S is a finite state space , A is a finite action space , p : S × A → ∆ ( S ) is the transition distribution ( where ∆ ( S ) is the probability simplex over S ) , r : S → R is the reward function , γ ∈ [ 0 , 1 ) is a discount factor , and µ ∈ ∆ ( S ) is the distribution over initial states . Note that the reward function is also frequently defined over state-action pairs ( s , a ) or triples ( s , a , s′ ) , but we restrict our analysis to state-based rewards for now . The goal of the agent is to maximize its expected return , or discounted cumulative reward ∑ t γ tr ( st ) . To simplify notation , we will frequently write r ( st ) : = rt and r ∈ R|S| as the vector of rewards for each state . The agent acts according to a stationary policy π : S → ∆ ( A ) . For finite MDPs , we can describe the expected transition probabilities under π using a |S| × |S| matrix Pπ such that Pπs , s′ = pπ ( s′|s ) : = ∑ a p ( s ′|s , a ) π ( a|s ) . Given π and a reward function r , the expected return is Qπr ( s , a ) = Eπ [ ∞∑ k=0 γkrt+k ∣∣∣st = s , at = a ] = Es′∼pπ ( ·|s ) [ rt + γQπr ( s′ , π ( s′ ) ) ] . ( 1 ) Qπr are called the state-action values or simply the Q-values of π . The expectation Eπ [ · ] is taken with respect to the randomness of both the policy and the transition dynamics . For simplicity of notation , from here onwards we will write expectations of the form Eπ [ ·|st = s , at = a ] as Eπ [ ·|st , at ] . This recursive form is called the Bellman equation , and it makes the process of estimating Qπr— termed policy evaluation—tractable via dynamic programming ( DP ; ( Bellman , 1957 ) ) . In particular , successive applications of the Bellman operator T πQ : = r + γPπQ are guaranteed to converge to the true value function Qπ for any initial real-valued |S| × |A| matrix Q . When the transition dynamics and reward function are unknown , temporal difference ( TD ) learning updates value estimates using a bootstrapped estimate of the Bellman target ( Sutton & Barto , 2018 ) . Given a transition sequence ( st , at , rt , st+1 ) and at+1 ∼ π ( ·|st+1 ) , Qπr ( st , at ) ← Qπr ( st , at ) + αδt , δt : = rt + γQπr ( st+1 , at+1 ) −Qπr ( st , at ) . ( 2 ) Once a policy has been evaluated , policy improvement identifies a new policy π′ such that Qπr ( s , a ) ≥ Qπ ′ r ( s , a ) , ∀ ( s , a ) ∈ Qπr ( s , a ) . Helpfully , such a policy can be defined as π′ ( s ) ∈ argmaxaQπr ( s , a ) . The successor representation The successor representation ( SR ; ( Dayan , 1993 ) ) is motivated by the idea that a state representation for policy evaluation should be dependent on the similarity of different paths under the current policy . The SR is defined as a policy ’ s expected discounted state occupancy , and for discrete state spaces can be stored in an |S| × |S| matrix Mπ , where Mπ ( s , s′ ) : = Eπ [ ∑ k γk1 ( st+k = s ′ ) ∣∣∣st ] = Eπ [ 1 ( st = s′ ) + γMπ ( st+1 , s′ ) ∣∣∣st ] , ( 3 ) where 1 ( · ) is the indicator function . The SR can also be conditioned on actions , i.e. , Mπ ( s , a , s′ ) : = Eπ [ ∑ k γ k 1 ( st+k = s ′ ) ∣∣∣st , at ] , and expressed in a vectorized format , we can write Mπ ( s ) : = Mπ ( s , · ) or Mπ ( s , a ) : = Mπ ( s , a , · ) . The recursion in Eq . ( 3 ) admits a TD error : δMt : = 1 ( st ) + γM π ( st+1 , π ( st+1 ) ) −Mπ ( st , at ) , ( 4 ) where 1 ( st ) is a one-hot state representation of length |S| . One useful property of the SR is that , once converged , it facilitates rapid policy evaluation for any reward function in a given environment : rTMπ ( s , a ) = rTEπ [ ∑ k γk1 ( st+k ) ∣∣∣st , at ] = Eπ [ ∑ k γkrt+k ∣∣∣st , at ] = Qπr ( s , a ) . ( 5 ) Fast transfer for multiple tasks In the real world , we often have to perform multiple tasks within a single environment . A simplified framework for this scenario is to consider a set of MDPsM that share every property ( i.e. , S , A , p , γ , µ ) except reward functions , where each task within this family is determined by a reward function r belonging to a setR . Extending the notions of policy evaluation and improvement to this multitask setting , we can define generalized policy evaluation ( GPE ) as the computation of the value function of a policy π on a set of tasks R. Similarly , generalized policy improvement ( GPI ) for a set of “ base ” policies Π is the definition of a policy π′ such that Qπ ′ r ( s , a ) ≥ sup π∈Π Qπr ( s , a ) ∀ ( s , a ) ∈ S ×A ( 6 ) for some r ∈ R. As hinted above , the SR offers a way to take advantage of this shared structure by decoupling the agent ’ s evaluation of its expected transition dynamics under a given policy from a single reward function . Rather than needing to directly estimate Qπ ∀π ∈ Π , Mπ only needs to be computed once , and given a new reward vector r , the agent can quickly peform GPE via Eq . ( 5 ) . As shown by Barreto et al . ( 2017a ) , GPE and GPI can be combined to define a new policy π′ via π′ ( s ) ∈ argmax a∈A max π∈Π rTMπ ( s , a ) . ( 7 ) For brevity , we will refer to this combined procedure of GPE and GPI simply as “ GPI '' , unless otherwise noted . The resulting policy π′ is guaranteed to perform at least as well as any individual π ∈ Π ( Barreto et al. , 2020 ) and is part of a larger class of policies termed set-improving policies which perform at least as well as any single policy in a given set ( Zahavy et al. , 2021 ) .

Summary: The paper proposes an alternative to the successor representation (SR) -- the first-occupancy representation (FR) -- which represents the expected time to first visitation of a state. The paper motivates the representation to be applicable in environments with non-Markovian rewards and show that the proposed FR can handle such cases. The authors present a well-executed comparison to SR and demonstrate improved exploration in MF settings and the ability to plan in a model-based setting. The authors also provide some theoretical analysis of the representation, most notably a convergence result on the Bellman operator for the FR, and planning optimality in a model-based setting.