A State Representation for Diminishing Rewards
Authors: Ted Moskovitz
Samo Hromadka
Ahmed Touati
Diana L Borsa
Maneesh Sahani
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
successor features
successor representation
neuroscience
TL;DR: We study the problem of diminishing marginal utility in reinforcement learning, deriving a novel state representation that is required for solving such problems and applying it in both tabular and continuous domains.

Abstract: A common setting in multitask reinforcement learning (RL) demands that an agent rapidly adapt to various stationary reward functions randomly sampled from a fixed distribution. In such situations, the successor representation (SR) is a popular framework which supports rapid policy evaluation by decoupling a policy's expected discounted, cumulative state occupancies from a specific reward function. However, in the natural world, sequential tasks are rarely independent, and instead reflect shifting priorities based on the availability and subjective perception of rewarding stimuli. Reflecting this disjunction, in this paper we study the phenomenon of diminishing marginal utility and introduce a novel state representation, the $\lambda$ representation ($\lambda$R) which, surprisingly, is required for policy evaluation in this setting and which generalizes the SR as well as several other state representations from the literature. We establish the $\lambda$R's formal properties and examine its normative advantages in the context of machine learning, as well as its usefulness for studying natural behaviors, particularly foraging.

A State Representation for Diminishing Rewards
Ted Moskovitz
Samo Hromadka Gatsby Unit, UCL
Gatsby Unit, UCL ted@gatsby. wcw.uccl..... ac.uk
samo www.hrooolo. 21@ucl.ac.oo
Ahmed Touati Meta atouati@meta, com
Diana Borsa Google DeepMind
Maneesh Sahani
Gatsby Unit, UCL maneesh@gatsbby ucl ac uk
borsa@google.com com
Abstract
A common setting in multitask reinforcement learning (RL) demands that an agent rapidly adapt to various stationary reward functions randomly sampled from fixed distribution. In such situations, the successor representation (SR) is a pop- ular framework which supports rapid policy evaluation by decoupling a policy's 's expected discounted, cumulative state occupancies from specific reward function. However, in the natural world, sequential tasks are rarely independent, and instead reflect shifting priorities based on the availability and subjective perception of re- warding stimuli. Reflecting this disjunction, in this paper we study the phenomenon of diminishing marginal utilityy and introduce a novel state representation, the À representation (AR) which, surprisingly, required for policy evaluation in this setting and which generalizes the SR well as several other state representations from the literature. We establish the AR's formal properties and examine nor- mative advantages in the context of machine learning, as well as its usefulness for studying natural behaviors, particularly foraging.
Introduction
The second ice cream cone rarely tastes as good the first, and once all the most accessible brambles have been picked,
the same investment of effort yields less fruit. In everyday life, the availability and our enjoyment of stimuli is sensitive to our past interactions with them. Thus, evaluate different courses of action and act accordingly, we might expect our brains form representations sensitive to the non-stationarity of rewards. Evidence in fields from behavioral economics [1, to neuroscience [3] supports this hypothesis. Surprisingly. however, most of reinforcement learning (RL) takes place under the assumptions of the Markov Decision Process [MDP; 4], where rewards and optimal decision-making remain stationary. Figure 1.1: Diminishing rewards.
Cumulative Reward
Constan: Utility
Diminishing Utlity
Stimulus Presentations
In this paper, seek bridge this gap by studying the phenomenon of diminishing marginal utility [5] in the context of RL. Diminishing marginal utility (DMU) is the subjective phenomenon by which repeated exposure a rewarding stimulus reduces the perceived utility one experiences. While DMU is thought to have its roots in the maintenance of homeostatic equilibrium (too much ice cream can result in stomach ache), it also manifests itself in domains in which the collected rewards are abstract, such economics (S10 vs. SO is perceived as bigger difference in value than $1,010 vs. $1,000), where it closely related to risk aversion [6, 7]. While DMU is well-studied in other fields, relatively few RL studies have explored diminishing reward functions [8, 9], and, to our
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

knowledge, none contain formal analysis of DMU within RL. Here, we seek to characterize both its importance and the challenge it poses for current RL approaches (Section 3).
Surprisingly, we find that evaluating policies under diminishing rewards requires agents to learn novel state representation which term the À representation (AR, Section 4). The XR gener- alizes several state representations from the RL literature: the successor representation [SR; 10], the first-occupancy representation [FR; 11], and the forward-buckward representation [FBR; 12], adapting them for non-stationary environments. Interestingly, despite the non-stationarity of the underlying reward functions, we show that the XR still admits Bellman recursion, allowing for efficient computation via dynamic programming (or approximate dynamic programming) and prove its convergence. We demonstrate the scalability of the XR large and continuous state spaces (Section 4.1,, show that it supports policy evaluation, improvement, and composition (Section 5), and show that the behavior induces consistent with optimal foraging theory (Section 6).
2 Preliminaries
Standard RL In standard RL, the goal of the agent is to act within its environment so as to maximize its discounted cumulative reward for some task T [13]. Typically. T is modeled discounted MDP, T = (S,A.p. A. r. Po), where is the state space, A is the set of available actions, p 8 A - P(S) is the transition kernel (where ((S) is the space of probability distributions over S). : S + R is the reward function, 7 E (0,1) 1) is a discount factor, and μl0 E PS)) is the distribution over initial states. Note that the reward function is also frequently defined over state- action pairs (s, a) or triples (s, a, s'), but to simplify notation we mostly consider state-based rewards. All of the following analysis, unless otherwise noted, can easily be extended to the (s)aa)oo or (s, a, s')- dependent cases. The goal of the agent is to maximize its expected return, or discounted cumulative reward )r(tr)) The role of the discount factor is twofold: from theoretical standpoint, ensures that this sum is finite for bounded rewards, and it induces myopia in the agent's 's behavior. To simplify notation, we will frequently write r(s;) ) 2 Tt and E ıR$i as the vector of rewards for each state. The agent acts according to a stationary policy can describe the expected transition probabilities under π using a |S| [S] matrix pr such that P. p" (s'[s) & En p(s' s, (alss. Given π and reward function r, the expected return associated with taking action in state s
S + P(A). For finite MDPs, we
(7, a) = Ex Σ y *ret 5,
Ex (-16) [r+ + (2, (s))))
(2.1)
The Q. are called the state-action values or simply the Q-values of π. The expectation Ex [-] 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; 8, a = al as E,, [-]st, The recursive form in Eq. (2.1) called the Beliman equation, and makes the process of estimating Q termed policy evaluation- tractable via dynamic programming. In particular, successive applications of the Bellman operator T _ + are guaranteed to converge the true value function Q" for any initial real-valued [ x [A] matrix Q. Once a policy has been evaluated, policy improvement identifies a new policy n' such that QT (s, a) > QT (s, as, V(s, a) E (s, a).
Helpfully, such policy can be defined N <'s) E argmax, Q7 (s,a).
Value Decomposition Often, it can be useful for an agent to evaluate the policies it's learned on new reward functions. In order to do so efficiently, it can make use of the successor representation [SR; 10], which decomposes the value function as follows:
(s) Ex,
Σ t k20
St
= r EE,
Σ A20
7 *1((()))
(2.2)
(M)
where 1(st) is one-hot representation of state st and M the SR, is an |S] |S] matrix such that
MT (s, s') A E.,
Σ x 1( =s²) = Ex 1(st s') + 7M "M
(2.3)
A20

---

Because the SR satisfies Bellman recursion, it can be learned in a similar manner to the value function, with the added benefit that if the transition kernel p is fixed, it can be used to instantly evaluate a policy on a task determined by any reward vector The SR can also be conditioned on actions, M (s, a, s'); ; in which case multiplication by produces the Q-values of x. The SR was originally motivated by the idea that a representation of state in the brain should be dependent on the similarity of future paths through the environment, and there is evidence that SR- Rike representations are present in the hippocampus [14]. A representation closely related to the SR is the first-occupancy representation [FR; IJ, which modifies the SR by only counting the first visit to each state:
F" (8) Em
Σ 1 ='' & {St...., 804k-13) 8t R20
(2.4)
The FR can be used in the same way as the SR, with the difference that the values it computes are predicated on ephemeral rewards those that are consumed or vanish after the first visit state.
Policy Composition If the agent has access a of policies = { and their associated SRs (or FRs) and is faced by a new task determined by reward vector r, it can instantly evaluate these policies by simply multiplying: {Q" a) M' a). This process is termed generalized policy evaluation [GPE; 15]. Similarly, generalized policy improvement (GPI) is defined as the identification of a new policy x' such that (s, a) 2 suP*em Q* (s.a) Vs, a E S x A. Combining both provides a way for an agent to efficiently compose its policies II that is, to combine them to produce a new policy without additional learning. This process, which we refer to as GPE+GPI, produces the following policy, which is guaranteed to perform at least as well as any policy I1 [16]:
T' T'(s) E argmax max r M (*,,a,,))
(2.5)
aEA TEll
3 Diminishing Marginal Utility
Problem Statement Motivated by DMU's importance in decision-making, our goal is to understand RL the context of the following class of non-Markov reward functions:
(s, ((s))())) ((s) E [0,
(3.1)
where n(s,&) 6) E N is the agent's visit count at up to time and F(s) describes the reward at the first visit to 8. A(s) therefore encodes the extent to which reward diminishes after each visit to 8. Note that for X(s) A = 1 we recover the usual stationary reward given by F, and so this family of rewards strictly generalizes the stationary Markovian rewards typically used in RL.
DMU is Challenging An immediate question when considering reward functions of this form is whether or not we can still define Bellman equation over the resulting value function. If this is the case, standard RL approaches still apply. However, the following result shows otherwise.
Lemma 3.1 (Impossibilityy; Informal). Given reward function of the form Eg. (3.1), it is impossible define a Bellman equation solely using the resulting value function and immediate reward.
We provide more precise statement, along with proofs for all theoretical results, in Appendix B. This result means that we can't It write the value function corresponding to rewards of the form Eq. 13.1) recursively only in terms of rewards and value in an analogous manner to Eq. (2.1). Nonetheless, we found that it is in fact possible to derive a recursive relationship in this setting, but only by positing a novel state representation that generalizes the SR and the FR, which term the À representation (AR). In the following sections, we define the XR, establish its formal properties, and demonstrate its necessity for RL problems with diminishing rewards.
4 The À Representation
Representation for DMU We now the derive the XR by decomposing the value function for rewards of the form Eq. (3.1) and show that it admits a Bellman recursion. To simplify notation,
3

---

use single À for all states, but the results below readily apply to non non-uniform A. We have
V (s) E
Σ * rr((sukk
= Ami((esaa)) (1)L+t)(s[=0) (4.1)
(e)
where we call (s) the À representation (AR), and (x( مe= 1 ($1+) 8). is the number of times state is visited from time up to but not including time k. Formally:
Definition 4.1 (AR). For an MDP with finite S and X [0, 1| the X representation is given by 8 such that
55) E
M) *1s s =
(4.2)
===0
where (s, k) A Ey=o (sits s) the number of times state s visited from time until time +
We can immediately see that for = 0, the XR recovers the FR (we take 0 = 1), and for À = 1, recovers the SR. For À (0, 1), the XR interpolates between the two, with higher values of reflecting greater persistence of reward in a given state or state-action pair and lower values of reflecting more ephemeral rewards.
The AR admits a recursive relationship:
E
(s(ss))) (
[1))+===)))----- =8
k=
[s: = 5 xmi(b'.1)
(1(1(((())) k-1 1(8 = =s))
k=
= 1s: = s' s')(1 + 7AE s041 rupi (7 (5(+1, + 7(1 - 1(st = =s')) (8
(4.3)
where (i) follows from ne(s', k)) na(s', ++++ ( k in Appendix A. Thus, we can define tractable Bellman operator:
1). 1). A more detailed derivation is provided
Definition 4.2 (AR Operator). Let E S|x(S| be an arbitrary real-valued matrix, and let G denote the AR Bellman operator for n, such that
9 A 1 (11 +>AP" + >1111 - -] p
(4.4)
where @ denotes elementwise multiplication and is the |S| x |S| identity matrix. In particular, for a stationary policy æ, 9 Tx
The following result establishes that successive applications of 9 converge to the XR. Proposition 4.1 (Convergence). Under the conditions assumed above, set (b() (1 A)T. For
k = =,1,2,... 1,2,..., suppose that (k-1) = (b)) Then
k+1
(k) 2%),
S
11 Xy
While such analysis fairly standard MDP theory. is noteworthy that the analysis extends to this case despite Lemma 3.1. That is, for ethologically relevant class of non-Markovian reward functions [17], is possible to define Markovian Bellman operator and prove that repeated applications of it converge to the desired representation. Furthermore, unlike in the stationary reward case, where decomposing the value function terms of the reward and the SR "optional" perform prediction or control, in this setting the structure of the problem requires that this representation be learned. To get an intuition for the AR consider the simple gridworld presented Fig. 4.1, where we visualize the AR for varying values of X. For X = 0, the representation recovers the FR, encoding the expected discount first occupancy of each state, while À increases, effective occupancy is accumulated accordingly states which are revisited. We offer more detailed visualization in Fig, F.2.

---

π
0.0 (FR)
À = 0.5
= 1.0 (SR)
Figure 4.1: The XR, interpolates between the FR and the SR. We visualize the XRs of the bottom left state for the depicted policy for À E {0.0, 0.5, 1.0].
4.1 Continuous State Spaces
When the state space is large or continuous, it becomes impossible to use a tabular representation, and we must turn function approximation. There are several ways to approach this, each with their own advantages and drawbacks.
Feature- Representations For the SR and the FR, compatibility with function approximation most commonly achieved by simply replacing the indicator functions their definitions with a base feature function Ø S + R2 create successor features [SFs; 15, 16] and first occupaccyy features [FFs; 11], respectively. The intuition this case is that ø for the SR and the FR is just a one-hot encoding of the state state-action pair), and for cases when |S] is too large, we can replace with a compressed representation. That is, we have the following definition
Definition 4.3 (SFs). Let Ø S + D be a base feature function. Then, the successor feaceue (SF) representation Fi S A + RD is defined as (p a) EE (((((=)))))))))))))))))))))())) St, at for all 8, E S A
One key fact to note, however, is due to this compression, all notion of state "occupancy" is lost and these representations instead measure feature accumulation. For any feasible À then, it is most natural define these representations using their recursive forms:
Definition 4.4 (AF). For X [0. 1] and bounded base features & : S 10, the X-feature (AF) representation of state s given by 2 such that
(x) ((+) ④ (1 + - + 7(1 ((x)) @ (1ap((1))))
(4.5)
In order to maintain their usefulness for policy evaluation, the main requirement of the base features that the reward should lie in their span. That is, a given feature function á is most useful for an associated set of reward functions R, given
R (r Fw R r(s,a) w (fsss, a) Vs, a S x A).
(4.6)
However, Barreto et al. [18] demonstrate that good performance can still be achieved for an arbitrary reward function as long as it's sufficiently close to some E
Set-Theoreths Formulations As noted above, computing expectations of accumulated abstract features is unsatisfying because it requires that lose the occupancy- -ased interpretation of these representations. also restricts the agent to reward functions which lie within R. An alternative approach extending the SR continuous MDPs that avoids this issue is the successor measure [SM; 19], which treats the distribution of future states as a measure over S:
*I((sL E X 8t s, s,a, a.
VX C S measurable,
(4.7)
which can be expressed in the discrete case M = - In the continuous case, matrices are replaced by their corresponding measures. Note that SFs can be recovered by integrating: PP (s = f (s, ds' dss'), We can define an analogous object for the AR as follows
T X) 2
(X..)) E E = s,
(4.8)
R 0
5

---

Atrue 0.5
Dynamic Programming
Tabular TD Learning
AF TD Learning
0.5
20
1.0
50
1800
0.0
20
1049
1550
Iterations
Episodes
Episodes
Figure 5.1: The XR is required for accurate policy evaluation. Policy evaluation of the policy depicted in Fig. 4.1 using dynamic programming. tabular TD learning, and XF TD learning produces the most accurate value estimates when using the XR with À = Atrue- Results are averaged over three random seeds. Shading indicates standard error.
where ny(X, A E5= 60.,, (X). However, this is not a measure because it fails to satisfy additivity for À < 1, i.e. for measurable disjoint sets A, B S S, (x A U B) < < A) + $5 B) (Lemma B.4). For this reason, we call Eq. (4.8) the operator (AO). We then minimize the following squared Bellman error loss for 2x (dropping sub/superscripts for concision):
( E -e [(p(S),s') 70( s 1)²
2E,, 2Berrplp((ss st]
(4.9)
+ + X)B " ((l (sx)))))))))))) ( 8, ())
where and are densities over S and (x(s) - (X())dia((H) in the discrete case, with P parametrized by neural network. = indicates a stop- gradient operation, i.e., a target network. A detailed derivation and discussion are given in Appendix H. While P can be any training distribution of transitions we can sample from, we require an analytic expression for μ.. Eq. (4.9) recovers the SM loss of Touati and Ollivier [12] when À = 1.
5 Policy Evaluation, Learning, and Composition under DMU
In the following sections, we experimentally validate the formal properties of the XR and explore its usefulness for solving RL problems with DMU, The majority of our experiments center on navigation tasks, as we believe this is the most natural setting for studying behavior under diminishing rewards. However, in Appendix we also explore potential for the XR's use other areas, such as continuous control, even when rewards do not diminish. There is also the inherent question of whether the agent has access to In a naturalistic context, À can be seen as an internal variable that the agent likely knows, especially if the agent has experienced the related stimulus before. Therefore, in subsequent experiments, treating À as a "given" can be taken to imply the agent has prior experience with the relevant stimulus. Further details for all experiments can be found in Appendix E.
5.1 Policy Evaluation
In Section 4, we showed that in order to perform policy evaluation under DMU, an agent is required to learn the AR. In our first experimental setting, we verify this analysis empirically for the policy depicted in Fig. 4.1 with a rewarded location in the state indicated by a g in Fig. 5. with Atrue = 0.5. We then compare the performance for agents using different values of across dynamic programming (DP), tabular TD learning, and XF TD learning with Laplacian features. For the latter two, we use linear function approximator with one-hot encoding of the state the base feature function. We then compute the Q-values using the AR with À {0.5, 1.0} (with X = 1 corresponding to the SR) and compare the resulting value estimates the true Q-values. Consistent with our theoretical analysis, Fig. 5.1 shows that the XR with x = Atrue is required to produce accurate value estimates.
5.2 Policy Learning
To demonstrate that XR useful in supporting policy improvement under diminishing rewards, implemented modified forms of Q-learning [20] (which we term Qx Qx-learning) and advantage

---

Aruee 0.5
Qu-Learning
RNN A2C-A
10
0-
100
400
590
2000
+000
6000
Episodes
Episodes
The XR required for strong performance. We apply a tabular 2-Jearning-style algorithm and deep actor-critic algorithm to policy optimization in the TwoRooms domain. The blue locations indicate reward, and the black triangle shows the agent's position. Results are averaged over three random seeds. Shading indicates standard error.
Figure 5.2:
actor-critic [A2C; 21] and applied them to the TwoRooms domain from the NeuroNav benchmark task set [22] (see Fig. 5.2). For a transition (St, at, St $t-1), we define the following operator:
T - 1s)) ④ (1 +
+ 711 - 1 (s:)) @ ((8++1, Cc-1),
(5.1)
where 4++1 = argmax,, Qx(s; [ a) argmax, r (st, a). This is an improvement operator for Qx. The results in Fig. 5.2 show that Qx Qx-learning outperforms standard Q-learning 1) for diminishing rewards, and that the "correct" À produces the best performance. To implement A2C with XR critie, we modified the standard TD target in a similar manner as follows:
(s;) r(s)) (V((4441) (A 1)wT()((ss))
(Ø(s;) @ (x(s+1)))),
(5.2)
where b were one-hot state encodings, and the policy, value function, and XF were output heads of a shared L.TT network [23]. Note this target is equivalent Definition 4.4 multiplied by the reward (derivation in Appendix E). Fig. 5.2 shows again that correct value targets lead to improved performance. Videos agent behavior can be found at Lambdarepresentatil github...... io.
5.3 Policy Composition
As we've seen, DMU problems of this form have an interesting property wherein solving one task requires the computation representation which on its own task agnostic. In the same way that the SR and FR facilitate generalization across reward functions, the XR facilitates generalization across reward functions with different rs. The following result shows that there is benefit having the "correct" X for given resource.
Theorem 5.1 (GPI). Let {M₃} S and M E M be set of tasks in an environment M and let Q denote the action-value function of an optimal policy of Mj when executed in M. Assume that the agent uses diminishing rate that may from the true environment diminishing rate Given estimates Q=: such that [[Q"5 - 11 S for all j. define
π (s) E argmax max Q (s,a).
Then,
2E + Alllrllo + 7(1 (1-r)r(s.a)) Q² 2 max Q (3,0)
1- xy
Note that for x= 1, we recover the original GPI bound due to Barreto et al. [18] with an additional term quantifying error accrued if incorrectly assuming À
Tabular Navigation We can see this result reflected empirically in Fig. 5.3, where we consider the following experimental set-up the classic FourRooms domain [24]. The agent is assumed to be given or have previously acquired four policies {To, 11, #2, л3) individually optimized to reach rewards located each of the four rooms of the environment. There are three reward locations {go, 91, 92 scattered across the rooms, each with its own initial reward and all with À 0.5. At the
7

---

Tabular GPI
Example Trajectories
25.0
Return
22.5
20.0
175
150
0.0
0.5
1.0
A= 0.5
À 0.0
Figure 5.3: Tabular GPL (Left) Average returns obtained by agents performing GPE+GPI using ARs with {0.0, 0.5, 1.0} over 50 episodes. Error bars indicate standard error. (Right) Sample trajectories, Agents with À set too high overstay in rewarding states, and those with À too low leave too early.
Example Observation
Pixel-based GPI
Feature Analysis
A
A
0.5
Figure 5.4: Pixel-Based GPI. Performance is strongest for agents using the correct À = 0.5. PCA on the learned features in each underlying environment state shows that the XFs capture the value- conditioned structure of the environment.
beginning of each episode, an initial state 80 sampled uniformly from the set of available states. An episode terminates either when the maximum reward remaining in any of the goal states is less than 0.1 or when the maximum number of steps H = is reached (when À = 1, the latter the only applicable condition). For each of the four policies, we learn ARs with À {0.0, 0.5, 1.0) using dynamic programming and record the returns obtained while performing GPE+GPI with each of these representations over 50 episodes. Bellman error curves for the XRs are shown in Fig. F.3, and demonstrate that convergence faster for lower A. In the left panel of Fig. 5.3, we can indeed see that using the correct (0.5) nets the highest returns. Example trajectories for each agent are shown in the remaining panels.
Pixel-Based Navigation We verified that the previous result is reflected in larger scales by repeating the experiment in partially-observed version of FourRooms in which the agent receives 128 x 128 RGB egocentric observations of the environment (Fig. 5.4, left) with H/ 50. In this case, the agent learns AFs for each policy for E {0.0, 0.5, 1.0}. where each XF is parameterized by feedforward convolutional network with the last seven previous frames stacked account for the partial observability. The base features were Laplacian eigenfunctions normalized to [0, 1]. which which were shown by Touati et al. [25] to perform the best of all base features for SFs across a range of environments including navigation.
6 Understanding Natural Behavior
Naturalistic environments often exhibit diminishing reward and give insight into animal behavior. The problem of foraging an environment with multiple diminishing food patches (i.c., reward states) is of interest in behavioral science [26-28]. The comerstone of foraging theory is the marginal value theorem [MVT; 28, 29], which states that the optimal time leave a patch when the patch's 's reward rate matches the average reward rate of the environment. However, the MVT does not describe which patch to move to once an agent leaves current patch. We show that the XO recovers MVT--ike behavior in discrete environments and improves upon the MVT by not only predicting when agents should leave rewarding patches, but also where they should go. Moreover, we provide scheme for learning À alongside the 0O using feedback from the environment.

---

Initial Reward
A 1..
A 7 0.1
Leammed X P
True
10
Epoch
Figure 6.1: OY trained via FB. a) À values of two agents in FourRooms, one which learns À and one which does not. b) Performance of the two agents from (a). Learning À improves performance. e) Reward structure and starting state of the asymmetric environment. d) Trajectory of an agent with A= 1. The optimal strategy is to reach the high reward state and exploit ad infinitum. e) Trajectory of an agent with À = 0.1. The optimal strategy is to exploit each reward state for a few time steps before moving to the next reward state.
To learn the ΛO, we take inspiration from the FBR [12] and use the following parametrization: "- (s, a, s') = F(s, a, =) ((s)) and n. (s) = argmax, F'ss, a, z) 2, where ç is a density with full support on S, e.g., uniform. We then optimize using the the loss in Eq. (4.9) under this parameterization (details in Appendix E). Given a reward function r - R at evaluation time, the agent acts according to argmax, Fss, a, =x) T zR. where zn E B(s)) Because the environment is non-stationary, 2R has to be re-computed at every time step. To emulate a more realistic foraging task, the agent learns À by minimizing the loss in Eq. (H.1) in parallel with the loss
((A) Es [1( S141) (r(se+1) Ar(ss))²
which provably recovers the correct value of À provided that p is sufficiently exploratory. In Appendix H.1 we provide experiments showing that using an incorrect value of X leads to poor performance on tabular tasks. In Fig. 6.1 show that the agent learns the correct value of X, increasing its performance. We illustrate the behavior of the XO in an asymmetric environment that has one large reward state on one side and many small reward states (with higher total reward) on the other. Different values of À lead to very different optimal foraging strategies, which the AO recovers and exhibits MVT-like behavior (see Appendix H.2 for more detailed analysis). Our hope is that the XR may provide a framework for new theoretical studies of foraging behavior and possibly mechanisms for posing new hypotheses. For example, an overly large may lead to overstaying in depleted patches, a frequently observed phenomenon [30].
7 Conclusion
Limitations Despite its advantages, there are several drawbacks to the representation which are a direct result of the challenge of the DMU setting. First, the AR is only useful for transfer across diminishing reward functions when the value of X at each state is consistent across tasks. In natural settings, this is fairly reasonable, as À can be thought of as encoding the type of the resource available each state (i.e., each resource has its own associated decay rate). Second, as noted in Section 4. 1, the AR does not admit measure- measure-theoretie theooetiic formulation, which makes it challenging to define principled, occupancy- based version compatible with continuous state spaces. Third, the AR is a prospective representation, and so while it is used to correctly evaluate a policy's future return under DMU, is not inherently memory-based and so performs this evaluation as the agent hasn't visited locations with diminishing reward before. Additional mechanisms (i.e., recurrence or frame-stacking) are necessary to account for previous visits. Finally, the XR is dependent on an episodic task setting for rewards to reset, as otherwise the agent would eventually consume all reward in the environment. An even more natural reward structure would include a mechanism for reward replenishment in addition to depletion. We describe several such candidates in Appendix J, but leave a more thorough investigation of this issue to future work.
In this work, we aimed to lay the groundwork for understanding policy evaluation, learning, and composition under diminishing rewards. To solve such problems, we introduced--and showed the necessity oh-the À represeniation, which generalizes the SR and FR. We demonstrated its usefulness for rapid policy evaluation and by extension, composition, as well as control. We believe the AR represents a useful step in the development of state representations for naturalistic environments.

---

References
[1] Daniel Kahneman and Amos Tversky. Prospect theory: An analysis of decision under risk. Econometrica, 47(2):263-292, 1979.
[2] Matthew Rabin. Risk aversion and expected-utiliy theory: calibration theorem. Econometrica, 68(5):1281-1292, 2000.
doi: 10.2307/2999450.
[3] Alex Pine, Ben Seymour, Jonathan P Roiser, Peter Bossaerts, Karl Friston, H Valerie Curran, and Raymond Dolan. Encoding of marginal utility across time in the human brain. Journal of Neuroscience, 29(30);9575995) 1975519511122 2009.
[4] Martin L. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley and Sons, 2010.
[5] Hermann Heinrich Gossen and Rudolph C Blitz. The laws of human relations and the rules of human action derived therefrom. Mit Press Cambridge, MA, 1983.
[6] Kenneth Arrow. The theory of risk-bearing: small and great risks. Journal of risk and uncertainty, 12:103-111, 1996.
[7] John W Pratt. Risk aversion in the small and the large. In Uncertainty economics, pages 59-79, Elsevier, 1978.
[8] Nathan Wispinski, Andrew Butcher, Kory Wallace Mathewson, Craig S Chapman, Matthew Botvinick, and Patrick M. Pilarski. Adaptive patch foraging in deep reinforcement learning agents. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: wwwprrvvweee..tttffooumrmdddaadd
[9] Sergey Shuvaev, Sarah Starosta, Duda Kvitsiani, Adam Kepecs, and Alexei Koulakov, R- learning in actor-critic model offers biologically relevant mechanism for sequential decision- making. Advances in neural information processing systems, 33:18872-18882, 2020.
[10] Peter Dayan. Improving generalization for temporal difference learning: The successor repre- sentation. Neural Computation, ((4):613-624, 1993. doi: 10.1 162/neco.
112./ecc..199.55.4444 1993.5.4.6133.
[11] Ted Moskovitz, Spencer R Wilson, and Maneesh Sahani. A first-ocupancyy representation for reinforcement learning. In International Conference Learning Representations, 2022. URL
https://openreviiee wwt.:/iiiwweeeefrnneeeeeeo Biw.orrmmerdd--222220120
[12] Ahmed Touati and Yann Ollivier. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021.
[13] Richard Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL http:://incommpeeeeeeess h/tp:////l.eaaaaaaaettootooooo// the-thhe--booooo 2nd. html.
[14] Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel J Gershman. The hippocampus as a predictive map. Nature neuroscience, 20(11):1643-165, 2017.
[15] Andre Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy Sciences, 1 117(48)30009 2020. ISSN 0027-8424. doi: 10.1073/pnas 19773/11.... URL https: / // //www.pmaa ppas. carnaa..ogg/coo
[16] André Barreto, Will Dabney, Rémi Munos, Jonathan Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017.
[17] Maor Gaon and Ronen Brafman. Reinforcement learning with non-markovian rewards. In Proceedings the AAAI conference on artificial intelligence, volume 34, pages 3980-3987, 2020.
10

---

[18] André Barreto, Diana Borsa, John Quan, Tom Schaul, David Silver, Matteo Hessel, Daniel Mankowitz. Augustin Zidek, and Rémi Munos. Transfer in deep reinforcement learning using successor features and generalised policy improvement, 2019.
[19] Léonard Blier and Yann Ollivier. The description length of deep learning models. Advances in Neural Information Processing Systems, 31, 2018.
[20]
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292, 1992.
[21] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforce- ment learning. In International conference machine learning, pages 1928-1937. PMLR. 2016.
[22] Arthur Juliani, Samuel Barnett, Brandon Davis, Margaret Sereno, and Ida Momennejad. Neuro- nav: a library for neurally-plausible reinforcement learning. arXiv preprint arXiv:2206.03312. 2022.
[23] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, (8):1735-1780, 1997.
[24] Richard S. Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction reinforcement learning. Artificial Intelligence, 112 (1):181211,, 1999. URL https://wwet http:///wwwwscieeeeeeeeeeee sciencedirect. com/ commscience///eeoooccocccoccoece httm//cience//rrcee//hhhhi $0004370299000521.
[25] Ahmed Touati, Jérémy Rapin, and Yann Ollivier.
Does zero-shot reinforcement learning exist? In The Eleventh International Conference on Learning Representations, 2023, URL
www.aarevieeeeeeeeeeeeee wtt.fformmmdddddd MYaappoPP OcQI.
[26] Benjamin Hayden, John M Pearson, and Michael Platt. Neuronal basis of sequential
foraging decisions in a patchy environment, Nature neuroscience, 14(7):933-939, 2011.
[27] Tehrim Yoon, Robert B Geary, Alaa Ahmed, and Reza Shadmehr. Control of movement vigor and decision making during foraging. Proceedings of the National Academy of Sciences, 115 (44):110466EE143)) 2018
[28] Anthony S Gabay and Matthew AJ Apps. Foraging optimally in social neuroscience: computa- tions and methodological considerations. Social Cognitive and Affective Neuroscience, 16(8): 782-794, 2021.
[29] Eric L. Charnov. Optimal foraging, the marginal value theorem. Theoretical Population Biology, 1(2):199-136, 1976. ISSN 0040-5809.
[30] Peter Nonacs. State dependent behavior and the marginal value theorem. Behavioral Ecology, 12(1):71-83, 2001. 2001.
[31] Ruosong Wang, Hanrui Zhang, Devendra Singh Chaplot, Denis Garagić, and Ruslan Salakhut- dinov. Planning with submodular objective functions. arXiv preprint arXiv:2010.11863. 2020.
[32] Manish Prajapat, Mojmir Mutný, Melanie N Zeilinger, and Andreas Krause. Submodular reinforcement learning. arXiv preprint arXiv:2307.13372. 2023.
[33] Tom Zahavy, Brendan O' O'Donoghue, Guillaume Desjardins, and Satinder Singh. Reward is enough for convex mdps, 2021.
[34] Eitan Altman. Constrained Markov decision processes, Routledge, 2021.
[35] Ted Moskovitz, Brendan O'Donoghue, Vivek Veeriah, Sebastian Flennerhag, Satinder Singh, and Tom Zahavy. Reload: Reinforcement learning with optimistic ascent-descent for last-iterate convergence in constrained mdps. International Conference on Machine Learning, pages 25303-25336. PMLR. 2023.
11

---

[36] Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides. Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, al. Behaviour suite for reinforcement learning. arXiv preprint arXiv: 190 003668, 2019.
[37] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization arXiv preprint urXiv:1607006400 2016.
[38] Martin Riedmiller. Neural fitted q iteration-first experiences with a data efficient neural reinforcement learning method. In Machine Learning: ECML 2005: 16th European Conference on Machine Learning, Porto, Portugal, October 3-7, 2005. Proceedings 16, pages 317-328. Springer, 2005.
[39] Peter Dayan and Geoffrey E Hinton. Feudal reinforcement learning. Advances in neural information processing systems, 5, 1992.
[40] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. Proceedings of the 35th International Conference on Machine Learning. ICML 2018, Stockholmsmässan Stockholm, Sweden, July 10-15, 2018, 2018.
[41] Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient wasserstein natural gradients for reinforcement learning. 2020.
[42] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor, 2018.
[43] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
[44] Ted Moskovitz, Jack Parker-Holder, Aldo Pacchiano, Michael Arbel, and Michael Jordan. Tacti- cal optimism and pessimism for deep reinforcement learning. Advances in Neural Information Processing Systems, 34:12849-12863, 2021.
[45] Nicolo Desa-Biancchi and Gábor Lugosi. Prediction, learning, and games. Cambridge university press, 2006.
12

---

Appendix: A State Representation for Diminishing Rewards GIFs of navigation agents can be found at Lemblarrpresentatioon.ggtt and in the supplementary material.
A Derivation of XR Recursion
We provide a step-by-step derivation of the AR recursion in Eq. (4.3):
0e
(8 E
A A = =5') 56){3c=
2
1 (st = 5') + Σ
mmi ()) 1(1++ =)))
- =
E
k=1
E
1 (s: = 5') + )me(o'2)
x )ct1ls 1 1(( 10054==511 =s)(
k=
0
(1l)
E
1 [(1 = s') + 11 (st = =')d Σ
x m1) 1 = =s)
k=1
+7(1 1(8: = =s'))
Σ
1 (b))) (- (81+H
- 5)) =
k=1
1(s 5') + 1ss 5') (7((411,88)
+ + - 1 = 5 ())eEett>p (8(41,18))
1(s s1)(1
7X x> (5:1 (5:11;1)))))
+ x1 1(s = s') E Estti D3
1 (91+1,
(A.1)
where (i) is because ne k)) = (s 1) + n++1(s', (s', k) and (ii) is because xma( ((r==s)) 1ss s')d (1 1(s = s))).
B Theoretical Analysis
Here, we provide proofs for the theoretical results in the main text.
Lemma 3.1 (Impossibility; Informal). Given reward function of the form Eq. (3.1), is impossible to define a Bellman equation solely using the resulting value function and immediate reward.
Before providing the formal statement and proof for Lemma 3.1, we introduce definition for Bellman operator.
Definition B.1 (Bellman Operator). A Bellman operator is contractive operator
RiS - RI that depends solely on V, one-step expectations under p", and learning hyperparameters (in our case 7 and X).
We can now give formal statement of Lemma 3.1:
13

---

Lemma B.1 (Impossibility; Formal). Assume that |S] > 1. Then, there does not exist Bellman operator T with fixed point V
Proof. Assume for a contradiction that T is a Bellman operator. By the Banach tixed-point theorem, V must be the unique fixed point of Hence, TV" must take on the following form (see the proof of Lemma 3.1 in Appendix B): for s ES, ((ss) V* (s') + 7(A ((s ( s). For the assumption to hold, there must exist a function f such that, for any s E S Ey Ex'op (11)) (7 5) = Ept(ts) (I, V", ),,))
(TV*)(s = -s)) +
Now, by definition,
(T(()= = Σ TT (((s,s')F(s) s's''s)
s'ES
F is vector in R Sl, so as long S > F- is non-trivial. Fix any F, V", s. Pick vector w E F²' \ {0} and define
(* = D7 5') + w(s')
for any 8,8' s E S. Note that
Σ
Σ (Ts,s')((5)) + Σ ((5))(()--Va = (s)),,
YES
NES
s'ES
as w 1 F. However,
Ex'ep' (11) (s = [p(ds1s)0 - + [d((((ssw(ss
SEES
s'ES
The final term evaluates to w(s). Because w + 0, there must exist some s such that w(s) # O. For this s, we have a single input (F, V 7, X, s) to f that corresponds to two distinct outputs: (s) (x'',, 8) and Eg (-1s) (s',sss)
Hence, f is a one-to-many mapping: for fixed input, there is more than one output. Therefore, f' is not a function, yielding a contradiction.
The following establishes 9 as a contraction.
Lemma B.2 (Contraction). Let G be the operator as defined in Definition 4.2 for some stationary policy π. Then for any two matrices Φ &' E [S (g) - ( s')l S (10(s, - D ssss
Proof. We have (gg Φ
- Φ
1 (11 + x) + (11 1) C p
-10 (11 + >>P= 2') 7 (11 - O"PP''')) ).oo)
7XP - 4') + 1(11 11 -
Pd $ (
I((1 11 ( 2')) 3.8¹
1
(
< ((7PT(0- - Φ0) 3,
=>((P=(0 ( $6) (')),,e' s.o"'
S 2 D') D'),,0/)
14

---

where (i) comes from using < 1 and simplifying.
Note that can actually get tighter contraction factor of xy for s ==s- this contractive property, we can prove its convergence with the use of the following lemma.
s'. s Given
Lemma B.3 (Max AR). The maximum possible value of Df (s, s') is
=( + 11 = 5)))
xy
Proof. For 8 = =',
(s) 1 + s (1se) (7 s).
This is just the standard SR recursion with discount factor Xy, so the maximum is
8
( 1-
(B.1)
k=
For # s', [s; =s) = 0, so
(7 5¹) s (181) D7 (7
Observe that (s s) 2 (s, s') for s' + 8, so the maximum is attained for 8t+1 = =s' We can then use the result for = s' to get
1 x (1-x)-
$*(s, s') =
(B.2)
Combining Eq. (B.1) and Eq. (B.2) yields the desired result.
Proposition 4.1 (Convergence). Under the conditions assumed above, set )(0) (I- For = 1,2,..., suppose that ((k+1) = -g Then
( -070)(0)..))
1- X7
Proof. Using the notation X
8,8' (8,,5') for matrix X:
( (k) 17),o// = (g) g (T s,,d
(g) - o,
(( s.
(B.3)
(s s¹)
(00)
gktl
Xy
where (i) is due to Lemma B.2, (ii) is because (0) (s, = 0 for + s', and (tii is due to Lemma B.3.
Lemma B.4 (Subadditivity). For any s E S, policy n, X E [0, 1), and disjoint measurable sets B S,
AUB) < (s, A) + (s B).
15

---

Proof. Note that for disjoint sets A, B, we have n, (AUB,)) = =,,AA + (n,BBk)
Hence, conditioned some policy π and St =,
X E AUB) = Xm((AA)xu(BB)) P(S+H E A) + (m(4A))m E B)
(AM(AA)P((SHA E A) + Aml(B.iiP(Sss E B),
where the first line follows from P(S++* E AU B) = P(s++* E A) + E B). Equality holds over all A, B, t, if and only if X = 1. Summing over l yields the result.
B.1 Proof of Theorem 5.1
We first prove two results, which rely throughout on the fact that 4x(s, S - for all 8, a, s', which follows from Lemma B.3. For simplicity, also assume throughout that all rewards are non-negative, but this assumption can easily be dropped by taking absolute values of rewards. The proofs presented here borrow ideas from those of [16].
Lemma B.5. Let {M}} S M M E M be a set tasks an environment M with diminishing rate X and let Q" denote the action-value function of an optimal policy of Mj when executed in M. Given estimates Q such that !|Q" - xo Se for all j, define
T (()) E argmax max (s.(sa))
Then,
2E + 7(1 - 1))(s,a)
(²(s,a) 2 max (s,
1-7
1- X7
where denotes the reward function of M.
Proof. Define Qmax a) = max, Q=i (s, a) and Qmax (s, a)
= == max, Q a). Let
T denote the Bellman operator of a policy v in task M. For all (s, a) E S x A and
16

---

all j.
@ a)) = =(ss, a) + 7 En(s'(s,a) a) (1=-))),,,,0 5) + mmax(s' πI (s'))
= as) + 7 Ep(s'ss,o) (A= s) + max amax (s', b)
r + Ep(s'(s,a) a) (1=1)))6,.)001() s) + max Qax (s', (2,b)) 7
r + E ((ss(ss a) (1 8) + m (s', (s'))) 76
r +7 Ep(s'(s,o) (
7E
ar(s,a) + (s (1-)))
-ye
+7(1 - 1)r(s,a) a)
7e 711 - -))
2 TT" Q/' (s, (5,(5))))))))))))))))))))
1-
7(1 - )))(s,a)
Q, (s, a) 7e-
1- A-
This holds for any j, so
Qmax (s, 2 max 023² (s) - -7E- 7(1 - -))ss))ss)))))
j
1- xy
max (s, a) - 7E 7(1 - a)r(s,a)
1 1-1 X7
7(1 - >)r(s,a) a)
2 Q max (s, a) - E - E-YE-
1- xy
Next, note that for any e E R,
T (Ôman (s, + cc= = -TmQmax (s,a) + y Σ [pss'ss,aj a)c
= -T max (s, + yc.
17

---

Putting everything together, and using the fact that T is monotonic and contractive, O7 (5, = lim (T 14xmmxsss ( a)
Y00 (
Qmax(s, (((17)) 7(1 - A)r(s,a)
k
Σ
lim
k0000
1 1 xy
1=0
1
E((1 +7) 7(1 - ))r(s,
Q max (8,
1 1-7 7
1 xy
- 1 1-7 (<(1 (c(1+)) +))+ 7(1 A)r(s, a)
Q max (5:a)
1- 1-
1
(2e
2E + 7(1 -1)()s))))))
2 Q (s,
1-
1- xy
This holds for every j, hence the result.
Lemma B.6. Let be any policy, X, Å E [0, 11, and Qx denote value function with respecting to diminishing rate x. Then,
11Q% -x S A ||||
7
Proof.
The proof follows from the definition of Q: for every (s, E ESX ESx
(xxss, a) - 27 a)l = E
( (mclseessh) r r(s(+t)
8t = 8, ag = a
k=
< E,
Σ
* ( {melo4ee) r(sett)
81 =, a, = a
k=
ne(
= E
y (r((()))) /) /a- {
Σ
8, = 0
5=
x
< A E
Σ
(((++b) 8, at =
A=
X 1||||||
Theorem 5.1 (GPI), Let {M1}s=1 C M and M E M be a set of tasks in an environment M and let Q denote the vetion-value function of an optimal policy of M, when executed in M. Assume that the agent uses diminishing rate that may differ from the true environment diminishing rate X. Given estimates QT,, such that I|Q": Qm S eor all j, define
T(s) E argmaa max (T)
Then,
1
2E + 1 - Åll|T|lo + 7(1 7(1 A)r(s,))
(=(8,a) > max
j
18

---

Proof. Let Q2 denote a value function with respect to diminishing constant X. We wish to bound
max
()(s,a),
i.e,, the value of the GPI policy with respect to the true compared to the maximum value of the constituent policies Tÿ used for GPI, which were used assuming X. By the triangle inequality,
max 2 (s,
(s, a) Q7 a) S max (s, (7 Q](s,a) a) + max 2 (s, a) max O (s,a))
I max Q (s, a) (x a) + max Q (s, a) o (s, as)
(1)
We bound (1) by Lemma B.5 and (2) by Lemma B.6 to get the result.
B.2 An Extension of Theorem 5.1
Inspired by [18], we prove an extension of Theorem 5.1:
Theorem B.1. Let M E M be a task in an environment M with true diminishing constant X. Suppose we perform GPI assuming a diminishing constant X:
Let {M})"= and M₃ be tasks in M and let Q₂' denote the action- value function of an optimal policy of Mj when executed in Mi.
Given estimates QTT such that ||Q, - 227 S eor all j, define (s) E argmax, max, (s,(sa
Let O7, and Q denote the action-valls functions of π and the M M-optimal policy n when executed in M, respectively. Then,
2 1IQT" 2e S 1-n
1 +e+ ||r Tilloo + min ||r; ryll00 + 1- X
d
where C is positive constant not depending on A:
C - 2|r rilloo + +2 mini ||r: rylloo + min (11r l|relloe) + min (IIralloo, ||rillo..... ||ral
7
Note that when X = 1, we recover Proposition 1 of [18] with an additional term quantifying error incurred by Â + X. The proof relies on two other technical lemmas, presented below.
Lemma B.7.
- 11 < ||- ri 1-7
min (1]r|/bc, ||r:lloo) + ||r rillxc (1 - 71 (1- X7)
+ 7111= A)
19

---

Proof. Define A == =1Q=* - -Q loo For any (s, E S x A,
1 (s, (8, - 0 (s.a)l as) = |r(s, a) + E Ep(s'ss,)) - (8.
- as) -7 p((ssss (s,a) (A_- Q7 ((
sir(s, a) ri(s, a))++ +? pp (s, Q, (s, a)l
+7(1 - 1) Σ p(s a) ((s,a)0 (8) 8) - (s (s'),
s ||r- 1:0100 + +74, + 7(1 A) \)||- r, क 20
The third term decomposes as
||r - -r; 10 S || - -r 1100 + ||ra= - r:
||/1b0 + ||r --|r--no 1ill00
1- xy
We could equivalently use the following decomposition:
||r - -,; 1100 I |ra - -r, 1l + ||,, - r,:b%i 1100 1000
|illoc ||| - rillo0
1 Ay
and so
min (I1rllso,
/|r:|l + +||r rilloo
||a=" - rb, 1l00 S
1 xy
The inequalities above hold for all s, a and so
min (11r l|rilloc) + |||r - rilloo
A < || ri + 7d + 7(1 - A)
1 - Xy
A |r 1r-rilles
min (117 ||rilloo) + ||| - 1|rl-riloo
+ 7111- 2)
1-7
(1 (1-) 1 Ay)
Hence the result.
Lemma B.8. For any policy n,
|r-rilla rillo
|| -r-lo
rilloo
||Q - 10 <
+ +x(1 - 1)
1 1-7
(1-7)(1-27)
20

---

Proof. Write A, == 1 - Q- Proceeding as in the previous lemma, for all
( x A, we have
[Q] (s, a)) r:(8, a) +
-r(s, a) -7 -
S a) - (s) + Σ
, (5'), 8)
+ 7 Σ
S ||r no + + +11111 rill
-
+ >
1 xy
A S rr 1:100 + 7(1 - A))lr- -iille-lle- + +74i
1- xy
||- 1ill00
+ 7(1 \)(l - 1ill00
1-7
- (1-7)(1-M)))) X7)
Finally, we prove Theorem B.1:
Proof of Theorem B. By the triangle inequality, | Qil S |1 - 2x + 11 - 2o
By Lemma B.6, the second term is bounded above by
A ||l|T|lo
1-7
The first term decomposes as follows (dropping the subscript on all action-value functions for clarity):
||QT - 10 S 1Q"-Q""" 11 of + +1 1
(1)
(2)
(3)
Applying Lemma B.5 to (2) (but with respect to M, rather than M), we have that for any j,
2E + 7(1 ))ri(s,a)
2 -7 a) < Q (s, a) - Q77 17,77777,,0) (s, + 1 1-7 7
1 1-d7
- 11 27 ll S 11 Q 1100 + 110, - Q,"² 110 1- 22 7(1- 1) (2.1)
1- 1
(2.2)
We bound (2.1) using Lemma B.7 and (2.2) using Lemma B.8 (but with respect to M, rather than M):
+ 110, - 1bc < 2||r: 12,
1-7 Y + >) min (11ralloc; (oo) + 2||r; ryllox
11 - 2, lloo l1o
(1 (1 X7)
We then apply Lemma B.7 to (1) and Lemma B.8 to (3) to get the result.
21

---

C An nth Occupancy Representation
To generalize the first occupancy representation to account for reward functions of this type, it's natural to consider an Nth occupancy representation-ting is, one which accumulates value only for the first N occupancies of one state s' starting from another state 8:
Definition C.1 (NR). For an MDP with finite S, the Nth-occupancy representation ( a policy is given by F E [J,N]]|s)]s]] such that
00
(s) 5))
E
#
(8t+1 = , # (fi St+j = S'.jE 10, 1}}) < N) st+)
=s',
St
K=O
(C.1)
Intuitively, such a representation sums the first N (discounted) occupancies of s' from time to t + & starting from s; We can also note that L7) is simply the FR and & (0) (s, s') = 0 Vs, s'. As with the FR and the SR, we can derive recursive relationship for the NR:
(T) s') 1ss s1)(1 + TE (N- (8 + +111111111111111 1(s = $'))ETT (N)
(st+1, s'), (C.2)
where the expectation is wrt p" (s.)) Once again, we can confirm that this is consistent with the FR by noting that for N 1, the NR recursion recovers the FR recursion. Crucially, we also recover the SR recursion in the limit as N - co:
lim (T) 1s: = =1) ( (Stt1, s')) 7(1 1(s s (0))) (8441,
Ni+0
1(; = s') + TE Tx) (81+1, s').
This is consistent with the intuition that the SR accumulates every (discounted) state occupancy in potentially infinite time horizon of experience. While Definition C. admits a recursive form which is consistent with our intuition, Eq. (C.2) reveals an inconvenient intractability: the Bellman target for (N) requires the availability of Φh T7 (N-1). This is challenge, because it means that we'd like to learn any NR for finite N > 1, the agent also must learn and store th (.).... ... (V-1)- Given these challenges, the question of how to learn a tractable general occupancy representation remains. From a neuroscientific perspective, fixed depletion amount is also incon- sistent with both behavioral observations and neural imaging [3], which indicate instead that utility disappears utility, rather than in proportion to the original utility. We address these theoretical and practical issues in the next section.
fixed rate in proportion to the current remaining
D Additional Related Work
Another important and relevant sub-field of reinforcement learning is work which studies non-stationary rewards. Perhaps most relevant, the setting of DMU can be seen as a special case of submodular planning and reward structures [31, 32]. [31] focus specifically on planning and not the form of diminishment we study, while [32] is a concurrent work which focuses on the general class of submodular
22

---

reward problems and introduces a policy-based, REINFORCE-like method which is necessarily on-policy. In contrast, we focus on a particular sub-class of problems especially relevant to natural behavior and introduce family of approaches which exploit this reward structure, are value-based (and which can be used to modify the critic in policy-based methods), and are compatible with off-policy learning. Other important areas include convex [33] and constrained [34, 35J MDPs. In these cases, non-stationarity is introduced by way of a primal-dual formulation of distinct problem classes into min-max games.
E
Further Experimental Details
E.1 Policy Evaluation
We perform policy evaluation for the policy shown in Fig. 4.1 on the 6 x 6 gridworld shown. The discount factor y was set to 0.9 for all experiments, which were run for H 10 steps per episode. The error metric was the mean squared error:
Qerror
(()"(s, - a(s,a))²
(E.1)
S
8,0
where Q" is the ground truth Q-values and Q is the estimate. Transitions are deterministic. For the dynamic programming result, we learned the XR using Eq. (4.3) for X {0.5, 1.0} and then measured the resulting values by multiplying the resulting AR by the associated reward vector r E f-1,0 0, 1}36, which was -1 in all wall states and +1 at the reward state g. We compared the results to the ground truth values. Dynamic programming was run until the maximum Bellman error across state-action pairs reduced below 5e-2. For the tabular TD learning result, we ran the policy for three episodes starting from every available (non-wall) state in the environment, and learned the AR for À {0.5, 1.0} above, but using the online TD update:
Φ (st, (1 a:) + ad;
5, = 1(s)) 2 (1 x a(+1)) + 7(1 - 1(s)) O T (S a(+1) - x at); where a,+1 ~ T(. s.+1). The learned Q-values were then computed in the same way as the dynamic programming case and compared to the ground truth. For the XF result, we first learned Laplacian eigenfunction base features as described in [25] from a uniform exploration policy and normalized them to the range [0, 1]. We parameterized the base feature network as a 2-layer MLP with ReLU activations and 16 units in the hidden layer. We then used the base features to learn the XFs as in the tabular case, but with the XF network parameterized as a three-layer MLP with 16 units in each of the hidden layers and ReLU activations. All networks were optimized using Adam with a learning rate of 3e-4. The tabular and neural network experiments were repeated for three random seeds, the former was run for 1,500 episodes and the latter for 2,000.
E.2 Policy Learning
We ran the experiments for Fig. 5.2 in a version of the TwoRooms environment from the NeuroNav benchmark [22] with reward modified to decay with a specified
23

---

Algorithm Online Tabular Qx Learning Update
1: Require: Current AR-values | current reward vector r()) observed
(St, de. St+1) tuple
2: Compute Qx-values 2 + (T()))
3: Select greedy action: a++1
argmax,e Qx (8+11,))
4: Update Φx::
@ (1+1) (s) (s,.
10 (s) + ad"h) where
2 ((0) 1(sr) o (1 + 110 (8411, +(111 (ss)) о 26 (4+41, - (sc, at).
5: Return updated (+++1)
Atrue 0.5 and discount factor 7 0.95. The initial rewards in the top right goal and the lower room goal locations were and the top left goal had initial reward 10. The observations in the neural network experiment were one-hot state indicators. The tabular Qx experiments run the algorithm in Algorithm for 500 episodes for À E {0.0, 0.5, 1.0], with Atrue set to 0.5, repeated for three random seeds. Experiments used a constant step size = 0.1. There were five possible actions: up, right, down, left, and stay. The recurrent A2C agents were based the implementation from the BSuite library [36] and were run for 7,500 episodes of maximum length H = 100 with > 0.99 using the Adam optimizer with learning rate 3e-4. The experiment was repeated for three random seeds. The RNN was an LSTM with 128 hidden units and three output heads: one for the policy, one for the value function, and one for the XF The base features were one-hot representations of the current state, 121 dildeiensional in this case.
E.3 Tabular GPI
The agent is assumed to be given or have previously acquired four policies {To, T1, n2; R3} individually optimized to reach rewards located each of the four rooms of the environment. There are three reward locations {90, 91, 92) scattered across the rooms, each with its own initial reward F = [5, 10, 5] and all with À = 0.5. At the beginning of each episode, initial state So is sampled uniformly from the set of available states. An episode terminates either when the maximum reward remaining in any of the goal states is less than 0. or when the maximum number of steps H 40 is reached. Empty states carry a reward of 0, encountering a wall gives reward of 1,, and the discount factor is set to 7 0.97.
For each of the four policies, we learn ARs with À equal to 0, 0.5, and 1.0 using standard dynamic programming (Bellman error curves plotted in Fig. F.3), and record the returns obtained while performing GPE+GPI with each of these representations over the course of 50 episodes. Bellman error curves for the XRs are In the left panel of Fig. 5.3, we can indeed see that using the correct À (0.5) nets the highest returns. Example trajectories for each of XR are shown in the remaining panels.

---

E.4 Pixel-Based GPI
In this case, the base policies II were identical those used in the tabular GPI experiments. First, we collected dataset consisting of 340 observation trajectories (Oo: O1...... OH-1) E O# with H 19 from each policy, totalling 6, 460 observa- tions. Raw observations were 128 x 128 x 3 and were converted to grayscale. The previous seven observations were stacked and used to train a Laplacian eigenfunction base feature network in the same way as [25]. For observations less than seven steps from the start of an episode, the remaining frames were filled in as all black observations (i.e., zeros). The network consisted of four convolutional layers with 32 filters with strides (2, 2, 2, 1), each followed by a ReLU nonlinearity. This was then flattened and passed through Layer Norm layer [37] and a tanh non- linearity before three fully fully connected layers, the first two with 64 units each and ReLU nonlinearities and the final, output layer with 50 units. The output was L2 normoaallized as in [25]. This network Φ O + R
D (with D = 50) was trained on the stacked observations for 10 epochs using the Adam optimizer and learning rate 64. To perform policy evaluation, the resulting features, evaluated on the dataset of stacked observations were collected into their own dataset of (s,, a++1, S1+1, a(+1) tuples, where S, - 01-6a- The "states" were normalized to be between 0 and I, and vector was fit to the actual associated rewards via linear regression on the complete dataset. The AF network was then trained using a form of neural fitted Q-iteration [FQI; 38] modified for policy evaluation with XFs (Algorithm 2). The architecture for the XF network was identical to the base feature network, with the exception that the hidden size of the fully connected layers was 128 and the output dimension was DIA| = 250. FQI was run for K = 20 outer loop iterations, with each inner loop supervised learning setting run for L 100 epochs on the current dataset. Supervised learning was done using Adam with learning rate 3e-4 and batch size B 64. Given the trained networks, GPI proceeded as in the tabular case, i.c.,
le-4 with batch size
at = argmax max T "p (St, a).
(E.2)
aaA
sell
50 episodes were run from random starting locations for H 50 steps and the returns measured. Learning curves for the base features and for XF fitting are shown in Fig. E. The XF curve measures the mean squared error as in Eq. (E.
The feature visualizations were created by performing PCA to reduce the average XF representations for observations at each state in the environment to 2D. Each point in the scatter plot represents the reduced representation on the zy plane, and is colored according to the X-conditioned value of the underlying state.
E.5
Continuous Control
X-SAC
See Appendix I for details.
E.6
Learning the 0O with FB
Training the OO with the FB parameterization proceeds in much the same way as in [12], but adjusted for a different norm and non-Markovian environment. We
25

---

Base Feature Learning
AF Learning
3500
3000
2500
MSE
2000
1500
1000
500
2
15
Epochs
FQI Iterations
Figure E.1: Learning curves for XF policy evaluation, Results are averaged over three runs, with shading indicating one unit of standard error.
Algorithm 2: Fitted Qx
1: Require: Dataset of base features {ø(s) ) decay rate X, discount factor 7 reward feature vector w E IRP, batch size B, learning rate a
2: Initialize AF yps parameters ((2) (we drop the subscript and superscript æ for concision)
3: for k 1. .....K do
4: // Stage I: Construct dataset
5: D - Ø
6: for (s, a) E S x do
7:
for (d) E S A do
8:
D+ DU
(s,a), w [0(s) @ (1 Xypoge (s', d'), + 7(1 o(s)) @ Pocks (s', (s',
(y(aa)
9:
end for end for
11:
// Stage 2: Supervised learning
Randomly initialize to
13:
for fare-1, 1,..- do
14:
Randomly shuffle for {((s, a), 1 E D do
15:
16:
Og + 0: a 2B Eb (yb w To (sb. az))
17: end for
18: end for
((++1) + Og
20: end for
summarize the learning procedure in Algorithm 3. The loss function L is derived in Appendix H, with the addition of the following regularizer:
(s) 11
This regularizer encourages B to be approximately orthonormal, which promotes identifiability of Fg and B, [12].
F
Additional Results
See surrounding sections.
26

---

Algorithm 3: XO Learning
1: Require: Probability distribution over RR, randomly initialized networks Fg. B learning rate !/, mini-batch size B, number of episodes F., number of epochs M, number of time steps per episode T, number of gradient steps N, regularization coefficient B. Polyak coefficient initial diminishing constant A, discount factor 7, exploratory policy greediness c, temperature
2: // Stage I: Unsupervised learning phase
3: D <-
4: for epoch L....... M do
5: for episode i = ... E do
6: Sample
7: Observe initial state $1
8: fort L,. do
9:
Select d; c greedy with respect to Fo(st, 0, =)
10:
Observe reward (s() and next state 8141
11:
D + DU {8t, at, (8t), 8441)
12:
end for
13:
end for
for = 1,.-,..... N do
15:
Sample minibatch {18, aj (8,), 844)))+1))1 jez C D of size |J|= B
Sample minibatch { HeJ D of size [J| B
Sample minibatch {8/))ee iid of size |J| B
17:
iid Sample minibatch {E)]je) of size |||
For every j E J, set TT (1185+1) = softmax (Fa- (8)+1:;;2)) 2///T) (/))
20:
w( -
Σ
Fo(s), (s( B -2 Σ Tsy (s)41, B TB. (5kk
2B²
j,kEJ
aEA
1
Folsj, aj. 23) (s))))
B
11
((8)) Fo(s;, Bu(ss) Σ
( B (8;)
B
JEJ
aEA
Σ J.kE
Bu(s;) Ba BL Ba (5x) B
(s (B(())))(())))))
B
JEJ
Update 0 and via one step of Adam on Sample minibatch {(8,, (3;), 8541, (5 size |J| B from D
22:
23:
(x(A) + 2B Eee [($]+1 8,) (15+1 - (r)(s;))²
24:
Update À via one step of Adam on LA end for
25:
26:
0 + +a
+ (1 a}Φ
27: w t aw + (1 a)w
28: end for
29: // Stage 2: Exploitation phase for a single episode with initial reward (s
30: zn < -Ees n(s (B)))
31: Observe initial state 81
32: for =,...., do
33:
+ arggax Fss 2R) T zR
Observe reward (s) and next state 8
35: 2R f- Esss
end for
27

---

71
112
50
Figure F.1: A simple grid and several policies.
Noise = 0.0
Noise = 0.2
Bellman Error
0.0
0.25
1.0
0.5
0.75
0.0
20
60
20
40
Iteration
Iteration
Figure F3: Convergence of dynamic programming FourRooms with and without stochastic transitions.
225
20.0
Return
17.5
125
100
1.0
0.5
0.0
1.00
Figure F.4: GPI with noisy transitions in FourRooms. To verify that performance was maintained even with stochastic transitions, we added 20% probability that a given action would result in a random transition to neighboring state. Results are consistent with Fig. 5.3, indicating the having the correct value of produces better performance.
G Advantage of the Correct A
Importantly, for GPE using the AR to work in this setting, the agent must either learn or be provided with the updated reward vector rx after each step/encounter with a rewarded state. This is because the XR is forward-looking in that it measures the (diminished) expected occupancies of states in the future without an explicit mechanism for remembering previous visits. For simplicity in this case, we provide this vector to the agent at each step step-though if we view such a multitask agent as simply as module carrying out the directives of a higher-level module or policy within a hierarchical framework as in, e.g., Feudal RL [39], the explicit provision of reward information is not unrealistic. Regardless, natural question in this case is whether there is actually any value in using the XR with the corret value of X in this
28

---

M
M
M
1.0
5g
5o
So
So
① 1.0
+
f12
113
1.0
1.0
5g
Sg
Sg
0.8
5o
① 0.88
+ fT2
+ 0.8
0.8
0.6
So
So
50
+ 711
+ f12
ft3
0.4
0.4
0.4
5gg
5g
5g
0.4
5o
So
+ 0.0
+
1T>
5g
0.2
5o
So
So
FM
F112
FTss
5g
Sg
Sg
0.0
50
Figure F.2. Visualizing the SR, the AR and the FR. We can see that the ΦT equivalent to the SR and og is equivalent to the FR, with intermediate values of À providing smooth transition between the two.
setting: If the agent is provided with the correct reward vector, then wouldn's 't policy evaluation work with any XR?
To see that this is not the case, consider the three-state toy MDP shown in Figure Fig. G.I. where F(s1) 10, F(s2) 6, r(su) 0, X(s 0, A(s2) 1.0, and 7 0.99. At time 0, the agent starts in Sg. Performing policy evaluation with A(s1) A(22) = (i.e., with the SR) would lead the agent to go left to s1. However,
29

---

=
=
a
50
$2
a
Figure G.1: A 3-state toy environment.
the reward would then disappear, and policy evaluation on the second step would lead it to then move right to so and then 82, where it would stay for the remainder of the episode. In contrast, performing PI with the correct values of À would lead the agent to go right to $2 and stay there. In the first two timesteps, the first policy nets total reward of 10 0 10, while the second policy nets 6 + 5.94 11.94. (The remaining decisions are identical between the two policies.) This is a clear example of the benefit of having the correct X, incorrect value estimation leads to suboptimal decisions even when the correct reward vector/function is provided at each step.
The À Operator H
To learn the Λo, we would like to define (s ds') & (P s')(dd') for some base policy μ. However, this would lead to a contradiction:
(s, =
(((ss,s') (((s,ds))((ss)
/ px ds ddd) = (s, +)+ + B)
for all disjoint A, B, contradicting Lemma B.4.
For now, we describe how to learn the Xo for discrete S, in which case we have (s, s') (s, (s'),, i.e., by learning & we learn weighted version of Φ. We define the following norm, inspired by Touati et al. [25]:
2
A E
8 s'reg
amp
(l(s))
where y is any density on S. In the case of finite S, we let be the uniform density. We then minimize the Bellman error for D with respect sub/superscripts on & and & for clarity):
|| 112 (dropping the
30

---

(( = 11ppu - (I (11" + AY pH) + 7(11" + +1) O Pd p))
Esprepstiou [(p(sr,)) 5')
1(s = =
(1S5)
+ 7(1 - 2) [ss 1) = 55)
((()+)111111111 8') 21
E (esti ((5++11,))) 6') ( (11a)) ((a)))(
uls')
# E Ex (10(s); (1(51+1,
Σ ((S') 80) 1 = 8))
2E s
MMS')
2y X)lE 1 En En(s')p(8)) 8 ()) 11 s')((')) 1(st = 8)
MsS') =
#e Es 1 [10(5),
- 76(8111 710 8'))² - 23 -2 22 st]]
S141, + 2 s) (tt ( (st+(()
Note that we recover the SM loss when À 1. Also, an interesting interpretation is that when the agent can never return to its previous state (i.e., ( St) = 0), then we also recover the SM loss, regardless of X. In this way, the above loss appears to "correct" for repeated state visits so that the measure only reflects the first visit.
[( = Es Bss. [F(s), 2) B(ss) (F (T(H)),, 2) B(s))) B(()
21 (F( at, an,))) =B
B B(x)]
+ 2y - 1) [u(st) at, 2) (((4))) 2) B [sst]
(H.1)
Even though the Xo is not measure, we can use the above loss to the continuous case, pretending as though we could take the Radon-Nikodym derivative ((
s(ds')
H.1 Experimental Results with the FB Parameterization
To show that knowing the correct value of À leads to improved performance, we trained OO with the FB parameterization on the FourRooms task of Fig. 5.3, but with each episode initialized at a random start state and with two random goal states. Average per-epoch reward is shown in Fig. H.2. We tested performance with Anee Augent {0.5, 1.0], where Aune denotes the true environment diminishing rate and >agent denotes the diminishing rate that the agent uses. For the purpose of illustration, we do not allow the agent to learn X. We see in Fig. H.2 that using the correct leads to significantly increased performance. In particular, the left plot shows that assuming À 1, i.c., using the SR, in diminishing environment can lead to highly suboptimal performance.
Hyperparameters used are given in Table (notation as in Algorithm 3).
31

---

Hyperparameter
Value
100
E
100
25
B
128
T
50
?
0.99
a
0.95
"
0.001
T
200
1
Table XO-FB hyperparameters.
Atrue 0.5
Atuu 1.0
Average Expert Cor-Epoch Reward
Average Expert Epoch Reward
37.5
25.0
Aujent
Aspert
Figure H.1: Performance of the XO-FB with two values of X. Results averaged over six seeds and 10 episodes per seed. Error bars indicate standard error.
H.2 XO and the Marginal Value Theorem
To study whether the agent's behavior similar to behavior predicted by the MVT, use very simple task with constant starting state and vary the distance between rewards (see Fig. H.l(a)). When an agent is in a reward state, we define an MVT-optimal leaving time as follows (similar to that of [8] but accounting for the non-stationarity of the reward).
Let R denote the average per-episode reward received by trained agent, r(s:) denote the reward received at time in a given episode, R3 = E r(s.) denote the total reward received until time in the episode, and let T be episode length. Then, on average, the agent should leave its current reward state at time if the next reward that it would receive by staying in 8t, i.e., Ar(st). is less than
R - R
T
In other words, the agent should lcave a reward state when its incoming reward falls below the diminished average per-step reward of the environment. We compute R by averaging reward received by trained agent over many episodes.
Previous studies have trained agents that assume stationary reward to perform foraging tasks, even though the reward in these tasks is non-stationary. These agents can still achieve good performance and MVT-like behavior [8]. However, because
32

---

1.0
THM PARCS Ralative Med
31000
Figure H.2: Analysis of MVT-like behavior of XO-FB. a) Three environments with equal start state and structure but different distances between reward states. b) Difference between the agent's 's leave times and MVT- MVT-predicted leave times for y 0.99, with discounting taken into account. The agent on average behaves similar to the discounted MVT. e) Difference between the agent's 's leave times and MVT-predicted leave times for 1.0, i.e., with discounting taken into account. The agent average behaves similar to the MVT. Results for (b) and (c) are averaged over three seeds. Error bars indicate standard error.
they target the standard RL objective
Em
k ((((++))
8t =
which requires y for convergence, optimal behavior is recovered only with respect to the discounted MVT, in which R (and in our case, R2) weights rewards by powers of 7 [8].
In Fig. H.1(b-c) we perform a similar analysis to that of [8] and show that, on average over multiple distances between rewards, XO-FB performs similarly to the discounted MVT for y = 0.99 and the standard MVT for y = 1.0. An advantage of the Oo is that it is finite for 7 1.0 provided that X < 1. Hence, opposed to previous work, we can recover the standard MVT without the need to adjust for discounting.
Hyperparameters used are given in Table (notation as in Algorithm 3).
I SAC
Mitigating Value Overestimation One well-known challenge in deep RL is that the use of function approximation to compute values is prone to overestimation. Standard approaches to mitigate this issue typically do so by using two value functions and either taking the minimum min, Q (s, a) to form the Bellman target for given (s, a) pair [40] or combining them in other ways [41]. However, creating multiple networks is expensive in both computation and memory. Instead, we hypothesized that it might be possible to address this issue by using X-based values. To test this idea, we modified the Soft Actor-Critic [SAC; 42] algorithm to compute XFsebased values by augmenting the soft value target TooftQ r, + (eVsoft((e)))) where given by the expression
E PaLINa(11a QQ50111, a441) + (A ( at) O ( at+1))
a log T (att1 (4t1)
33

---

A derivation as well as pseudocode for the modified loss is provided in Appendix E.5. Observe that for À = 1, we recover the standard SAC value target, corresponding to an assumed stationary reward. We apply this modified SAC algorithm, which we term X-SAC to feature-based Mujoco continuous control tasks within OpenAl Gym [43]. We found that concatenating the raw state and action observations Φi = [st, a. and normalizing them to [0, 1] make effective regressors to the reward. That is, we compute base features as
0f 3 mink, 3
max; Φ ming of
where b indexes (st, at) within batch. Let X E [0, Bxx be the concatenated matrix of features for a batch, where dim(S) + dim(A). Then,
w, = (X X
where here denotes the vector of rewards from the batch. In addition to using fixed À value, ideally we'd like an agent to adaptively update À to achieve the best balance of optimism and pessimism in its value estimates, Following [44], we frame this decision as a multi-armed bandit problem, discretizing À into three possible values {0, 0.5, 1.0] representing the arms of the bandit. At the start of each episode, random value of X is sampled from these arms and used in the value function update. The probability of each arm is updated using the Exponentially Weighted Average Forecasting algorithm [45], which modulates the probabilities in proportion to a feedback score. As in [44], we use the difference in cumulative (undiscounted) reward between the current episode and the previous one l - 1 as this feedback signal: R₂ R That is, the probability of selecting a given value of increases if performance is improving and decreases if it's decreasing. We use identical settings for the bandit algorithm as in [44]. We call this variant X-SAC.
We plot the results for SAC with two critics (as is standard), SAC with one critic, SAC with single critic trained with XF-based values ("r-SAC" denotes SAC trained with fixed À = 2)), and X-SAC trained on the HalfCheetah-v2 and Hopper-v2 Mujoco environments. All experiments were repeated over eight random seeds. HalfCheetah-v2 was found by [44] to support "optimistic" value estimates in that even without pessimism to reduce overestimation was possible to perform well. Consistent with this, we found that single-critic SAC matched the performance of standard SAC, did 1-SAC (which amounts training a standard value function with the auxiliary task of SF prediction). Fixing lower values of À performed poorly, indicating that over-pess is harmful in this environment. However, X-SAC eventually manages to learn to set À = and matches the final performance of the best fixed algorithms. Similarly, in [44] it was observed that strong performance in Hopper-v2 was associated with pessimistic value estimates. Consistent with this, X-SAC learns to select lower values of X, again matching the performance of SAC while only requiring one critic and significantly reducing the required FLOPS Fig. I.2. We consider these results to be very preliminary, and hope to perform more experiments on other environments. We also believe A-SAC could be improved by using the difference between the current episode's total reward and the average of the total rewards from previous episodes Re (E - 1) Di=l R. as a more
34

---

stable feedback signal for the bandit. There is also non-stationarity in the base features due to the per-batch normalization, which could also likely be improved. Hyperparameters are described in Table 2.
Hyperparameter
Value
Collection Steps
1000
Random Action Steps
10000
Network Hidden Layers
256:256
Learning Rate
3 10-1
Optimizer
Adam
Replay Buffer Size
1 x 106
Action Limit
[-1,1] 1]
Exponential Moving Avg. Parameters
5 x 10-3
(Critic Englate:Eviirommen Step) Ratio
(Policy Update:Fnnvirommntt Step) Ratio
Has Target Policy?
No
Expected Entropy Target
(.m(
Policy Log- Variance Limits
20, 2]
Bandit Learning Rate
0.1 {0, 0.0,
Options*
Table 2: Hyperparameters for SAC Mujoco experiments. *Only applicable to X-SAC
HaifCheetah-v2
Hopper-v2
À Choice
12500
2500
10000
Test Test Reward
2000
0.6
1500
2500
1000
500
0.2
HallCheetah-x22
Hopper-v2
-2500
Timestops
Timestops
T Timesteps
SAC critics)
0-SAC
0.5-SAC
Figure I.1: X-SAC (1 critic) matches the performance of SAC (2 critics) by adapting online.
HalfCheetah-v2
Hopper-v2
2600
13800
SNC crtess)
Final Test Reward
2400
13600
2200
13400
1800
FLOPS
FLOPS
Figure L.2: X-SAC matches performance while saving in computational cost.
35

---

Replenishing Rewards
We list below a few candidate reward replenishment schemes, which are visualized in Fig. J.1.
Time elapsed rewards
r(s,t) = F(s);
where m(s, 1) is the time elapsed since the last visit to state 8:
m(s t) _ t max
Sttj stts 8, 0 05 St 1t
Due to the max term in mss t), the corresponding representation
Amle.))/mls.)) 1 (St1k= s') =
E
A=
does not admit closed-fomm recursion. However, empirically tested version of this type of reward with Qx-leaninng in the TwoRooms environment, modified so that the exponent on À is n n(s, t0.()))))))) (m,1mms ()). This was done so that reward replenishes at a slow rate, reducing the deviation from the standard diminishing setting. Episode length was capped at H 100. All other settings are identical to the Qx experiment described in Appendix E. Results are presented in Fig. J.2 and GIF is included in the supplementary material.
Eligibility trace rewards
Σ x 1 =s) F(s):
r(s,t)= = 1
1-
(1
Xa)
where A A E [0, 1] are diminishment and replenishment constants, respectively.
Denoting the corresponding representation by S2T i.c.,
(1- (1-
(8, = E
Ao) Σ k-
1 =sy) 1 1
= s') =
K=
1=0
we obtain the following recursion:
(2, = (s = s')(Aa 77r(1 - Ad) A sn st (-16) M
+ y E (11)) (1s)) (51+1,
where M,,,, denotes the successor representation of π with discount factor 7Ap. This representation could be learned by alternating TD learning between Ω and M We leave this to future work.
Total time rewards
r(s,t) A (ssf)) (p(st))(((s),
where Xa, Ar E [0, 1] are diminishment and replenishment constants, respectively. The corresponding representation is
p. (T = E
1 X
ne Ak- nils (8E1k s') 81 = 8
36

---

Time Elapsed
Ellgibllity Trace
50
130
10
Roleara
Reward
14
20
54
30
Time
Time
Time
Figure J.1: Visualizing three different replenishment schemes. For all schemes, F(s) 1 and visits to are at t = 2, 5. (Left) The time elapsed reward with À 0.5; (Middle) The eligibility trace reward with Ar Xd 0.5; (Right) The total time reward with Xa 0.5, 2, 0.9.
Q^ Time Elapsed Replenishing Rewards
25
Return
10
5
=
0.5
1.0
100
200
300
400
500
Episodes
Figure J.2: Performance on TwoRooms with replenishing rewards. Return is averaged over five runs, with shading indicating one unit of standard error.
which satisfies the following recursion:
pa (s, = (ss = 5') + 7 = s')+1 +1
- 1(ss s1) 1 s)) + 1(s s Pa (s (st1,
While neither the reward nor the representation are guaranteed be finite, p could be learned via TD updates capped at suitable upper bound.
K
vs. Y
We now briefly discuss the interaction between the temporal discount factor y commonly used in RL and the diminishing utility rate X. The key distinction between the two is that all rewards decay in value every time step with respect to y, regardless of whether state is visited or not. With A, however, decay is specific to each state (or (s, a) pair) and only occurs when the agent visits that state. Thus, decays reward in a global manner which is independent of the agent's behavior, and À decays reward in a local manner which dependent on the agent's behavior. In combination, they have the beneficial effect of accelerating convergence in dynamic programming (Fig. F.3). This indicates the potential for the use of higher discount factors in practice, as paired with a decay factor X, similar (or faster) convergence rates could be observed even as agents are able to act with a longer effective temporal horizon.
37

---

L Compute Resources
The XF-based experiments shown were run on a single NVIDIA GeForce GTX 1080 GPU. The recurrent A2C experiments took roughly 30 minutes, base feature learning for policy composition took approximately 45 minutes, XF learning for policy composition took approximately 10 hours, and the SAC experiments took approximately hours per run. The XF training required roughly 30GB of memory due to the size of the dataset. All experiments in Section 6 and Appendix H were run on a single RTX5000 GPU and each training and evaluation run took about 30 minutes. All other experiments were run on 2020 MacBook Air laptop 1.1 GHz Quad-Core Intel Core i5 CPU and took less than one hour to train.
M Broader Impact Statement
We consider this work be primarily of theoretical nature pertaining sequential decision-making primarily in the context of natural intelligence. While it may have applications for more efficient training of artificial RL agents, it is hard to predict long-term societal impacts.
38


Reviews:
[{"review_id": "xsMwnz6oYu", "reviewer": "Reviewer_7pTK", "summary": "The authors explore diminishing marginal utility in the context of reinforcement learning. Specifically, they study reinforcement learning when reward obtained at a state diminishes with each visit to that state (following a particular mathematical form). The authors show that, under such diminishing rewards, agent policies can be evaluated by using a novel state representation introduced in the paper. This representation generalises earlier state representations in the literature, such as the successor representation.", "questions": "I do not have any particular questions but, as I noted above, I do not find the paper particularly well-motivated in its present form. Why and how is diminishing marginal utility, in the form studied in the paper, useful for reinforcement learning agents? \n\nI have read the rebuttal and thank the authors for their additional thoughts on the subject.", "limitations": "Yes, the limitations have been discussed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originally: The concept of diminishing rewards is well known and explored in other fields. The paper presents new results within the context of reinforcement learning. \n\nQuality: The authors provide a thoughtful and useful analysis under the assumptions they make. This includes the discussion of how the proposed representation relates to existing state representations in RL. The recursion they have identified is interesting and potentially useful.\n\nThe paper is well structured. The writing is clear.", "weaknesses": "The problem is not strongly motivated: why is studying diminishing marginal utility within the context of RL, with the specific assumptions made in this paper, important? Reasoning that goes beyond \"people have diminishing marginal utility\" would be useful. Why and how is this concept, in the form studied in the paper, useful for reinforcement learning agents? \n\nThe writing should clearly distinguish between two cases: (1) an objective quantity in the environment (e.g., available food) gets smaller, (2) the subjective value attached by the agent to an outside quantity gets smaller (e.g., the second ice cream cone does not taste as good). \n\nIn Section 6, the authors write that naturalistic environments often exhibit diminishing reward and give the example of foraging with diminishing food patches. I see a discrepancy between this example and the problem addressed in the rest of the paper. In the foraging example, the available food decreases because other agents consume it, and the amount of reduction should generally be a function of the time that passed between visits to the site. In contrast, the core assumption in the paper is that the reduction in reward in a given state is a function of how many times the agent itself visited that particular state, regardless of how much time passed since its last visit to that state. \n\nAdditional comments: Forward-backward representation should be explained in Section 2 (Preliminaries).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your detailed review and careful consideration of our paper! We are glad that you appreciated the originality of the work, found it to be of high quality, well-structured, and clearly written. We aim to address your concerns below:\n\n- Motivation: We believe that there are two central questions that can be asked regarding the usefulness of DMU and RL. (1) Does incorporating DMU into standard RL paradigms/agents improve their performance on problems of interest in machine learning? We believe the following three areas are of particular importance in answering this question: (i) faster convergence for policy evaluation with lambda < 1 (Proposition 4.1 and Attachment Figure L1), reducing value function overestimation (e.g., $\\lambda$-SAC, described in Appendix H and with additional results presented in Attachment Figures L3 and L4), (iii) greater understanding of non-stationary rewards in RL, and (iv) the set-theoretic formulation of the FR (lambda = 0)--suitable for continuous spaces–is novel. The FR facilitates navigation along shortest paths between goal states, which is of interest in robotics and other areas of machine learning. 2) Does modeling DMU within an RL framework provide new avenues of study for understanding natural behavior? This is the question for which we are most hopeful to provide an affirmative answer. The SR and FR are already of interest to neuroscientists [1,2,3] for the connections between the ways they encode spatial information and hippocampal activity. By generalizing these representations (and the problem settings for which they are suited) to DMU, we believe that the $\\lambda$R can guide both further study in this area and offer new connections to other phenomena associated with DMU (e.g., foraging). Previous work on the SR and FR has already been grounded within the formalism of RL. When neuroscientists use these representations, then, they have clarity on what it is they can formally claim regarding behavior. We also believe these results may be of interest to economists interested in DMU. Our primary goal in this area is to provide the same type of formally sound framework for DMU as has been established for other problem areas, so that the $\\lambda$R can act as a solid foundation on which to build. We view the experimental results in Section 6 as a proof of concept for the $\\lambda$R as a useful tool in this setting. We will emphasize these motivations for the work in the paper.\n\n- The Nature of Diminishment: This is indeed a very good point, and one which we should emphasize and explore in more depth in the paper! It is especially interesting because in natural foraging, both the objective depletion of resources in the environment and subjective diminishment of perceived utility play a role in behavior (e.g., an animal may leave a berry patch before all berries have been consumed if another resource’s appeal supersedes the consumption of more fruit). One especially important point that you bring up is that of utility replenishment–i.e., that utility should be a function of the amount of time since the agent has last visited a particular state rather than just how many times it’s visited that state overall. As we note in our discussion of limitations in Section 7, the current framework we present models this replenishment process as being linked to the episodic nature of the environment, but in an ideal case, replenishment would take place gradually in a continuing environment. We explore this case in a fair amount of detail in Appendix I, providing some possible representations in line with such a framework as well as preliminary simulations. We agree this is the natural next step for studying DMU and foraging within RL! \n\n- FB Representation Background: Yes, that is a good point–we will make the necessary adjustments.\n\nThank you very much once again for the helpful feedback! We hope we have adequately addressed your concerns. If not, we’re more than happy to continue the discussion!\n\n[1] https://www.nature.com/articles/nn.4650\n\n[2] https://www.jneurosci.org/content/38/33/7193 \n\n[3] https://www.cell.com/neuron/pdf/S0896-6273(23)00230-1.pdf \n\n"}, {"review_id": "PRY8bOu1Kj", "reviewer": "Reviewer_jk48", "summary": "The paper considers the case of non-Markov rewards, where rewards at states decay over time (motivated by the diminishing marginal utility phenomenon). While successor representations (SR) have been used in standard MDPs, the decaying of rewards means that they cannot be used here, and so the paper develops a \"successor-like\" representation that can handle this case, which generalises SR and similar other representations. Theoretical results ensure that these representations can be learned using a Bellman backup (as in the case of SR), which can be used to evaluate and improve policies. Further extensions to the function approximation case (where features are required) are also demonstrated, and bounds related to composition (similar to the successor features bound) are presented. Results in tabular and pixel-based domains showcase how the representation can be used with model-free RL algorithms to solve this class of tasks, where taking into account the decaying rewards (instead of assuming rewards are stationary) is shown to be beneficial. \n", "questions": "On line 291, it states that there is not yet a measure-theoretic version of the representation. Having said that, while there may not be strict theoretical justification, would there be any issue in practice with applying this approach to continuous state/action domains?\n\nI'm slightly confused about the composition in the Four Rooms domain. There are four base policies but three goals. Is each policy to reach the centre of each room (and in three of the rooms, a goal exists)?", "limitations": "The limitations are well described in Section 7, including many that I would not have even considered. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper considers an interesting extension to the typical MDPs considered in reinforcement learning. The motivation regarding natural behaviours is an interesting one, but I think on a more fundamental level, any approach that attacks decision problems in RL that aren't Markovian in some way is of interest to the field.\n\nThis paper packs a lot into it - it develops many different ideas, ranging from the $\\lambda$ representation in the tabular case to the function approximation \"feature\" case. It also shows how these can be used to evaluate and learn policies, compose existing value functions, and be applied in a deep RL context. \n\nAside from the extensive theory presented, the paper also does a good job of relating the approach to existing work in successor representations and demonstrating under what conditions that approach here simplifies to existing work. This helps to unify these otherwise disparate methods.", "weaknesses": "As presented, the experiments demonstrate that the method can be applied in both tabular and function approximation settings, that using the right value for $\\lambda$ is necessary to achieve the best performance, and that composition can be achieved just as in the SR/SF case. While these experiments confirm the utility of the representation and go hand in hand with the theory, it would have been nice to see the performance of the approach compared to different baselines. This could answer questions like how it compares to SR/FR/SM under a variety of conditions: when $\\lambda < 1$ or when $\\lambda =1$ or when states are only ever visited once (because it's continuous/very high dimensional).\n\nThe pixel domain was helpful to showcase that the approach can be applied in high-dimensional domains like this, but it would have been nice to see a more \"compositional\" domain as has been demonstrated in previous work to investigate composition (e.g. the domains where the agent collects objects of different colours/shapes with different priorities) [1,2,3]. \n\nI realise space is an issue, but it would have been helpful to see a very brief discussion on related work that focuses on non-Markov rewards (outside of the context of successor representations). e.g. [4]\n\nMinor:\n\nSome of the drawbacks of composition with SF are also inherited by this approach. For instance, the bound in Theorem 5.1 is pretty loose (see [2]), but I understand the point is to show it's the same bound with an extra term that is positive only when the $\\lambda$ estimates are wrong, so it's not very important\n\nWhile I found the motivation for considering these forms of non-Markovian rewards interesting, it may be the case that many RL practitioners who are not interested in modelling animal/real-world behaviour will not adopt this since for many practical applications (or popular benchmarks like Atari or MuJoCo), such a formulation is not natural. \n\nSome of the contractions used made it hard to parse sentences (e.g. L74). I'd recommend replacing these with the full words.\n\n[1] Barreto, André, et al. \"Fast reinforcement learning with generalised policy updates.\" Proceedings of the National Academy of Sciences 117.48 (2020): 30079-30087.\n\n[2] Nangue Tasse, Geraud, et al. \"Generalisation in lifelong reinforcement learning through logical composition.\" International Conference on Learning Representations. 2022.\n\n[3] Alver, Safa, and Doina Precup. \"Constructing a Good Behavior Basis for Transfer using Generalised Policy Updates.\" International Conference on Learning Representations. 2022.\n\n[4] Gaon, Maor, and Ronen Brafman. \"Reinforcement learning with non-markovian rewards.\" Proceedings of the AAAI conference on artificial intelligence. Vol. 34. No. 04. 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your detailed feedback, and we’re glad that you liked the paper! We hope to address your concerns below: \n\n- Baselines: We completely agree, and identifying appropriate baselines was a challenge for us, simply because (as far as we know) DMU has not been addressed in an RL context before. Experiments where we set $\\lambda = 1$ are meant to compare our approach to standard RL, but if you have any specific suggestions for additional baselines, we’d very much appreciate any pointers!\n\n- Compositional Domains: Thank you for pointing this out! We agree, one direction that we’d like to push in the future is to explore both more compositional tasks (e.g., environments where the agent has to pick up a key and proceed to a locked door) as well as novel ways to use the $\\lambda$R to compose policies (e.g., the FR can be used to perform a form of shortest path planning over a set of base policies [1]). \n\n- Discussion of Non-Markov Rewards: Yes, we certainly agree–thank you for pointing this out! We will add a section in the Appendix with a more detailed discussion. \n\n- Looseness of the Bound: Yes, indeed it is a rather loose bound, but as you note, our hope was to draw the connection with the previously established bound for GPI :). We have slightly tightened the bound recently, replacing the $L_1(r)$ term with $L_\\infty(r)$.\n\n- Interest to the Broader Community: We believe this is an important point, and one that we’ve given careful consideration. It is certainly the case that the most natural area of application for DMU is in studying natural behavior. In that lens, we hope this paper establishes a formal framework for DMU within RL that can be applied in the future more extensively within this context. However, we also believe that addressing non-stationarity in environmental rewards is also an important topic for the design of artificial agents. If agents are designed with the implicit assumption that rewards are ever-present, it could lead to value overestimation and suboptimal behavior. While this theme is present in all of our experiments, we’d like to highlight $\\lambda$-SAC (Appendix H and Attachment Figures L3 and L4) as an interesting example case. \n\n- Contractions: Thank you! We will update the paper accordingly. \n\n- Q1 (Measure/Set-Theoretic Formulation in Practice): Good question! The only continuous-domain experiments that we performed were using $lambda$-features (Fig. 5.4), but we do not anticipate any issues with using the $\\lambda$-operator in practice, even though it isn’t entirely theoretically justified. Continuous-domain experiments with the $\\lambda$O are an exciting direction for future work that we intend to look at, especially because we could adjust FR planning (as mentioned above and detailed in [1]) to be compatible with the first-occupancy operator (i.e. $\\lambda$O with $\\lambda=0$) to construct a continuous-domain planning scheme. We can mention this in the paper.\n\n- Q2 (FourRooms): Yes, that’s correct! We’ll make it more clear in the text. \n\n\n[1] https://openreview.net/forum?id=JBAZe2yN6Ub\n"}, {"review_id": "R9M1FYfokS", "reviewer": "Reviewer_2Ebj", "summary": "This paper studies the phenomenon of diminishing marginal utility in RL, where the a state-based reward r(s) decays as the agent visits it more often. To solve this problem setting, this work introduced a novel state representation, named the λ representation (λR). The author showed that we fail to define a Bellman equation in the λR setting and introduced a recursive relationship.", "questions": "- Since exploration is the natural setting for  λR, why there is no comparisons to different exploration baselines, i.e., RND, state entropy, pseudo-count ...?\n- I am curious of how will the λR perform in more complex pixel-based tasks.", "limitations": "- The application of the proposed method seems to be narrow, and there is a lack of comparisons to SOTA exploration baselines.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The probelm setting of DMU in RL is interesting.\n- The writing is clear and easy to follow.\n- The paper provides theoretical analysis of the λ representations.", "weaknesses": "- Typos and grammar error: line 156 (when when ...), line 156 (Instead of keeping ...), line 177 (The mechanism)\n- Missing baselines: the current experiments do not compare to other SOTA exploration baselines with intrinsic rewards. So it's hard to evaluate the effectiveness of the proposed methods.\n- Experiments for continuous control is only conducted on Halfcheetah, a toy example that most RL baselines can easily solve.  Experiments on more complex tasks, i.e., humanoid-run, hopper-hop, acrobot-swingup, fish-swim, from dm_control would be more persuasive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you very much for your detailed review and feedback! We are glad that you appreciated the problem setting, found the writing to be clear, and valued our theoretical analysis. We aim to address your concerns below: \n\n- Typos: Thank you for pointing these out! We’ll correct them in the updated paper. \n\n- Baselines: We think there may be a misunderstanding of our experiments. While both the FR and the SR have been used as exploration bonuses, we don’t examine the $\\lambda$R’s application to exploration in this paper, but rather its use for accurate policy evaluation, control, and composition in settings where agents experience diminishing marginal utility from repeated exposure to rewarding stimuli. We completely agree that had we used the $\\lambda$R as an exploration bonus that these would be valuable baselines to compare against, but they aren’t as applicable in the experimental settings we do consider, which are primarily focused on establishing the usefulness of the $\\lambda$R for supporting policy evaluation, control, and composition for problems with DMU. We completely agree that exploration is another interesting application area–we believe it is out of the scope of the current work, but would absolutely like to pursue it in the future. \n\n- Continuous Control Domains: Yes, this is a good point! We have added Hopper-v2 as an additional domain (Attachment Figures L3 and L4) and found that $\\lambda$-SAC again matches SAC’s performance while being much less computationally demanding. Please see the general response above for a more detailed description.\n\n- Q1 (Exploration): See response above re: baselines. We are primarily concerned with problem settings involving diminishing marginal utility and show that in order to do RL using dynamic programming in such settings, the $\\lambda$R is required. We didn’t offer comparisons to exploration baselines because we didn’t perform any experiments focused on exploration, but rather focused on establishing the $\\lambda$R as useful for evaluation, control, and composition for settings with DMU.\n\n- Q2 (Further Pixel-Based Experiments): We certainly agree! We think the fact that the $\\lambda$F was effective at facilitating pixel-based navigation is an encouraging first step, as successfully performing with 128 x 128 RGB observations is a good sign that the approach scales well. We hope to pursue further experimentation in the future! \n\n\nThank you very much once again for your comments and questions! We hope our response has provided clarification and addressed your concerns. If that’s the case, we would greatly appreciate it if you would consider raising your score. If not, we would be more than happy to continue discussion!\n\n\n"}, {"review_id": "bbAEosvSdc", "reviewer": "Reviewer_GnnZ", "summary": "The authors study a problem of DMU in RL, e.g. if visiting the same state may result in a smaller reward. They introduce a \\lambda R representation that considers a particular form of diminishing rewards and provides convergence results. They extend it to continuous domains and perform preliminary experimental studies.", "questions": "Can you mention the assumptions behind Proposition 4.1? It is counterintuitive for me. $\\lambda=0$ (FR) should be a much harder problem as compared to $\\lamda=1$ (SR). However, error bound is better for $\\lambda=0$?\n\nI would suggest authors add an algorithm block for a high-level understanding of the algorithm.\n\nI did not fully understand the experiments.\n- Are your environments deterministic?\n- In fig 5.2 b), why does \\lambda =1 perform well? I would expect it to stay at one of the blue cells forever since it assumes that it gets the same reward always, and on the plot, it shall have an exponential decay.\n- Is there an assumption of finite state-action spaces, or can you extend it to continuous dynamics as well?\n\nTypos: line 268 the the, 156 when when,", "limitations": "no societal negative impact", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Typically, such problems having non-markovian rewards are hard, but by having assumptions of \\lamda rewards representation, the authors are able to recover some near-optimality convergence results.\n\n- Extension to the continuous domain", "weaknesses": "I think the assumption of decay with \\lambda is quite strong (very restrictive), and it's hard to imagine practical problems satisfying this. E.g. Authors have motivated the problem with examples of ice cream, finance and neuroscience, but it is hard(impossible) that a decay of visiting again is exactly a factor of 0<\\lambda<1. Arguably a more interesting setting to consider would be:\n- \\lambda R for is a lower or upper bound on the true rewards that one may earn after visiting the states multiple times, i.e., the reward will decay at least by that amount or at max by a certain amount if visited again. \n- I think another interesting setting will be if \\lambda changes with state visitation frequency and thus needs to be estimated.\n\nBoth these settings will consider quite a large part of DMU-type problems.\n\nLemma 3.1 is perhaps not the right way to impose Impossibility/hardness results, and the proof needs to be rigorous. I think it can be interpreted in multiple ways, and also, the proof considers a specific format in which the reward is written and a specific definition of the value function. For the lemma statement to be true, the proof shall eradicate all the ways in which value function and rewards can be defined. For e.g. consider a reward function $r(s) = F(\\tau \\cup s) - F(\\tau)$, where $F(\\tau)$ is a general set function, $\\tau$ is the agent's trajectory. $F(\\tau)$ denotes the cumulative value earned for a trajectory $\\tau$ as per the reward function (eqn 3.1). This can be written as Bellman recursion with $r(s_i) = F(s_{0:i}) - F(s_{0:i-1})$ (Imagine $\\tau = \\{s_0, s_1, s_2 ..... \\}$). I would suggest not having it as a lemma but just in the text mentioning it is difficult to solve it with the Bellman equation directly.\n\nExperiments: I think authors should consider more diverse reward functions, e.g., coverage functions, experiment design, and informative path planning objectives. For e.g., consider a camera-type measurement model, and if we observe the same region again, I will not get much information (there can be some noise model to establish \\lamda type of decay). Similarly, in experiment design/ informative planning objectives: Information gained on the nearby states will decrease if I gather/collect measurements at my current location.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your detailed review and helpful feedback! We are glad that you appreciated the convergence results and the extension of the $\\lambda$R to continuous domains. We aim to address your concerns below: \n\n- Assumptions: This is a thought-provoking point! We completely agree that examining modifications of the specific DMU assumptions represent an interesting avenue for future work. We’d like to note that the neuroscience, foraging, and economics works studying DMU that we cite in the paper all use a fixed value for the diminishing rate (and in fact use the equivalent of a fixed value across states, an assumption which our framework can relax). Previous work [1,2] that applies standard RL methods to DMU problems also assumes a constant diminishment rate. In our work we show that it is possible to construct principled and tractable dynamic programming algorithms to solve DMU problems in RL and use these algorithms to achieve better performance than standard methods. Therefore, we believe that our work is an important first step towards RL problems with more general reward diminishment schemes, both theoretically and empirically. Additionally, we present the continuous control results in Attachment Figures L3 and L4 as an interesting example case where the “true” $\\lambda$ is actually 1, but picking $\\lambda < 1$ can actually result in a performance gain due to the tendency towards value overestimation in deep RL. In this case, $\\lambda$-SAC estimates $\\lambda$ online to maximize performance. More detail can be found in the general response above and in Appendix H. \n\n\n- Lemma 3.1: Thank you for alerting us to this! We agree that the lemma is unclear in its current form. Perhaps a more rigorous statement would be the following: “There does not exist a contractive operator $T^\\pi$  that has fixed point $V^\\pi$ and does not depend on $\\Phi_\\lambda^\\pi$.” A sketch of the proof would be as follows: By the Banach fixed point theorem, $V^\\pi$ is a unique fixed point, and so $T^\\pi$ must be of the form given in the proof of Lemma 3.1. Hence, $T^\\pi$ has a term containing $\\Phi_\\lambda^\\pi$. For $T^\\pi$ to not depend on $\\Phi_\\lambda^\\pi$, we would need a function $f$ such that $\\Phi_\\lambda^\\pi = f(\\bar r, V^\\pi)$. However, because $V^\\pi = \\bar r^\\top \\Phi_\\lambda^\\pi$, such a function cannot exist. In this sense the $\\lambda$R is necessary for this problem setting, unlike the Markovian case in which using the SR is optional. We are happy to rewrite this as a remark rather than a lemma if you think it would be more appropriate. Regarding your concerns about our definition of reward and value being too narrow–as noted above, (1) our approach is consistent with studies of DMU in other fields, and (2) we believe theoretical results for specific settings are useful starting points for more general results in the future. Your suggestion to define reward as a difference of set-valued functions is interesting–thank you! However, can you please clarify your final point about writing a Bellman recursion for $r(s_i) = F(s_{0:i}) - F(s_{0:i-1})$?\n\n- Alternative Reward Functions: Thank you, this is an interesting point! In this paper, we’re primarily concerned with a specific setting as a starting point for studying DMU using RL. We believe a detailed examination of these alternative reward functions is beyond the scope of our work, but they represent an exciting opportunity for future study, and we will add a discussion to the paper. \n\n- Q1 (Proposition 4.1): Thank you for highlighting this! We agree that one of the most interesting, and perhaps broadly valuable, aspects of the $\\lambda$R (and by extension, for $\\lambda=0$, the FR) is that convergence is faster for lower values of $\\lambda$. The assumptions are exactly those stated in the text, and the proofs can be found in Appendix B. We also show empirically in Figure E3/Attachment Figure L1 that dynamic programming indeed converges faster for lower $\\lambda$. Intuitively, a lower $\\lambda$ can be seen to speed convergence in a similar way to a lower discount factor in standard RL. \n\n- Q2 (Algorithm Block): Good point! Indeed, there are algorithm pseudocode blocks for $Q_\\lambda$-Learning (Algorithm 1), Fitted $Q_\\lambda$ Iteration (Algorithm 2), and $\\lambda$O Forward-Backward Learning (Algorithm 3) in the Appendix, but can certainly add more details. \n\n- Q3 (Experiments): (i)  While the experiments presented in the initial submission all use deterministic environments, the theoretical analysis does not depend on this property, and we present new policy evaluation and composition results demonstrating that strong performance is maintained with stochastic dynamics in Attachment Figures L1 and L2. (ii) This is a subtle detail! It’s because, as we note in the conclusion, all of the representations we consider are prospective in nature and therefore don’t naturally account for previous visits. The $\\lambda=1$ policy moves on because the amount of reward is decreasing–that is, $\\vec r(s)$ is going down. The source of its suboptimality is that at every step using $\\lambda = 1$ assumes that the current level of reward will persist, while using the correct $\\lambda$ accounts for future decay. (iii) The $\\lambda$F and $\\lambda$O representations we introduce (Section 4.1, Appendix H, and Figures 5.4 and H.1) are designed to handle the continuous case. \n\n- Typos: Thank you for pointing these out! We’ll fix them.  \n\nThank you very much once again for your comments and questions! We will certainly integrate your suggestions into the paper. We hope our response has provided clarification and addressed your concerns. If that’s the case, we would greatly appreciate it if you would consider raising your score. If not, we would be more than happy to continue discussion!\n\n[1] https://openreview.net/forum?id=a0T3nOP9sB \n\n[2] https://proceedings.neurips.cc/paper/2020/hash/da97f65bd113e490a5fab20c4a69f586-Abstract.html \n\n\n"}]

Author rebuttal:

We’d like to sincerely thank all reviewers for their time and for their helpful feedback and suggestions for the paper. We believe that incorporating this feedback will make the paper stronger, and we look forward to a constructive discussion! We are glad that reviewers consistently appreciated the originality of the paper, the theoretical analysis provided, the consideration of both tabular and deep RL instantiations of our approach, and found the writing to be clear. Overall, we view this as exploratory work—DMU is a prominent factor in human and animal decision-making. RL is likely the most dominant approach to modeling sequential decision-making in machine learning and a common approach for doing so in neuroscience and other fields, yet until now (as far as we are aware) has not contained a framework that accounts for (and formally characterizes) DMU. The goal of this paper is to provide a basic approach for closing this gap.

We have also attached a pdf file with several additional experimental results:

- In Figure L1, we demonstrate the accelerated convergence of dynamic programming—in this case, policy evaluation—for lower values of $\lambda$ both with and without stochastic transition dynamics in the environment (in this case, FourRooms). This result augments the one in Figure E3.

- In Figure L2, we show that the usefulness of the $\lambda$R for supporting policy composition is maintained with stochastic transition dynamics.

- In Figure L3, we augment the results in Figure H1 in the following ways: (1) we add five additional random seeds to our HalfCheetah-v2 results, and (2) we repeat our experiment on Hopper-v2, a more challenging environment (also for eight seeds total). We found that, consistent with our initial results, $\lambda$-SAC is able to match the performance of SAC while only using a single critic (rather than two). It does this by using an online bandit approach to select values of $\lambda$ used for training the critic that result in higher reward (see Appendix H for more details). In the rightmost panel, we can see how the choice of $\lambda$ changes over the course of learning for each environment. Interestingly, in [1], it was found that strong performance in HalfCheetah-v2 is associated with optimistic/high value estimates, while strong performance in Hopper-v2 is associated with pessimistic/lower value estimates (i.e., it benefits more from the pessimistic evaluation performed when SAC takes the minimum across its two critics). Consistent with this finding, we observed that $\lambda$-SAC learns to select high values of $\lambda$ on HalfCheetah-v2, and lower values on Hopper-v2. We believe this is a promising example of a possible area of benefit that the $\lambda$R can bring even to standard RL problem settings where the reward is stationary. 

- In Figure L4, we plot the average final performance of SAC and $\lambda$-SAC versus the FLOPS used by each agent (calculated using the `flopth` Python package) as a means of emphasizing that $\lambda$-SAC is able to match SAC’s performance while substantially saving on computational cost (~32% for each task). 

We aim to address specific concerns in the responses below. In these responses, when we refer to the additional figures described above, we will name them as “Attachment Figures.” Thank you all once again!

[1] https://openreview.net/forum?id=a4WgjcLeZIn 



Meta-review: Paper looks at diminishing marginal utility and how this is factored into an MDP setting, which has been done before (as the authors state) but has novelty here with in terms of RL applications by generalizing the successor representation. Some of the experimental setups arguably could have used more baselines. The pixel-based experiments showcase the authors' method quite well. Reviewer questions mostly answered.
Decision: Accept (poster)