Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 REINFORCEMENT LEARNING PRELIMINARIES
text: 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).

heading: 3 THE FIRST-OCCUPANCY REPRESENTATION
text: While the SR encodes states via total occupancy, this may not always be ideal. If a task lacks a time limit but terminates once the agent reaches a pre-defined goal, or if reward in a given state is consumed or made otherwise unavailable once encountered, a more useful representation would instead measure the duration until a policy is expected to reach states the first time. Such natural problems emphasize the importance of the first occupancyand motivate the first-occupancy representation (FR). Definition 3.1. For an MDP with finite S, the first-occupancy representation (FR) for a policy π Fπ ∈ [0, 1]|S|×|S| is given by
Fπ(s, s′) := Eπ [ ∞∑ k=0 γk1(st+k = s ′, s′ /∈ {st:t+k}) ∣∣∣st] , (8) where {st:t+k} = {st, st+1, . . . , st+k−1}, with the convention that {st:t+0} = ∅.
That is, as the indicator equals 1 iff st+k = s′ and time t + k is the first occasion on which the agent has occupied s′ since time t, Fπ(s, s′) gives the expected discount at the time the policy first reaches s′ starting from s. The idea of learning policy-dependent distances to target states has a long history in RL (Kaelbling, 1993; Pong et al., 2018; Hartikainen et al., 2020). Previous methods don’t learn these distances as state representations and measure distance in the space of time steps, rather than discount factors. A more thorough discussion can be found in Appendix A.8. We can write a recursive relationship for the FR (derivation in Appendix A.1):
Fπ(s, s′) = Est+1∼pπ(·|s) [ 1(st = s ′) + γ(1− 1(st = s′))Fπ(st+1, s′) ∣∣∣st] (9)
This recursion implies the following Bellman operator, analogous to the one used for policy evaluation:
Definition 3.2 (FR Operator). Let F ∈ R|S|×|S| be an arbitrary real-valued matrix. Then let Gπ denote the Bellman operator for the FR, such that
GπF = I|S| + γ(11T − I|S|)PπF, (10)
where 1 is the length-|S| vector of all ones. In particular, for a stationary policy π, GπFπ = Fπ .
The following result establishes Gπ as a contraction, with the proof provided in Appendix A.4. Proposition 3.1 (Contraction). Let Gπ be the operator as defined in Definition 3.2 for some stationary policy π. Then for any two matrices F, F ′ ∈ R|S|×|S|,
|GπF (s, s′)− GπF ′(s, s′)| ≤ γ|F (s, s′)− F ′(s, s′)|, (11)
with the difference equal to zero for s = s′.
This implies the following convergence property of Gπ . Proposition 3.2 (Convergence). Under the conditions assumed above, set F (0) = I|S|. For k = 0, 1, . . . , suppose F (k+1) = GπF (k). Then
|F (k)(s, s′)− Fπ(s, s′)| < γk (12)
for s 6= s′ with the difference for s = s′ equal to zero ∀k.
Therefore, repeated applications of the FR Bellman operator GkF → Fπ as k → ∞. When the transition matrix Pπ is unknown, the FR can instead be updated through the following TD error:
δFt = 1(st = s ′) + γ(1− 1(st = s′))Fπ(st+1, s′)− Fπ(st, s′). (13)
In all following experiments, the FR is learned via TD updates, rather than via dynamic programming. To gain intuition for the FR, we can imagine a 2D environment with start state s0, a rewarded state sg , and deterministic transitions (Fig. 1a). One policy, π1, reaches sg slowly, but after first encountering it, re-enters sg infinitely often (Fig. 1b). A second policy, π2, reaches sg quickly but never occupies it again (Fig. 1c). In this setting, because π1 re-enters sg multiple times, despite arriving there more slowly than π2, Mπ1(s0, sg) > Mπ2(s0, sg), but because the FR only counts the first occupancy of a given state, Fπ1(s0, sg) < Fπ2(s0, sg). The FR thus reflects a policy’s path length between states.
Policy evaluation and improvement with the FR Like the SR, we can quickly perform policy evaluation with the FR. Crucially, however, the FR induces the following value function:
rTFπ(s, a) = Eπ [∑ k γkrFt+k ∣∣∣st, at] := QπrF (s, a), (14) where rF : S → R is a reward function such that rF (st+k) = r(st+k) if st+k /∈ {st:t+k} and 0 otherwise. In other words, multiplying any reward vector by the FR results in the value function for a corresponding task with non-Markovian reward structure in which the agent obtains rewards from states only once. Policy improvement can then be performed with respect to QπrF as normal. This structure is a very common feature of real-world tasks, such as foraging for food or reaching to a target. Accordingly, there is a rich history of studying tasks with this kind of structure, termed non-Markovian reward decision processes (NMRDPs; (Bacchus et al., 1996; Peshkin et al., 2001;
Littman et al., 2017; Gaon & Brafman, 2020)). Helpfully, all NMRDPs can be converted into an equivalent MDP with an appropriate transformation of the state space (Bacchus et al., 1996).
Most appproaches in this family attempt to be generalists, learning an appropriate state transformation and encoding it using some form of logical calculus or finite state automaton (Bacchus et al., 1996; Littman et al., 2017; Gaon & Brafman, 2020). While it would technically be possible to learn or construct the transformation required to account for the non-Markovian nature of rF , it would be exponentially expensive in |S|, as every path would need to account for the first occupancy of each state along the path. That is, |S| bits would need to be added to each successive state in the trajectory. Crucially, the FR has the added advantage of being task-agnostic, in that for any reward function r in a given environment, the FR can immediately perform policy evaluation for the corresponding rF .
Infinite state spaces A natural question when extending the FR to real-world scenarios is how it can be generalized to settings where |S| is either impractically large or infinite1. In these cases, the SR is reframed as successor features ψπ (SFs; (Kulkarni et al., 2016; Barreto et al., 2017b)), where the dth SF is defined as ψπd (s) := Eπ [∑∞ k=0 γ kφd(st+k) ∣∣st = s], where d = 1, . . . , D and φ : S → RD is a base feature function. The base features φ(·) are typically defined so that a linear combination predicts immediate reward (i.e., wTφ(s) = r(s) for some w ∈ RD), and there are a number of approaches to learning them (Kulkarni et al., 2016; Barreto et al., 2018; Ma et al., 2020). A natural extension to continuous S for the FR would be to define a first-occupancy feature (FF) representation ϕπ , where the dth FF is given by
ϕπd (s) := Eπ [ ∞∑ k=0 γk1(φd(st+k) ≥ θd, {φd(st′)}t′=t:t+k < θd) ∣∣∣st = s]
= 1(φd(st) ≥ θd) + γ(1− 1(φd(st) ≥ θd))Est+1∼pπ [ϕπd (st+1)] (15)
where θd is a threshold value for the dth feature. The indicator equals 1 only if st+k is the first state whose feature embedding exceeds the threshold. Note that this representation recovers the FR when the feature function is a one-hot state encoding and the thresholds {θd}Dd=1 are all 1.

heading: 4 EXPERIMENTS
text: We now demonstrate the broad applicability of the FR, as well as highlight its unique properties with respect to the SR. We focus on 4 areas: exploration, unsupervised RL, planning, and animal behavior.

heading: 4.1 THE FR AS AN EXPLORATION BONUS
text: Table 1: Exploration results. ± values denote 1 SE across 100 trials.
method RIVERSWIM SIXARMS
SARSA + FR 1, 547, 243± 34, 050 11, 9149± 42, 942 SARSA + SR 1, 197, 075± 36, 999 1, 025, 750± 49, 095 SARSA 25, 075± 1, 224 376, 655± 8, 449
Intuitively, representations which encode state visitation should be useful measures of exploration. Machado et al. (2020) proposed the SR as a way to encourage exploration in tasks with sparse or deceptive rewards. Specifically, the SR is used as a bonus in on-policy learning with Sarsa (Rummery & Niranjan, 1994):
δt = rt + β
‖Mπ(st)‖1 + γQπ(st+1, π(st+1))−Qπ(st, at), (16)
where β ∈ R controls the size of the exploration bonus. Machado et al. (2020) show that during learning, ‖Mπ(s)‖1 can act as a proxy for the state-visit count n(s), with ‖Mπ(s)‖−11 awarding a progressively lower bonus for every consecutive visit of s. In the limit as t → ∞, however, ‖Mπ(s)‖−11 → 1 − γ ∀s as π stabilizes, regardless of whether π has effectively explored. To encourage exploration, we’d like for a bonus to maintain its effectiveness over time. In contrast to ‖Mπ(s)‖1, 1 ≤ ‖Fπ(s)‖1 ≤ κ|S| := 1−γ |S|+1
1−γ , where κ|S| > 1 for |S| ≥ 1. Note that ‖Fπ(s)‖1 = κ|S| only if π reaches all states in as many steps. Because ‖Fπ‖1 only grows when new states or shorter paths are discovered, we can instead augment Sarsa as follows:
δt = rt + β‖Fπ(st)‖1 + γQπ(st+1, π(st+1))−Qπ(st, at). (17) 1Recent work (Blier et al., 2021; Touati & Ollivier, 2021) has shown ways of extending the SR to continuous
S with the need for base features. We leave this as an interesting avenue for future work.
We tested our approach on the RIVERSWIM and SIXARMS problems (Strehl & Littman, 2008), two hard-exploration tasks from the PAC-MDP literature. In both tasks, visualized in Appendix Fig. 8, the transition dynamics push the agent towards small rewards in easy to reach states, with greater reward available in harder to reach states. In both cases, we ran Sarsa, Sarsa with an SR bonus (SARSA + SR) and Sarsa with an FR bonus (SARSA + FR) for 5,000 time steps with an -greedy policy. The results are listed in Table 1, where we can see that the FR results in an added benefit over the SR. It’s also important to note that the maximum bonus κ|S| has another useful property, in that it scales exponentially in |S|. This is desirable, because as the number of states grows, exploration frequently becomes more difficult. To measure this factor empirically, we tested the same approaches with the same settings on a modified RIVERSWIM, RIVERSWIM-N, with N = {6, 12, 24} states, finding that SARSA + FR was more robust to the increased exploration difficulty (see Appendix Table 2 and Appendix A.3 for results and more details). We also tested whether these results extend to the function approximation setting, comparing a DQN-style model (Mnih et al., 2015) using the SF and FF as explorating bonuses on the challenging DEEPSEA task (Osband et al., 2020), finding that the advantage of the FF is conserved. See Appendix A.3 for details. Developing further understanding of the relationship between the FR and exploration represents an interesting topic for future work.

heading: 4.2 UNSUPERVISED RL WITH THE FF
text: We demonstrate the usefulness of the FF in the unsupervised pre-training RL (URL) setting, a paradigm which has gained popularity recently as a possible solution to the high sample complexity of deep RL algorithms (Liu & Abbeel, 2021; Gregor et al., 2016; Eysenbach et al., 2018; Sharma et al., 2020). In URL, the agent first explores an environment without extrinsic reward with the objective of learning a useful representation which then enables rapid fine-tuning to a test task.
Continuous MountainCar We first demonstrate that if the test task is non-Markovian, the SR can produce misleading value estimates. To do so, we use a modified version of the continuous MountainCar task (Brockman et al., 2016) (Fig. 2(a)). The agent pre-trains for 20,000 time steps in a rewardless environment, during which it learns FFs or SFs for a set of policies Π which swing back and forth with a fixed acceleration or “power.” (details in Appendix A.3). During fine-tuning, the agent must identify the policy π ∈ Π which reaches a randomly sampled goal location as quickly as possible. We use radial basis functions as the base features φd(·) with fixed FF thresholds θd = 0.7. Fig. 2(b) plots the FF and SF values versus policy power from the start state for two different goal locations. The low-power policies require more time to gain momentum up the hill, but the policies which maximize the SF values slow down around the goal locations, dramatically increasing their total “occupancy” of that area. In contrast, high-powered policies reach the goal locations for the first time much sooner, and so the policies with the highest FF values have higher powers. In the
test phase, the agent fits the reward vector w ∈ RD by minimizing ∑ t ||rt − wTφ(st)||2, as is standard in the SF literature (Barreto et al., 2017b;a; 2018; 2020; Zahavy et al., 2021). The agent then follows the set-max policy (SMP; (Zahavy et al., 2021)), which selects the policy in Π which has the highest expected value across starting states: πSMP ∈ argmaxπ∈Π Es0∼µ [V π(s0)], where V π(s0) = wTϕπ(s0) (with ϕπ replaced by ψπ for SF-based value estimates). Fig. 2(c) shows both the estimated (Vest) and true (Vtrue) values of the SMPs selected using both the SF and FF, along with the value of the optimal policy V ∗. We can see that the accumulation of the SFs results in a significant overestimation in value, as well as a suboptimal policy. The FF estimates are nearly matched to the true values of the selected policies for each goal location and achieve nearly optimal performance.
Robotic reaching To test whether these results translate to high-dimensional problems, we applied the FF to the 6-DoF JACO robotic arm environment from Laskin et al. (2021). In this domain, the arm must quickly reach to different locations and perform simple object manipulations (Fig. 4(a)). We modify the Active Pretraining with Successor features (APS; (Liu & Abbeel, 2021)) URL algorithm, which leverages a nonparametric entropy maximization objective in conjunction with SFs during pre-training
(Hansen et al., 2020) in order to learn useful and adaptable behaviors. Our modification is to replace the SFs with FFs, resulting in Active Pre-training with First-occupancy features (APF), using the same intuition motivating the MountainCar experiments: cumulative features are misleading when downstream tasks benefit from quickly reaching a desired goal, in this case, the object. Here, the agent is first trained for 1e6 time steps using the aforementioned intrinsic reward objective before being applied to a specified reaching task. As additional baselines, we compare against two variants of dynamic distance learning - unsupervised (DDLUS and DDLUS-G; (Hartikainen et al., 2020)). Details can be found in Appendix A.3. We found that the FF accelerated convergence (Fig. 4(b)).

heading: 4.3 PLANNING WITH THE FR
text: While SRs effectively encode task-agnostic, pre-compiled environment models, they cannot be directly used for multi-task model-based planning. GPI is only able to select actions based on a onestep lookahead, which may result in suboptimal behavior. One simple situation that highlights such a scenario is depicted in Fig. 3. As before, there are start and goal states in a simple room (Fig. 3(a)), but here there are three policies comprising Π = {π1, π2, π3} (Fig. 3(b-d)). GPI selects π1 because it is the only policy that reaches the goal sg within one step of the start s0: π1 = maxπ∈Π rTMπ(s0, sg) (Fig. 3(e)). (Note that using GPI with the FR instead would also lead to this choice.) To gain further intuition for the FR versus the SR, we plot the representations for the policies in Appendix Fig. 17. However, the optimal strategy using the policies in Π is instead to move right using π2 and up using π3. How can the FR be used to find such a sequence?
Intuitively, starting in a state s, this strategy is grounded in following one policy until a certain state s′, which we refer to as a subgoal, and then switching to a different policy in the base set. Because the FR effectively encodes the shortest path between each pair of states (s, s′) for a given policy, the agent can elect to follow the policy π ∈ Π with the greatest value of Fπ(s, s′), then switch to another policy and repeat the process until reaching a desired state. The resulting approach is related to the hierarchical options framework (Sutton et al., 1999; Sutton & Barto, 2018), where the planning process effectively converts the base policies into options with—as we show—optimal termination conditions. For a more detailed discussion, see Appendix A.7.
More formally, we can construct a DP algorithm to solve for the optimal sequence of planning policies πF and subgoals sF . Denoting by Γk(s) the total discount of the full trajectory from s to sg at step k of the procedure, we jointly optimize over policies π and subgoals s′ for each state s:
Γk+1(s) = max π∈Π,s′∈S
Fπ(s, s′)Γk(s ′), with πFk+1(s), s F k+1(s) = argmax π∈Π,s′∈S Fπ(s, s′)Γk(s ′).
Intuitively, the product Fπ(s, s′)Γk(s′) can be interpreted as the expected discount of the plan consisting of following π from s to s′, then the current best (shortest-path) plan from s′ to sg. Note
that it is this property of the FR which allows such planning: multiplying total occupancies, as would be done with the SR, is not well-defined. The full procedure, which we refer to as FR-planning (FRP), is given by Alg. 1 in Appendix A.2. Appendix Fig. 13 depicts the resulting policies πF and subgoals sF obtained from running FRP on the example in Fig. 3. The following result shows that under certain assumptions, FRP finds the shortest path to a given goal (proof in Appendix A.4).
Proposition 4.1 (Planning optimality). Consider a deterministic, finite MDP with a single goal state sg, and a base policy set Π composed of policies π : S → A. We make the following coverage assumption: there exists a sequence of policies that reaches sg from a given start state s0. Then Alg. 1 converges so that Γ(s0) = γL ? π , where L?π is the shortest path length from s0 to sg using π ∈ Π.
Performance and computational cost We can see that each iteration of the planning algorithm adds at most one new subgoal to the planned trajectory from each state to the goal, with convergence when no policy switches can be made that reduce the number of steps required. If there are K iterations, the overall computational complexity of FRP is O(K|Π||S|2). The worst-case complexity occurs when the policy must switch at every state en route to the target—K is upper-bounded by the the number of states along the shortest path to the goal. In contrast, running value iteration (VI; (Sutton & Barto, 2018)) for N iterations is O(N |A||S|2). Given the true transition matrix P and reward vector r, VI will also converge to the shortest path to a specified goal state, but FRP converges more quickly than VI whenever K|Π| < N |A|. To achieve an error between the estimated value function and the value function of the optimal policy, VI requires N ≥ 1(1−γ) log 2 (1−γ)2 (Puterman, 1994), which for γ = 0.95, = 0.1, e.g., gives N ≥ 180 iterations2. To test convergence rates in practice, we applied FRP, VI, and GPI using the FR to the classic FOURROOMS environment (Sutton et al., 1999) on a modified task in which agents start in the bottom left corner and move to a randomly located goal state. Once the goal is reached, a new goal is randomly sampled in a different location until the episode is terminated after 75 time steps. For GPI and FRP, we use four base policies which each only take one action: {up, down, left, right}, with their FRs learned by TD learning. We ran each algorithm for 100 episodes, with the results plotted in Fig. 5(a). Note that here GPI is equivalent to FRP with K = 0 iterations. To see this, observe that when there is a single goal sg such that only r(sg) > 0, the policy selected by GPI is
πGPI(s) ∈ argmax π∈Π rTFπ(s) = argmax π∈Π r(sg)F π(s, sg) = argmax π∈Π Fπ(s, sg). (18)
When there are n goal states with equal reward, finding the ordering of the goals that results in the shortest expected path is in general O(n!) (see Appendix A.6). Due to the nature of the base policies used above, the number of subgoals on any path is equal to the number of turns the agent must take from its curent state to the goal, which for this environment is three. We can then see that FRP reaches the optimal performance obtained by the converged VI after K = 3 iterations (Fig. 5(a)). In contrast, for the same number of iterations, VI performs far worse. This planning process must be repeated each time a new goal is sampled, so that the computational benefits of FRP versus traditional DP methods compound for each new reward vector. Example FRP trajectories between goals for K = 0 and K = 3 iterations are plotted in Fig. 5(b). Finally, to test FRP’s robustness to stochasticity, we added transition noise to the FOURROOMS task. That is, the agent moves to a random adjacent state with probability regardless of action. We compared FRP to converged VI for increasing , with the
2Other DP methods like policy iteration (PI), which is strongly polynomial, converge more quickly than VI, but in the example above, PI still needs N ≥ log(1/((1− γ) ))/(1− γ) = 106 (Ye, 2011), for instance.
results plotted in Fig. 5(c), where we can see that FRP matches the performance of VI across noise levels. It’s important to note that this ability to adaptively interpolate between MF and MB behavior, based on the value of K, is a unique capability of the FR compared to the SR3. The same FRs can be combined using DP to plan for one task or for GPI on the next.
4.4 ESCAPE BEHAVIOR
In prey species such as mice, escaping from threats using efficient paths to shelter is critical for survival (Lima & Dill, 1990). Recent work studying the strategies employed by mice when fleeing threatening stimuli in an arena containing a barrier has indicated that, rather than use an explicit cognitive map, mice instead appear to memorize a sequence of subgoals to plan efficient routes to shelter (Shamash et al., 2021). When first threatened, most animals ran along a direct path and into the barrier. Over subsequent identical trials spanning 20 minutes of exploration, threat-stimulus presentation, and escape, mice learned to navigate directly to the edge of the wall before switching direction towards the
shelter (Fig. 6). Follow-up control experiments suggest that mice acquire persistent spatial memories of subgoal locations for efficient escapes. We model this task and demonstrate that FRP induces behavior consistent with these results.
We model the initial escape trial by an agent with a partially learned FR, leading to a suboptimal escape plan leading directly to the barrier. Upon hitting the barrier, sensory input prompts rapid re-planning to navigate around the obstacle. The FR is then updated and the escape plan is recomputed, simulating subsequent periods of exploration during which the mouse presumably memorizes subgoals. We find a similar pattern of behavior to that of mice (Fig. 6(b)). See Appendix A.3 for experimental details.
We do not claim that this is the exact process by which mice are able to efficiently learn escape behavior. Rather, we demonstrate that the FR facilitates behavior that is consistent with our understanding of animal learning in tasks which demand efficient planning. Given the recent evidence in support of SR-like representations in the brain (Stachenfeld et al., 2017; Momennejad et al., 2017), we are optimistic about the possibility of neural encodings of FR-like representations as well. We also re-emphasize that this type of rapid shortest-path planning is not possible with the SR.

heading: 5 CONCLUSION
text: In this work, we have introduced the FR, an alternative to the SR which encodes the expected path length between states for a given policy. We explored its basic formal properties, its use as an exploration bonus, and its usefulness for unsupervised representation learning in environments with an ethologically important type of non-Markovian reward structure. We then demonstrated that, unlike the SR, the FR supports a form of efficient planning which induces similar behaviors to those observed in mice escaping from perceived threats. As with any new approach, there are limitations. However, we believe that these limitations represent opportunities for future work. From a theoretical perspective, it will be important to more precisely understand FRP in stochastic environments. For the FF, we have limited understanding of the effect of feature choice on performance, especially in high dimensions. FRP is naturally restricted to discrete state spaces, and it could be interesting to explore approximations or its use in partially-observable MDPs with real-valued observations and discrete latents (e.g., (Vértes & Sahani, 2019; Du et al., 2019)). Further exploration of FRP’s connections to hierarchical methods like options would be valuable. Finally, it would be informative to test hypotheses of FR-like representations in the brain. We hope this research direction will inspire advancements on representations that can support efficient behavior in realistic settings.
3We’d like to stress that this claim of uniqueness is only with respect to the SR. Previous work also explores the use of MF methods to support MB learning (e.g., Pong et al. (2018))
Reproducibility statement Experiments: All experiments are described in detail in Section 4 in the main text and in Appendix A.3. We have attached code for the tabular experiments, and in Appendix A.3 we provide a link to the base implementation for APS as well as an explanation of where to find the hyperparameters and modifications made to produce APF. Theoretical results: Theoretical results, including assumptions and proofs, are provided in Appendix A.1 and Appendix A.4.

heading: A APPENDIX

heading: A.1 FR RECURSION
text: For clarity, we provide the derivation of the recursive form of the FR below:
Fπ(s, s′) = Eπ [ ∞∑ k=0 γk1(st+k = s ′, s′ /∈ {st:t+k}) ∣∣∣st]
= Eπ [ 1(st = s
′, s′ /∈ ∅) + ∞∑ k=1 γk1(st+k = s ′, s′ /∈ {st:t+k})
∣∣∣st]
= Eπ [ 1(st = s ′) +
∞∑ k=1 γk1(st+k = s ′, st 6= s′, s′ /∈ {st+1:t+k}) ∣∣∣st] = Est+1∼pπ(·|s) [ 1(st = s ′) + γ1(st 6= s′)Fπ(st+1, s′) ∣∣∣st]
= Est+1∼pπ(·|s) [ 1(st = s ′) + γ(1− 1(st = s′))Fπ(st+1, s′) ∣∣∣st]
(19)

heading: A.2 FRP ALGORITHM
text: We present the full algorithm for FR planning (FRP) below.
Algorithm 1: FR Planning (FRP)
1: input: goal state sg , base policies Π = {π1, . . . , πn} and FRs {Fπ1 , . . . , Fπn}. 2: // initialize discounts-to-goal Γ 3: Γ0(s)← −∞ ∀s ∈ S 4: for s ∈ S do 5: Γ1(s)← maxπ∈Π Fπ(s, sg) 6: πF1 (s), sF1 (s)← argmaxπ∈Π Fπ(s, sg), sg 7: end for 8: // iteratively refine Γ 9: k ← 1
10: while ∃s ∈ S such that Γk(s) > Γk−1(s) do 11: for s ∈ S do 12: Γk+1(s)← maxπ∈Π,s′∈S Fπ(s, s′)Γk(s′) 13: πFk+1(s), s F k+1(s)← argmaxπ∈Π,s′∈S Fπ(s, s′)Γk(s′) 14: end for 15: k ← k + 1 16: end while 17: return πF , sF

heading: A.3 ADDITIONAL EXPERIMENTAL DETAILS
text: All experiments except for the robotic reaching experiment were performed on a single 8-core CPU. The robotic reaching experiment was performed using four Nvidia Quadro RTX 5000 GPUs.
Tabular exploration We reuse all hyperparameter settings from Machado et al. (2020) in both the RIVERSWIM and SIXARMS environments, with the only difference being a lower value for β, the exploration bonus coefficient, for the FR, as the bonuses given by the FR are generally larger. The hyperparameters are {α, η, γSR/FR, β, , η}, which are the Sarsa learning rate, the SR/FR learning rate, the SR/FR discount factor, the exploration bonus coefficient, and the probability of taking a random action in the -greedy policy. For RIVERSWIM, these values were {0.25, 0.01, 0.95, 100/50, 0.1}, respectively, and for SIXARMS they were {0.1, 0.01, 0.99, 100/50, 0.01}. For the RIVERSWIM-N task, we chose N = {6, 12, 24} as default RIVERSWIM has N = 6 states, and we chose to successively double the problem size. As N increased, the number of unrewarded central states was multiplied (with the same transition structure), while the endpoints remained the same. It’s also worth noting that β could be manually adjusted upwards to compensate for the SR bonus’ invariance to problem size, though this would require a longer hyperparameter search generally, which we believe is less preferable to a bonus which naturally scales.

heading: RIVERSWIM-N SARSA + FR SARSA + SR SARSA
text: Exploration with function approximation In order to test the usefulness of the FR/FF in a function approximation setting, we use a similar approach to Machado et al. (2020). That is, we train a modified DQN agent (Mnih et al., 2015) using an architecture inspired by Machado et al. (2020) and Oh et al. (2015) (see Fig. 9), such that the base feature representation φ(s) is an intermediate layer of the network. Like the standard DQN, the architecture outputs an |A|-length vector of predicted Q-values for the current state, trained off-policy using minibatches of transition tuples
{(sit, ait, rit, sit+1)}Bi=1 (where B is the minibatch size) to minimize the squared Bellman error
LQ = B∑ i=1 ‖rit + γmax a Q−(s i t+1, a)−Q(sit, ait)‖22 (20)
via gradient descent (the subscript − on the target Q-values indicates that gradients do not flow through it). Unlike the standard DQN, the network has two additional output heads. The first is a reconstruction head which, given the base feature representation of a state st and an embedding of the subsequent action at produces a prediction of the following state ŝt+1. It’s trained to minimize the reconstruction loss
Ls = B∑ i=1 ‖st+1 − ŝt+1‖22. (21)
The final output head of the network is an FF/SF prediction, trained using the squared FF error in the former case:
Lϕ = B∑ i=1 ‖φ̃(st) + γ(1− φ̃(st))ϕ−(st+1)− ϕ(st)‖22, (22)
and the squared SF error in the latter
Lψ = B∑ i=1 ‖φ(st) + γψ−(st+1)− ψ(st)‖22. (23)
For the FF, the features φt are passed through a sigmoid function to compress them in the range [0, 1] and then thresholded at 0.75. The total loss is then given as a weighted combination
L = wQLQ + wsLs + wXLX , (24)
where X ∈ {ϕ,ψ} and wQ, ws, wX ∈ R are fixed weights. As in Machado et al. (2020), gradients from LQ and Ls, but not LX , are permitted to flow through to φ. Thus, the base features are trained to be both reward-predictive and to carry information about the environment transition dynamics. The norm of the FF/SF vector is then used to compute an intrinsic exploration bonus to the task reward in the same manner as in the tabular setting.
To test this model, we chose the DEEPSEA task from the Behavior Suite (bsuite; (Osband et al., 2020)) set of benchmark tasks. DEEPSEA is a challenging exploration task set up in an N ×N grid (Fig. 10, top left). At the beginning of each episode, the agent starts at the top left of the grid. Each time step, the agent descends one level, and can choose to move either right or left. The episode ends after N steps, when the agent reaches the bottom level. There is a small negative reward of −0.01 if the agent moves right, but if the agent moves right N times in a row, there is a large reward +1
located at the bottom right of environment—this is the only policy which nets the agent a positive reward. In the bsuite framework, the agent is separately trained on increasing problem sizes N = 5, 6, 7, . . . , 50 for 10,000 episodes each, with the final score the proportion of N for which the agent reached an average regret of less than 0.9 faster than 2N episodes.
We tested the standard DQN, DQN + SF (Machado et al., 2020), and DQN + FF agents on this task, with training hyperparameters described in Table 3. For all models, the network consisted entirely of fully-connected layers, with φ(st) being a 2-layer MLP with 64 units per layer, Q(st, ·) being a linear function of φ(st) with |A| = 2 units, ŝt+1 consisting of a 64-unit layer followed by an N2-unit output layer (the action embedding is 64-dimensional as well), and the SFs/FFs also a 64-dimensional linear layer over φ(st). The agent is trained using -greedy action selection. Our DQN implementation was coded in JAX (Bradbury et al., 2018) and based off that of Osband et al. (2020). To select the values of β, ws, and wX (wQ was always kept at 1) we performed a sweep over β ∈ {0.01, 0.05, 0.1} and wX , wQ ∈ {0.001, 0.1, 1, 10, 100, 1000}, choosing the best-performing values for each method.

heading: HYPERPARAMETER DQN + FF DQN + SF DQN
text: Results are presented in Fig. 10. We can see that DQN+FF significantly outperforms the other methods (top right), with the intuition from the tabular experiments—particularly RIVERSWIMN—carrying over into the function approximation setting. That is, as N increases, the norm of the SF approaches its asymptotic value regardless of the degree of exploration. In contrast, for the FF, the maximum bonus scales with the problem size. This enables the bonus to remain effective in environments with larger state spaces. We hope to investigate this approach and its theoretical properties further in future work.
MountainCar experiment In our version of the task, the feature representations are learned in a rewardless environment, and at test time the reward may be located at any location along the righthand hill. We evaluate the performance of a set of policies Π = {πi} with a constant magnitude of acceleration and which accelerate in the opposite direction from their current displacement when at rest and in the direction of their current velocity otherwise (see Python code below for details). That is, each πi will swing back and forth along the track with a fixed power coefficient ai. For each possible reward location along the righthand hill, then, the best policy from among this set is the one whose degree of acceleration is such that it arrives at the reward location in the fewest time steps. There is a natural tradeoff–too little acceleration and the cart will not reach the required height. Too much, and time will be wasted traveling too far up the lefthand hill and then reversing momentum back to the right.
We hypothesized that the FF would be a natural representation for this task, as it would not count the repeated state visits each policy experiences as it swings back and forth to gain momentum.
We defined a set of policies with acceleration magnitudes |ai| = 0.1i for i = 1, . . . , 9, and learned both their SFs and FFs via TD learning on an "empty" environment without rewards over the course of 100 episodes, each consisting of 200 time steps, with the SFs using just the simple RBF feature functions without thresholds. Python code for the policy class is shown below.
1 class FixedPolicy: 2 3 def __init__(self, a): 4 # set fixed acceleration/power 5 self.a = a 6 7 def get_action(self, pos, vel): 8
9 if vel == 0: 10 # if stopped, accelerate to the opposite end of the environment 11 action = -sign(pos) * self.a 12 else: 13 # otherwise, continue in the current direction of motion 14 action = sign(vel) * self.a 15 16 return action
We repeated this process for 20 runs, with the plots in Fig. 7 showing the means and standard deviations of the TD learning curves across runs. For the FFs, the thresholds were constant across features at θd = θ = 0.7. Because of the nature of the environment, all of the policies spent a
significant portion of time coasting back and forth between the hills, causing their SFs to accumulate in magnitude each time states were revisited.
Given the learned representations, we then tested them by using them as features for policy evaluation in different tasks, with each task containing a different rewarded/absorbing state. Note that a crucial factor is that the representations were learned in the environment without absorbing states. This is natural, as in the real world reward may arise in the environment anywhere, and we’d like a representation that can be effective for any potential goal location.
Fig. 11 shows the value of the FF and SF at the start state for different policies with fixed power and for different feature dimensions (number of basis functions) in continuous MountainCar. The results show that the policies for which the FF is highest is closer in power to the optimal policy than for the policies at which the SF is greatest across all but the coarsest feature dimensionalities. This provides an indication of the robustness of the FF to the choice of feature dimensionality.
Robotic reaching experiment We used the custom JACO domain as well as the APS base code from Laskin et al. (2021), located at this link: https://anonymous.4open.science/r/ urlb/README.md. Both the critic and actor networks were parameterized by 3-layer MLPs with ReLU nonlinearities and 1,024 hidden units. Observations were 55-dimensional with 10-dimensional features φ(·). For all other implementation details, including learning rates, optimizers, etc. see the above link. All hyperparameters and network settings are kept constant from those provided in the linked .yaml files. All experiments were repeated for 10 random seeds.
We now describe each training phase. Pre-training: Agents were trained for 1M time steps on the rewardless JACO domain by maximizing the intrinstic reward
rintrinsic(s, a, s′) = rexploit(s, a, s′) + rexplore(s, a, s′)
= wTφ(s) + log 1 + 1 k ∑ h(j)∈Nk(φ(s′)) ‖φ(s′)− φ(s′)(j)‖nhnh , (25) where w ∈ RD, D = 10 is a random reward vector sampled from a standard Gaussian distribution and the righthand term is a particle-based estimate of the state-based feature entropy, with Nk(·) denoting the k nearest-neighbors (see Liu & Abbeel (2021) for details). In standard APS, this reward is used to train the successor features by constructing the bootstrapped target
yAPS = rintrinsic(s, a, s′) + γwTψ(st+1, a ′, w), (26)
where a′ = argmaxaw Tψ(s′, a, w). For the FF, we make the following modification:
yAPF = yAPF−exploit + yexplore (27)
= wTφ̃(s)︸ ︷︷ ︸ :=rF (s) +γwT(1− φ̃(s))V (st+1) + rexplore(s, a, s′) + γV (st+1) (28)
= rF (s) + rexplore + γ [ wT1− rF + 1 ] V (st+1), (29)
with V (st+1) = maxa′ wTϕ(st+1, a′, w), φ̃(·) the thresholded base features, such that φ̃d(st) = 1(φ(st) ≥ θd, {φd(st′)}t′=0:t < θd), 1 is the D-length vector of ones, and φ(st) ∈ [0, 1]. Interestingly, if the features are kept in the range [0, 1], φ(·) can be thought of encoding a form of “soft” first feature occupancy, rather than the hard threshold given by the indicator function.
The agent is then trained with the off-policy deep deterministic policy gradient (DDPG; (Lillicrap et al., 2019)) algorithm, where given a stored replay buffer of transitions D = {(st, at, rt, st+1)}, the (SF/FF) critic Qω (with Q formed from either the SF or FF and ω being the parameters) is trained to minimize the squared Bellman loss
LQ(ω,D) = E(st,at,rt,st+1)∼D [ (yr −Qω(st, at))2 ] , (30)
where in the pre-training phase yr ∈ {yAPS, yAPF} (the target parameters are an exponential moving average of the weights—gradients do not flow through them). The deterministic actor πθ is trained using the derministic policy gradient loss:
Lπ(θ,D) = Est∼D [Qφ(st, πθ(st))] . (31)
Fine-tuning: After pre-training, the agent is fine-tuned on the target task, REACHTOPLEFT, where the learning proceeds exactly as in the pre-training phase, but instead of intrinsic reward, the agent is given the task reward—that is, yr = rtaskt + γV (st+1). We performed this task-specific training for an additional 1M steps.
As additional baselines, we implemented two versions of the dynamic distance learning method of Hartikainen et al. (2020), the original DDLUS and an additional variant DDLUS-G. In standard DDLUS, the goals g ∈ S for the pre-training phase are generated according to
g? ∈ argmax g∈D dπ(s0, g),
where D is a stored set of trajectories. In Hartikainen et al. (2020), this is useful as a mechanism for encouraging the agent to learn skills which move the agent as far as possible from the start state. However, since the object for manipulation in the fine-tuning phase of the JACO task is not typically especially far from the initial point, we also tested DDLUS-G, which samples goals in the same manner as APS and APF, via
g ∼ N (0, I),
to ensure a more fair comparison. All other hyperparameters match those of Hartikainen et al. (2020), with the exception that the base agent is DDPG, implemented using the same framework as APS and APF, rather than SAC (Haarnoja et al., 2018). For a more detailed discussion of DDL and its relationship to the FR/FF, see Appendix A.8.
In future work, it would be interesting to explore the interaction of the FF with other off-policy algorithms (Haarnoja et al., 2018; Fujimoto et al., 2018; Moskovitz et al., 2021) and whether on-policy learning (e.g., with (Schulman et al., 2017; Kakade, 2002; Williams, 1992; Moskovitz et al., 2020; Hartikainen et al., 2020)) has different effects.
FourRoom experiments The FourRoom environment we used was defined on an 11×11 gridworld in which the agent started in the bottom left corner and moved to a known goal state. The action space was A = {up,right,down,left} with four base policies each corresponding to one of the basic actions. TD Learning curves for the base policies are depicted in Fig. 12. Once reaching the goal, a new goal was uniformly randomly sampled from the non-walled states. At each time step, the agent received as state input only the index of the next square it would occupy. Each achieved goal netted a reward of +50, hitting a wall incurred a penalty of −1 and kept the agent in the same

heading: TD learning curves, FourRooms
text: place, and every other action resulted in 0 reward. There were 75 time steps per episode—the agent had to reach as many goals as possible within that limit. The discount factor γ was 0.95, and the FR learning rate was 0.05. In order to learn accurate FRs for each policy, each policy was run for multiple start states in the environment for 50 episodes prior to training. FRP (for different values of K), GPI, and VI were each run for 100 episodes. VI was given the true transition matrix and reward vector in each case. In the stochastic case, for each level of transition noise = 0.0, 0.1, 0.2, . . . , 1.0, both VI and FRP were run to convergence (≈ 180 iterations for VI, 3 iterations for FRP) and then tested for 100 episodes.

heading: 1,2 1,2 1,2 1,2

heading: 1,2 1,2 1,2 1,2

heading: 2,2 2,2 1,2 2,2
text: Escape experiments The escape experiments were modeled as a 25 × 25 gridworld with eight available actions,
A = {up,right,down,left,up-right,down-right,down-left,up-left} (32)
and a barrier one-half the width of the grid optionally present in the center of the room. The discount factor γ was 0.99, and the FR learning rate was 0.05. In test settings, the agent started in the top central state of the grid, with a single goal state directly opposite. At each time step, the agent receives only a number corresponding to the index of its current state as input. The base policy set Π consisted of eight policies, one for each action in A. Escape trials had a maximum of 50 steps, with termination if the agent reached the goal state. In (Shamash et al., 2021), mice were allowed to explore the room starting from the shelter location. Accordingly, during “exploration phases,” the agent started in the goal state and randomly selected a base policy at each time step, after which it updated the corresponding FR. Each exploration phase consisted of 5 episodes—a sample trajectory is shown in Fig. 14(a). After the first exploration phase, the agent started from the canonical start state and ran FRP to reach the goal state. Because most of its experience was in the lower half of the grid, the FRs for the upper half were incomplete (Fig. 14(b)), and we hypothesized that in this case, the mouse should either i) default to a policy which would take it to the shelter in the area of the room which it knew well (the down policy) or ii) default to a policy which would simply take it away from the threat (again the down policy). During the first escape trial, the agent selects the down policy repeatedly, continuing to update its FRs during the testing phase. Upon reaching the wall and getting stuck, the FR for the down policy is eventually updated enough that re-planning with FRP produces a path around the barrier. After updating its FR during the first escape trial and during another exploration period, the FRs for the upper half of the grid are more accurate (Fig. 14(c)) and running FRP again from the start state produces a faster path around the barrier on the second escape trial. TD learning curves for the experiment (repeated with the SR for completeness) are plotted in Fig. 15.
For completeness, we repeated this experiment with the SR. In this case, the planning algorithm is ill-defined for K > 0, so we default to GPI (K = 0). As expected, without the barrier, the down policy is selected and the goal is reached (Fig. 16,left). However, when there is a barrier, while the SR updates when the agent hits it (Fig. 12), since there is no single policy that can reach the shelter, GPI fails to find a path around the barrier (Fig. 16,middle,right).

heading: A.4 ADDITIONAL PROOFS
text: Below are the proofs for Proposition 3.1 and Proposition 3.2, which are restated below.
Proposition 3.1 (Contraction). Let Gπ be the operator as defined in Definition 3.2 for some stationary policy π. Then for any two matrices F, F ′ ∈ R|S|×|S|,
|GπF (s, s′)− GπF ′(s, s′)| ≤ γ|F (s, s′)− F ′(s, s′)|, (33)
with the difference equal to zero for s = s′.
Proof. For s 6= s′ we have
|(Gπ − GπF ′)s,s′ | = γ|(PπF − PπF ′)s,s′ | = γ|Pπ(F − F ′)s,s′ | ≤ γ|(F − F ′)s,s′ |,
where we use the notation Xs,s′ to mean X(s, s′), and the inequality is due to the fact that every element of Pπ(F − F ′) is a convex average of F − F ′. For s = s′, we trivially have |(GπF − GπF ′)s,s| = |1− 1| = 0.
Proposition A.1 (Convergence). Under the conditions assumed above, set F (0) = I|S|. For k = 0, 1, . . . , suppose F (k+1) = GπF (k). Then
|F (k)(s, s′)− Fπ(s, s′)| < γk (34)
for s 6= s′ with the difference for s = s′ equal to zero ∀k.
Proof. We have, for s 6= s′ and using the notation Xs,s′ = X(s, s′) for a matrix X ,
|(F (k) − Fπ)s,s′ | = |(GkF (0) − GkFπ)s,s′ | ≤ γk|(F (0) − Fπ)s,s′ | (Proposition 3.1) = γkFπ(s, s′) < γk (Fπ(s, s′) ∈ [0, 1)).
(35)
For s = s′, |(F (k) − Fπ)s,s| = |1− 1| = 0 ∀k.
Below is the proof of Proposition 4.1, which is restated below.
Proposition 4.1 (Planning optimality). Consider a deterministic, finite MDP with a single goal state sg, and a policy set Π composed of policies π : S → A. We make the following coverage assumption, there exists some sequence of policies that reaches sg from a given start state s0. Under these conditions, Alg. 1 converges such that Γ(s0) = γL ? π , where L?π is the shortest path length from s0 to sg using π ∈ Π.
Proof. Since the MDP is deterministic, we use a deterministic transition function ρ : S × A → S. We proceed by induction on L?π .
Base case: L?π = 1 If L?π = 1, s0 must be one step from sg. The coverage assumption guarantees that ∃π ∈ Π such that ρ(s0, π(s0)) = sg. Note also that when both the MDP and policies in Π are deterministic, Fπ(s, s′) = γLπ , where Lπ is the number of steps from s to s′ under π, and we use the abuse of notation Lπ =∞ if π does not reach s′ from s.
Then following Alg. 1,
Γ1(s0) = max π∈Π
Fπ(s0, sg) = γ (guaranteed by coverage of Π)
Γ1(sg) = max π∈Π
Fπ(sg, sg) = 1 (by definition of Fπ).
Moreover,
Γ2(s0) = max π∈Π,s′∈S
Fπ(s0, s ′)Γ1(s ′)
= max π∈Π {Fπ(s0, s0)Γ1(s0), Fπ(s0, sg)Γ1(sg)}
= max{1 · γ, γ · 1} = γ.
Then Γ2(s) = Γ1(s) ∀s and Alg. 1 terminates. Thus, Γ(s0) = γ = γL ? π and the base case holds.
Induction step: Assume Proposition 4.1 holds for L?π = L Given the induction assumption, we now need to show that Proposition 4.1 holds for L?π = L+ 1. By the induction and coverage assumptions, there must exist at least one state within one step of sg that the agent can reach in L steps, such that the discount for this state or states is γL. Moreover, the coverage assumption guarantees that ∃π ∈ Π such that for at least one such state sL, ρ(sL, π(sL)) = sg . Then this problem reduces to the base case—that is, Alg. 1 will select the policy π ∈ Π that transitions directly from sL to sg—and the proof is complete.

heading: A.5 EXPLICIT PLANNING
text: Below we describe a procedure for constructing an explicit plan using πF and sF .
Algorithm 2: ConstructPlan
1: input: goal state sg , planning policy πF , subgoals sF 2: Λ← [] . init. plan 3: s← s0 . begin at start state 4: while s 6= sg do 5: Λ.append((πF (s), sF (s))) . add policy-subgoal pair for current state to plan 6: s← sF (s) 7: end while 8: return Λ

heading: A.6 FRP WITH MULITPLE GOALS
text: Here we consider the application of FRP to environments with multiple goals {g1, . . . , gn}. To find the shortest path between them given the base policy set Π, we first run FRP for each possible goal state in {gi}, yielding an expected discount matrix ΓΠ ∈ [0, 1]|S|×n, such that ΓΠ(s, gi) is the expected discount of the shortest path from state s to goal gi. We denote by gσ = [s0, σ(g1), σ(g2), . . . , σ(gn)] a specific ordering of the goals in {gi} starting in s0. The expected discount of a sequence of goals is then
Ξ(gσ) = n∏ i=1 ΓΠ(gσ(i− 1), gσ(i)), (36)
with the optimal goal ordering gσ∗ given by
gσ∗ = argmax gσ∈G Ξ(gσ), (37)
where G is the set of all possible permutations of {gi}, of size n!. This is related to a form of the travelling salesman problem. Fortunately, in most settings we don’t expect the number of goals n to be particularly large.

heading: A.7 CONNECTIONS TO OPTIONS
text: The options framework (Sutton et al., 1999) is a method for temporal abstraction in RL, wherein an option ω is defined as a tuple (πω, τω), where πω is a policy and τω ∈ ∆(S) is a state-dependent termination distribution (or function, if deterministic). Executing an option at time t entails sampling an action at ∼ πω(·|st) and ceasing execution of πω at time t + 1 with probability τω(st+1). The use of options enlarges an MDP’s action space, whereby a higher-level policy selects among basic, low-level actions and options.
Options are connected to FRP (Alg. 1) in that by outputting a set of policies and associated subgoals {(πF , sF )}, FRP effectively converts each base policy to an option with a deterministic termination function, i.e., the agent will follow πF until terminating at sF . One of the key difficulties in the options literature is how to learn the best options to add to the available action set. For the class of problems considered in this paper, FRP then provides a framework for generating optimal (in the sense of finding the fastest path to a goal) options from a set of standard policies, subject to the fulfillment of the coverage assumption. Importantly, the associated FRs (which can be learned via simple TD updating) can be reused across tasks, so that FRP can re-derive optimal options for a new goal.
FRP can also be seen as related to the work of Silver & Ciosek (2012), which demonstrates that value iteration performed on top of a set of task-specific options converges more quickly than value iteration performed on the default state space of the MDP. One critical difference to note is that the FR/FRP is transferable to any MDP with shared transition dynamics. While value iteration on a set of options for a given MDP is more efficient than value iteration performed directly on the underlying MDP, this process must be repeated every time the reward function changes. However, the FR enables an implicit representation of the transition dynamics to be cached and reused.

heading: A.8 FURTHER CONNECTIONS TO RELATED WORK
text: We now describe the connection between the FR/FF and several related approaches in the literature. Table 4 summarizes a high-level view of these connections.
The Dynamic Distance Function The dynamic distance function (DDF; (Hartikainen et al., 2020)) is defined as
dπ(s, s′) = Eπ [ j−1∑ k=0 γkc(st+k, st+k+1) ∣∣∣st = s, sj = s′] , (38)
where c : S × S → R is a local cost function. In practice, c(st+k, st+k+1) = γ = 1, giving
dπ(s, s′) = Eπ [ j−1∑ k=0 1 ∣∣∣st = s, sj = s′] . (39)
n practice, dπ(s, s′) is parameterized via a neural network with parameters ψ trained on-policy from full trajectories τ using the loss
L(ψ) = 1 2 Eτ∼D,i∼[0,T ],j∼[i,T ](dπψ(si, sj)− (j − i))2, (40)
whereD is a buffer of stored trajectories. On a downstream navigation task, the agent policy is trained to minimize
Lπ(φ) = Eπ [ ∞∑ t=0 γtdπψ(st, g), ] , (41)
where g ∈ S is a goal state and φ are the policy parameters. There are several important differences between the DDF and the FR, both in theory and practical application.
This form of the DDF, as discussed by Hartikainen et al. (2020), is also naturally restricted to onpolicy learning from full trajectories, and cannot be updated off-policy and/or via one-step temporal difference learning. A natural additional consequence is that it is only applicable to finite-horizon MDPs. This confinement to finite horizons is significant because it leads to the conditioning problem described by Hartikainen et al. (2020). This problem occurs because the DDF is conditioned on the policy successfully reaching the goal state—when this doesn’t happen, it can lead to significant value estimation errors.
A second difference is that the policy is trained, via Eq. (41), to minimize the cumulative discounted distance to the goal rather than via greedy distance minimization.
Another difference is that Eqs. (38) and (39) don’t explicitly require time step j to be the first time the agent enters s′, although this may have been the authors’ intention. In fact, this definition is more closely related to the SR than the FR (with equality in the infinite horizon setting), but the implementation of the DDF in practice is more closely related to the FR. It’s also important to note that when γ = 1, the number of steps k appears linearly within the expectation rather than exponentially (this is important, as in general γE[k] 6= E[γk]. Another significant difference, then, is that the DDF as presented cannot be in a meaningful sense be considered a state representation, but rather a function mapping pairs of states to expected distances. That is, dπ(s, ·) has no meaningful semantics when implemented as a mapping dπ : S × S → R and trained using Eq. (40). In practice, it is used to support policy optimization, rather than evaluation. In contrast, the FF representation, like the SF representation, maps a single state to a fixed length vectorvalued encoding, ϕπ : S → Rd. This difference has significant implications for the applications of the DDF. In particular, it is unclear how the DDF could be used as an exploration bonus in the manner of the FF/SF, and the scalar value would make it challenging to implement the type of parallel updates useful for efficient planning.
Dynamic Goal Learning Another related approach is that of dynamic goal (DG) learning (Kaelbling, 1993), a method for optimal control related to Q-learning. The optimal DG function G? : S ×A× S → R is defined recursively for a goal state g ∈ S as
G?(s, a, g) = 1 + Es′∼P (·|s,a) [
min a′∈A
G?(s′, a′, g) ] , (42)
where G?(g, a, g) := 0. There are several important differences between this approach and the FR. First, as mentioned above, DG learning is a method for optimal control, rather than policy evaluation. That is, D?(s, a, g) converges to the expected number of steps for the optimal policy for reaching g starting from s. It cannot be reused for policy evaluation for a policy π 6= π?, and also implicitly assumes that g is reachable in finite time. DG learning is thus susceptible to the same conditioning problem discussion above for the DDF, and is only studied by Kaelbling (1993) in a gridworld environment. A second difference is that the DG function scales linearly with k, the number of steps for the optimal policy to reach g, while the FR scales exponentially at a rate of γ. This is important, as we can note (with a slight abuse of notation) in general that γEπ [k] 6= Eπ [ γk ] —it is nontrivial to recover the FR for the optimal policy from the DG function (and vice-versa).
Temporal Difference Models Temporal difference models (TDMs; (Pong et al., 2018)) are another related approach. TDMs are an optimal control method motivated by the observation that goalconditioned Q-functions can be used to construct an implicit dynamics model when the discount factor γ = 0. For γ > 0, the authors introduce a horizon variable τ to interpolate between modelbased and model-free learning. Like the FR, this formulation allows agent to consider multiple
goals in parallel. As in DG learning, and unlike the FR, TDMs represent a method for optimal control, not evaluation. The emphasis is on efficiently improving a policy, rather than learning a state representation which can be leveraged for multiple uses (e.g., policy evaluation or as an exploration bonus). Another difference from the FR is that, like the DDF and DG learning, TDMs rely on a finite-horizon setting.
A.9 FR VS. SR VISUALIZATION

Review summaries:
This paper proposes a new notion of state representation, the first-occupancy representation (FR), inspired by the Successor Representation. It is motivated by situations where the rewards are non-Markovian. Similarly to the SR, the FR can be learnt by TD learning. The usefulness of the FR is demonstrated on exploration, unsupervised RL, planning and escape behaviour tasks.  

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.

The paper proposes a novel representation of the dynamics of an environment that is independent of the reward structure encoding a particular sequential decision task in it, and thus can be reused to speed up the computation of policies for multiple tasks in this environment. Unlike successor representations, which compute the likelihood of occupying a particular state at any time in the future, when starting from a given initial state, the proposed first-occupancy representation encodes the likelihood of reaching a particular state for the first time. Because of this, the novel representation effectively represents the expected path length between all pairs of states when following a particular policy. The benefits of the proposed representation are illustrated in several decision problems with non-Markovian reward structure. It is also argued that animals might use similar representations of their environment to effectively find escape routes in a short amount of time.