A F IRST -O CCUPANCY R EPRESENTATION FOR R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Both animals and artiﬁcial agents beneﬁt from state representations that support rapid transfer of learning across tasks and which enable them to efﬁciently traverse their environments to reach rewarding states.
The successor representation (SR), which measures the expected cumulative, discounted state occupancy under a ﬁxed policy, enables efﬁcient transfer to different reward structures in an otherwise constant Markovian environment and has been hypothesized to underlie aspects of biological behavior and neural activity.
However, in the real world, rewards may move or only be available for consumption once, may shift location, or agents may simply aim to reach goal states as rapidly as possible without the constraint of artiﬁcially imposed task horizons.
In such cases, the most behaviorally-relevant representation would carry information about when the agent was likely to ﬁrst reach states of interest, rather than how often it should expect to visit them over a potentially inﬁnite time span.
To reﬂect such demands, we introduce the ﬁrstoccupancy representation (FR), which measures the expected temporal discount to the ﬁrst time a state is accessed.
We demonstrate that the FR facilitates exploration, the selection of efﬁcient paths to desired states, allows the agent, under certain conditions, to plan provably optimal trajectories deﬁned by a sequence of subgoals, and induces similar behavior to animals avoiding threatening stimuli.
I NTRODUCTION
In order to maximize reward, both animals and machines must quickly make decisions in uncertain environments with rapidly changing reward structure.
Often, the strategies these agents employ are categorized as either model-free (MF) or model-based (MB) (Sutton & Barto, 2018).
In the former, the optimal action in each state is identiﬁed 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 efﬁcient but ﬂexible planning.
One such approach, based on the successor representation (SR) (Dayan, 1993), has been the subject of recent interest in the context of both biological (Stachenfeld et al., 2017; Gershman, 2018; Momennejad et al., 2017;
Vértes & Sahani, 2019;
Behrens et al., 2018) and machine (Kulkarni et al., 2016; Barreto et al., 2017b;a; 2018;
Machado et al., 2020; Ma et al., 2020; Madarasz & Behrens, 2019) learning.
The SR associates with each state and policy of action a measure of the expected rate of future occupancy of all states if that policy were to be followed indeﬁnitely.
This cached representation can be acquired through experience in much the same way as MF methods and provides some of the ﬂexibility of MB behaviour at reduced computational cost.
Importantly, the SR makes it possible to rapidly evaluate the expected return of each available policy in an otherwise unchanging environment, provided that the reward distribution remains consistent.
However, these requirements limit the applicabilty of the SR.
In the real world, rewards are frequently non-Markovian.
They may be depleted by consumption, frequently only being available on the ﬁrst entry to each state.
Internal goals for control—say, to pick up a particular object—need to be achieved as rapidly as possible, but only one at a time.
Furthermore, while a collection of SRs for different policies makes it possible to select the best amongst them, or to improve upon them all by considering the best immediate policy-dependent state values (Barreto et al., 2018), this capacity still falls far short of the power of planning within complete models of the environment.
Here, we propose a different form of representation in which the information cached is appropriate for achieving ephemeral rewards and for planning complex combinations of policies.
Both features arise from considering the expected time at which other states will be ﬁrst accessed by following the available policies.
We refer to this as a ﬁrst-occupancy representation (FR).
The shift from expected rate of future occupancy (SR) to delay to ﬁrst occupancy makes it possible to handle settings where the underlying environment remains stationary, but reward availability is not Markovian.
Our primary goal in this paper is to formally introduce the FR and to highlight the breadth of settings in which it offers a compelling alternative to the SR, including, but not limited to: exploration, unsupervised RL, planning, and modeling animal behavior.
R EINFORCEMENT L EARNING P RELIMINARIES
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-deﬁned environment.
Typically, task T is modelled as a ﬁnite Markov decision process (MDP; (Puterman, 2010)), T = ( S , A , p, r, γ, µ ) , where S is a ﬁnite state space, A is a ﬁnite action space, p : S ×
A → ∆( S ) is the transition distribution (where ∆( S ) is the probability simplex over S ), r : S → R is the reward function, γ ∈ [0 , 1) is a discount factor, and µ ∈ ∆( S ) is the distribution over initial states.
Note that the reward function is also frequently deﬁned over state-action pairs ( s, a ) or triples ( s, a, s (cid:48) ) , 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 (cid:80) t γ t r ( s t ) .
To simplify notation, we will frequently write r ( s t ) :
= r t and r ∈ R |S| as the vector of rewards for each state.
The agent acts according to a stationary policy π : S → ∆( A ) .
For ﬁnite MDPs, we can describe the expected transition probabilities under π using a |S| × |S| matrix P π such that P πs,s (cid:48)
=
p π ( s (cid:48) | s ) :
= (cid:80) a p ( s (cid:48) | s, a ) π ( a | s ) .
Given π and a reward function r , the expected return is
[Equation]
Q πr are called the state-action values or simply the Q -values of π .
The expectation E π [ · ] is taken with respect to the randomness of both the policy and the transition dynamics.
For simplicity of notation, from here onwards we will write expectations of the form E π [ ·| s t = s, a t = a ] as E π [ ·| s t , a t ] .
This recursive form is called the Bellman equation , and it makes the process of estimating Q πr — termed policy evaluation —tractable via dynamic programming (DP; (Bellman, 1957)).
In particular, successive applications of the Bellman operator T π Q := r + γP π Q are guaranteed to converge to the true value function Q π for any initial real-valued |S| × |A| matrix Q .
When the transition dynamics and reward function are unknown, temporal difference (TD) learning updates value estimates using a bootstrapped estimate of the Bellman target (Sutton & Barto, 2018).
Given a transition sequence ( s t , a t , r t , s t +1 ) and a t +1 ∼ π ( ·| s t [Equation]
[Equation]
Once a policy has been evaluated, policy improvement identiﬁes a new policy π (cid:48) such that Q πr ( s, a ) ≥
Q π (cid:48) r ( s, a ) , ∀ ( s, a ) ∈
Q πr ( s, a ) .
Helpfully, such a policy can be deﬁned as π (cid:48)
( s ) ∈ argmax a Q πr ( s, a ) .
The successor representation The successor representation (SR; (Dayan, 1993)) is motivated by the idea that a state representation for policy evaluation should be dependent on the similarity of different paths under the current policy.
The SR is deﬁned as a policy’s expected discounted state occupancy, and for discrete state spaces can be stored in an |S| × |S| matrix M π , where
[Equation]
where 1 ( · ) is the indicator function.
The SR can also be conditioned on actions, i.e., M π ( s, a, s (cid:48) ) :
=
E π (cid:104)(cid:80) k
γ k 1 ( s t + k = s (cid:48) )
(cid:12)(cid:12)(cid:12) s t , a t (cid:105) , and expressed in a vectorized format, we can write M π ( s ) : =
M π ( s, · ) or M π ( s, a ) :
= M π ( s
[Equation]
a, · ) .
The recursion in Eq.
(3) admits a TD error:
[Equation]
where 1 ( s t ) is a one-hot state representation of length |S| .
One useful property of the SR is that, once converged, it facilitates rapid policy evaluation for any reward function in a given environment:
[Equation]
Fast transfer for multiple tasks In the real world, we often have to perform multiple tasks within a single environment.
A simpliﬁed framework for this scenario is to consider a set of MDPs M that share every property (i.e., S , A , p, γ, µ ) except reward functions, where each task within this family is determined by a reward function r belonging to a set R .
Extending the notions of policy evaluation and improvement to this multitask setting, we can deﬁne generalized policy evaluation (GPE) as the computation of the value function of a policy π on a set of tasks R .
Similarly, generalized policy improvement (GPI) for a set of “base” policies Π is the deﬁnition of a policy π (cid:48)
such that
[Equation]
for some r ∈ R .
As hinted above, the SR offers a way to take advantage of this shared structure by decoupling the agent’s evaluation of its expected transition dynamics under a given policy from a single reward function.
Rather than needing to directly estimate Q π ∀ π ∈ Π , M π only needs to be computed once, and given a new reward vector r , the agent can quickly peform GPE via Eq.
As shown by Barreto et al. (2017a), GPE and GPI can be combined to deﬁne a new policy π (cid:48) via
[Equation]
For brevity, we will refer to this combined procedure of GPE and GPI simply as “GPI", unless otherwise noted.
The resulting policy π (cid:48) is guaranteed to perform at least as well as
any individual π ∈ Π
(Barreto et al., 2020) and is part of a larger class of policies termed set-improving policies which perform at least as well as any single policy in a given set (Zahavy et al., 2021).
T HE F IRST -O CCUPANCY R EPRESENTATION
While the SR encodes states via total occupancy, this may not always be ideal.
If a task lacks a time limit but terminates once the agent reaches a pre-deﬁned goal, or if reward in a given state is consumed or made otherwise unavailable once encountered, a more useful representation would instead measure the duration until a policy is expected to reach states the ﬁrst time.
Such natural problems emphasize the importance of the ﬁrst occupancyand motivate the ﬁrst-occupancy representation (FR).
Deﬁnition 3.1.
For an MDP with ﬁnite S , the ﬁrst-occupancy representation (FR) for a policy π
F π ∈ [0 , 1] |S|×|S| is given by
[Equation]
where { s t : t + k } = { s t , s t +1 , . . .
, s t + k − 1 } , with the convention that { s t : t +0 } = ∅ .
That is, as the indicator equals 1 iff s t + k = s (cid:48) and time t + k is the ﬁrst occasion on which the agent has occupied s (cid:48) since time t , F π ( s, s (cid:48) ) gives the expected discount at the time the policy ﬁrst reaches s (cid:48) starting from s . The idea of learning policy-dependent distances to target states has a long history in RL (Kaelbling, 1993; Pong et al., 2018; Hartikainen et al., 2020). Previous methods don’t learn these distances as state representations and measure distance in the space of time steps, rather than discount factors.
A more thorough discussion can be found in Appendix A.8. We can write a recursive relationship for the FR (derivation in Appendix A.1):
[Equation]
This recursion implies the following Bellman operator, analogous to the one used for policy evaluation: [Figure] ss g a 2D grid (b) (c)
M ( s 0 , s g ) F ( s 0 , s g ) (d) 0.
1ss g
GPI trajectories
[Figure] Deﬁnition 3.2 (FR Operator) .
Let F ∈ R |S|×|S| be an arbitrary real-valued matrix.
Then let G π denote the Bellman operator for the FR, such that
[Equation]
where 1 is the length- |S| vector of all ones.
In particular, for a stationary policy π , G π F π = F π .
The following result establishes G π as a contraction, with the proof provided in Appendix A.4.
Proposition 3.1 (Contraction) .
Let G π be the operator as deﬁned in Deﬁnition 3.2 for some stationary policy π .
Then for any two matrices F, F (cid:48) ∈ R |S|×|S| ,
[Equation]
(11) with the difference equal to zero for s = s (cid:48) .
This implies the following convergence property of G π .
Proposition 3.2 (Convergence) .
Under the conditions assumed above, set F (0)
=
I |S| .
For k = 0 , 1 , . . .
, suppose F ( k +1) =
G π F ( k ) .
Then
[Equation]
for s (cid:54) =
s (cid:48) with the difference for s = s (cid:48) equal to zero ∀ k .
Therefore, repeated applications of the FR Bellman operator G k F → F π as k → ∞ .
When the transition matrix P π is unknown, the FR can instead be updated through the following TD error:
[Equation]
In all following experiments, the FR is learned via TD updates, rather than via dynamic programming.
To gain intuition for the FR, we can imagine a 2D environment with start state s 0 , a rewarded state s g , and deterministic transitions (Fig. 1a).
One policy, π 1 , reaches s g slowly, but after ﬁrst encountering it, re-enters s g inﬁnitely often (Fig. 1b).
A second policy, π 2 , reaches s g quickly but never occupies it again (Fig. 1c).
In this setting, because π 1 re-enters s g multiple times, despite arriving there more slowly than π 2 , M π 1 ( s 0 , s g ) >
M π 2 ( s 0 , s g ) , but because the FR only counts the ﬁrst occupancy of a given state, F π 1 ( s 0 , s g )
<
F π 2 ( s 0 , s g ) .
The FR thus reﬂects a policy’s path length between states.
Policy evaluation and improvement with the FR Like the SR, we can quickly perform policy evaluation with the FR.
Crucially, however, the FR induces the following value function:
[Equation]
where r F : S → R is a reward function such that r F ( s t + k ) =
r ( s t + k ) if s t + k / ∈ { s t : t + k } andotherwise.
In other words, multiplying any reward vector by the FR results in the value function for a corresponding task with non-Markovian reward structure in which the agent obtains rewards from states only once. Policy improvement can then be performed with respect to Q πr F as normal.
This structure is a very common feature of real-world tasks, such as foraging for food or reaching to a target.
Accordingly, there is a rich history of studying tasks with this kind of structure, termed non-Markovian reward decision processes (NMRDPs; (Bacchus et al., 1996; Peshkin et al., 2001; Littman et al., 2017; Gaon & Brafman, 2020)).
Helpfully, all NMRDPs can be converted into an equivalent MDP with an appropriate transformation of the state space (Bacchus et al., 1996).
Most appproaches in this family attempt to be generalists, learning an appropriate state transformation and encoding it using some form of logical calculus or ﬁnite 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 F , it would be exponentially expensive in |S| , as every path would need to account for the ﬁrst occupancy of each state along the path.
That is, |S| bits would need to be added to each successive state in the trajectory.
Crucially, the FR has the added advantage of being task-agnostic, in that for any reward function r in a given environment, the FR can immediately perform policy evaluation for the corresponding r F .
Inﬁnite state spaces A natural question when extending the FR to real-world scenarios is how it can be generalized to settings where |S| is either impractically large or inﬁnite 1 .
In these cases, the SR is reframed as successor features ψ π (SFs; (Kulkarni et al., 2016; Barreto et al., 2017b)), where the d th SF is deﬁned as ψ πd ( s ) :
=
E π (cid:2)(cid:80) ∞ k =0 γ k φ d ( s t + k ) (cid:12)(cid:12) s t
= s (cid:3) , where d = 1 , . . .
, D and φ : S → R D is a base feature function.
The base features φ ( · ) are typically deﬁned so that a linear combination predicts immediate reward (i.e., w T φ ( s )
=
r ( s ) for some w ∈ R D ), 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 deﬁne a ﬁrst-occupancy feature (FF) representation ϕ π , where the d th FF is given by
[Equation]
where θ d is a threshold value for the d th feature.
The indicator equals 1 only if s t + k is the ﬁrst state whose feature embedding exceeds the threshold.
Note that this representation recovers the FR when the feature function is a one-hot state encoding and the thresholds { θ d } Dd =1 are all 1.
E XPERIMENTS
We now demonstrate the broad applicability of the FR, as well as highlight its unique properties with respect to the SR.
We focus on 4 areas: exploration, unsupervised RL, planning, and animal behavior
. 4. T HE FR AS AN E XPLORATION B ONUS
[Table]
Intuitively, representations which encode state visitation should be useful measures of exploration.
Machado et al.
(2020) proposed the SR as a way to encourage exploration in tasks with sparse or deceptive rewards.
Speciﬁcally, the SR is used as a bonus in on-policy learning with Sarsa (Rummery & Niranjan, 1994): [Equation]
[Equation]
where β ∈ R controls the size of the exploration bonus.
Machado et al.
(2020) show that during learning, (cid:107) M π ( s ) (cid:107) 1 can act as a proxy for the state-visit count n ( s ) , with (cid:107) M π ( s ) (cid:107) − 1 1 awarding a progressively lower bonus for every consecutive visit of s .
In the limit as t → ∞ , however, (cid:107) M π ( s ) (cid:107) − 1→ 1 − γ ∀
s as π stabilizes, regardless of whether π has effectively explored.
To encourage exploration, we’d like for a bonus to maintain its effectiveness over time.
In contrast to (cid:107) M π ( s ) (cid:107) 1 , 1 ≤ (cid:107) F π ( s ) (cid:107) 1 ≤ κ |S| := 1 − γ |S| +1 1 − γ , where κ |S| > 1 for |S| ≥ 1 .
Note that (cid:107) F π ( s ) (cid:107) 1 = κ |S| only if π reaches all states in as many steps.
Because (cid:107) F π (cid:107) 1 only grows when new states or shorter paths are discovered, we can instead augment Sarsa as follows:
[Equation]
[Figure]
[Figure]
) (c) (d) (e)
(a) ss g
A Grid (b) (c) (d) (e) ss g
GPI (f) ss g
FR-planning We tested our approach on the R IVER S WIM and
S IX
A RMS problems (Strehl & Littman, 2008), two hard-exploration tasks from the PAC-MDP literature.
In both tasks, visualized in Appendix Fig. 8, the transition dynamics push the agent towards small rewards in easy to reach states, with greater reward available in harder to reach states.
In both cases, we ran Sarsa, Sarsa with an SR bonus (S ARSA +
SR) and Sarsa with an FR bonus (S ARSA + FR) for 5,000 time steps with an (cid:15) -greedy policy.
The results are listed in Table 1, where we can see that the FR results in an added beneﬁt over the SR.
It’s also important to note that the maximum bonus κ |S| has another useful property, in that it scales exponentially in |S| .
This is desirable, because as the number of states grows, exploration frequently becomes more difﬁcult.
To measure this factor empirically, we tested the same approaches with the same settings on a modiﬁed R IVER S WIM , R IVER S WIM -N, with N = { 6 , 12 , 24 } states, ﬁnding that S ARSA + FR was more robust to the increased exploration difﬁculty (see Appendix Tableand Appendix A.3 for results and more details). We also tested whether these results extend to the function approximation setting, comparing a DQN-style model (Mnih et al., 2015) using the SF and
FF as explorating bonuses on the challenging D EEP S EA task (Osband et al., 2020), ﬁnding that the advantage of the FF is conserved.
See Appendix A.3 for details.
Developing further understanding of the relationship between the FR and exploration represents an interesting topic for future work.
4. U NSUPERVISED RL WITH THE FF
We demonstrate the usefulness of the FF in the unsupervised pre-training RL (URL) setting, a paradigm which has gained popularity recently as a possible solution to the high sample complexity of deep RL algorithms (Liu & Abbeel, 2021; Gregor et al., 2016;
Eysenbach et al., 2018; Sharma et al., 2020).
In URL, the agent ﬁrst explores an environment without extrinsic reward with the objective of learning a useful representation which then enables rapid ﬁne-tuning to a test task.
Continuous MountainCar
We ﬁrst demonstrate that if the test task is non-Markovian, the SR can produce misleading value estimates.
To do so, we use a modiﬁed version of the continuous
MountainCar task (Brockman et al., 2016) (Fig. 2(a)).
The agent pre-trains for 20,000 time steps in a rewardless environment, during which it learns FFs or SFs for a set of policies Π which swing back and forth with a ﬁxed acceleration or “power.”
(details in Appendix A.3).
During ﬁne-tuning, the agent must identify the policy π ∈ Π which reaches a randomly sampled goal location as quickly as possible.
We use radial basis functions as the base features φ d ( · ) with ﬁxed FF thresholds θ d = 0 .
Fig. 2(b) plots the FF and SF values versus policy power from the start state for two different goal locations.
The low-power policies require more time to gain momentum up the hill, but the policies which maximize the SF values slow down around the goal locations, dramatically increasing their total “occupancy” of that area.
In contrast, high-powered policies reach the goal locations for the ﬁrst time much sooner, and so the policies with the highest FF values have higher powers.
In the test phase, the agent ﬁts the reward vector w ∈ R D by minimizing (cid:80) t
|| r t − w T φ ( s t ) ||
2 , as is standard in the SF literature (Barreto et al., 2017b;a; 2018; 2020;
Zahavy et al., 2021).
The agent then follows the set-max policy (SMP; (Zahavy et al., 2021)), which selects the policy in Π which has the highest expected value across starting states: π SMP ∈ argmax
π ∈ Π E s 0 ∼ µ
[ V π ( s 0 )] , where
V π ( s 0 ) =
w T ϕ π ( s 0 ) (with ϕ π replaced by ψ π for SF-based value estimates).
Fig. 2(c) shows both the estimated ( V est ) and true ( V true ) 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 signiﬁcant overestimation in value, as well as a suboptimal policy.
The FF estimates are nearly matched to the true values of the selected policies for each goal location and achieve nearly optimal performance.
[Figure] Robotic reaching To test whether these results translate to high-dimensional problems, we applied the FF to the 6-DoF J ACO robotic arm environment from Laskin et al.
In this domain, the arm must quickly reach to different locations and perform simple object manipulations (Fig. 4(a)).
We modify the Active Pretraining with Successor features (APS; (Liu &
Abbeel, 2021)) URL algorithm, which leverages a nonparametric entropy maximization objective in conjunction with SFs during pre-training (Hansen et al., 2020) in order to learn useful and adaptable behaviors.
Our modiﬁcation is to replace the SFs with FFs, resulting in Active Pre-training with First-occupancy features (APF), using the same intuition motivating the MountainCar experiments: cumulative features are misleading when downstream tasks beneﬁt from quickly reaching a desired goal, in this case, the object.
Here, the agent is ﬁrst trained for 1e6 time steps using the aforementioned intrinsic reward objective before being applied to a speciﬁed reaching task.
As additional baselines, we compare against two variants of dynamic distance learning - unsupervised (DDLUS and DDLUS-G; (Hartikainen et al., 2020)).
Details can be found in Appendix A.3. We found that the FF accelerated convergence (Fig. 4(b)).
4. P LANNING 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 onestep lookahead, which may result in suboptimal behavior.
One simple situation that highlights such a scenario is depicted in Fig.
As before, there are start and goal states in a simple room (Fig. 3(a)), but here there are three policies comprising Π
= { π 1 , π 2 , π 3 } (Fig. 3(b-d)).
GPI selects π 1 because it is the only policy that reaches the goal s g within one step of the start s 0 : π 1 = max π ∈ Π r T M π ( s 0 , s g )
(Fig. 3(e)).
(Note that using GPI with the FR instead would also lead to this choice.)
To gain further intuition for the FR versus the SR, we plot the representations for the policies in Appendix Fig. 17.
However, the optimal strategy using the policies in Π is instead to move right using π 2 and up using π 3 .
How can the FR be used to ﬁnd such a sequence?
Intuitively, starting in a state s , this strategy is grounded in following one policy until a certain state s (cid:48) , which we refer to as a subgoal , and then switching to a different policy in the base set.
Because the FR effectively encodes the shortest path between each pair of states ( s, s (cid:48) ) for a given policy, the agent can elect to follow the policy π ∈ Π with the greatest value of F π ( s, s (cid:48) ) , then switch to another policy and repeat the process until reaching a desired state.
The resulting approach is related to the hierarchical options framework (Sutton et al., 1999; Sutton & Barto, 2018), where the planning process effectively converts the base policies into options with—as we show—optimal termination conditions.
For a more detailed discussion, see Appendix A.7.
More formally, we can construct a DP algorithm to solve for the optimal sequence of planning policies π F and subgoals s F .
Denoting by Γ k ( s ) the total discount of the full trajectory from s to s g at step k of the procedure, we jointly optimize over policies π and subgoals s (cid:48) for each state s :
[Equation]
with
[Equation]
Intuitively, the product F π ( s, s (cid:48) )
Γ k ( s (cid:48) ) can be interpreted as the expected discount of the plan consisting of following π from s to s (cid:48) , then the current best (shortest-path) plan from s (cid:48) to s g . Note [Figure] r e t u r n
K = 0 iterations ss g s g s g s g s g
K = 3 iterations ss g s g s g s g s g (a) (b) FRP VI GPI
VI-converged 0.
1.0 transition noise 04080r e t u r n
VI GPI
FRP (c) 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-deﬁned.
The full procedure, which we refer to as FR-planning (FRP), is given by Alg. 1 in Appendix A.2.
Appendix Fig.
13 depicts the resulting policies π F and subgoals s F obtained from running FRP on the example in Fig.
The following result shows that under certain assumptions, FRP ﬁnds the shortest path to a given goal (proof in Appendix A.4).
Proposition 4.1 (Planning optimality) .
Consider a deterministic, ﬁnite MDP with a single goal state s g , and a base policy set Π composed of policies π : S → A .
We make the following coverage assumption: there exists a sequence of policies that reaches s g from a given start state s 0 .
Then Alg.
1 converges so that Γ( s 0 ) =
γ L (cid:63)π , where L (cid:63)π is the shortest path length from s 0 to s g using π ∈ Π .
Performance and computational cost We can see that each iteration of the planning algorithm adds at most one new subgoal to the planned trajectory from each state to the goal, with convergence when no policy switches can be made that reduce the number of steps required.
If there are K iterations, the overall computational complexity of FRP is O
( K | Π ||S| 2 ) .
The worst-case complexity occurs when the policy must switch at every state en route to the target— K is upper-bounded by the the number of states along the shortest path to the goal.
In contrast, running value iteration (VI; (Sutton & Barto, 2018)) for N iterations is O ( N |A||S| 2 ) .
[-13: π Fk +
-]
[-( s ) , s Fk +
-]
Given the true transition matrix P and reward vector r , VI will also converge to the shortest path to a speciﬁed goal state, but FRP converges more quickly than VI whenever K | Π | < N |A| .
To achieve an error (cid:15) between the estimated value function and the value function of the optimal policy, VI requires N ≥
[Equation]
(Puterman, 1994), which for γ = 0 .
95 , (cid:15) = 0 .
1 , e.g., gives N ≥ 180 iterations 2 .
To test convergence rates in practice, we applied FRP, VI, and GPI using the FR to the classic F OUR R OOMS environment (Sutton et al., 1999) on a modiﬁed task in which agents start in the bottom left corner and move to a randomly located goal state.
Once the goal is reached, a new goal is randomly sampled in a different location until the episode is terminated after 75 time steps.
For GPI and FRP, we use four base policies which each only take one action: { up , down , left , right } , with their FRs learned by TD learning.
We ran each algorithm for 100 episodes, with the results plotted in Fig. 5(a).
Note that here GPI is equivalent to FRP with K = 0 iterations.
To see this, observe that when there is a single goal s g such that only r ( s g ) > 0 , the policy selected by GPI is
[Equation]
When there are n goal states with equal reward, ﬁnding the ordering of the goals that results in the shortest expected path is in general O ( n !)
(see Appendix A.6).
Due to the nature of the base policies used above, the number of subgoals on any path is equal to the number of turns the agent must take from its curent state to the goal, which for this environment is three.
We can then see that FRP reaches the optimal performance obtained by the converged VI after K = 3 iterations (Fig. 5(a)).
In contrast, for the same number of iterations, VI performs far worse.
This planning process must be repeated each time a new goal is sampled, so that the computational beneﬁts of FRP versus traditional DP methods compound for each new reward vector.
Example FRP trajectories between goals for K =and K = 3 iterations are plotted in Fig. 5(b).
Finally, to test FRP’s robustness to stochasticity, we added transition noise (cid:15) to the F OUR R OOMS task.
That is, the agent moves to a random adjacent state with probability (cid:15) regardless of action.
We compared FRP to converged VI for increasing (cid:15) , with the results plotted in Fig. 5(c), where we can see that FRP matches the performance of VI across noise levels.
It’s important to note that this ability to adaptively interpolate between MF and MB behavior, based on the value of K , is a unique capability of the FR compared to the SR 3 .
The same FRs can be combined using DP to plan for one task or for GPI on the next.
4. E SCAPE BEHAVIOR
ggg no barrier ggg barrier trialggg barrier trial
S 10 cm S peed ( c m / s ) Trial Trial
S 1 - edge d i r e c t i on 0.
S
[Figure] 0 - s he l t e r d i r e c t i on [Figure] In prey species such as mice, escaping from threats using efﬁcient paths to shelter is critical for survival (Lima & Dill, 1990).
Recent work studying the strategies employed by mice when ﬂeeing threatening stimuli in an arena containing a barrier has indicated that, rather than use an explicit cognitive map, mice instead appear to memorize a sequence of subgoals to plan efﬁcient routes to shelter (Shamash et al., 2021).
When ﬁrst threatened, most animals ran along a direct path and into the barrier.
Over subsequent identical trials spanning 20 minutes of exploration, threat-stimulus presentation, and escape, mice learned to navigate directly to the edge of the wall before switching direction towards the shelter (Fig. 6).
Follow-up control experiments suggest that mice acquire persistent spatial memories of subgoal locations for efﬁcient escapes.
We model this task and demonstrate that FRP induces behavior consistent with these results.
We model the initial escape trial by an agent with a partially learned FR, leading to a suboptimal escape plan leading directly to the barrier.
Upon hitting the barrier, sensory input prompts rapid re-planning to navigate around the obstacle.
The FR is then updated and the escape plan is recomputed, simulating subsequent periods of exploration during which the mouse presumably memorizes subgoals.
We ﬁnd a similar pattern of behavior to that of mice (Fig. 6(b)).
See Appendix A.3 for experimental details.
We do not claim that this is the exact process by which mice are able to efﬁciently learn escape behavior.
Rather, we demonstrate that the FR facilitates behavior that is consistent with our understanding of animal learning in tasks which demand efﬁcient planning.
Given the recent evidence in support of SR-like representations in the brain (Stachenfeld et al., 2017; Momennejad et al., 2017), we are optimistic about the possibility of neural encodings of FR-like representations as well.
We also re-emphasize that this type of rapid shortest-path planning is not possible with the SR.
C ONCLUSION
In this work, we have introduced the FR, an alternative to the SR which encodes the expected path length between states for a given policy.
We explored its basic formal properties, its use as an exploration bonus, and its usefulness for unsupervised representation learning in environments with an ethologically important type of non-Markovian reward structure.
We then demonstrated that, unlike the SR, the FR supports a form of efﬁcient planning which induces similar behaviors to those observed in mice escaping from perceived threats.
As with any new approach, there are limitations.
However, we believe that these limitations represent opportunities for future work.
From a theoretical perspective, it will be important to more precisely understand FRP in stochastic environments.
For the FF, we have limited understanding of the effect of feature choice on performance, especially in high dimensions.
FRP is naturally restricted to discrete state spaces, and it could be interesting to explore approximations or its use in partially-observable MDPs with real-valued observations and discrete latents (e.g., (Vértes & Sahani, 2019; Du et al., 2019)).
Further exploration of FRP’s connections to hierarchical methods like options would be valuable.
Finally, it would be informative to test hypotheses of FR-like representations in the brain.
We hope this research direction will inspire advancements on representations that can support efﬁcient behavior in realistic settings
Reproducibility statement Experiments: All experiments are described in detail in Sectionin the main text and in Appendix A.3.
We have attached code for the tabular experiments, and in
Appendix A.3 we provide a link to the base implementation for APS as well as an explanation of where to ﬁnd the hyperparameters and modiﬁcations made to produce APF.
Theoretical results: Theoretical results, including assumptions and proofs, are provided in Appendix A.1 and Appendix A.
A A PPENDIX A. FR RECURSION
For clarity
[Equation]
we provide the derivation of the recursive form of the FR below:
[Equation]
s t (cid:105) (19) A. FRP A LGORITHM We present the full algorithm for FR planning (FRP) below.
Planning (FRP)
A. A DDITIONAL E XPERIMENTAL D ETAILS
All experiments except for the robotic reaching experiment were performed on a single 8-core CPU.
The robotic reaching experiment was performed using four Nvidia Quadro RTX 5000 GPUs.
Tabular exploration We reuse all hyperparameter settings from Machado et al.
(2020) in both the R IVER S WIM and S IX
A RMS environments, with the only difference being a lower value for β , the exploration bonus coefﬁcient, for the FR, as the bonuses given by the FR are generally larger.
The hyperparameters are { α, η, γ SR / FR , β, (cid:15), η } , which are the Sarsa learning rate, the SR/FR learning rate, the SR/FR discount factor, the exploration bonus coefﬁcient, and the probability of taking a random action in these were respectively, they were { 0 . 1 , 0 . 01 , 0 . 99 , 100 / 50 , 0 . 01 } .
For the R IVER S WIM -N task, we chose N = { 6 , 12 , 24 } as default R IVER S WIM has N = 6 states, and we chose to successively double the problem size.
As N increased, the number of unrewarded central states was multiplied (with the same transition structure), while the endpoints remained the same.
It’s also worth noting that β could be manually adjusted upwards to compensate for the SR bonus’ invariance to problem size, though this would require a longer hyperparameter search generally, which we believe is less preferable to a bonus which naturally scales. 1 SE across 100 trials. R IVER S WIM -N S ARSA + FR
S ARSA + SR
S ARSA
N [Table] [Table]
Exploration with function approximation In order to test the usefulness of the FR/FF in a function approximation setting, we use a similar approach to Machado et al.
That is, we train a modiﬁed DQN agent (Mnih et al., 2015) using an architecture inspired by Machado et al.
and Oh et al.
(see Fig. 9), such that the base feature representation φ ( s ) is an intermediate layer of the network.
Like the standard DQN, the architecture outputs an |A| -length vector of predicted Q -values for the current state, trained off-policy using minibatches of transition tuples { ( s it , a it , r it , s it +1 ) } Bi =1 (where B is the minibatch size) to minimize the squared Bellman error
[Equation]
20) via gradient descent (the subscript − on the target Q -values indicates that gradients do not ﬂow through it).
Unlike the standard DQN, the network has two additional output heads.
The ﬁrst is a reconstruction head which, given the base feature representation of a state s t and an embedding of the subsequent action a t produces a prediction of the following state
ˆ s t +1 .
It’s trained to minimize the reconstruction loss
[Equation]
The ﬁnal output head of the network is an FF/SF prediction, trained using the squared FF error in the former case:
[Equation]
22) and the squared SF error in the latter
[Equation]
For the FF, the features φ t are passed through a sigmoid function to compress them in the range [0 , 1] and then thresholded at 0 .
The total loss is then given as a weighted combination
[Equation]
(24) where X ∈ { ϕ, ψ } and w Q , w s , w X ∈ R are ﬁxed weights.
As in Machado et al.
(2020), gradients from L Q and L s , but not L X , are permitted to ﬂow through to φ .
Thus, the base features are trained to be both reward-predictive and to carry information about the environment transition dynamics.
The norm of the FF/SF vector is then used to compute an intrinsic exploration bonus to the task reward in the same manner as in the tabular setting.
To test this model, we chose the D EEP S EA task from the Behavior Suite ( bsuite ; (Osband et al., 2020)) set of benchmark tasks.
D EEP S EA is a challenging exploration task set up in an N × N grid (Fig.
10, top left).
At the beginning of each episode, the agent starts at the top left of the grid.
Each time step, the agent descends one level, and can choose to move either right or left.
The episode ends after N steps, when the agent reaches the bottom level.
There is a small negative reward of − 0 .if
the agent moves right, but if the agent moves right N times in a row, there is a large reward +
[Figure]
v
[Figure]
r
[Figure]
e
[Figure]
i s
[Figure]
i c r
[Figure]
t
[Figure]
n
[Figure]
located at the bottom right of environment—this is the only policy which nets the agent a positive reward.
In the bsuite framework, the agent is separately trained on increasing problem sizes
N = 5 , 6 , 7 , . . .
, 50 for 10,000 episodes each, with the ﬁnal score the proportion of N for which the agent reached an average regret of less than 0 .
9 faster than 2 N episodes.
We tested the standard DQN, DQN + SF (Machado et al., 2020), and DQN + FF agents on this task, with training hyperparameters described in Table 3.
For all models, the network consisted entirely of fully-connected layers, with φ ( s t ) being a 2-layer MLP with 64 units per layer, Q ( s t , · ) being a linear function of φ ( s t ) with |A| = 2 units, ˆ s t +1 consisting of a 64-unit layer followed by an N 2 -unit output layer (the action embedding is 64-dimensional as well), and the SFs/FFs also a 64-dimensional linear layer over φ ( s t ) .
The agent is trained using (cid:15) -greedy action selection.
Our
DQN implementation was coded in JAX (Bradbury et al., 2018) and based off that of Osband et al.
To select the values of β , w s , and w X ( w Q was always kept at 1) we performed a sweep over β ∈ { 0 . 01 , 0 . 05 , 0 . 1 } and w X , w Q ∈ { 0 . 001 , 0 . 1 , 1 , 10 , 100 , 1000 } , choosing the best-performing values for each method.
[Table]
Results are presented in Fig.
We can see that DQN+FF signiﬁcantly outperforms the other methods (top right), with the intuition from the tabular experiments—particularly R IVER S WIMN—carrying over into the function approximation setting.
That is, as N increases, the norm of the SF approaches its asymptotic value regardless of the degree of exploration.
In contrast, for the
FF, the maximum bonus scales with the problem size.
This enables the bonus to remain effective in environments with larger state spaces.
We hope to investigate this approach and its theoretical properties further in future work.
MountainCar experiment In our version of the task, the feature representations are learned in a rewardless environment, and at test time the reward may be located at any location along the righthand hill.
We evaluate the performance of a set of policies Π
= { π i } with a constant magnitude of acceleration and which accelerate in the opposite direction from their current displacement when at rest and in the direction of their current velocity otherwise (see Python code below for details).
That is, each π i will swing back and forth along the track with a ﬁxed power coefﬁcient a i .
For each possible reward location along the righthand hill, then, the best policy from among this set is the one whose degree of acceleration is such that it arrives at the reward location in the fewest time steps.
There is a natural tradeoff–too little acceleration and the cart will not reach the required height.
Too much, and time will be wasted traveling too far up the lefthand hill and then reversing momentum back to the right.
We hypothesized that the FF would be a natural representation for this task, as it would not count the repeated state visits each policy experiences as it swings back and forth to gain momentum.
We deﬁned a set of policies with acceleration magnitudes | a i | = 0 .
1 i for i = 1 , . . .
, 9 , and learned both their SFs and FFs via TD learning on an "empty" environment without rewards over the course of 100 episodes, each consisting of 200 time steps, with the SFs using just the simple RBF feature functions without thresholds.
Python code for the policy class is shown below.
[-else: # otherwise, continue in the current direction of motion action = sign(vel) * self.a 
-]
return action
We repeated this process for 20 runs, with the plots in Fig.
7 showing the means and standard deviations of the TD learning curves across runs.
For the FFs, the thresholds were constant across features at θ d = θ = 0 .
Because of the nature of the environment, all of the policies spent a signiﬁcant portion of time coasting back and forth between the hills, causing their SFs to accumulate in magnitude each time states were revisited.
Given the learned representations, we then tested them by using them as features for policy evaluation in different tasks, with each task containing a different rewarded/absorbing state.
Note that a crucial factor is that the representations were learned in the environment without absorbing states.
This is natural, as in the real world reward may arise in the environment anywhere, and we’d like a representation that can be effective for any potential goal location.
[Table]
li
[Table]
d
[Table]
[Figure]
Fig.
11 shows the value of the FF and SF at the start state for different policies with ﬁxed power and for different feature dimensions (number of basis functions) in continuous MountainCar.
The results show that the policies for which the FF is highest is closer in power to the optimal policy than for the policies at which the SF is greatest across all but the coarsest feature dimensionalities.
This provides an indication of the robustness of the FF to the choice of feature dimensionality.
Robotic reaching experiment We used the custom J ACO domain as well as the APS base code from Laskin et al.
(2021), located at this link: https://anonymous.4open.science/r/ urlb/README.md .
Both the critic and actor networks were parameterized by 3-layer MLPs with
ReLU nonlinearities and 1,024 hidden units.
Observations were 55-dimensional with 10-dimensional features φ ( · ) .
For all other implementation details, including learning rates, optimizers, etc. see the above link.
All hyperparameters and network settings are kept constant from those provided in the linked .yaml ﬁles.
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 J ACO domain by maximizing the intrinstic reward
[Equation]
where w ∈ R D , D = 10 is a random reward vector sampled from a standard Gaussian distribution and the righthand term is a particle-based estimate of the state-based feature entropy, with N k ( · ) denoting the k nearest-neighbors (see Liu & Abbeel (2021) for details).
In standard APS, this reward is used to train the successor features by constructing the bootstrapped target
[Equation] [Equation]
where a (cid:48)
= argmax a w T ψ ( s (cid:48) , a, w ) .
For the FF, we make the following modiﬁcation:
[Equation]
with V ( s t +1 )
= max a (cid:48) w T ϕ
( s t +1 , a (cid:48) , w ) , ˜ φ ( · ) the thresholded base features, such that ˜ φ d ( s t ) =( φ ( s t ) ≥ θ d , { φ d ( s t (cid:48) )
} t (cid:48)
=
: t < θ d ) , 1 is the D -length vector of ones, and φ ( s t ) ∈ [0 , 1] .
Interestingly, if the features are kept in the range [0 , 1] , φ ( · ) can be thought of encoding a form of “soft” ﬁrst feature occupancy, rather than the hard threshold given by the indicator function.
The agent is then trained with the off-policy deep deterministic policy gradient (DDPG; (Lillicrap et al., 2019)) algorithm, where given a stored replay buffer of transitions D = { ( s t , a t , r t , s t +1 ) } , the (SF/FF) critic Q ω (with Q formed from either the SF or FF and ω being the parameters) is trained to minimize the squared Bellman loss
[Equation]
where in the pre-training phase y r ∈ { y APS , y APF } (the target parameters are an exponential moving average of the weights—gradients do not ﬂow through them).
The deterministic actor π θ is trained using the derministic policy gradient loss:
[Equation]
Fine-tuning:
After pre-training, the agent is ﬁne-tuned on the target task, R EACH T OP L EFT , where the learning proceeds exactly as in the pre-training phase, but instead of intrinsic reward, the agent is given the task reward—that is, y r = r task t + γV ( s t +1 ) .
We performed this task-speciﬁc training for an additional 1M steps.
As additional baselines, we implemented two versions of the dynamic distance learning method of Hartikainen et al. (2020), the original DDLUS and an additional variant DDLUS-G. In standard
DDLUS, the goals g ∈ S for the pre-training phase are generated according to
[Equation]
where D is a stored set of trajectories.
In Hartikainen et al.
(2020), this is useful as a mechanism for encouraging the agent to learn skills which move the agent as far as possible from the start state.
However, since the object for manipulation in the ﬁne-tuning phase of the J ACO task is not typically especially far from the initial point, we also tested DDLUS-G, which samples goals in the same manner as APS and APF, via g ∼ N (0 , I ) , to ensure a more fair comparison.
All other hyperparameters match those of Hartikainen et al.
(2020), with the exception that the base agent is DDPG, implemented using the same framework as APS and APF, rather than SAC (Haarnoja et al., 2018).
For a more detailed discussion of DDL and its relationship to the FR/FF, see Appendix A.8.
In future work, it would be interesting to explore the interaction of the FF with other off-policy algorithms (Haarnoja et al., 2018; Fujimoto et al., 2018; Moskovitz et al., 2021) and whether on-policy learning (e.g., with (Schulman et al., 2017; Kakade, 2002; Williams, 1992; Moskovitz et al., 2020;
Hartikainen et al., 2020)) has different effects.
experiments The FourRoom environment we used was deﬁned on an 11 × 11 gridworld in which the agent started in the bottom left corner and moved to a known goal state.
The action space was A = { up , right , down , left } with four base policies each corresponding to one of the basic actions.
TD Learning curves for the base policies are depicted in Fig.
Once reaching the goal, a new goal was uniformly randomly sampled from the non-walled states.
At each time step, the agent received as state input only the index of the next square it would occupy.
Each achieved goal netted a reward of +50 , hitting a wall incurred a penalty of − 1 and kept the agent in the same [Figure] place, and every other action resulted in 0 reward.
There were 75 time steps per episode—the agent had to reach as many goals as possible within that limit.
The discount factor γ was 0 .
95 , and the
FR learning rate was 0 .
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 (cid:15) = 0 . 0 , 0 . 1 , 0 .
, 1 . 0 , both VI and FRP were run to convergence ( ≈ 180 iterations for VI, 3 iterations for FRP) and then tested for 100 episodes. [Figure] [Figure]
The escape experiments were modeled as a 25 × 25 gridworld with eight available actions,
[Equation] [Equation]
F R : || F ||exploration phase   escape trial     exploration phase    5 episodes
S R : ||
M ||  40 time steps    5 episodes and a barrier one-half the width of the grid optionally present in the center of the room.
The discount factor γ was 0 . 99 , and the FR learning rate was 0 .
In test settings, the agent started in the top central state of the grid, with a single goal state directly opposite.
At each time step, the agent receives only a number corresponding to the index of its current state as input.
The base policy set Π consisted of eight policies, one for each action in A .
Escape trials had a maximum of 50 steps, with termination if the agent reached the goal state.
In (Shamash et al., 2021), mice were allowed to explore the room starting from the shelter location.
Accordingly, during “exploration phases,” the agent started in the goal state and randomly selected a base policy at each time step, after which it updated the corresponding FR.
Each exploration phase consisted of 5 episodes—a sample trajectory is shown in Fig.
14(a).
After the ﬁrst exploration phase, the agent started from the canonical start state and ran FRP to reach the goal state.
Because most of its experience was in the lower half of the grid, the FRs for the upper half were incomplete (Fig. 14(b)), and we hypothesized that in this case, the mouse should either i) default to a policy which would take it to the shelter in the area of the room which it knew well (the down policy) or ii) default to a policy which would simply take it away from the threat (again the down policy).
During the ﬁrst escape trial, the agent selects the down policy repeatedly, continuing to update its FRs during the testing phase.
Upon reaching the wall and getting stuck, the FR for the down policy is eventually updated enough that re-planning with FRP produces a path around the barrier.
After updating its FR during the ﬁrst escape trial and during another exploration period, the FRs for the upper half of the grid are more accurate (Fig. 14(c)) and running FRP again from the start state produces a faster path around the barrier on the second escape trial.
TD learning curves for the experiment (repeated with the SR for completeness) are plotted in
Fig.
For completeness, we repeated this experiment with the SR.
In this case, the planning algorithm is ill-deﬁned for K > 0 , so we default to GPI ( K = 0 ).
As expected, without the barrier, the down policy is selected and the goal is reached (Fig. 16,left).
However, when there is a barrier, while the
SR updates when the agent hits it (Fig. 12), since there is no single policy that can reach the shelter,
GPI fails to ﬁnd a path around the barrier (Fig.
16,middle,right).
[Figure]
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 deﬁned in Deﬁnition 3.2 for some stationary policy π .
Then for any two matrices F, F (cid:48) ∈ R |S|×|S|
[Equation]
with the difference equal to zero for s = s (cid:48) .
Proof.
For s (cid:54) =
s (cid:48)
we have
[Equation]
where we use the notation X s,s (cid:48) to mean X
( s, s (cid:48) ) , and the inequality is due to the fact that every element of P π ( F − F (cid:48) ) is a convex average of F − F (cid:48) .
For s = s (cid:48) , we trivially have | ( G π F −
G π F (cid:48) ) s,s | = | 1 − 1 | = 0 .
Proposition A.1 (Convergence) .
Under the conditions assumed above, set F (0)
=
I |S| .
For k = 0 , 1 , . . .
, suppose F ( k +1) =
G π F ( k ) .
Then
[Equation]
s [Equation] s [Equation]
| < [Equation] 34) for s (cid:54) = s (cid:48) with the difference for s = s (cid:48) equal to zero ∀ k .
Proof.
We have, for s (cid:54) =
s (cid:48) and using the notation X s,s (cid:48)
=
X ( s, s (cid:48) ) for a matrix X ,
[Equation]
F [Equation] For s = s (cid:48) , | ( F ( k ) − F π ) s,s | = | 1 − 1 | = 0 ∀ k .
Below is the proof of Proposition 4.1, which is restated below.
Proposition 4.1 (Planning optimality) .
Consider a deterministic, ﬁnite MDP with a single goal state s g , and a policy set Π composed of policies π : S → A .
We make the following coverage assumption, there exists some sequence of policies that reaches s g from a given start state s 0 .
Under these conditions, Alg. 1 converges such that Γ( s 0 ) =
γ L (cid:63)π , where L (cid:63)π is the shortest path length from s 0 to s g using π ∈ Π .
Proof.
Since the MDP is deterministic, we use a deterministic transition function ρ : S ×
A → S .
We proceed by induction on L (cid:63) π .
Base case: L (cid:63)π =
If L (cid:63)π = 1 , s 0 must be one step from s g .
The coverage assumption guarantees that ∃ π ∈ Π such that ρ ( s 0 , π ( s 0 ))
=
s g .
Note also that when both the MDP and policies in Π are deterministic,
F π ( s, s (cid:48) ) =
γ L π , where L π is the number of steps from s to s (cid:48) under π , and we use the abuse of notation L π = ∞ if π does not reach s (cid:48) from s .
Then following Alg. 1,
[Equation]
guaranteed by coverage of Π
[Equation]
by deﬁnition of F π ) .
Moreover,
[Equation]
Then Γ 2 ( s ) =
Γ 1 ( s ) ∀ s and Alg.
1 terminates.
Thus, Γ( s 0 ) =
γ = γ L (cid:63)π and the base case holds.
Induction step: Assume Proposition 4.1 holds for L (cid:63)π =
L
Given the induction assumption, we now need to show that Proposition 4.1 holds for L (cid:63)π =
L + 1 .
By the induction and coverage assumptions, there must exist at least one state within one step of s g that the agent can reach in L steps, such that the discount for this state or states is γ L .
Moreover, the coverage assumption guarantees that ∃ π ∈ Π such that for at least one such state s L , ρ ( s L , π ( s L )) =
s g .
Then this problem reduces to the base case—that is, Alg. 1 will select the policy π ∈ Π that transitions directly from s L to s g —and the proof is complete.
A. E XPLICIT P LANNING
Below we describe a procedure for constructing an explicit plan using π F and s F .
A. FRP WITH M ULITPLE G OALS
Here we consider the application of FRP to environments with multiple goals { g 1 , . .
. , g n } .
To ﬁnd the shortest path between them given the base policy set Π , we ﬁrst run FRP for each possible goal state in { g i } , yielding an expected discount matrix Γ Π ∈ [0 , 1] |S|× n , such that Γ Π ( s, g i ) is the expected discount of the shortest path from state s to goal g i .
We denote by g σ =
[ s 0 , σ ( g 1 ) , σ ( g 2 ) , . .
, σ ( g n )] a speciﬁc ordering of the goals in { g i } starting in s 0 .
The expected discount of a sequence of goals is then
[Equation]
) with the optimal goal ordering g σ ∗ given by
[Equation]
37) where G is the set of all possible permutations of { g i } , of size n ! .
This is related to a form of the travelling salesman problem.
Fortunately, in most settings we don’t expect the number of goals n to be particularly large.
A. C ONNECTIONS TO OPTIONS
The options framework (Sutton et al., 1999) is a method for temporal abstraction in RL, wherein an option ω is deﬁned as a tuple ( π ω , τ ω ) , where π ω is a policy and τ ω ∈ ∆( S ) is a state-dependent termination distribution (or function, if deterministic).
Executing an option at time t entails sampling an action a t ∼ π ω ( ·| s t ) and ceasing execution of π ω at time t + 1 with probability τ ω ( s t +1 ) .
The use of options enlarges an MDP’s action space, whereby a higher-level policy selects among basic, low-level actions and options.
Options are connected to FRP (Alg. 1) in that by outputting a set of policies and associated subgoals { ( π F , s F ) } , FRP effectively converts each base policy to an option with a deterministic termination function, i.e., the agent will follow π F until terminating at s F .
One of the key difﬁculties in the options literature is how to learn the best options to add to the available action set.
For the class of problems considered in this paper, FRP then provides a framework for generating optimal (in the sense of ﬁnding the fastest path to a goal) options from a set of standard policies, subject to the fulﬁllment of the coverage assumption.
Importantly, the associated FRs (which can be learned via simple TD updating) can be reused across tasks, so that FRP can re-derive optimal options for a new goal.
FRP can also be seen as related to the work of Silver & Ciosek (2012), which demonstrates that value iteration performed on top of a set of task-speciﬁc options converges more quickly than value iteration performed on the default state space of the MDP.
One critical difference to note is that the
FR/FRP is transferable to any MDP with shared transition dynamics.
While value iteration on a set of options for a given MDP is more efﬁcient 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. F URTHER CONNECTIONS TO RELATED WORK
We now describe the connection between the FR/FF and several related approaches in the literature.
The Dynamic Distance Function The dynamic distance function (DDF; (Hartikainen et al., 2020)) is deﬁned as
[Equation]
where c : S × S → R is a local cost function.
In practice, c ( s t + k , s t + k +1 )
= γ = 1 , giving
[Equation]
n practice, d π ( s, s (cid:48) ) is parameterized via a neural network with parameters ψ trained on-policy from full trajectories τ using the loss
[Equation]
where D is a buffer of stored trajectories.
On a downstream navigation task, the agent policy is trained to minimize
[Equation]
[Equation]
where g ∈ S is a goal state and φ are the policy parameters.
There are several important differences between the DDF and the FR, both in theory and practical application.
This form of the DDF, as discussed by Hartikainen et al.
(2020), is also naturally restricted to onpolicy learning from full trajectories, and cannot be updated off-policy and/or via one-step temporal difference learning.
A natural additional consequence is that it is only applicable to ﬁnite-horizon
MDPs.
This conﬁnement to ﬁnite horizons is signiﬁcant because it leads to the conditioning problem described by Hartikainen et al.
This problem occurs because the DDF is conditioned on the policy successfully reaching the goal state—when this doesn’t happen, it can lead to signiﬁcant value estimation errors.
A second difference is that the policy is trained, via Eq.
(41), to minimize the cumulative discounted distance to the goal rather than via greedy distance minimization.
Another difference is that Eqs.
(38) and (39) don’t explicitly require time step j to be the ﬁrst time the agent enters s (cid:48) , although this may have been the authors’ intention.
In fact, this deﬁnition is more closely related to the SR than the FR (with equality in the inﬁnite horizon setting), but the implementation of the DDF in practice is more closely related to the FR.
It’s also important to note that when γ = 1 , the number of steps k appears linearly within the expectation rather than exponentially (this is important, as in general γ E [ k ] (cid:54) =
E [ γ k ] .
Another signiﬁcant 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.
In practice, it is used to support policy optimization, rather than evaluation.
In contrast, the FF representation, like the SF representation, maps a single state to a ﬁxed length vectorvalued encoding, ϕ π : S → R d .
This difference has signiﬁcant implications for the applications of the DDF.
In particular, it is unclear how the DDF could be used as an exploration bonus in the manner of the FF/SF, and the scalar value would make it challenging to implement the type of parallel updates useful for efﬁcient planning.
Dynamic Goal Learning
Another related approach is that of dynamic goal (DG) learning (Kaelbling, 1993), a method for optimal control related to Q -learning.
The optimal DG function
G (cid:63) : S ×
A × S → R is deﬁned recursively for a goal state g ∈ S as
[Equation]
where G (cid:63) ( g, a, g ) := 0 .
There are several important differences between this approach and the FR.
First, as mentioned above, DG learning is a method for optimal control, rather than policy evaluation .
That is, D (cid:63) ( s, a, g ) converges to the expected number of steps for the optimal policy for reaching g starting from s .
It cannot be reused for policy evaluation for a policy π (cid:54) =
π (cid:63) , and also implicitly assumes that g is reachable in ﬁnite time.
DG learning is thus susceptible to the same conditioning problem discussion above for the DDF, and is only studied by Kaelbling (1993) in a gridworld environment.
A second difference is that the DG function scales linearly with k , the number of steps for the optimal policy to reach g , while the FR scales exponentially at a rate of γ .
This is important, as we can note (with a slight abuse of notation) in general that γ E π [ k ] (cid:54) =
E π (cid:2) γ k (cid:3) —it is nontrivial to recover the FR for the optimal policy from the DG function (and vice-versa).
Temporal Difference Models Temporal difference models (TDMs; (Pong et al., 2018)) are another related approach. TDMs are an optimal control method motivated by the observation that goalconditioned Q -functions can be used to construct an implicit dynamics model when the discount factor γ = 0 . For γ > 0 , the authors introduce a horizon variable τ to interpolate between modelbased and model-free learning. Like the FR, this formulation allows agent to consider multiple goals in parallel.
As in DG learning, and unlike the FR, TDMs represent a method for optimal control, not evaluation.
The emphasis is on efﬁciently improving a policy, rather than learning a state representation which can be leveraged for multiple uses (e.g., policy evaluation or as an exploration bonus).
Another difference from the FR is that, like the DDF and DG learning, TDMs rely on a ﬁnite-horizon setting.
A. FR VS . SR V ISUALIZATION
ss g
Mss g
Mss g
Mss g
Fss g
Fss g
F0.