A First-Occupancy Representation for Reinforcement Learning
Authors: Ted Moskovitz
Spencer R Wilson
Maneesh Sahani
Venue: ICLR 2022 Poster (2022)
Keywords: successor representation
successor features
generalized policy improvement
GPI

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 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 first-occupancy 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.

Published as a conference paper at ICLR 2022.
A FIRST-OCCUPANCY REPRESENTATION FOR REINFORCEMENT LEARNING
Ted Moskovitz
Spencer R. Wilson?
Maneesh Sahani'
1 Gatsby Unit, UCL
2 Sainsbury Wellcome Centre, UCL
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 reach rewarding states. The successor representation (SR), which measures the expected cumulative, discounted state occupancy under a fixed policy, enables efficient transfer 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 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 behaviorall--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 first- occupancy representation (FR), which measures the expected temporal discount the first time 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.
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 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 associates with each state and policy of action measure of the expected rate of future occupancy of all states that policy were to be followed indefinitely. This cached representation can be acquired through experience in much the same way 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 transition distribution remains consistent.
However, these requirements limit the applicability 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 a state. Internal goals for control-san say, to pick up a particular object-need to be achieved as rapidly as possible, but only once at a time.
"Correspondence: teddgat.sby uuc1 ac.ukk

---

Published as a conference paper at ICLR 2022.
Furthermore, while 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 al., 2018), this capacity still falls far short of the power of planning within complete models of the environment.
Here, we propose 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 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 offers a compelling alternative to the SR, including, but not limited to: exploration, unsupervised RI,, planning, and modeling animal behavior.
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. We model a task T as finite Markov decision process (MDP; (Puterman, 2010)), T = AS,A, p, 7, n),, where is a finite state space, A is finite action space, : S x A + ASS) is the transition distribution (where (SS) is the probability simplex over S), r
is the reward function, E [0, 1) is discount factor, and μ E A(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, 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 Σ, )rre()) To simplify notation, we will frequently write r(s:) == тı and r E Rlss as the vector of rewards for each state. The agent acts according
describe the expected transition probabilities under æ using a |S| x [S] matrix pr such that pr a, = p" (s'ls) =
stationary policy π S + (AA). For finite MDPs,
En p(s'|s (a)s). Given π and a reward function r, the expected return is
(s, = E, Σ
Arru s, = E, Euioptin [Tr + (2, ,')l d'wa(11e))
(I)
kk=
QF are called the state-action values or simply the Q-values of ा. The expectation ,r, [-] 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,, [l| = s, ag = a) as E,, H|s, a.). This recursive form is called the Bellman equation, and it makes the process of estimating Q tertee policy -uulattion-tacttbble via dynamic programming (DP; Bellman, 1957). In particular, successive applications of the Bellman operator T Q ;= r 7 Q are guaranteed to converge to the true value function Q for any initial realevalued S x |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 transition sequence (St, de, rt, $141) and ac+1 T-1(5142),
O7 (st, d+) + (s., + axd,, d, Once policy has been evaluated, policy improvement identifies a new policy x' such that Q a) 2 QF' a), V(s, a) (7, a). Helpfully, such policy can be defined n' (s) E argmax, Q7 (s,a).
(s a(+1) Q (st, az).
(2)
The successor representation The successor representation (SR: (Dayan, 1993)) is motivated by the idea that state representation for policy evaluation should be dependent the similarity of different paths under the current policy. The SR is a policy's 's expected cumulative discounted state occupancy, and for discrete state spaces be stored in an |S] x |S| matrix M where
= -Er
* 1 1(8+4* ( =50
[(sr Er
+ 7M" 7 (s (3)
2

---

Published as a conference paper at ICLR 2022.
where 1(-) is the indicator function. The SR can also be conditioned on actions, i.e,, M (s, = E Σ 7 (s = s') St, az) and expressed in vectorized format, we can write (*)s) = M (s. or (s, = M (s, a, -). The recursion in Eq. (3) admits a TD error:
8, - (s.) ( ((+++1)) (*
where (s;) 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 given environment:
(s,a) = F'EE Σ (((1(s4))) Sc, 0t
Exx
(s,
(5)
Fast transfer for multiple tasks In the real world, often have to perform multiple tasks within a single environment. A simplified framework for this scenario is to consider a set of MDPs M that share every property (i.e., A, p: 7: ()) except reward functions, where each task within this family determined by a reward function r belonging to a set R. 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 policy π on set of tasks R. Similarly, generalized policy improvement (GPI) for set of "base" policies IT is the definition a policy n' such that
º (s,a) 2 sup O (s, a) v(s,a) E ESX A rell
(6)
for some E R. As hinted above, the SR offers way to take advantage of this shared structure by decoupling the agent's 's evaluation of its expected transition dynamics under a given policy from single reward function. Rather than needing to directly estimate Q" Vir E II, M only needs to be computed once, and given 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 new policy x' via
(s) E argmax max r M (8,
aEA xEll
For brevity, we will refer this combined procedure of GPE and GPI simply "GPI". unless otherwise noted. The resulting policy n' is guaranteed to perform at least as well as any individual II (Barreto et al., 2020) and is part of 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).
3 THE FIRST-OCCUPANCY REPRESENTATION
While the SR encodes states via total occupancy, this may not always be ideal. If task lacks a time limit but terminates once the agent reaches a pre-defined goal, or if reward in given state is consumed or made otherwise unavailable once encountered, a more useful representation would instead measure the duration until policy is expected to reach states the first time. Such natural problems emphasize the importance of the first occupancy and motivate the first-occupancy representation (FR).
Definition 3.1. For MDP with finite S, the first-occupancy representation (FR) for policy F" E [o, 1] || 5) is given by
8
F (" (s) = ES
Σ
x *1(514 = 8',8' (sne+*))
where {8uE4R] = {St, S14k with the convention that {Sta=0} = Ø.
That is, as the indicator equals iff $ s' and time t + k is the first occasion which the agent has occupied s' since time f, F (s, (s, s') gives the expected discount at the time the policy first reaches s' starting from The idea of learning policy-dependent distances to target states has long history in RL (Kaelbling, 1993; Pong et al., 2018; Hartikainen et al., 2020). However, 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 recursive relationship for the FR (derivation Appendix A.1):
Fd (s, Emstspp (11)) 1(st s') + 1 ((: = F))) (S s')
(9)
This recursion implies the following Bellman operator, analogous to the one used for policy evaluation:
3

---

Published as a conference paper at ICLR 2022.
2D grid
GPI trajectories
56
5g)
(b)
(c)
SR
(e)
FR
Figure 1: The is higher for shorter paths. (a-c) A 2D gridworld and fixed policies. (d) The FR from so to Sy is higher for 22, but the SR is lower. (e) SR-GPI with the SR picks mi while FR-GPI selects m2.
Definition 3.2 (FR Operator). Let F E R|S|x|SI be an arbitrary real-valued matrix. Then let G denote the Bellman operator for the FR, such that
G" F = Fsst + -111' F,
(10)
where 1 is the length-|S] vector of all ones. In particular, for stationary policy n,, F F
F",
The following result establishes G as a contraction, with the proof provided in Appendix A.4. Proposition 3.1 (Contraction). Let G" the operator as defined in Definition 3.2 for some stationary policy т. Then for any two matrices F, F, E R|S|x|S|
F(", - F FFFF(8 =F((s&')
(11)
with the difference equal to zero for 8 8'.
This implies the following convergence property of G*.
Proposition 3.2 (Convergence). Under the conditions assumed above, set F(0) -(ss For k 0, o,..... suppose p(k+1) = G p(k), Then
F (s, < <h
(12)
for + s' with the difference for s = s' equal to zero vk.
Therefore, repeated applications the FR Bellman operator F + F as + cc. When the transition matrix pa is unknown, the FR can instead be updated through the following TD error:
8f = 1(: = 5') >1 1(st S")) (8t41, 1 (st.
(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 2D environment with start state 8p, a rewarded state 8g- and deterministic transitions (Fig. 1a). One policy, Vi, reaches sy slowly, but after first encountering it, re-enters 8, infinitely often (Fig. 1b). A second policy, n2, reaches 8y quickly but never occupies again (Fig. Ic). In this setting, because गi re-enters sy multiple times, despite arriving there more slowly than 72, M (so, 8q) =2 (so, 8g), but because the FR only counts the first occupancy of a given state, a (Su, 8g) F (Sg. Sy). The FR thus reflects a policy's 's path length between states.
Policy evaluation and improvement with the FR As with the SR, evaluation with the FR. Crucially, however, the FR induces the following value function:
quickly perform policy
(T = Er Σ x
8t, az
O (s,
(14)
where ,.F: S -> IR is reward function such that f (St+k) = r(st+t) if 8t+h & (SES+*) 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 performed with respect to Q,, as normal. This structure is very common feature of real-worhe tasks, such as foraging for food or reaching to target. Accordingly, there is a rich history of studying tasks with this kind of structure, termed non-Markovian reward decision processes (NMRDPs; (Peshkin al., 2001; Littman al., 2017;
4

---

Published as a conference paper at ICLR 2022.
Gaon & Brafman, 2020; Ringstrom & Schrater, 2019)). Helpfully, all NMRDPs be converted into an equivalent MDP with an appropriate transformation of the state space (Bacchus et al., 1996).
Most approaches in this family attempt to be generalists, learning an appropriate state transformation and encoding 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 r 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 in a given environment, the FR can immediately perform policy evaluation for the corresponding r".
Infinite state spaces natural question when extending the FR to real-world scenarios is how can be generalized to settings where |S] is either impractically large or infinite'. In these cases, the SR reframed as successor features 13th (SFs; Kulkarni et al., 2016; Barreto et al., 2017b), where the dth s]. where d = 1, ..... D and Φ S - RD is a base feature function. The base features o(-) are typically defined so that linear combination predicts immediate reward (i.e., w (T((s)
SF is defined as v5 (s) = ER, -_k= )
r(s) for some w RRD), 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 ffrst--ccupancy feature (FF) representation sp" where the dth FF is given by
( Ex
Σ 7 (((a(())4k) 2 Oa, ( < <0a)
(15)
= 2 0²) 711 (0b((s:) 2 0a) E (pā (8
where 0 is threshold value for the dth feature. The indicator equals only s is the first state whose feature embedding exceeds the threshold. Note that this representation recovers the FR when the feature function a one-hot state encoding and the thresholds {0a] are all
EXPERIMENTS
We now demonstrate the broad applicability of the FR, and highlight ways its properties differ from those of the SR. We focus on 4 areas: exploration, unsupervised RL, planning, and animal behavior.
4.1 THE FR AS AN EXPLORATION BONUS
Intuitively, representations which encode state visitation should be useful measures of exploration. Machado al. (2020) proposed the SR as a way encourage exploration in tasks with sparse or deceptive rewards. Specifically, the SR used as a bonus on-policy learning with Sarsa (Rummery & Niranjan, 1994):
Table 1: Exploration results. + values denote
SE across 100 trials.
SARSA + FR
1,547, 243 4 34,
11. 9149 42, .422
SARSA SR
1,197 00. 1. 36, 999
1,025, 750 1 29, 095
SARSA
075 1,424
376, 655 1 8,49
B
8. +
[M)(ll + (0+1(1111
- Q" (8t,0t),
(16)
where B E E controls the size of the exploration bonus. Machado et al. (2020) show that during learning, ||M (s)]l can act as a proxy for the state-visit count n((s., with | (s)ili+ awarding progressively lower bonus for every consecutive visit of s. In the limit as t + oc, however, + (+)*((
Vs as π stabilizes, regardless of whether π has effectively explored. To encourage exploration, we'd like for bonus maintain its effectiveness over time. In contrast
18++ where Ki > 1 for [S] 2 1. Note that || (s)]l1 Kiss only if reaches all states in as many steps. Because || | 1 only grows when new states or shorter paths are discovered, we can instead augment the reward as follows:
to (s]l, 1 < (s)ll < Kilss
1-7
Tt - r. + 31| (s))ll
(17)
Recent 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.
5

---

Published as a conference paper at ICLR 2022.
Vite
Vil Yem
Vih Vi
Figure 2: The FF facilitates accurate policy evaluation and selection. Shading denotes SE over 20 seeds.
Grid
R
GPI
FR-planning
(8)
(b)
(c)
(d)
Figure 3: The FR enables efficient planning. (a-d) A 2D gridworld with start (so) and goal(sg) states, along with three fixed policies. (c) GPI follows T1. (f) Planning with the FR enables the construction of shorter path.
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 easy to reach states, with greater reward available in harder to reach states. In both cases, SR) and Sarsa with an FR bonus (SARSA + FR) for 5,000 time steps with an e-greedy policy. The results are listed in Table where
Sarsa, Sarsa with SR bonus (SARSA +
that the FR results in added benefit It's also important to note that the maximum bonus H S) has another useful property, that it scales exponentially S). This is desirable, because as the number of states grows, exploration frequently becomes more difficult. To measure this factor empirically, DOVERSWIM-NN with = {6, 12, 24} states, finding
the SR.
tested the same approaches with the
same settings on a modified RIVERSWIM, robust the increased exploration difficulty (see Appendix Table
that SARSA FR and Appendix A.3 for results and more details). We also tested whether these results extend to the function approximation setting, comparing DQN-style model (Mnih et al., 2015) using the SF and FF as exploration bonuses on the challenging DEEPSBA task (Osband et al., 2020), finding that the advantage of the FF 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.
4.2 UNSUPERVISED RL WITH THE FF
We demonstrate the usefulness the FF in the unsupervised pre-training RL (URL) setting, paradigm which has gained popularity recently possible solution to the high sample complexity of deep RL algorithms (Liu & Abbeel, 2021; Gregor et al., 2016; Eysenbach et al., 2018; Sharma al., 2020). In URL., the agent first explores environment without extrinsic reward with the objective of learning useful representation which then enables rapid fine-tuning to 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 modified version of the continuous MountainCar task (Brockman al., 2016) (Fig. 2(a)). The agent pre-trains for 20,000 time steps in rewardless environment, during which learns FF's or SFs for set of policies II which swing back and forth with fixed acceleration or "power." (details Appendix A.3). During fine-tuning, the agent must identify the policy use radial basis functions as the base features da(-) with fixed thresholds 0g = 0.7.
II which reaches randomly sampled goal location the fastest.
Fig. 2(b.d) 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-poweree policies reach the goal locations for the first time much sooner, and the policies with the highest FF values have higher powers. In

---

Published as a conference paper at ICLR 2022.
the test phase, the agent fits the reward vector R by minimizing E, r+ w wT ( 1 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 II which has the highest expected value across starting states: π SMP E arggax xen Exg (8n),, where V (so) = w w (so) (with s" replaced by e for SF-based value estimates). Fig. 2(c,e) shows both the estimated (Vest) and true (Virue) 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 suboptimal policy. The FF estimates are nearly matched the true values of the selected policies for each goal location and achieve nearly optimal performance.
Robotic reaching To test whether these re- sults 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 manipula- tions (Fig. 4(a)). We modify the Active Pre- training with Successor features (APS; Liu & Abbeel, 2021) URL algorithm, which leverages a nonparametric entropy maximization objec- tive 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- Pre-training with First-occupancy features (APF), using the same intuition motivating the MountainCar ex experiments: cumulative features are misleading when downstream tasks benefit from quickly reaching desired goal, in this case, the object. Here, the agent first trained for 106 time steps using the aforementioned intrinsic reward objective before being applied to specified reaching task. As additional baselines, 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)).
APS
DDLUS-G
DDLUS
64
staps
(a)
(b)
Figure 4: APF accelerates convergence in robotic reaching. Shading denotes SE over 10 seeds.
4.3 PLANNING WITH THE FR
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 one- step 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())), but here there are three policies comprising II = [11. T2, A3} Tz (Fig. 3(b-d)). GPI selects λı because it the only policy that reaches the goal s, within one step of the start sp: क max r MM" 8y) (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 11 is instead move right using T22 and up using T3. How can the FR be used to find such a sequence?
Intuitively, starting in state s, this strategy is grounded in following one policy until 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. for given policy, the agent can elect to follow the policy π E II with the greatest value of (8 s"), then switch to another policy and repeat the process until reaching 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 more detailed discussion, see Appendix A.7.
More formally, we can construct DP algorithm to solve for the optimal sequence of planning policies TF and subgoals s Denoting by ('u(s) the total discount of the full trajectory from s to Sy at step k of the procedure, we jointly optimize over policies π and subgoals s' for each state s:
I [s) max TEIL,YESS
F (s, (')'s',,, with Tf) sf = argmax (s, (')''''). xell,s'S
Intuitively, the product (T "',,s')(&(s (s'')) can be interpreted as the expected discount of the plan consisting of first following from to s', then the current best (shortest-path) plan from s' to Sy-
7

---

Published as a conference paper at ICLR 2022.
GM
Y
GPI
return
K
0 iterations
K
iterations
14
08
iterations
transition noise
(a)
(b)
(c)
Figure 5: FRP interpolates between GPI and model-based DP. Shading represents SE.
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 plaanmingg (FRP), is given by Alg. 1 in Appendix A.2. Appendix Fig. 13 depicts the resulting policies x 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 given goal (proof in Appendix A.4).
Proposition 4.1 (Planning optimality). Consider deterministic, finite MDP with a single goal state Sg. and a base policy set 11 composed of policies a S + A. We make the following coverage assumption: there exists a sequence of policies that reaches 8y from a given start state Su- Then Alg.
I converges that T (s0) = where 1 is the shortest path length from So 8g using и E TT.
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 [1| S)²). The worst-case complexity occurs when the policy must switch at every state route to the target- K is opper-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 [||SS P). 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 [I| < NIAA To achieve an error e between the estimated value function and the value function of the optimal policy, VI requires N 2 (1-7) log (1-7)æc (Puterman,
2
1994), which for 7 = 0.95, e = 0.1, e.g., gives 2 180 iterations 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, new goal is randomly sampled in 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, ri ght,, 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 FRP with K- 0 iterations. To see this, observe that when there is a single goal so such that only r(sg)= > 0, the policy selected by GPI
π (P)
argmaxr FT (s) arggmax ( (5,)) (s, argmax F Sg)- #ETI
(18)
#El
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 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 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 = 3 iterations are plotted in Fig. 5(b). Finally, to test FRP's s robustness to stochasticity, we added transition noise to the FOUR task. That is, the agent moves to random adjacent state with probability regardless of action. We compared FRP to converged VI for increasing c, with the
after K iterations (Fig. 5(a)). In contrast,
Other DP methods like policy iteration (PI), which is strongly polynomial, converge more quickly than VI, but in the example above, Pl still needs N 2 log(1 - 711))/(1 - 7) 106 (Ye, 201 20, for instance.
8

---

Published as a conference paper at ICLR 2022.
results plotted in Fig. 5(c), where levels. It's important to note that this ability to adaptively interpolate between MF and MB behavior, based on the value of K, is unique capability of the FR compared to the SR² The same FRs can be combined using DP to plan for one task or for GPI on the next.
that FRP matches the performance of VI across noise
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 contain- ing barrier has indicated that, rather than use an explicit cognitive map, mice instead appear to memorize sequence of subgoals to plan ef- ficient routes to shelter (Shamash et al., 2021).
(a)
open field
wall obstacle
Saat
Made
(b)
barrier
barrier trial
barrier trial
When first threatened, most animals ran along direct path and into the barrier. Over subsequent identical trials spanning 20 minutes of explo- ration, streal-stimulu 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.
Figure 6: FRP induces realistic escape behavior.
We model the initial escape trial by an agent with partially learned FR, leading 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 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, demonstrate that the FR facilitates behavior that is consistent with under- standing of animal learning in tasks which demand efficient planning. Given the recent evidence in support of SR-like representations in the brain (Stachenfeld al., 2017: Momennejad 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 thortest-path planning not possible with the SR.
CONCLUSION
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 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 approach, there 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 MDPs with real-valued observations
to explore approximations or its use in partially-observahh 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.
We'd like stress that this claim of uniqueness is only with respect to the SR. Previous work also explores the of MF methods to support MB learning (c.g., Pong et al. (2018))

---

Published as a conference paper at ICLR 2022.
Reproducibility statement Experiments: All experiments described in detail in Section 4 in the main text and in Appendix A.3. We have attached code for the tabular experiments (also available github woot/teddoskooiite first occupancy), y and in Appendix A.3 provide a link to the base implementation for APS as well as an explanation of where to find the hyperparameters and modifications made produce APF. Theoretical results: Theoretical results, including assumptions and proofs, are provided in Appendix A.l and Appendix A.4.
REFERENCES
Fahiem Bacchus, Craig Boutilier, and Adam Grove. Rewarding behaviors. In Proceedings of the Thirteenth National Conference on Artificial Intelligence Volume 2, AAAI'9,, pp. 1160-1167. AAAI Press, 1996.
Andre Barreto, Will Dabney, Remi Munos, Jonathan J Hunt. Tom Schaul, Hado P van Has- selt, and David Silver. Successor features for transfer in reinforcement learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 7a. URL https: http://pooceeeengeeee httpepp..cccpppppeeee//eeeepee/e wccpppppe//0111111 0000////ee///////// 35Ccb081a666 2eliaa16611222222 515555235235dd. 9b8c05 Paper.po
Andre Barreto, Will Dabney, Remi Munos, Jonathan Hunt, Tom Schaul, Hado van Has- selt, and David Silver. Successor features for transfer in reinforcement learning. In Guyon, V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- neurips
ciates, Inc., 2017b. URL ptattpa:/pprcceeedeeee https/eeee..ccccppppe////000000 wwtpppppppp./////ee.e.pe pdf.
Andre Barreto, Diana Borsa, John Quan, Tom Schaul, David Silver, Matteo Hessel, Daniel Mankowitz, Augustin Zidek, and Remi Munos. Transfer in deep reinforcement learning using successor features and generalised policy improvement. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 501-510. PMLR, 10-15 Jul 2018. URL http: htstp::/pprcceeiinig mlr. press/v8 /ttp/bbarroooooooooooooooooooole 8a html.
Andre Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy of Sciences, 117(48):30079-3008) 2020. ISSN 0027-8424. doi: 00.1073//paa.1907370117 URL https: //www. pnas. content///////eeeeeeeeecoeeetc /117 / /4///////////
Timothy E.J. Behrens, Timothy H. Muller, James C.R. Whittington, Shirley Mark, Alon B. Baram, Kimberly I.. Stachenfeld, and 7eb Kurth-Nelson. What is cognitive map? organizing knowl- edge for ficxible behavior. Neuron, 100(2)4400500, 2018. doi: http://doi.org/10.0116// 2euron.2018.100 URL https:// http:///ennneeeeeeeeooce/eeeeeee pi1/50896627318305666
Richard Bellman. Dynamic Programming. Dover Publications, 1957.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent
values: A mathematical viewpoint, 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnh and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
httt://githubs
jax.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation, 5(4):613-622 1993. doi: 10. 0116//ecco.199) 1993.5.4.613.
10

---

Published as a conference paper at ICLR 2022.
Simon S. Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding, 2019.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function, 2018.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässsan, Stockholm, Sweden, July 10-15 15, 2018, 2018.
Maor Gaon and Ronen Brafman. Reinforcement learning with mon-markovian rewards. Proceedings of the AAAI Conference on Artificial Intelligence, 34(04):3980-3388) Apr. 2020. doi: 10.1609/ aaai, mhai.v34i0k v34i04.5814. URL hww.::tts///ooto wwwp:ossssaaaaaa... http:://///....aaaaroooooo/iieie wtrrg//inex..e http://///aaaatttccl///eeeeee 5814.
Samuel J. Gershman. The successor representation: Its computational logic and neural substrates.
Journal of Neuroscience, 38(33):7193-7200 7200, 2018. doi: 101523/NEUROSCLO051.1B URL https::///// // /wwww jneurosei. org. http:///o/t/t///ee/ee/ee/eee /ttpet/////////////
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control, 2016.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor, 2018.
Steven Hansen, Will Dabney, Andre Barreto, David Warde-Farley, Tom Van de Wiele, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. In International Confer- ence on Learning Representations, 2020. URL htwph:tt/////c...ttt http::///ppreeeee.ecmm net/ forum? dd= BJeAHKrYDS.
/
Kristian Hartikainen, Xinyang Geng, Tuomas Haarnoja, and Sergey Levine. Dynamical distance learn- ing for semi-supervised and unsupervised skill discovery. In International Conference Learning Representations, 2020. URL https: / http:///wppeeeeoo/oe/e/ee/ net forum fieerummmmi-Hi1 11laa
Leslie Pack Kaelbling. Learning to achieve goals. In Proceedings of IJCAI-93, pp. 1094-1098. Morgan Kaufmann, 1993.
Sham M Kakade. A natural policy gradient. In Advances in neural information processing systems, pp. 1531-1538, 2002.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Tejas D. Kulkarni, Ardavan Saeedi, Simanta Gautam, and Samuel J. Gershman. Deep successor reinforcement learning, 2016.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. Urlb: Unsupervised reinforcement learning benchmark. In NeurlPS 2021 Datasets and Benchmarks Track, Submitted, 2021. URL https: : /oppnreevvvvcccc ew net/ 1
Timothy Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning, 2019.
Steven Lima and Larry Dill. Behavioral decisions made under the risk of predation: A review and prospectus. Canadian Journal of Zoology-revue Canadienne De Zoologie CAN J ZOOL, 68: 619-640, 04 1990. doi: 10. 10.1999970--09.
Michael L. Littman, Ufuk Topcu, Jie Fu, Charles Isbell, Min Wen, and James MacGlashan. inniioomment-ndeeenddde Task Specifications via GLTL. arXiv e-prints, 2017.
Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 Proceedings of Machine Learning Research, pp. 6736-676 6747. PMLR, 18-24 Jul 2021. URL https: proceedd ngs ml wtt.r/...ess..////eeeee http///19////// uulb html.

---

Published as a conference paper at ICLR 2022.
Chen Ma, Dylan R. Ashley, Junfeng Wen, and Yoshua Bengio. Universal successor features for transfer reinforcement learning. 2020.
Marlos C. Machado, Marc G. Bellemare, and Michael Bowling. Count-based exploration with the successor representation. Proceedings of the AAAI Conference on Artificial Intelligence, 34(04): 5125-5133, Apr. 2020. doi: www.6///aaa...444....... index. phpp AAA I aetaicilee/eee...... wwwccleeviiee wwwe///w5. 5955.
URL https: aaai.org/
Tamas Madarasz and Tim Behrens. Better transfer learning with inferred successor maps. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL https: ptttpa://prooeeedings. http://g.....rr.p.p..pppppp////a httpscccppppee119.33 9/file/ 2
Paper.....
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-333 2015. URL hitps: /doiiooorgggooo org/ /1/g///..00 110100.00 1038/0..///.. http/nauuuuee.....
I. Momennejad, E. M. Russek, J. H. Cheong, M. M. Botvinick, N. D. Daw, and S. Gershman. The successor representation human reinforcement learning. Nature Human Behaviour, 1(9): 680-692, 2017.
Ted Moskovitz, Michael Arbel, Ferene Huszar, and Arthur Gretton. Efficient wasserstein natural gradients for reinforcement learning. 2020.
Ted Moskovitz, Jack Parker-Holder, Aldo Pacchiano, Michael Arbel, and Michael I. Jordan. Tactical optimism and pessimism for deep reinforcement learning, 2021.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard Lewis, and Satinder Singh. Action-conditional video prediction using deep networks in atari games, 2015.
Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado Van Hasselt. Behaviour suite for reinforcement learning. International Conference Learning Representations, 2020. URL https /oppnreeivveeeeeeee ew. nww. ennet/foruum forum? id=rygf-kSYwtt
Leonid Peshkin, Nicolas Meuleau, and Leslie Pack Kaelbling. Learning policies with external memory, 2001.
Vitchyr Pong, Shixiang Gu Murtaza Dalal, and Sergey Levine. Temporal difference models: Model- free deep RL for model-based control. In International Conference Learning Representations, 2018. URL https: /
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 1994.
Martin Puterman. Markov decision processes: discrete stochastic dynamic programming. John
Wiley and Sons, 2010.
Thomas J. Ringstrom and Paul R. Schrater. Constraint satisfaction propagation: Non-stationary policy synthesis for temporal logic planning. CoRR, abs/1901. 10405, 2019. URL http:///rxix: org/abs/11 9011.... 10405.
G. Rummery and Mahesan Niranjan. On-line q-learning using connectionist systems. Technical Report CUED/F-INFENG/Y 166, 11 1994.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017.
12

---

Published as a conference paper at ICLR 2022.
Shamash, Sarah F. Olesen, Panagiota lordanidou, Dario Campagner, Banerjee Nabhojit, and Philip
Tiago Branco. Mice learn multi-step routes by memorizing subgoal locations. bioRxiv, 2021. doi: 10.1101/2020.08.1900 19....9..555.. 19.256867. URL https:://// http//wwwiioo..o www.:wwww.ivv..oo org/ httgg//onnent//errr///ee 2021 /05/ / 08/ / 2020. 08 256867.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. In International Conference on Learning Representations, 2020.
URL battps://ppenreview..t Eew./ wotoruniid-HggLLeH. id-HJgLZR4KYA.
David Silver and Kamil Ciosek. Compositional planning using optimal option models, 2012.
Kimberly L. Stachenfeld, Matthew M Botvinick, and Samuel J Gershman. The hippocampus predictive map. Nature Neuroxcience, 20(11):1643-1653, 2017. URL https:t///i..ooooo 1038/nmg ....00..
Alexander L. Strehl and Michael L. Littman. An analysis of model-based interval estimation for markov decision processes. Journal of Computer and System Sciences, 740$):1309-1331, 2008. ISSN 0022-0000. doi: http://cw.......ea..aa.////////. www.
URL https: Learn-
wtw.darretocooeeeeeeeoooooooooee ing Theory 2005.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL http: httt:://nccompleteeeeee. wttpa/a/ttttetoootteeeeetteeeeee -bokkoh---dddd....oo html.
Richard S. Sutton, Doina Precup. and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction reinforcement learning. Artificial Intelligence, 112(1): 181-211, 1999, URL https: //ww://wwwwww/w../ www.ecedireeeeneeeoooeeecoeoeeee com/ http://cccc//aaiee/iiii.i. S0C04370299000555
Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards, 2021.
Eszter Vértes and Maneesh Sahani. A neurally plausible model learns successor representations in partially observable environments. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Cur- ran Associates, Inc., 2019. URL https: http:::::///ppppeeppessooo neurips. cc/paper/2019/ -aapeedddddd--apeee..oororaoae
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, ((3-4):29-2254) 1992.
Yinyu Ye. The simplex and policy-iteration methods are strongly polynomial for the markov decision problem with fixed discount rate. Mathematics Operations Research, 36(4):593-60) 2011. URL https:: / /oo. .rr 1287 /moor..... 1110. 051 6.
Tom Zahavy, Avinatan Hassidim, Haim Kaplan, and Yishay Mansour. Planning in hierarchical reinforcement learning: Guarantees for using local policies. CoRR, abs/1902.10140. 2019. URL
http:://rxiin
1190//111111..
Tom Zahavy, Andre Barreto, Daniel J Mankowitz, Shaobo Hou, Brendan O'Donoghue, Iurüi Kemaev, and Satinder Singh. Discovering a set of policies for the worst case reward. International Conference Learning Representations, 2021. URL https: /openrevieww net forum? id=PUkhNz6ddy5.
13

---

Published as a conference paper at ICLR 2022.
A APPENDIX
A.1 FR RECURSION
For clarity, we provide the derivation of the recursive form of the FR below:
00
F" (s.
ET
1 (s = 8' s' $ (sLc+))) St
k==0
Eg
[(s = s' 5 222) +
* *(1++ = s'ss' (sucts)) s1
k=
(19)
E
1(s, = s') Σ
1st 11(44 = =8 F s' $ (sen
k=1
[(ss = 5') + 7 # 5 ( 5 s1
(1
[(st= = =s) + +1 - (s: = F (Sc+1, F
A.2 FRP ALGORITHM
We present the full algorithm for FR planning (FRP) below.
Algorithm I: FR Planning (FRP)
I: input: goal state 8,, base policies { TnY and FRs {F F
2: // initialize discounts to-goal
3: Pa(s) + -00 Vs E S
4: for E S do
5: (s) max+en (s)
6: (s), sf (s) + argmax F* (8) Sq
7: end for
8: // iteratively refine r
& 1-11
10: while As E S such that Ts(s) > 1=-1(8) do
11: for s E do
12:
Ts)1(s) max F" (s. (sss')
13:
F (s), + argmax Frt (1))))))))))))
14: end for
15: lo 2-6 ++
16: end while
17: return x²
m
114
50 episodes
25
30 episodes
Figure FF and SF learning curves for continuous MountainCar. Results averaged over runs. Shading represents one standard deviation.
14

---

Published as a conference paper at ICLR 2022.
A.3 ADDITIONAL EXPERIMENTAL DETAILS
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.
(1: 0.6;; 0)
(1: 0.0.0 0)
(1: O.G, 6)
(1: 0.6; 0J
(1: 0.3: 0
01:0
2
3
10000)
(9:
1: 0) (1: 0:1:
(0:
(0;
(0:
(0:
(1: 0.1: 0)
(1: 0.1: 6)
(1: 0.1: 0)
(1:
RiverSwim
(0-3 1: 50)
4-5
(3: 900)
(0:
0.03: oy
(0
(1: 0.15, S0
(4: 1: 1600)
(1: 1: 195)
(4: 0.3; 01
2-5
(0-4: 00
0.1: 3
(9: 0000)
(2: 500)
(9-1, 3-5 1;
SixArms
Figure 8: Tabular environments for exploration. The tuples marking each transition denote (action id(s); probability; reward). In RiverSwim, the agent starts in either state 1 or state with equal probability, while for SixArms the agent always starts in state 0.
Tabular exploration We reuse all hyperparameter settings from Machado et al. (2020) in both the RIVERSWIM and SIXARMS environments, with the only difference being lower value for 3, the exploration bonus coefficient, for the FR, as the bonuses given by the FR are generally larger. The hy- perparameters are {a, 17, /SR/ 3, c, n), 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 e-greedy policy. For RIVERSWIM, these values were (0.25, 0.01, 0.95, 100/50 /50 0.1 }. respectively, and for SIXARME they were (0.1, 0.01, 0.99, 100/50 50, 0.01}. For the RIVER SWIM-N task, we chose N = {6. 12, 24} as default RIVERS WIM has N = 6 states, and we chose to suc- cessively 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 3 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 less preferable a bonus which naturally scales.
Table 2: RIVERSWIM-N results. + values denote SE across 100 trials.
RIVERSWIM-N
SARSA + FR
SARSA SR
SARSA
N
1,547, 243 :: 34, 050
1, 197, 075 399999
25, 075 +: 224
N 12
1,497, 937 : 29, 291
714. 797 + 34,4,557 574
14, 590 +: 3, 145
24
962,376 + 33, 325
511,511-
20, 580
11,950 950 + 2, 643
Exploration with function approximation In order to test the usefulness of the FR/FF function approximation setting, we use 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 intermediate layer of the network. Like the standard DQN, the architecture outputs an |Al-length vector of predicted Q-values for the current state, trained off off-po using minibatches of transition tuples
15

---

Published as a conference paper at ICLR 2022.
$ s++
81-
Φ
(s..)) -)
(s)))
Figure 9: Network architecture for DQN + FF and DQN + SF
(6 ob s121) (where is the minibatch size) to minimize the squared Bellman error
LQ= = - 1rt max Q (8(+1, Q(st. aB)
(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 reconstruction head which, given the base feature representation state 8, and an embedding of the subsequent action a, produces prediction of the following state 8 It's trained to minimize the reconstruction loss
C, Σ
| 81+1l12.
(21)
The final output head of the network an FF/SF prediction, trained using the squared FF error in the former case:
B
Cp = Σ 1lo(s) + ( ((s))) -
(22)
and the squared SF error in the latter
Cw Σ ((((s)) 7804(1+1) ((s))1
(23)
For the FF, the features é, passed through sigmoid function compress them in the range [0, 1] and then thresholded at 0.75. The total loss is then given as a weighted combination
L =w@L w.C, wx
(24)
where X E {yp, and wo, w., wx E IR are fixed weights. As in Machado et al. (2020), gradients from Co and L., but not Lx. are permitted flow through to Φb Thus, the base features are trained to be both reward-predictin and to carry information about the environment transition dynamics. The norm of the FF/SF vector then used compute an intrinsic exploration bonus to the task reward in the same manner 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 challenging exploration task set up in 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 choose move either right 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 row, there is large reward +1
16

---

Published as a conference paper at ICLR 2022.
bsulte score for deep_sed sea experiment
2
3
N
0.70
2
agent
DGN
0.80
DON
3
025
N
DGK
agent
0.30
average
35.00 5000
7500
10000
35000
5000
7500
10600
2500
5000
7.00
10000
Figure 10: Exploration with function approximation. (Top left) Visualization of the DeepSea environment, credit to Osband al. (2020). (Top right) DQN FF signficantly outperforms standard DQN and DQN SF. (Bottom) Different runs across problem sizes.
located at the bottom right of environment- this is the only policy which the agent positive reward. In the bsuite framework, the agent is separately trained on increasing problem sizes N 5, 7, 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 2 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 fully-connected layers, with ((x)) being a 2-layer MIP with 64 units per layer, Q(+t) being linear function of Φ(st) with |A| = units, 8141 consisting of 64-unit layer followed by an 2 output layer (the action embedding is 64-dimensional as well), and the SFs/FFs also 64-diimensional linear layer over ((x)). The agent is trained using e-greedy action selection. Our DQN implementation was coded JAX (Bradbury et al., 2018) and based off that of Osband et al. (2020). To select the values of B, w., and wxx (wQ was always kept E {0.01, 0.05, 0.13 and wxx wg E (0.001, 0.1, 1, 10, 100, 1000), choosing the best pesst-eeforming values for each method.
performed sweep
17

---

Published as a conference paper at ICLR 2022.
Table 3: Hyperparameter settings for the DEEPSEA experiment
HYPERFARAMETER
DQN FF
DQN SF
DQN
optimizer
Adam (Kingma & Ba, 2017)
Adam
Adam
leaming rate
0.001
0.001
0.001
0.05
0.01
wig: Wa,
(1, 100, 1000)
(1, 0.0000 1000)
B
32
32
32
replay buffer size
10,000
10,000
10,000
target update period
4
7
0.99
0.99
0.05
0.05
0.05
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 N-carrying over into the function approximation setting. That is, N increases, the norm of the SF approaches its asymptotic value regardless of the degree exploration. 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.
-aatennntttttttlulalllyly RIVERSWIM,
MountainCan Car experiment In our version of the task, the feature representations are learned in a rewardless environment, and test time the reward may be located at any location along the righthand hill. We evaluate the performance of set of policies 11 = {m. 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 ny will swing back and forth along the track with a fixed power coefficient a;. 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 natural tradeoff-to 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 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
swings back and forth to gain momentum.
We defined set of policies with acceleration magnitudes |a.| = 0.1i for = 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.
class FixedPalicy:
def iiiit_ttss (self, a):
# set fixed acceleration/power self.a
def get_action pos, vel):
if vel 0:
10
if stopped, accelerate to the opposite end of the environment action = spos((po)) * self.a
11
12
else:
otherwise, continue in the current direction of notion action sign (vel) self.a
14
15
16
return action
We repeated this process for 20 runs, with the plots in Fig. showing the means and standard deviations of the TD learning curves across runs. For the FFs, the thresholds were constant across features at 0g 0.7. Because of the nature of the environment, all of the policies spent
18

---

Published as a conference paper at ICLR 2022.
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, then tested them by using them features for policy evaluation in different tasks, with each task containing a different rewardec//bssorbnng state. Note that a crucial factor is that the representations were learned in the environment without absorbing states. This 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.
D-44
D
D
M
M
Figure 11: The FF is robust to feature dimensionality. FF and SF representation strengths for difference feature dimensionalities between the start and goal locations for an example goal in continuous MountainCar. The vertical dashed line marks the power of the optimal policy. We can see that for all but the coarsest feature representation, the FF is highest for the policy closest to the optimal.
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 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 this link: https: / anonymous 4open wereensscieneee scissceieeesee / /r/ urlb/ BAAb/REDMME... 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
p intrinsic (s)a,s') = r explait (s, (ss,,,0 s') + explore (s,
(25)
= w Tdss) + log 1
Σ
llØ(s') ((s))()1m))
where w E RR D = 10 is a random reward vector sampled from a standard Gaussian distribution and the righthand term is particle-based estimate of the state-based feature entropy, with (-) denoting the & nearest-neighborg (see Liu & Abbeel (2021) for details). In standard APS, this reward used train the successor features by constructing the bootstrapped target
APS intrinsscc (s, a, 8') + ww (1 w),
(26)
19

---

Published as a conference paper at ICLR 2022.
where d' a w TT (s'' a, w). For the FF, make the following modification:
yAPF y xxPP-exploi + exxllore
(27)
= w ((s) + tw (1- (1- (L)) + eexpoor (s, a, a,,ss) s¹) + (V(4)))
(28)
(x)())
= r (s) + tr explore + [ww[ 1----- + ((s(11),
(29)
with ((8++1) max, w "( d', w) a(.) the thresholded base features, such that (a())) = (0(x)) 2 Os, (o)((())))ee
02). is the D-length vector of ones, and d(x)) E [o. 1]. Interest-
ingly, if the features are kept in the range [0, 1), o(-) 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-pollicy deep deterministic policy gradient (DDPG; (Lillicrap et al., 2019)) algorithm, where given stored replay buffer of transitions D {( d, Pt, 8(+1) })]; the (SF/FF) critic Qw (with Q formed from either the SF or FF and w being the parameters) is trained minimize the squared Bellman loss
Lolw, D) = (e [(V" - Qw (8t,,))))] )+)
(30)
where in the pre-training phase y" Ays . (PF] (the target parameters are an exponential moving average of the weigght-grrddinns do not flow through them). The deterministic actor xo is trained using the derministic policy gradient loss:
Eg (2 (o(s)))
(31)
Fine-tuming: After pre- pre-training, the agent is fine-tuned on the target task, REACHTOPLEPT where the learning proceeds exactly in the pre-training phase, but instead of intrinsic reward, the agent given the task reward-thar is, y" task + (S++1)). We performed this task-specific training for an additional IM steps.
As additional baselines, we implemented two versions of the dynamic distance learning method of Hartikainen et (2020), the original DDLUS and an additional variant DDLUS-G In standard DDLUS, the goals g S for the pre- pre-training phase are generated according to
g" E argmax d (so.s))
JED
where 'D is stored set of trajectories. In Hartikainen et al. (2020), this is useful as a mechanism for encouraging the agent 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 the same manner APS and APF, via
(O 7).
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 APS and APF, rather than SAC (Haarnoja et al., 2018). For a more detailed discussion of DDL and its relationship 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 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 al., 2020)) has different effects.
FourRoom experiments The FourRoom environment we used was defined on an 11 x 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, new goal 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 reward of +50, hitting wall incurred penalty of -1 and kept the agent in the same
20

---

Published as a conference paper at ICLR 2022.
learning FourRooms
PH4
neight
redown
n-
2B1
Figure 12: FOURROOMS learning curves. FOURROOMS base policies learning curves (average L2 norm of TD errors over 10 runs; shaded area is one standard deviation); top row is for FRs. bottom is for SRs.
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 y 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 convergence ( 180 iterations for VI, 3 iterations for FRP) and then tested for 100 episodes.
(s)
sf(s)
0,0
0,1
0,2
1,0
1,2
1.3
2.2
2,0
2.0
2.1
3.2
3.0
Figure 13: Implicit planning output. (Left) The plamning policies n (s) that the agent will elect to follow in each state en route to the goal (see Fig. 3(a)). Arrows denote the action taken by the chosen policy in each state. (Middle) The (row, column) subgoals for each state s (s). (Right) The state space S, for reference.
exploration
pre-update
post-update
post-update
(a)
(b)
(c)
(d)
Figure 14: Exploration and escape (a) A sample trajectory from the "exploration phase" starting from the shelter. (b) Because the agent starts from the shelter during exploration, the first time is tested starting from the top the grid, its FR for the down policy for that state is still initialization. (c) After updating its FR during testing and further exploration, the FR for the down policy from the start state accurate, stopping at the barrier. (d) We can see that we were to use the SR instead, value in the state above the wall would accumulate when it gets stuck.
Escape experiments The escape experiments were modeled as a 25 x 25 gridworld with eight available actions.
A fup, right, down, left, htgpl//cw.........r.rrrroooooore
(32)
21

---

Published as a conference paper at ICLR 2022.
exploration phase
escape trial
exploration phase 2
SRc 116
apisocisss
steps
Figure 15: Escape learning curves. Learning curves (norms of TD errors) for the first exploration phase, the first escape trial, and the second exploration phase for the "down" policy. The vertical dotted lines in the escape trial mark the time step at which the agent encounters the barrier. This causes temporary jump in the TD errors, as representation learning did reflect the wall this point. The top row consists of FR results and the bottom row is from SRs, averaged over 10 runs. The shading represents one standard deviation.
and barrier one-half the width of the grid optionally present in the center of the room. The discount factor y 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 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 base policy at each time step, after which it updated the corresponding FR. Each exploration phase consisted shown in Fig. 14()). 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 hypothesized that in this case, the mouse should either i) default to policy which would take to the shelter the area the room which it knew well (the down policy) or ii) default to a policy which would simply take away from the threat (again the down policy). During the first escape trial, the agent selects the down policy repeatedly, continuing update FRs during the testing phase. Upon reaching the wall and getting stuck, the FR for the down policy is eventually updated enough that re- p-lanninng 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 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.
episodes- sample trajectory
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 barrier, while the SR updates when the agent hits (Fig. 12). since there is no single policy that can reach the shelter, GPI fails to find path around the barrier (Fig. 16, riiii.
22

---

Published as a conference paper at ICLR 2022.
no barrier (SR)
barrier trial (SR)
barrier trial 2 (SR)
Figure 16: An SR cannot effectively escape under the same conditions an FR agent.
A.4 ADDITIONAL PROOPS
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 Definition 3.2 for some stationary policy æ. Then for any two matrices F, F,F' S|
Fg,
(')ss'))
(33)
with the difference equal to zero for s == s'.
Proof. For o s we have
1(g -""F ~(P=FF=
Pa - F').,>l ( ( -F'),.'
where we use the notation Xe.e' to mean ((,,s') s'), and the inequality due to the fact that every element of (" F') is convex average of F F. For 8', we trivially have ( F 9" F')s,ss =
= O.
17
Proposition A.1 (Convergence). Under the conditions assumed above, set F(O) T(s). For k 0, suppose p(++1) G p(k), Then
(s,s') F 7t
(34)
for s s with the difference for s s' equal zero vk.
Proof. We have, for 8 + s' and using the notation X 88) for matrix X,
[F F ),,rr = 1(g p(0) gg* (),''),,r))
- (1((1))))) - 1*)*))a))) ),,0r| (Proposition 3.1) (* < th (FF"), E ())
(35)
For s', ((F())
PT) )o.oo = vk..
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 8g. and policy set II composed of policies и S A. We make the following coverage assumption, there exists some sequence of policies that reaches 8n from a given start state 86- Under these conditions, Alg. converges such that I(sn) = = where L is the shortest path length from 8o to 8g using æ E
Proof. Since the MDP is deterministic, We proceed by induction on L
deterministic transition function p p:S A S
Base case: T- =
If L 8o must be step from Sy- The coverage assumption guarantees that 3r II such that p(Sp. x (So)) = $g. Note also that when both the MDP and policies in are deterministic, F* F"ss.''= - - where L is the number steps from s to under n, and we use the abuse of notation L if does reach s' from
23

---

Published as a conference paper at ICLR 2022.
Then following Alg. 1,
T (so)
max P (50; 8g)
(guaranteed by coverage of II)
REll
(1
max F (Sg: $g) = (by definition of F*).
xEIl
Moreover,
(2))
max F (s).s')((
ES
max F 80): Fd (s0, (5)))
=Ell
maxx1 97,111
Then (2(8) (8) Vs and Alg. terminates. Thus, r(s0) - 7- 7 and the base case holds.
Induction step: Assume Proposition 4.1 holds for L = L
Given the induction assumption, we now need show that Proposition 4.1 holds for L; L By the induction and coverage assumptions, there must exist at least one state within one step of sn that the agent can reach in L steps, such that the discount for this state or states is 7-. Moreover, the coverage assumption guarantees that 3r E II such that for at least one such state 8L p(8L, (+L)) 89-
Then this problem reduces to the base case-that is, Alg. will select the policy æ E IT that transitions directly from sg Sy -andnd the proof is complete.
A.5 EXPLICIT PLANNING
Below we describe procedure for constructing an explicit plan using T" and s",
Algorithm 2: ConstructPian
input: goal state 8y, planning policy #F, subgoals F
2: A k-1
> init. plan begin at start state
3: s + 50
4: while + 8y do
5: A.append(( ( (8), (s)))
add policy-subgoal pair for current state to plan
6: 8 +- 8f(s)
7: end while
8: return
A.6 FRP WITH MULITPLE GOALS
Here we consider the application of FRP to environments with multiple goals fgi...., 9n.. To find the shortest path between them given the base policy set II, we first run FRP for each possible goal state in {gi]. yielding an expected discount matrix I'll [0, 1]ls x such that (s,gg) is the expected discount of the shortest path from state to goal 9:- We denote by 9. = (so, (gg)) (g))), a (g.)) ) a specific ordering of the goals in {g:} starting sg. The expected discount of sequence of goals then
12
E(ga)
II i=1
r (go 11, (i)),
(36)
with the optimal goal ordering 9o- given by
Jo-
argmax E(go),
(37)
E.
where is the set all possible permutations of {9:}. of size n!. This is related to a form of the travelling salesman problem, and refer the reader to Zahavy et al. (2019) for formal investigation of the use of local policies to construct shortest paths. Fortunately, in most settings we don't expect the number of goals to be particularly large.
24

---

Published as a conference paper at ICLR 2022.
Table 4: Overview of basic points of comparison between the FR/FF, SR/SF (Dayan, 1993; Barreto ct al., 2017b), DDL (Hartikainen et al., 2020), TDMs (Pong et al., 2018), and DG (Kaelbling, 1993).
FR/FF
SR/SF
DDL
TDM
DG
On- v. Off-policy
Both
Both
On-policy
Off-policy
Off-policy
Eval. v. Control
Eval.
Eval.
Eval.
Control
Control
Finite v. Inf. Horizon
Both
Both
Finite Horizon
Finite Horizon
Finite Horizon
State representation?
Yes
Yes
No
No
No
A.7 CONNECTIONS TO OPTIONS
The options framework (Sutton et al., 1999) is method for temporal abstraction in RL, wherein an option w is defined as a tuple (Tw: Tw), where Tw a policy and тu E A(S) a state-dependent termination distribution (or function, if deterministic). Executing an option at time entails sampling an action ag ~ (-Is;) and ceasing exccution of π. at time t + 1 with probability Tu ($1+1). The of options enlarges MDP's action space, whereby higher-level policy selects among basic, low-level actions and options.
Options connected to FRP (Alg. 1) in that by outputting FRP effectively converts each base policy to an option with deterministic termination function, i.e., the agent will follow nF 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 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 new goal.
of policies and associated subgoals
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 transfcrable 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.
A.8 FURTHER CONNECTIONS TO RELATED WORK
We now describe the connection between the FR/FF and several related approaches in the literature.
Table 4 summarizes high-level view of these connections.
The Dynamic Distance Function is defined as
The dynamic distance function (DDF; (Hartikainen et al., 2020)
d3 (s, En
Σ
(( =8,; =
(38)
k==
where c: S S is a local cost function. In practice, $14b-1) 1, giving
d= (8, 8')
Er
Σ 8t = 8, s
(39)
n practice ds, parameterized via a neural network with parameters vp trained on- policolcy from full trajectories using the loss
((w) - B... - -
(40)
25

---

Published as a conference paper at ICLR 2022.
where D is buffer of stored trajectories. On a downstream navigation task, the agent policy is trained to minimize
00
(0) = Ex
r (sts
g)
(41)
1=0
where g S is 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 on- policy learning from full trajectories, and cannot be updated off- 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 finite horizons significant because it leads 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, can lead significant value estimation errors.
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 y = 1, the number of steps & appears linearly within the expectation rather than exponentially (this is important, as in general 7 [A + E[>*1.
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, used to support policy optimization, rather than evaluation. In contrast, the FF representation, like the SF representation, maps single state to a fixed length vector- valued encoding, : S + Rd. This difference has significant implications for the applications of the DDF. In particular, it is unelear 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 (Kael- bling, 1993), method for optimal control related Q-learning. The optimal DG function G* : A x S + IR is defined recursively for goal state g E S as
(*,,,,,)) = + (g/ min G* (*)
(42)
a'EA
where G'(g- a, == There 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, (s,s, a, g) converges to the expected number of steps for the optimal policy for reaching starting from s. It cannot be reused for policy evaluation for policy π * x², 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 gridworld environment. A second difference is that the DG function scales linearly with k, the number of steps for the optimal policy reach g, while the FR scales exponentially at rate of This is important, as we can note (with slight abuse of notation) in general that 7 o Ex [ it is nontrivial to recover the 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 goal- conditioned Q-fuccioons can be used to construct an implicit dynamics model when the discount factor y = 0. For y > 0, the authors introduce a horizon variable r to interpolate between model- based and model-free learning. Like the FR, this formulation allows agent to consider multiple
26

---

Published as a conference paper at ICLR 2022.
goals in parallel. As in DG learning, and unlike the FR, TDMs represent 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 bonus). Another difference from the FR is that, like the DDF and DG learning, TDMM rely on finite-horizon setting.
exploration
A.9 FR VS. SR VISUALIZATION
M
M
M
1.0
50
Se
06
F2
F3
04
-00
Figure 17: SR vs. FR visualization The SRs and FRs from the start state for the policies in Fig. Fig. For the SRs (Fig. 17, top row), we can see that states that are revisited (or in which the policy simply stays) are more highly weighted, while for the FRs (Fig. 17, bottom row), the magnitude of F(s0: s') is higher for states s' that are closer along the path taken by policy.
27


Reviews:
[{"review_id": "tkSb2qhQMe", "reviewer": "Reviewer_SXgZ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed representation (FR) seems to have clear advantages in comparison to approaches that convert MDPs with non-Markovian reward structure to such with Markovian one. The use of the first-occupancy values as an exploration bonus results in much more efficient exploration. The authors also demonstrate that FR can be used effectively for unsupervised pre-training RL with non-Markovian rewards, where the successor representations tends to produce incorrect estimates of the value function. The use of FR for planning is also demonstrated on a high-dimensional continuous control problem involving a robotic arm. All of these illustrate well the advantages of the proposed representation, and overall, I find this representation highly original.\n\nOne limitation of this general approach is probably the assumption that a fixed, and relatively small number of policies will be provided in advance, and the solution of the sequential decision problem can be expressed by switching between these policies. It is not clear if and under what conditions this assumption is justified. Could the (sub-)optimality of this approach be analyzed theoretically or empirically, by computing the true optimal policy and comparing with its performance? Section 3 talks about \"Policy evaluation and improvement with the FR\", and although the use of FR for fast policy evaluation is explained very well, I do not see any discussion about policy improvement there.\n\nA minor typo, on page 4, last paragraph: \"automoton\" -> \"automaton\".\n \n ", "summary_of_the_paper": "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.", "summary_of_the_review": "I find the proposed novel representation highly original, and ICLR is the right venue for its publication. It is a clear advance of the state of the art in computing intermediate representations that cache the system dynamics in a format that is more suitable for policy evaluation than the basic transfer function of the MDP. Its advantages over successor representations for the case of non-Markovian rewards are clearly explained and supported by empirical evaluations. For these reasons, I recommend acceptance. It is also apparent that the proposed approach has some limitations, and the authors readily acknowledge this. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "mxAVwC0aOt", "reviewer": "Reviewer_epv4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Short Version\nOverall, I think this is a very nice paper that presents a neat little idea with a lot of potential impact in the community. Successor features have resurfaced in the RL community in the past few years, and the proposed \"FF\" features improve upon them in a very specific, but evident way. My biggest concern is the somewhat limited empirical evaluations and broader comparisons that bring out the strengths of the method.  \n[Full Disclosure: I have not checked the accuracy of the theoretical claims thoroughly; they seem appropriate on a quick glance.]\n\n---\n## Strengths\n1. The writing is very succinct and the authors make a great case of comparing against the SR in every way possible (and do so fairly).\n2. The paper is well-written in terms of reproducibility and the authors make plenty of information available about the experiments and algorithm to reproduce the results. The theoretical results are also discussed elaborately in the appendix.\n3. The suite of capabilities considered, and demonstrated, is very complete and makes a good case for \"FR > SR\" -- exploration bonus, unsupervised pretraining and planning. Particularly, the section on planning with FR makes a good case for their strengths and could be made more prominent in the contributions/introduction, in my opinion. The paper makes a great case for the _completeness_ of a first-occupancy representation.\n4. The parallels to animal behavior, although a bit far-fetched to me, are very nice inclusions in the article. I appreciate the candor of the authors in presenting the facts (\"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 under- standing of animal learning\") and not overfitting their findings to tall claims about understanding the animal brain. The findings are interesting in their own right and add value to the paper.\n\n\n---\n## Weaknesses\n1. While the paper does a great job at comparing against the vanilla successor representation and shows (quite clearly) that FR > SR in tabular domains, my biggest concern is that the results don't go too far beyond that. In most realistic domains, the infinite state-space version (FF) would be more relevant and there is almost no empirical analysis of how the method performs in more complex domains/tasks, beyond Section 4.2. The authors also do not compare against a broader suite of methods in the recent years that use successor features to do more interesting control and planning problems; a more thorough empirical evaluation will greatly increase the impact of the work and serve as a strong baseline for future works to present results against. As a suggestion, I would recommend [a recent paper by Janner et al.](https://arxiv.org/abs/2010.14496) which conducts thorough empirical evaluation of a model similar to the SR, and can provide some pointers towards more tasks to consider for empirical analysis.\n2. While motivating the need for a _better_ representation, the authors frequently make claims about how the problem of non-Markovian rewards is widespread and _natural_, e.g. \"Such natural problems emphasize the importance of FR...\". However, the text fails to provide a concrete example for the reader to conceptualize the claims and contributions of the paper. As a suggestion on the presentation, I would urge the authors to present a concrete example that motivates the need for the first-occupancy early on, or a toy problem that can isolate the issue, for clarity.", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper presents a simple idea which is supported by ample theoretical analysis and (some) experiments. Comparison to prior work, while not exhaustive, is convincing. I vote to accept the paper but strongly urge the authors to improve the empirical analysis for the overall impact of the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "mjIbIbzeVvK", "reviewer": "Reviewer_Ugir", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well-written and well-motivated. The idea is simple but is a nice extension of the concept of the SR. It also seems promising on the different tasks considered by the authors. In particular, the case of non markovian rewards is studied in section 4.2.\n\nHow does the FR relate to the dynamic distance function introduced in https://arxiv.org/pdf/1907.08225.pdf ?\n\nIn Section 4.2, could you provide a reference for  “a bonus to maintain its effectiveness as time progresses in order to prevent the rate of policy improvement from exponentially decaying“? Unlike the analogy of SR / pseudo-count, it s still a bit unclear what the motivation for the FR is. Testing the benefit of the FR on harder exploration tasks would make the paper stronger I think (e.g., Machado et al 2020 tested their approach on Atari games)\n\nRegarding infinite state spaces, the authors might want to note that recent works has shown that the SR can be extended the continuous state spaces when viewed as a measure without the need to rely on basis functions. See https://arxiv.org/abs/2101.07123 https://arxiv.org/abs/2103.07945\n\nMinor points: Typos: take advantage *of*, theoeretical ", "summary_of_the_paper": "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. \n", "summary_of_the_review": "I think this paper provides a nice contribution to the topic of representation learning in RL. The authors provided theoretical results for their FR and demonstrated its benefits on 4 different tasks. Although it would be nice to confirm these findings on harder tasks e.g., in the case of section 4.1, I believe the variety of tasks considered by the authors still shows some promise which is why I recommend an accept for this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "eKu80ww9S_M", "reviewer": "Reviewer_aieJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper proposes a general-purpose modification of the successor representation that has a number of nice properties: in particular, it can be used for shortest path search and can be stitched together to form a multi-step plan in the spirit of options. The paper also shows demonstrations of its usefulness in unsupervised RL and exploration. The evidence for each individual application is somewhat preliminary, which I think is fine for a paper that intends primarily to introduce a new prediction problem and outline potential applications. The most compelling idea to me is the planning (FRP) procedure, which naturally fits into other DP-style methods and multi-step planning frameworks.\n\nThe main issue I can find is that the proposed first-occupancy representation might not be as new as the paper describes. As early as [Kaelbling 1993](http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.51.3077) there was the so-called \"dynamic goal learning\" (DG) algorithm, which uses TD to predict the minimum number of timesteps required to reach a goal. The FR is in the space of discounts and DG in the space of timesteps, but at least for deterministic policies and dynamics it seems like there would be a straightforward conversion between the two ($FR = \\gamma^{DG}$?).\n\nFor a more modern take on the same problem, [dynamical distance learning (DDL)](https://arxiv.org/abs/1907.08225) also estimates the expected number of timesteps required to reach a goal state, and [temporal difference models (TDMs)](https://arxiv.org/abs/1802.09081) estimate how close a policy will come to a goal state in a fixed number of timesteps. \n\nIt is not necessarily an issue that there are papers from the last few years investigating a similar line of inquiry; the paper is still a thorough description of this idea and has enough in it that it is not subsumed by these prior works. But since the space has already been scoped out a bit, the framing of this paper as a first introduction of the idea might not be the right one, and raises the bar slightly in terms of experimental evaluation. It would be valuable to compare it to the related approaches that have already been scaled to higher dimensions (like DDL and TDMs).\n\n**Minor**\n1. \"to take advantage this shared\" --> \"advantage of\"\n2. In Equation 8, should the condition read $s_t = s$?\n", "summary_of_the_paper": "This paper presents the first-occupancy representation, a sort of non-Markovian analogue of the successor representation that carries information about the time required to reach a state for the first time rather than the full discounted occupancy of a policy.", "summary_of_the_review": "The paper gives a thorough description of a variant of the successor representation and its potential application in a few different settings. Its main issue is that it frames itself as a sort of preliminary description of an idea, with a somewhat limited evaluation in any one setting, but very similar ideas have already been shown to scale to the types of problems commonly studied in contemporary deep RL. Including a comparison to these ideas (both experimentally and also in the framing) would improve the paper and situate it better in context.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a first-occupancy representation for reinforcement learning problems, with potential benefits on problems with non-stationary rewards.  The representation is defined analogously to the successor representations, but captures the expected discounted time to first arrive at a state instead of measuring discounted visitations.  The paper develops the idea and illustrates some uses for exploration, unsupervised RL, and non-stationary reward functions (for example when food rewards are consumed).

The reviews brought forward a number of related older ideas in the literature, where several aspects of the method have been previously developed.  These include dynamic goal learning, option conditional predictions, general value functions, dynamical distance learning, and temporal difference models.  However, from the author response and ensuing discussion, the exact form of the proposed representation has not been studied for the purposes presented in this paper.  The reviewers appreciated the utility of this representation for problems with non-Markovian rewards, in particular that “the use of the first-occupancy values as an exploration bonus results in much more efficient exploration”.  Multiple reviewers commented on the desire for a stronger empirical evaluation, but they were satisfied with the contribution of the paper.

The reviewers arrived at a consensus that the paper contributes a new representation for RL problems with non-stationary rewards, with two reviewers strongly convinced and none opposed.  The paper is therefore accepted.
Decision: Accept (Poster)