Reward-Aware Proto-Representations in Reinforcement Learning
Authors: Hon Tik Tse
Siddarth Chandrasekar
Marlos C. Machado
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Representation Learning
Proto-Representation
Default Representation
Successor Representation
Primary area: reinforcement_learning
TL;DR: We lay the theoretical foundation underlying the default representation, a reward-aware proto-representation, and empirically demonstrate its benefits over the successor representation in tabular settings.

Abstract: In recent years, the successor representation (SR) has attracted increasing attention in reinforcement learning (RL), and it has been used to address some of its key challenges, such as exploration, credit assignment, and generalization. The SR can be seen as representing the underlying credit assignment structure of the environment by implicitly encoding its induced transition dynamics. However, the SR is reward-agnostic. In this paper, we discuss a similar representation that also takes into account the reward dynamics of the problem. We study the default representation (DR), a recently proposed representation with limited theoretical (and empirical) analysis. Here, we lay some of the theoretical foundation underlying the DR in the tabular case by (1) deriving dynamic programming and (2) temporal-difference methods to learn the DR, (3) characterizing the basis for the vector space of the DR, and (4) formally extending the DR to the function approximation case through default features. Empirically, we analyze the benefits of the DR in many of the settings in which the SR has been applied, including (1) reward shaping, (2) option discovery, (3) exploration, and (4) transfer learning. Our results show that, compared to the SR, the DR gives rise to qualitatively different, reward-aware behaviour and quantitatively better performance in several settings.

Reward-Aware Proto-Representations in Reinforcement Learning
Hon Tik Tse
Siddarth Chandrasekar
Marlos C. Machado
University of Alberta, Alberta Machine Intelligence Institute (Amii)
*Canada CIFAR AI Chair
(hontik, siddart2, machado) Qual)ertaaaaaaaaa
Abstract
In recent years, the successor representation (SR) has attracted increasing attention in reinforcement learning (RL), and it has been used to address some of its key challenges, such as exploration, credit assignment, and generalization. The SR can be seen as representing the underlying credit assignment structure of the environment by implicitly encoding its induced transition dynamics. However, the SR is reward-agnostic. In this paper, we discuss a similar representation that also takes into account the reward dynamics of the problem. We study the default representation (DR), a recently proposed representation with limited theoretical (and empirical) analysis. Here, we lay some the theoretical foundation underlying the DR in the tabular case by (1) deriving dynamic programming and (2) temporal- difference methods to learn the DR, (3) characterizing the basis for the vector space of the DR, and (4) formally extending the DR to the function approximation case through default features. Empirically, we analyze the benefits of the DR in many of the settings in which the SR has been applied, including (1) reward shaping, (2) option discovery, (3) exploration, and (4) transfer learning. Our results show that, compared to the SR, the DR gives rise to qualitatively different, reward-aware behaviour and quantitatively better performance in several settings.
Introduction
Learning appropriate representations is a key challenge in reinforcement learning (RL). The successor representation (SR) [10], which represents states as the expected discounted number visits to successor states, is particularly promising idea. It has been shown to be effective distance metric
for reward shaping [4 49 promising inductive bias for temporally-extender exploration (21,22),
and effective representation for credit assignment (27 23), generalization (20), and zero-shot
RL 47 48 6 44 I. Furthermore, recent evidence suggests that the SR is good computational
model for explaining decision-making
28 and neural activity in the brain [41 40
The underlying idea behind the SR is incorporate the temporal aspect of the problem into the representation so that two states are considered similar if they lead to similar future outcomes. Similar ideas permeate other concepts such as proto-value functions [27 26 and slow feature refer to such temporal representations that implicitly capture the environment dynamics, with the SR the most prominent only captures the transition dynamics of the environment, and does not take rewards into account. In this paper, we demonstrate the benefits of learning a proto-representation that also takes rewards into account.
analysis [39 (53). In this work,
the term proto-representation
of them. However, the SR (and others [29] 30 12 52
One candidate for reward-aware proto-representations is the default representation (DR) Piray and Daw, in neuroscience study, applied the DR mainly to replanning tasks and explaining
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

different cognitive phenomena, such as habits and cognitive control. Nevertheless, we claim that the DR is an interesting concept beyond neuroscience, due to its relationship with the SR and its reward-awareness. In this context, many aspects of the DR have yet to be explored. There are no efficient, general algorithms for learning the DR incrementally and online with a linear cost, akin to lemporal-difference (TD) learning [4). Additionally, given the successful applications of the cigenvectors of the SR for reward shaping [54 49), and lemporally-extendee exploration [25], the eigenvectors of the DR have seen limited investigation. Finally, the behavior of the DR for a reward function not constant over non-terminal states is underexplored.
In this paper, we (1) derive methods for learning the DR using dynamic programming (DP) and temporal-differene
(TD) learning [42), and prove the convergence for the DP method, (2) characterize the basis for the vector space of the DR, also describing the settings in which the DR and the SR have the same eigenvectors, and (3) define default features, taking first step towards using the DR with function approximation.
Furthermore, we empirically investigate whether the DR provides benefits over the SR. While prior work (34] 3 has mainly focused on applying the DR to explain phenomena in neuroscience, here we consider applications more common in computational RL. We demonstrate that: (1) in environments with low-reward regions to be avoided, using the DR for reward shaping achieves superior performance over the SR, (2) using the DR for online eigenoption discovery [25 the DR exhibits reward-aware exploratory behavior and obtains higher rewards than the SR over the course of exploration, (3) the DR, similar to the SR, can used for count-based exploration, and (4) default features enable efficient transfer learning when the rewards at terminal states change.
2 Background
In this paper, we use lowercase symbols (e.g., r, to denote functions, calligraphic font (e.g. A) to denote sets, bold lowercase symbols (e.g.. r, to denote vectors, and bold uppercase symbols (e.g., V Z) to denote matrices. We index the jith entry of a matrix A by A (1,j).
2.1 Reinforcement Learning and the Successor Representation
In the standard RL framework, the environment is formulated as MDP (8,A, r. p, ? where 8 is the state space, A is the action space, r : 8 + (or r 8 x A -> RR is the reward function, p: 8 x A - (88) is the transition function, and 7 E [0, 1) is the discount factor. At every time step t. the agent observes a state, Sr, and selects an action, Ar. The environment then transitions to a next state, S ~ p(-|S,, At), and the agent receives a reward, Rg r(S:) (or R, = r(S,, At)). The agent's goal is to learn a policy π. S + A(A) that maximizes the expected discounted return E, ( R'ee. Ra).
In this framework, the successor representation [SR; [0) represents states the expected discounted number of visits to their successor states. Given policy æ, the SR, T E TR|8 (88) is defined as
V (s. = Ex Σ
y (( Su =
where 1 is the indicator function. Furthermore, denoting the transition probability matrix induced by
p and x by p the SR can be computed in closed form by Ei=o > '((P''' (I 11-7 -I--P >Px)--1 Note
that in this work, we set the rows corresponding to terminal states to be all zeros for p Denoting
the set of trajectories from s to s' by JJ-ss', we can express the entries of the SR [5] as
T (s,
Σ P (T))))((T)
(2)
TET
where (T(T) is the probability of following 7 under ग,, and (7()) denotes the number of steps in т.
Finally, the SR can be learned using temporal-difference learning 42 10 for all j E 8:
T TT' (St, < 4 (St.)) + a [1 (ss--)) + + (SH4,, 1 (St..))
(3)
2.2 Linearly Solvable Markov Decision Processes
The default representation (DR) is defined in the framework of linearly solvable MDPs 45 46 simplification of MDPs in which the optimal value function can be expressed by a linear equation. In this work, building on the work by Piray and Daw [34], define linearly solvable MDPs as

---

(s TT), where S is the state space, A is the action space, r 8 + R is the reward function, p SxA - +S) is the transition function, and Ra is the default policy that assigns non non- probabilities to all state-action pairs. It is standard in linearly solvable MDPs to assume no discounting of rewards, i.e. = 1. We also make the assumption that r(s) < 0 for all non-terminal 8. In order to make MDPs linearly solvable, the agent is tasked not only with maximizing the rewards, but also with not deviating too much from the default policy. Deviating from the default policy incurs a cost to the agent, causing the agent to also receive a penalty À KL(p" (-) [S)) at every time step t, where (L(ullv) denotes Kullback-Leibhe
(KL) divergence between u and v, p" (s'ls) denotes the probability of transi- tioning to s' from under policy म, and À determines the relative importance of the deviation cost.
The DR, introduced under this formulation, is a reward-aware proto-representation that encodes the expected rewards for visiting successor states. Let R be the vector of rewards at all states, and Pd be the transition probability matrix induced by Ta (i.e., (=*(s, s') is the probability of transition from to s' under TT), the DR, denoted by Z. can be computed in closed form by
Z = (diag(exp(-r//))) - PPd
where exp denotes element-wiss wise exponentiation.
Importantly, the DR can be used to retrieve the optimal value function with set of linear equations 34]. Let N, T be the set of indices of non-terminal and terminal states, and 8 and Sy denote the set of non-terminal and terminal states. We have
(5) where viy is the vector of optimal values for non- con-terminal states, and PNT E ||8x|X[8r] denotes the matrix of transition probabilities. Note that after learning the DR, assuming access to the transition dynamics, it is possible to directly compute the optimal values for any new configurations of terminal state rewards, fact that has been used in the past perform transfer learning (34]]]]]
cxp(vM/A) ZNN PN exp(rr/A),
3 Reward-Aware Proto-Representationn
In the previous section, we presented the DR in a much more limited way than when discussing the SR. This is not by accident; up to this point, the DR has not been that well developed. Little is written about the impact of environment rewards on the DR, how those impact the space the DR spans, and we still do not have a state-action formulation of the DR. We address these issues in this section.
3.1 The Default Representation as a Generalization of the Successor Representation
In problems in which the reward signal is the same across the whole state space, the DR spans the same space spanned by the SR. However, in settings in which the rewards received by the agent vary throughout an episode (see Figure ID, the DR captures the expected rewards obtained when traveling between two states, instead of the number of transitions required to travel between states, as in the SR.
We first formalize the idea that the DR and SR span the same space in Theorem 3. settings in which the agent does not have access to reward signal. Prior work had been limited to comparing the analytical form the SR and DR (3).
Theorem 3.1. Suppose both the SR and DR are computed with respect to the same policy, i.e., π Td. When the reward func- tion is constant and negative, ie.,, r(s) r(s') < 0 Vss s' E 8, the SR and DR share the same set of eigenvectors. Furthermore, when the SR and DR are symmetric, the i-th eigenvectors of the SR and DR are equivalent, and the i-th eigenvalues of the SR (HSRR,) )) and DR (HDRR)) are related as follows:
Figure 1: Episodic envs. adapted to incorporate negative rewards. Clockwise: 1) grid task [10].
2) four rooms (3), 3) grid room [49]. and 4) grid maze [49].
Start state in blue. The agent re- ceives -1 reward at every time step unless steps on red tiles (-20 re- ward) or reaches the goal in green (0 reward).
Tsse:=-(((rre (HTR. -ex (-r(s)))) + t +
piRRi =
(6)
where y E (0, 1) is the discount factor of the SR, r(s) is the state reward, and X is the relative importance of the deviation cost of the DR.
3

---

3R DR SR DR SR DR
3R
DR
FE
E
Figure 2: Top eigenvectors of the SR and the DR the environments shown in Figure the logarithm of the DR for better visualization due to very different magnitudes.
1 We report
Proof sketch. We write the DR as PI - P-]-1, where F exp(-r(s)/ and perform algebraic
manipulations on the expression Ze = Xe to obtain eigenvectors of 7P* J_ Then, we show that the eigenvalues of the SR and DR are related by function that is monotonically increasing over the range DR eigenvalues. The complete proof is available Appendix [3.2
This result is particularly interesting because the eigenvectors of the SR have been shown to be useful in settings such as reward shaping and option discovery. The top eigenvector of the SR captures temporal distance between states, and such distance has been used for reward shaping as it leads In option
better performance than distances that are based on the coordinates of states 44 54
discovery, it has been shown that eigenoptions, which are obtained by training policy to maximize reward signal induced by the eigenvectors of the SR, allow the agent to reach less explored areas of the state space, facilitating exploration [25]. We will revisit these applications in Section 6
In problems in which the reward function is not constant, the DR is able capture such information.
While this is evident by the reward vector, r, in Equation 4 we can re-write the DR to match the closed-form of the SR (Equation 2 to make such relationship clearer:
Z(s,8')= = Σ p*4(T) exx(((())))
(7)
TET
where r(r) is the sum of rewards obtained along trajectory (see derivation in Appendix B.1
By writing Equation 7 similarly the SR's, we see that both the SR and the DR compute weighted average of statistic over all trajectories from to s', where the weights correspond to the probabilities the trajectories under the policies w.r.t which the SR and DR are computed. For the SR, the statistic captures the number of transitions required to travel between states, while for the DR, it captures the rewards obtained when travelling between states.
Figure 2 illustrates the eigenvectors of the DR and SR in the environments depicted in Figure I to help the reader develop an intuition of the similarities and differences between these proto- representations. While the top eigenvectors of the SR reflect the proximity of states in terms of transitions, the top eigenvectors of the DR not only capture the transition dynamics of the environment, but they also accurately reflect the positions of low-reward regions in the environment.
3.2 Paate-Action-Depennee
Rewards
The form of the DR discussed far is defined for reward functions that depend only the state, and cannot be applied when the reward function depends on state-action pairs, which is prevalent in RL.
To extend the DR to state-action pairs, need to extend linearly solvable MDPs to state-action pairs, which was first presented the work by Ringstrom et al. [37]. However, the DR was not discussed in this work, so provide our derivation in Appendix B.3 Interestingly, the resulting cquation looks similar to Equation 4 Using the overbar for vectors and matrices the state-action setting, c.g., P instead of P. the DR for state-action-dependde rewards is:
Ż diag (exp(-r/X)) PT
8
where and F E |||/|| is the vector of state-action rewards, and Pm e [(8A411 the matrix of transition probabilities between state-action pairs under ma-
By learning the DR for the state-action space, the optimal Q-values can now be computed by
cxp(āv/A) /X) ZNN PT NT exp(Fr/A),
(9)

---

where q E ||8||4 is the vector of optimal state-action values. Furthermore, the optimal policy can be computed simply as
Ta(als) exp(("")))"") x (a(ss) - E. Ta cxp(q"ss.a'))) (sx)
(10)
As we show in Section the DR under this formulation, combined with default features, unlike prior work, allows direct computation of optimal policies when terminal rewards change, enabling efficient transfer without assuming access to transition dynamics 34
General Algorithms for Estimating the Default Representation
We now introduce efficient and general algorithms for learning the DR through dynamic programming and TD learning. Before the results we outline below, there were two update mechanisms for the DR: one based the Woodbury matrix identity (34. which is limited to the tabular case; and TD learning method recently introduced in a pre-print that is limited to settings in which the reward function is constant over non-terminal states (3) We refer the reader back to the previous section to emphasize how different the DR is in more general settings, with rewards that vary across states.
4.1 Dynamic Programming
We first present method to learn the DR by dynamic programming (DP), and prove it converges. The update rule for the DP algorithm and the convergence result are presented the theorem below.
Theorem 4.1. Let SN denote the set of non-terminal states. Assume r(s) 0 Vs &S.. Let R diag ( (ex((-r/)), where is the vector all state rewards. Let Zo R The update rule
Z
(11)
converges to the DR, that is, limg-> + Z₁ Z.
Proof sketch. We recursively expand the RHS of Equation and look at its closed-form in the limit, which we show is a Neumann series. We conclude the proof by leveraging the convergence of the Neumann series and plugging back into the derived limit. The complete proof is in Appendix B.4
4.2 Temporal Difference Learning
To learn the DR by TD learning, we build on the DP method and start with Equation state 8, for all j. Z( j), can be learned using DP
For any
(( exp (r(s)/A) (1(=))) + (x)p))(()()))) [Zs' :))):
(12)
which requires access to the transition probabilities under the default policy.
In TD learning, we do not assume access to these probabilities, but we can approximate the above equation using samples, estimating the expectation over next states s' with sampled next state. Then, given a transition (s, as, r, r.)') s') sampled under the default policy, we can update Z(s,j) as
[ss - Z(8, + a Z(s.
(13)
where is the step size and Y the bootstrapping target. For non-terminal states, Y = exp(r/A)(1 (-=i) Z(s while for terminal states, Y = exp(r/A)I (s=j)- Naturally, we can leverage importance sampling 36 if the default policy is different from the behaviour policy.
5 Extending The DR Beyond Classical Tabular RL
We have been using the tabular case to formalize the ideas we are introducing; now we extend them be- yond this setting. Given Equation B it is trivial to formalize the parameterization of the DR such that Z(s. s'; 0) x Zis, s'), for some parameters In this section, we go beyond that, defining the concept of default features, which is semantically meaningful when talking about the DR in terms of fea- tures. We also define a vrotoreppresenttiion for the maximum entropy RI. framework. However, since this paper focuses the DR, we defer the discussion of this proto-representation to Appendix B.5
5

---

Similar to successor features [SFs: 2). we derive here a decomposition of the value function into features and weights such that these features capture the dynamics of the environment. Different from SFs, default features also consider non-terminal rewards as part of the environment dynamics.
Recall can the DR obtain the optimal value function with set of linear equations (Eq. 5
exp(v SopvvN/M /A) NNN P NT exp(rt/A).
(14)
Inspired by the decomposition performed with SFs, we define default features (DFs) by decomposing the terminal rewards in this equation. We assume that the exponential of the terminal state rewards can be computed as exp (r(s)/A) = ((s) w, where ((s), w E IR are the features of a terminal state, 8, and the weights of the reward function, respectively. Note that this decomposition is not restrictive, since we can fully recover any terminal state rewards if (fs) = exp (r(s)/A).
Under this decomposition, the above equation can be written
exp(v /A) ZNNP" NT /))(rr/A) ZNN
PM NT Dw,
(15)
where E RRIriidd the matrix of terminal state features. We define the product ZNN PTd + as the DFs matrix, and each row <(s) x P NT D),, to be the DFs of non-terminal state. Then, the optimal value for each non-terminal state can be computed ((s) w. Note that such form of optimal value computation is only possible when only the rewards at terminal states change. The reason is that, when only the terminal state rewards change, C(s) stays the same for all non-terminal 8. However, when the rewards for non-terminal states change, ((s) changes since ((s) = ZNN P $ and NNN depends
con-terminal state rewards.
Building on Theorem we can learn the DFs using TD learning from a transition (s, 0, T; sampled under the default policy:
((8) C(8) (exp((e/)x(()))))) -
(16)
where a is the step size, and we define ((s) = ((s) for terminal states, s. DFs can be thought of as features propagating from the terminal states all non-terminal states in a manner that respects the reward and transition dynamics. The definition of DFs can be easily extended to the case of sere--actiondeeeeneee rewards, following the formulation in Section 3.2]
Apart from enabling function approximation, another benefit of learning the DFs over the DR is that it does not require access the transition dynamics, P N7. retrieve optimal values. This is rewards, since we can directly
especially useful when learning the DFs for state-action-dependee retrieve the optimal policy using the corresponding state-action values (see Equation I0). After learning the DFs in this setting, even without access to environment dynamics, we can efficiently compute the optimal policy under any terminal state reward configuration, enabling transfer learning. as shown in Section 6 Note that while DFs enable directly computing the optimal policy, they are limited to the scenario where only the terminal state rewards change. On the other hand, while SFs can only compute a policy as good as the ones used to compute the SFs, they can be applied when any state rewards change, allowing more flexibility than DFs.
Experiments
rrotorreprssnntttions, like the SR, have been used in RL for reward shaping [54 49]. option
discovery 22 25], count-based exploration [24], and transfer [2], among others. We now revisit
these settings to assess the impact of using reward-aware representations.
6.1 Reward Shaping
we first evaluate the effectiveness of the DR distance metric. In reward shaping experiments, our goal is assess whether the DR serve as a useful shaping signal, capturing the geometry of the state space, and whether it not only guides the agent toward the goal state but also helps it avoid negative rewards.
Drawing inspiration from results with the SR 54 49
Assuming e is the top eigenvector of the SR, prior work [54] 49 has used the shaping reward
A: (e( (e(Sgou)) ) e( e(*--1) 1)²= where e(s) denotes the entry for state s. Here, we instead use the top eigenvector e of the DR for potential-basee reward shaping (31), defining the shaping reward as
c re(s++1) e(st). We compare approach (DR-pot) with three basclines: 1) potential-base

---

Grid Task
Four Rooms
Grid Room
Grid
DR pot
SRooot
Return
-200
-100
-500
-500
-400
-200
1000
-1000
2.5
7.5
2.5
15
45
50
100
Stops (x101)
Steps (x10')
Stops (x10')
Steos (x10%)
Figure 3: The avg. undiscounted return over 50 runs for potential- reward shaping using the DR (DR-pot) the SR (SR-pot), the prior approach using the SR (SR-prior) [49], and no shaping (ns) in the environments shown in Figure The shaded area indicates 95% confidence interval.
shaping with the SR's 's top eigenvector (SR-pot), 2) the distance-based SR method from 49 (SR- (SR-pri), and 3) no shaping (ns). Both the SR and DR are computed with respect to the uniform random policy,
We carry out the experiments in the environments shown in Figure ! Importantly, the shortest path from the start state to the goal state passes through low-reward regions and is sub-optimal, so the agent needs to learn to avoid low-reward regions in the environment to achieve the optimal return. For the reward shaping approaches, we train Q-learning [51 agent using a convex combination of the original environment reward, and the shaping reward, ft. resulting in the expression
(1-,)r,+ + fř,, where 3 E [0, 1] is hyperparameter. Note that we assume access to the eigenvectors of the SR and DR prior to training the agent with the potential-based reward. Future work can explore learning the eigenvectors and the policy simultaneously. For the no shaping baseline, we simply train the Q-learning agent using the original environment reward. We use 7 0.99, 0.05 for e-greedy exploration, À = 1.3 for the DR, and perform a grid search over the Q-learning's 's step size ((0.1, 0.3, L.0) and 3 ([0.25, 0.5, 0,75, 1.0).. We run 20 seeds for each hyperparameter setting, and after identifying the best hyperparameters, re-run 50 seeds to avoid maximization bias.
Figure 3 shows the undiscounted returns of our approach and the baselines in the environments shown in Figure []] First, we observe that compared prior methods, potential-based reward shaping is a better way of utilizing the eigenvectors of the proto-representations.. Second, we observe that DR-pot performs significantly better than any of the baselines, including the SR-based approaches. This is because, shown in Figure 2 the top eigenvector of the DR captures the location of low-reward regions, and is capable of guiding the agent along the optimal path that does not cross them, whereas the top eigenvector of the SR fails to do so, and simply leads the agent along the shortest sub-optimal path to the goal. As per our result in Section the performance using the SR and the DR is very similar in the absence of low-reward regions (sec Appendix C.2)
6.2 Option Discovery
The options framework 335 43) allows the agent to interact with the environment in a temporally- extended manner. Prior work (25) demonstrated that the SR can be used for discovering options. In particular, their iterative online eigenoption discovery method using the SR, called covering eigenoptions (CEO), greatly reduces the average number of steps required to visit every state, and enables exploring the state space more uniformly. Furthermore, CEO combined with deep learning has shown promising results and outperformed state-of-the-an baselines in a wide variety of settings [8]. Here, we consider using the DR for iterative online eigenoption discovery.
Our approach, reward-aware covering eigenoptions (RACE), is similar to CEO but learns the DR instead of the SR, using its top eigenvector define the intrinsic reward. In RACE, at every iteration, the agent collects samples to learn the DR, the cigenvector of which is then used to compute an cigenoption. The eigenoption will be used in the next RACE iteration collect samples. This forms cycle of using eigenoptions to improve the learned representation, and using the learned representation to then refine eigenoptions. The full algorithm and hyperparameters are Appendix C33 We compare RACE with CÉO and uniform random walk in the environments from Figure L using 100-step episodes with no terminal states. To evaluate exploration, we measure the percentage of states visited; to assess risk awareness, we track the average reward per time step. Our goal is to test whether RACE enables reward-aware exploration. The results are depicted in Figure 4 where each point represents the average performance of one hyperparameter setting over 10 seeds.
7

---

Grid Task
Four Rooms
Grid Room
Grid Maze
-1.0
-1.0
Avg. Reward
RACE
CEO
RW
-1.4
0.90
0.95
0.8
0.9
0.25
0.50
0.75
0.4
0.6
0.8
State Visitation %
State Visitation %
State Visitation
State Visitation
Grid Room
Grid Maze
Grid Room (L)
Grid Maze (L)
-70
-70
140
RACE+O
1000
Return
-80
CEO+O
-B0
160
-90
90
-150
- 180
-100
200
200
0
50
100
0
50 100
50 100
50
100
of Iterations
of lterations
of Iterations
of Iterations
Figure 4: Top row: Average reward vs. state visitation percentage for various hyperparameter settings of iterative online eigenoption discovery via CEO, RACE, and random walk (RW) in the environments from Figure | For reference, solid dots mark settings with the highest visitation. Bottom row: Undiscounted return RACE+Q, CEO+Q. and QL (baseline), averaged over 50 seeds. Rightmost environments are shown in Figure 10 Shaded areas indicate 95% confidence intervals.
RACE is better than CEO at avoiding low-reward regions. Interestingly, because of its reward- awareness, RACE has a slightly lower average state visitation percentage than CEO. The agent sometimes needs to take detours to visit states that are obstructed by low-reward regions, thus requiring larger number of steps to visit other states (see Figure 2 in Appendix C.3). Again, consequence of our Theorem is that CEO and RACE behave very similarly environments with constant rewards across the whole state space (see Figure in Appendix C.3 Finally, RACE's exploratory behaviour also leads to faster learning when accumulating rewards. Results combining iterative online eigenoption discovery with Q-learning (51) are shown in Figure Appendix C.3 provides details about these experiments. Interestingly, RACE+Q performs much better than CEO+Q because CEO does not distinguish between high-reward and low-reward paths, and can reach the goal state following a sub-optimal one, causing offline Q-learning to first learn the sub-optimal path.
6.3 Count-Basee Exploration
Machado et al. (24] have shown that the norm the SR, while is being learned, encodes state-visitation counts. While reward-aware rep- resentations are somewhat at odds with pure exploration, given their risk-averse nature, can still ask whether the norm the DR can be used as density model for pseudo- counts [4]. Here, we evaluate the use of the E, norm of the DR as an exploration bonus. Specifically, the intrinsic reward is defined (s, a)) = 3 log( ZZa.; l|2), where is a scaling factor.
Table 1: Model-free count-based exploration. Values in thousands (x (x10*), 95% conf. intervals shown in parentheses.
Environment
Sarsa
+SR
+DR
RIVERSWIM
25 (0.8)
1,206 (566)
2,964 (252)
SIXARMS
265 (157)
1,066 (2,708)
3,518 (4.571)
We follow Machado et al.'s (24] experimental setup, and we consider two traditional exploration problems: Riverswim and SixArms. In Table we report the performance (total undiscounted return) of Sarsa [3 with e-greedy exploration, Sarsa SR [24] and Sarsa DR. Details about hyperameters are available in Appendix C.4 All results are averaged over 100 independent runs.
Our results suggest that the norm of the DR can indeed work a density function since its use has led to performance orders of magnitude better than random exploration. Maybe surprisingly, Sarsa+DD DR even outperforms, with statistical significance, Sarsa+SR in RiverSwim. This is an interesting result, as, to learn the optimal policy, the agent needs to overcome its aversion to negative rewards, which is encoding in its representation; but maybe this is precisely why it is somewhat more effective, such representation allows it to learn faster. This interplay is an interesting topic of future work.

---

6.4 Transfer Learning
The DR has been applied in the transfer learning setting to compute the optimal policy when the rewards at terminal states change [3]. However, such computations require access to the transition dynamics of the environment. With our extension of the DR to seree-acttio-eependdee
rewards and default features, it is now possible to compute the optimal policy without requiring access to the transition dynamics. We compare our approach with successor features [SFs; 2] which allow efficient transfer when the rewards of any states change. Note, however, that in this work we only consider the setting when only the rewards at terminal states change, since this is the setting that DFs can be applied to. The comparison of DFs and SFs under this setting does not serve as evidence that SFs is not effective for the setting it is designed for.
We perform experiments in the environment shown in Figure [5] consisting of four goal states. We describe our procedure for learning and using the SFs and DFs for transfer in Appendix C.5 as well as the features Φ used. Figure shows the performance of the DFs, the SEs computed with respect to different numbers of policies, and the optimal performance obtained by training optimal policy under each new terminal state reward configuration using Q Q-learning [51]. As previously mentioned in Section [S] in this setting, DFs is extremely effective because it directly computes the optimal policy for new terminal state reward configurations, while SFs only computes policy that is at least as good as the policies used to compute the SFs. Note, however, that DFs computes an optimal policy in the linearly solvable MDP setting that balances the reward function and the cost of deviating from the default policy.
Cosmuliion netum (x10%)
Optimal
25
50
75
100
Terminal Reward Configuration
Figure 5: Left: Four rooms with multiple goals. Right: Cumulative return across new terminal state reward configurations. Curves are averaged over 50 runs. The shaded area shows 95% conf. interval.
7 Related Work
Throughout this paper, we have discussed the relevant related work for both the theoretical development and the empirical evaluation of our work. A line of work we did discuss that might appear relevant for learning representations that capture both reward and transition dynamics is based on bisimulation metrics [13] 77 which quantify behavioural similarity between pairs of states by comparing their immediate rewards and distributions over the following states. These works enforee the bisimulation metric structure on the representation space by encouraging the difference between state representations to approximate the bisimulation metric [56] & 55 17. However, such approaches learn a distance distance-preserving mapping, where behaviorally similar states are mapped to nearby representations. This is in contrast to the DR or the SR, which encodes the environment dynamics directly in a semantically meaningful manner (see Eq. Eq. and 2). As an immediate result, the DR supports planning (see 5) as well other use cases discussed here, such exploration and option discovery, in a manner not possible with bisimulation ceeiiion-metic-bbaeeeee representations.
Conclusion
In this work, we advance the theoretical foundation of the default representation (DR), reward-aware nuto-reeresennttiion. Our contributions include: (1) deriving dynamic programming and TD methods to learn the DR, (2) analyzing its eigenspectrum, and (3) extending it to function approximation via default features. Empirically, we show that the DR, by incorporating reward structure, yields qualitatively different behavior than the SR across tasks such reward shaping, option discovery, exploration, and transfer.
While the DR provides benefits in environments where being reward-awar is important, there are some limitations. First, naively computing the DR and its eigendecomposition can lead to numerical instabilities due to taking the exponential of negative rewards, discussed in Appendix C. This could be problematic, especially in more complicated environments with long horizons. Second, as show in Figures 6 and 8 in Appendix C in environments in which all non-terminal state rewards

---

are the same, the DR will perform as well as the SR. Given the numerical issues associated with the DR. the SR may be better choice in these environments. Third, while the DR captures the environment dynamics more fully, this comes at the expense of flexibility. While the SR, which reward-agnostic, can allow transfer learning when any state rewards change, the DR, which reward-aware, can only perform transfer in settings where only the terminal state rewards change.
Finally, this work lays a foundation for the DR potential middle ground between reward-agnostic and risk-sensitive approaches. While we focused on the tabular case for clarity and to be able to evaluate the DR across range of use cases where reward- -arard-ggooticc proto-representations have been applied, restricting the scope to this setting remains limitation of this work. Future work should explore the use of the DR in more complex settings, starting with adapting neural network methods developed for approximating the SR [e.g.. 19 14 9
Acknowledgments and Disclosure of Funding
We thank Brett Daley for useful discussions and feedback. This research was supported in part by the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada CIFAR AI Chair Program, and Alberta Innovates. It was also enabled in part by computational resources provided by the Digital Research Alliance of Canada.
References
[1] S. Agarwal, H. Sikchi, Stone, and A. Zhang. Proto successor measure: Representing the behavior space of an RL agent. In International Conference on Machine Learning (ICML), 2025.
[2] A. Barreto, W. Dabney, R. Munos, J. Hunt, T. Schaul, D. Silver, and H. van Hasselt. Successor features for transfer in reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[3] A. Bazarjani and P. Piray. Efficient learning of predictive maps for flexible planning. PsyArXiv, 2025. doi: 10.31234/osf.o/ak577 URL osf. ho/ppr/ppiitts/ppaarr//////aaaaa
[4] M. G. Bellemare, S. Srinivasan, G. Ostrovski, T. Schaul, D. Saxton, and R. Munos. Unifying count-based exploration and intrinsic motivation. In Advances in Neural Information Processing Systems (NeurIPS), 2016.
[5] I. Blier, C. Tallee, and Y. Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. CoRR, (tbs/2111.0122, 2021.
[6] D. Borsa, A. Barreto, J. Quan, D. Mankowitz, H. van Hasselt, R. Munos, D. Silver, and Schaul. Universal successor features approximators. International Conference on Learning Representations (ICLR), 2019.
[7] P. S. Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. In AAAI Conference on Artificial Intelligence (AAAI), 2020.
[8] P. S. Castro, T. Kastner, P. Panangaden, and M. Rowland. MICo: Improved representations via sampling-basee state similarity for Markov decision processes. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[9] R. Chua, A. Ghosh, C. Kaplanis, B. A. Richards, and D. Precup. Learning successor features the simple way. In Advances Neural Information Processing Systems (NeurIPS), 2024.
[10] Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural Computation, 5(4):613-624, 1993.
[11] K. Dvijotham and Todorov. Inverse optimal control with linearly-solvable MDPs. In International Conference Machine Learning (ICML), 2010. 2010.
10

---

[12] J. Farebrother, J. Greaves, R. Agarwal, C. Lan, R. Goroshin, P. S. Castro, and M. G. Bellemare. Proto-value networks: Scaling representation learning with auxiliary tasks. In International Conference Learning Representations (ICIR), 2023.
[13] N. Ferns, P. Panangaden, and D. Precup. Metries for finite Markov decision processes. In Conference on Uncertainty in Artificial Intelligence (UAI), 2004.
[14] D. Gomez, M. Bowling, and M. C. Machado. Proper Laplacian representation learning. In International Conference on Learning Representations (ICLR), 2024.
[15] T. Haarnoja, H. Tang, P. Abbeel, and S. Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning (ICML), 2017.
[16] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Leurning (ICML), 2018.
[17] M. Kemertas and T. Mamnentado-Ammstronn Towards robust bisimulation metric learning. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[18] M. Klissarov and M. C. Machado. Deep Laplacian-based options for temporally-extended exploration. In International Conference on Machine Learning (ICML), 2023.
[19] T. D. Kulkarni, A. Saeedi, S. Gautam, and S. J. Gershman. Deep successor reinforcement learning. CoRR, abs/16066000 02396, 2016.
[20] C. Le Lan, S. Tu, A. Oberman, R. Agarwal, and M. G. Bellemare. On the generalization of representations reinforcement learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
[21] M. C. Machado. Efficient exploration in reinforcement learning through time-based representa- fions. PhD thesis, University of Alberta, Canada, 2019.
[22] M. C. Machado, M. G. Bellemare, and M. Bowling. A Laplacian framework for option discovery in reinforcement learning. In International Conference on Machine Learning (ICML), 2017.
[23] M. C. Machado, C. Rosenbaum, X. Guo, M. Liu, G. Tesauro, and M. Camphell. Eigenoption discovery through the deep successor representation. International Conference on Learning Representations (ICLR), 2018
[24] M. C. Machado, M. G. Bellemare, and M. Bowling. Count-based exploration with the successor representation. In AAAI Conference on Artificial Intelligence (AAAI), 2020.
[25] M. C. Machado, A. Barreto, D. Precup, and M. Bowling. Temporal abstraction in reinforcement learning with the successor representation. Journal of Machine Learning Research, 24(80): 1-69, 2023.
[26] S. Mahadevan. Proto-value functions: Conference on Machine learning (ICML), 2005.
Developmental reinforcement learning. In International
[27] S. Mahadevan and M. Maggioni. Proto-value functions: A Laplacian framework for learning representation and control in Markov decision processes. Journal of Machine Learning Research, (110)2169-22311 2007.
[28) I. Momennejad, F. M. Russek, J. H. Cheong, M. M. Botvinick, N. D. Daw, and S. J. Gershman. The successor representation in human reinforcement learning. Nature Human Behaviour, 1(9): 680-692, 2017.
[29] T. Moskovitz, S. R. Wilson, and M. Sahani. A first-ocuppancy representation for reinforcement learning. In International Conference on Learning Representations (ICLR), 2022.
[30] Moskovitz, S. Hromadka, A. Touati, Borsa, and M. Sahani. A state representation for diminishing rewards. Advances in Neural Information Processing Systems (NeurIPS), 2023,
11

---

[31] A. Y. Ng, D. Harada, and S. Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In International Conference on Machine Learning (ICML), 1999.
[32] J. Peters, K. Mulling, and Y. Altun. Relative entropy policy search. In AAAI Conference Artificial Intelligence (AAAT), 2010.
[33] S. U. Pillai, T. Suel, and S. Cha. The Perron-Frobeniu theorem: some of its applications. LEEE Signal Processing Magazine, 22(2):62 75, 2005.
[34] P. Piray and N. D. Daw. Linear reinforcement learning in planning, grid fields, and cognitive control. Nature Communications, 12(1):4942, 2021.
[35] D. Precup. Temporal abstraction in reinforcement learning. PhD thesis, University of Mas- sachusetts Amherst, 2000.
[36] D. Precup, R. Sutton, and S. Singh. Eligibility traces for off aff-pooicy policy evaluation. In International Conference on Machine Learning (ICML), 2000.
[37] T. J. Ringstrom, M. Hasanbeig, and A. Abate. Goal kernel planning: Linearly-solvable non- Markovian policies for logical tasks with goal-conditioned options. CoRR, abs/2007.02553 2025.
[38] G. A. Rummery and M. Niranjan. On-line Q-Learning using connectionist systems. CUED/F- INFENG/TR 166, Cambridge University Engineering Department, 1994.
[39] H. Sprekeler. On the relation of slow feature analysis and Laplacian eigenmaps. Neural Computation, 23(12):328733002, 201
[40] K. Stachenfeld, M. M. Botvinick, and S. J. Gershman. Design principles of the hippocampal cognitive map. In Advances in Neural Information Processing Systems (NeurIPS), 2014.
[41] K. L. Stachenfeld, M. M. Botvinick, and S. J. Gershman. The hippocampus as a predictive map. Nature Neuroscience, 20(11):164331653, 2017.
[42] R. S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 9-44, 1988.
[43] R. S. Sutton, D. Precup, and S. Singh. Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial Intelligence, 1 112(1-2)1811-211 1999.
[44] A. Tirinzoni, A. Touati, J. Farebrother, M. Guzek, A. Kanervisto. Y. Xu, A. Lazaric, and M. Pirotta. Zero-shot bodle-body humanoid control via behavioral foundation models. In International Conference on Learning Representations (ICLR), 2025.
[45] E. Todorov. sooeall Markov decision problems. In Advances in Neural Information Processing Systems (NeurIPS), 2006.
the [46] E. Todorov. Efficient computation of optimal actions. Proceedings the National Academy of Sciences, 106(28); 1478-11483 2009.
[47] A. Touati and Y. Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information Processing Systems (NeurIPS), 2021.
[48] A. Touati, J. Rapin, and Y. Ollivier. Does zero-shot reinforcement learning exist? In Interna- lional Conference on Learning Representations (ICLR), 2023.
[49] K. Wang, K. Zhou, Q Zhang, J. Shao, B. Hooi, and J. Feng. Towards better Laplacian representation in reinforcement learning with generalized graph drawing. In International Conference on Machine Learning (ICML), 2021.
[50] K. Wang. Zhou, J. Feng, B. Hooi, and X. Wang. Reachability Laplacian representation in reinforcement learning. In International Conference on Machine Learning (ICML), 2023.
[51] C. Watkins. Learning from delayed rewards. PhD thesis, University of Cambridge, 1989.
12

---

[52] H. Wiltzer, J. Farebrother, A. Gretton, Y. Tang, A. Barreto, W. Dabney, M. G. Bellemare, and M. Rowland. distributional analogue the successor representation. In International Conference Machine Learning (ICML), 2024.
[53] L. Wiskott and T. J. Sejnowski. Slow feature analysis: Unsupervised learning of invariances. Neural Computation, 14(4))715-770, 2002.
[54] Y. Wu, G. Tucker, and O. Nachum. The Laplacian in RL: Learning representations with efficient approximations. In International Conference Learning Representations (ICLR), 2019.
[55] H. Zang, Li, and M. Wang. Simsr: Simple distance-based state representations for deep reinforcement learning. In AAAI conference on Artificial Intelligence (AAAI), 2022.
[56] A. Zhang, R. Т. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant repre- sentations for reinforcement learning without reconstruction. In International Conference on Learning Representations (ICLR), 2021.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction are supported by theoretical or empirical evidence.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in the Conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification:
We provide complete proofs or derivation for our theoretical results in Ap-
pendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: provide sufficient details of our algorithms and baselines for reproduction in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide source code in the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://niis.co cc/
hhblltt/guuddes/oddeelbbm.ssss wwt.:/uabmiiiiooooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips.co. ccippi.cc guwdes/ddeeeee CooeeSumeeeeeeeeeeeeee ssi onyylicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide description of experimental settings sufficient for understanding in the main text, and provide additional details, e.g., hyperparameters and hyperparameter search procedures, in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report 95% confidence intervals or visualize all independent runs wherever possible, and discuss statistical significance in Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We discuss compute resources Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: This paper conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
As explained in Appendix A this paper does not have direct societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the creators of the code used in the paper properly credited, and license information is mentioned.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We follow the NeurIPS code submission guidelines for our submitted code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs have only been used for editing and formatting purposes.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy battps://neurre /0tps///ecccoooeeeeee////2////// for what should or should be described.
20

---

A
Broader Impact
This work investigates novel form of representation in reinforcement learning-a foundational area within the field of artificial intelligence. Given the fundamental and theoretical nature of this research, does not pose any foresecable negative societal impacts. As with most projects focused on core scientific inquiry in AI, the associated risks are minimal, and no significant mitigation strategies required.
B Theory
In this section, we present the detailed derivation of the theoretical results in the main text.
B.1
Derivation of Equations andd] 7
We present the full derivation of Equations 2 and 7
B.1.1 Derivation of Equation 7
Equation states that the (si, 85)-th entry of the DR be written as:
Z (s), =
(T exp (r(r)/A).
(17)
TEP
Proof. Recall that the analytical form of the DR 34
|iag (exp(-r/A)) Pmo
To
simplify notation, let = diag (exp( (r/ We then have
Z (R P
(18)
RR P p=2) )== II
(19)
RZ pT4 =
(20)
RZ + P
(21)
Z = R -1 + R '
(22)
From the above equation, we know that
Z(5;, 5;) = exp (r(s))/A)4( + Ep (st, s'xp exp (r(ss)/A)Z(ss'.)
(23)
the right-hand side of which can be repeatedly expanded to give
Z(8;, 85) = exp (r(s:) (,==e,}
P (s;, exp exp (r(ss) + r(s')((/)))))))) (( (w-s;)
a (st 5') Σ (s ("', exp
exp (r(s.) + rrs's) + +r(s")) /2) A(m))s))
(24)
In the above equation, it can be seen that the expression enumerates over trajectories from state 8; to state sy with different lengths. The first row corresponds (containing one state), which is the simple case when the agent starts and ends in sj without taking an action. The second row enumerates all trajectories from s; to 8, with length 2, and so on.
single trajectory with length
The above equation can be written in the following more compact form:
Z(s,,
Σ (T) exp (r(T)/A), TET..
(25)
where Ja.- is the set of trajectories from state s; to state 8,. P= is the probability of following trajectory under the policy wa. and r(r) is the sum of rewards obtained in trajectory D
21

---

B.1.2 Derivation of Equation 2
Equation 2 appeared in the work by Blier et al. (5], but the derivation was not shown. For completeness, we provide the derivation here, which similar to that for Equation 7 states that the (8;, 4th entry of the SR can be written as:
Equation
L (st. 8y)
(T(T))(())
(26)
TEJ,, 12g
Proof. Recall that the analytical form of the SR
[0] is t (I (--PP))1
We know that
y (I pp")
(27)
(I - y P y =
(28)
4 I+ u
(29)
from which we know that
L (s; 83) = 1 ( + Σ pd s' u (s',8))
(30)
Expanding the hight-hand side, we have
4 (8i, 83) = [0:=5}}
(31)
Ps (5 (7)
(32)
Σ PT (8 8))
Σ
pd (5',8")77 (s"=a,)
(33)
+...
(34)
which can be written as
V (sts 85)
Σ
P PT (T)))())
(35)
TEJ..-ry
where Jsy is the set of trajectories from state to state sj. (T((T is the probability of following trajectory T under the policy π, and 17(T) is the number of steps it takes to visit sy from 8; in trajectory T. Note that when (s:) = (sy), nnT]= = O.
B.2 Theorem 3.1
Theorem 3.1. Suppose both the SR and DR are computed with respect to the same policy, i.e. Nd- When the reward function is constant and negative, iee, r(s) = r(s') < 0 Vss s' E 8, the SR and DR share the same set of eigenvectors. Furthermore, when the SR and DR are symmetric, the i-th eigenvectors of the SR and DR are equivalent, and the i-th eigenvalues of the SR (USR,) and DR ()) are related as follows:
JARR, ( (MDRi exp ( - (())/)) ((03//)+7-11) +
(36)
where 7 E (0,1) the discount factor of the SR, r(s) is the state reward, and X is the relative importance of the deviation cost of the DR.
Proof. Let μ.: e; be the i-th eigenvalue and eigenvector of the DR. Assuming that the reward function is constant, and letting F = exp rs))/A), the DR can be written as [FI p=] Then, we have:
FI -=] Mili
(37)
FI - pd le Le
(38)
P e, = 7k, 1 F)e
(39)
(1-pP)ee >e = 7/4
+ 7 11)
(40)
(I 1-7pt)-te [>65 1 8
+ "1] -1]- e;
(41)
b.c = 7 14 -1 exp ( - (-r(s)/A) + Ci
(42)
22

---

We have thus shown that the SR and DR share the same set of eigenvectors. When the SR and DR are symmetric, their eigenvalues are real, and we can further show that the orders of the eigenvectors of the SR and DR by their corresponding eigenvalues are identical.
From Eq. 42 we know that the eigenvalues of the SR and DR related by the following function:
PSR ((uDR) by( exp( +
(43)
Taking the derivative, we have
('UDD) + -xp((r()))]] When 7 exp(r(s)/A), f' ("DD) is always positive, so the orders of the eigenvectors of the SR and DR are identical. When 7 & exp(r(s)/A), J'(DRR) > O except at the vertical asymptote at JDR -xpj We now show that all eigenvalues of the DR lie on the same side of the vertical asymptote, and so the orders of eigenvectors are still identical. We do so by first deriving lower and upper bounds on the eigenvalues of the DR.
(44)
First, show that all eigenvalues of the DR are positive. We start with the inverse of the DR, Z-1 = diag(exp( r/)) - PMd Consider a row of Z corresponding to a non-terminal state. We know that Z (i, exp(-r(s)/A)
(i, Ej IZ- (i,j), For a row i corresponding to a terminal state, Z (1, > Ejf --, easily holds. Then, by the Gershgorin circle theorem, we know that this matrix has all positive eigenvalues, and thus, all eigenvalues of the DR are positive.
P (i, >
Next, we derive an upper bound on the eigenvalues of the DR. Also applying the Gershgorin circle theorem on the inverse of the DR, we know that the set of eigenvalues of the DR lies within the union of the Gershgorin discs. For the i-th row of the inverse of the DR, if corresponds to non-terminal state, we have disc centered at exp(-r(s)) state, the disc centered at exp(-r(s)/A with radius O. is then not hard to see that the minimum of the union of the discs is simply exp(- exp(-r(s)/A) 1. Since the lower bound on the eigenvalues of the inverse of the DR is exp(-r(s)) A) - 1, the upper bound on the eigenvalues of the DR is
Pm (i, 8) with radius - P (i, i). For terminal
expl (xp(-r(s/AA)-T'
We now know that all eigenvalues of the DR lie in the range (0, (xpp-r(3)/AA-i] We now show that the vertical asymptote of f. if it exists, always lies outside of this range. We focus the following two cases: 1) 7 < exp(r(s)/M) A),, and 2) 7 > exp(r(s)/A).
First, when exp(r(s)/A), it can be easily shown that the vertical asymptote 1--1-exp--(aa) 0. Since all eigenvalues positive, they all lie on the right hand side of the vertical asymptote. For the remaining case, we first have that
7 <1
(45)
7- >
(46)
gー1 (pp(-r(s)/A) exp(-r(s)/A).
(47)
Since y > cxp(r(s)/2) and we assume r(s) < 0, both sides of the inequality are negative. Then,
1
(48)
7- exp( exp(-r(s)/M)
1- 1--exp(-r(s))
(49)
7-1 - (01-ep((-(ss/A)
exp(-r(s)/1) -1
Since all eigenvalues are less than or equal to exp( ( they all lie on the left hand side of the vertical asymptote. We have thus shown that when a vertical asymptote exists, it lies outside of the range of eigenvalues, and the orders of the eigenvectors of the SR and DR are still identical. O
B.3 Extension of the DR to State-Action- Reward
In this subsection, we derive the DR for the case when the reward function depends on both the state and action. The extension of linearly solvable MDPs to the state-action case was first presented by
23

---

Ringstrom et al. [37. However, since (1) the DR was not defined in this work, and (2) there are minor differences in the formulation, derivation, and notation, we present the full derivation below nonethe- less. Note, also, that the derivation bears similarity to relative entropy policy search (32), which constrains the KI. divergence between an observed data distribution and that generated by the policy.
For this setting, instead of formulating the deviation cost as KL mulate the deviation cost from the default policy as KL (. A (.. Ar)), where (" (A) denotes the distribution over the next state-action pairs given the current state- action pair (St. Ar) and the policy т. Then, every time step t, the agent receives reward of F(Se. A,) r(S, At) XKL (s At)ll"" (t, (S,, where À > 0 determines the relative importance of the deviation cost.
(-15)), we for-
Let T be the random variable denoting the final time step of an episode. The return starting from time is then
G = + ((f.))))---((S---AA---. +r(sr Ar),
(50)
where we allow the agent to select action A, at the terminal state St obtain final reward AT),
The state-action value function defined as q² (s, a) E, [G s, a]. For terminal states, (s(ss, a) is simply r(s, a). For non-terminal states, we have
(s,
E,, [ s,
(51)
-E Ar) + 17 A,
(52)
= + Ex +  (|
= s, = a)
(53)
-F(s,a) + Σ Epr(s, (
A+41) + ....
(54)
f(ST-1, AT-I &,
8, A. = a, S.+1 = 8 Act1 =d]
(55)
= =((s,a)+ + Σ p (s E, [ Ac41) + +...
(56)
d'aa
Ar-) (rr.A.A)5111 Ar 5 Ar1 = d'l
(57)
=((s,a)+ + Σ Edd
(58)
= =(ss a) + E. (ols.)) a').
(59)
Then, the optimal ataten value function for non-terminal states satisfies
as = max
((s,a) Eg (v (olan)) (0'(s'
(60)
((()a)) -KL (1s (.,1s,) + Eg (o (s',
= max "
(61)
In ('(6' a) pne (s', p"dssd''ssa pra(s', a) exp (5 (s', )')\x)
g a(66,09)
=r(s,a) + max
a' upp (v))),))
(62)
A
}
= =r(s, a) + max
-XE oa (v(lsa)
(63)
The expectation term above resembles a KI. divergence, but the denominator of the fraction is not properly normalized. Define the normalization factor C Ex al pd (5', a'Is, a) exp (4' d')/A).
24

---

The above equation then becomes
p p' (64)
(5..))r(ss,0) + max
XB, a (vla,)) In
pad (s'',''sss, (x
=r(s,a) + A max { - AKL (p" (11(111,)))lpp"" exp (q' (.-)A)//))
(65)
= a)) AlnC,
(66)
where the last equality follows from the fact that the minimum of the KL divergence is 0. Note that in order to make the KL divergence O, the optimal policy needs to satisfy
r² (s', d' (s) x p="( a'|s, a) exp (4 (d'))/))))) (a') (ss (s,0) x ma(a') Is (p(s (psss1,,,,, exp exp
(67)
(68)
(a (s') x na(a'Is') exp exp (g' ('),)))A))
(69)
Then, given the optimal acteon value function and the default policy, it is straightforward to retrieve the optimal policy, ;²:
Ta(als) exp (4 (s,
(als) E. Ta(a'|s) exp (s'(s))))AA)
(70)
Substituting C' back to Equation 66 we have
((s,a)=r(s,0) Aln Σ (' (x)((())d))))
(71)
0)..))/)+ In 110
(72)
s'aa
Σ (s exp (f exp (s)(a))) exp exp (r(8,a)/A)
(73)
We represent the above equation in matrix form. Let q RIS|AA| be the vector of optimal state-action values, R |8|4||| be the vector of state-action rewards, and P || [|||a|] be the matrix of transition probabilities between state-action pairs under the default policy, i.e., pra s'a') ( d' a). Furthermore, let N, T be the set of indices of non-terminal and terminal state-action pairs respectively. The above equation, defined for non-terminal states, can then be written
exp(âN/ /A) diag((ep((xx/A))(PPP)) N'N exp( + P NT exp(ar/A)) diaglexp((Ex/A)) (xx)xxxx T N' exp /X) + T (xp((T/A)
(74)
(75)
diag(exp(-Ew/A)) P
N' exp(qN /A) TH exp(Fr/A)
(76)
exp(qx/A) [diag(exp(-FF]
P NN
TT exp(Fr/A)
(77)
The DR for state-action-deeeeddett rewp:n/dddddeeneeeewwarsssssss is then ddia((ixp( states, and [diag(exp(-F/A) - - for all states.
- T N'
for non-terminal
B.4 Convergence Proof for Dynamic Programming Algorithm
Theorem 4.1. Let Sx denote the set of non-terminal states. Assume r(s) < Vs Sx. Let
R diag(exp(-r/X)) where is the vector of all slute rewards. Let Zn R R-11 The update rule
Z++1 = -R 1--+R-1 - Z₁
(78)
converges to the DR, that is, limg-+o Z₁ Z.
25

---

Proof. Recursively expanding the right hand side of Equation 78 we have
k = R (R
(pT
R
+
R
p
)
R
(79)
1+++
Σ R
R
(80)
Taking the limit, we have
20
lim Z₁- = k0000
Σ
R
pruje
R
(81)
Note that [> R R [ is a Neumann series. We now establish the convergence of this
Neumann series. Since Eg' p*a 8') equal to for every non-terminal 8, and for all terminal have
R 'p'd
1bx max exp (r(s)/A) Σ (p))(s,s))
max exp (r(s)/A) sENN
(82)
sE8
where the final inequality follows from our assumption that r(s) < Equation 82 is sufficient condition for the convergence of the Neumann series, so we know that this senes converges to RI- (pa)-1 Plugging this back, we have
lim Zμ = (I R PPma)= -RR-
(83)
k0000
R '(R - P=*]]] 'R-1
(84)
[RR (R - p=u]]
because B A'A'' B-'A-' =
(AB)) -1
(85)
diag (exp(-r/A)) -pad Z.
(86)
B.5 Proto-Repressntatiin in Maximum Entropy RL
Similar to how the SR is a proto-representation derived the standard RL formulation, and the DR in the linearly solvable MDPs, we can derive prote-represnntations for other formulations. Specifically, due its popularity, and for completeness, we introduce proto-representatioe in the maximum entropy (MaxEnt) RL framework [16 15 We call it maximum entropy representation (MER).
In our MaxEnt RL formulation, at each time step I, the agent receives both the reward r(;)) and en- tropy bonus XH (p (-(Sa)). where pn (-IS;) denotes transition probabilities over successor states given policy म, and X > 0 controls the weight of the entropy term. Assuming 7 = 1, we define the MER as: Definition 4.1. Let r be the vector of state rewards, and A be the adjacency matrix, i.e., A(s, 8') is equal to if s' can be reached in one step from s, and O otherwise. The MER is defined as
M diag ( (exp(-r)/A) A
(87)
The MER differs from the DR (Equation 4 in that it uses the adjacency matrix rather than the default policy's 's transition probabilities. This connection arises because maximum entropy RL and linearly solvable MDPs are closely related as can be transformed into the other [1I] As result, most of our contributions readily extend to the MER. Preliminary analyses showed that the MER also behaves similarly to the DR, we focus our experiments the DR. Still, the MER and DR are distinct rooorseepessnn with different formulations, assumptions, and value functions, which lead to different optimal policies. While we did not observe substantial differences in the settings we considered, exploring them in other regimes remains an open direction.
We now provide the derivation of the MER. To derive the MER for deterministic transitions, we can simply formulate the entropy as (-((.-)).). However, to derive the MER for stochastic transitions, need to formulate the entropy slightly differently. We assume that at every time step t, the
26

---

agent receives a reward of F(S)) r(t)) + ()((p=(-SS)) where p= (-1S) denotes the transition probabilities over the successor states of S, given policy π, and X 0 determines the relative importance of the entropy term.
Let T be random variable denoting the final time step of an episode. The return starting from time t is then
G. ((S)) +... + F(S--1) + +
(88)
Note that similar to prior work 46 34 allow the agent to obtain a final reward of r(Sr) the terminal state, Sr.
The value function is defined as (=(s) = E, (G.SS = =s] For terminal states, v² is simply r(s). For non-ter non-terminal states, we have
v=(s) E,, [r =
(89)
F.,, (F(S.) ( (750))++---((5))+++++)++++) s, 8
(90)
-(() Σ
(91)
E = -((s) +75 = sS
(92)
= =((s) + (d'(s)))
(93)
= =(s)) + (1 [b (5').
(94)
The optimal value function for non-terminal states satisfies
v" (s) max F(s) Eg'o (110) (1s)) [s']]
(95)
max π
(r(s) + XI((p"((15) + E (u
(96)
= r(s) + max
AE (115) Imp" + (os
(97)
p² (s'(s)
= r(s) + max
- -AB
In exp(u*(s')//)
(98)
The expectation term above resembles KL divergence, but exp (5*())/X) is not normalized distribution. To normalize it, we define C(s) Eves, exp (v" (s')/A). where S, [ (s'(ss) > 0} denotes the set of successor states state s. We then have
In (((((s))/C(ss) exp (1- (1)))))))/(()
(s) = r(s) + max
AE, AB
(99)
(=((sss)
= =r(s) + max
- -ABerppr(1e)
+ AlnC(())
(100)
exp (1)(s))))//((
=r(s) + Aln C(s) + +max
- -AKLL (--(11s)]] exp (v" (()))/(()))
(101)
= r((s) + C(s),
(102)
where the final equality follows from the fact that the minimum value of KL divergence is 0.
Then, we have that
+ Aln Σ exp (v*(s))AA)
(103)
SE
:. (s))/A-((s)) r(s) + In
Σ exp (u*(s)/A)
(104)
exp (xp (b)/2) = exp (r(s)/A)
Σ
exp (v (d(())/)))
(105)
27

---

Now, we express the above equation in a matrix form. Let v be the vector of optimal state values, be the vector of state rewards, and A be adjacency matrix, i.e. A(s, s') = 1 (s',,)) Furthermore, let N, T be the set of indices of non-terminal and terminal states respectively. The above equation, defined for non-terminal states, can be written as
exp(VN/ (A) (lag(exp(rw)/))(AAm)) eXp(vN/A) + ANT exp(vr/A))
(106)
diag (exp(_rr)/ exp(vx//) ANN ANNexppwVN/N + NT exp(rt/A)
(107)
diag ( (exp( -rpp()))/AA) - ANN] exp(vx/A) = ANT exp(rr/A)
(108)
exp(vw/A)
diag (exp(-NN)/\) ANN]
ANT exp(rt/A)
(109)
We define the MER for non-terminal states as diag (exp(-rr)// ANN] A more general
definition of the MER for all states is ddag (exp(-r)/A) -A]
C More Experiment Details
We provide additional experiment details omitted from the main text due to space constraints.
C.1 Numerical Considerations
Reward shaping and option discovery experiments involve the top eigenvector of the DR. We describe numerical considerations when performing eigendecomposition of the DR. First, in practice, we perform eigendecomposition of the symmetrized DR, Sym(Z) = (Z Z )/2, to ensure real eigenvalues and eigenvectors. Second, as the DR involves exponentiating negative rewards (see 7. the magnitude of the DR entries, especially the off-diagonal entries, can become very small. To mitigate the resulting numerical issues, we use the library python-fling improved precision comes at the cost of increased runtime. As a larger has the effect of reducing the magnitude of negative trajectory returns in the DR (see Eq. 7). can alleviate numerical instabilities. In our experiments, we initially started with À = 1, and slowly increased À by 0.1 until we settled on 1.3, which, paired with the use of python- ffint, resulted in no numerical issues.
Note, however, that the
Finally, the magnitude of the top eigenvector entries can be very small. Therefore, in practice, we use the logarithm of the top eigenvector in place of the top eigenvector. We now show that under mild assumptions, the top eigenvector of the DR is positive, allowing us to take the logarithm. We first present mild assumption:
Assumption C.1. There only one start state, so. and it possible to reach any state from so under the default policy, i.e., - PT (Tau-t)) > 0 for all s E S.
To avoid introducing unnecessary bias, we use the uniform random policy as the default policy in all of our experiments, which is the standard practice [34 3 We also perform experiments in grid world environments, where it is possible to reach any state from the start state. Under these conditions, Assumption callly holds. We now present the following proposition:
Proposition C.2. Under Assumption
C.I the top eigenvector of Sym(Z) (Z) is positive.
Proof. We can rearrange the rows of Z so that the first row corresponds to the start state, 80- By Assumption C.1 and Eq. 7 the first of Z has all positive entries, i.e., Z(sn, s) > 0 Vss E 8, while all other entries are non-negative. After symmetrization. Sym(Z) (Z) has positive first row and first column. It is not hard see, then, that Sym(Z)² has all positive entries, and Sym (Z) a primitive matrix (33. Finally, by the Perron's 's theorem [3 the largest eigenvalue of Sym(Z) is positive, and the corresponding eigenvector is positive.
and
Proposition C.2 applies the case when the DR is computed in closed form. We now extend the proposition to the case when the DR is learned by TD learning. When learning the DR using TD
hutps://gthhbb.ccm//iiiiiihhhhop flint, MIT license
28

---

Four Rooms
Room
0
DR pot
SP-pot
Return
-50
SR-pri
-100
-300
-500
-100
-200
-1000
-1000
2.5
7.5
50
100
Stops (x101)
Steps (X201)
Scops (x10')
Stoos (x10%)
Figure 6: The average undiscounted return of potential- reward shaping using the DR (DR-pot) and the SR (SR-pot), the prior distance-bait based reward shaping using the SR (SR -R-prioo)) [49]. and no shaping (ns) over 50 independent runs in the variations of the environments shown in Figure without low-reward regions. The shaded area indicates 95% confidence interval.
the agent interacts with the environment and learn the DR in an incremental, online manner. Let Sy be the set of states visited by the agent. this setting, the top eigenvector refers to the top eigenvector of the sub-matrix of the symmetrized DR corresponding to the states in Sy.
learning (Eq. 13
Proposition C.3. Assume that there only one start state so- Let D be a dataset containing the transitions collected by an agent starting from Su and following the default policy. Importantly, the transitions are stored in the order in which they were collected. Initialize Z as the identity matrix. When learning the DR by TD learning (Eq. 13) with step size E (0, 1), backward sweep through DD guarantees that the top eigenvector of Symm(vv is positive, where Zvv is the sub-matrix of Z that corresponds to the states in Sy.
Proof. We can rearrange the rows of Zvv so that the first row corresponds to the start state, 8o- The dataset, D. consists of one more trajectories starting from so. When iterating through trajectory in backward manner, for the row in Zvv corresponding to state, 8, all the entries corresponding states visited after in the trajectory will be updated to have positive values. is then hard to see that after iterating through all transitions, the first row of Zyv has all positive entries, while the remaining rows are non negaatvve After symmetrization. Sym(Zvv) has positive first row and first column. It is then easy to see that Sym (Zvv) is a primitive matrix and, by Perron's 's theorem B3 its top eigenvector is positive.
For MDPs with an initial distribution, Po, over multiple start states, we can define new start state that transitions to these start states under po, no matter what action taken. The reward for this new start state also does not affect the optimal policy, and only needs to be negative (see Theorem 4.1). can then be applied.
Proposition C.2 and C.3
In practice, when learning the DR using TD learning, for example, in RACE (Algorithm T. we initialize the DR the identity matrix, and perform least one backward sweep through the dataset of collected transitions. We then compute the eigendecomposition for Sym (Zvv), and take the logarithm of the resulting top eigenvector. To project this transformed eigenvector in R|8v| back to we simply augment this vector with zeros for states in 8y.
R 8|
C.2 Reward Shaping
Figure 6 shows reward-shaping results in the variations of the environments shown in Figure D without low-reward regions. In these environments, we do not observe a significant performance difference between potential-based reward shaping using the DR and the SR. This is because in these environments, all states apart from the terminal state have the same reward. Although Theorem 3.] cannot be directly applied since the terminal state has different reward (0) than the rest of the states (-1), it is reasonable to expect that the DR and the SR have similar eigenvectors in this setting, leading to similar results.
C,3 Option Discovery
In Algorithm T we present the full algorithm of RACE, an instance of the framework of representattion-drivee option discovery [ROD; 25 that relies the DR instead of the SR for option
29

---

Algorithm Reward-Aware Covering Eigenoptions (RACE)
Input: a;
Step size for learning the DR Relative importance of the deviation cost the DR
A;
Miaam
Number of iterations through collected dataset to learn the DR
Noption
> Number of latest eigenoptions to keep
Poption
Probability of sampling option instead of primitive action > Step size and discount factor for learning the options' policies Number of interactions with the environment each ROD iteration b Number of iterations of the ROD cycle
Co. 7o
N Nseps
Viter
D
0
Ω
Z
for i <- to Niter do
Collect samples
D Dour -- 0
for j - 0 Natops do
With prob. 1- Poption randomly sample, uniformly, a primitive action a; otherwise uniformly sample an option w from Ω if primitive action was sampled then state s, take action a and observe state s' and reward r
D Denr + Dear (s, a, r,,,
Append transition to current dataset
clse
while option not terminated do
In state 8, take action a sampled under the option policy, and observe state s' and reward Deur <- D euml (s, a, r,
Append transition to current dataset
end while
end if end for
b Learn the default representation
Update Z by Nieam sweeps through Deur using Equation 13 with step size and deviation importance À
Learn eigenoption
D +
DI|Deuur
Compute the top eigenvector of
Use Q-learning to compute an cigenoption w' using e, D, ag. and 70
Add the learned eigenoption the set of eigenoptions
Add w' to Ω and remove the oldest option from Ω if [S2] > Noption
end for
discovery. Note that by learning the SR instead of the DR Algorithm I we recover CEO 25 and therefore omit the corresponding pseudocode for brevity. In this version of RACE, given the complexity of performing importance sampling with options, we treat transitions generated from options the same those generated from the default policy, and use both for learning the DR.
We share the following hyperparameters for RACE and CEO: ag = 0.1, "o = 0.99, NNacp = 100, Niter 50 for grid task and four rooms, and Niter 120 for grid room and grid maze. We use 7 = 0.99 for learning the SR, 1.3 for learning the DR, and sweep over the remaining hyperparameters, described in Table 2 We perform 10 independent runs for each hyperparameter setting.
In Figure 4 we used solid dots to highlight the hyperparameter settings that lead to the highest state visitation percentage. We show the learning curves these hyperparameter settings in Figure 7 We can see that RACE maintains a similar state visitation percentage as CEO, but obtains much higher rewards when exploring the environments. This is because RACE, being reward-aware, avoids low- reward regions when exploring the environments. The random walk, on the other hand, struggles to
30

---

Table 2: Hyperparameter search for eigenoption discovery.
Name
Description
Values
Poption
Probability of selecting an option
[0.01, 0.05, 0.1]
Niam
Number of iterations through the collected dataset to learn the SR/DR
[1, 10, 100]
Step size for learning the SR/DR
[0.01, 0.03, 0.1]
N Noption
Number of latest options keep
1000]
Grid
Four Rooms
Room
Maze
State Vistatoon N
1.00
czo
RACE
1.00
1.00
1.00
RW
0.50
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0-
0.
01
0.
Cumutative Rowards
-4
-5
-4-
-8-
100
-8
-12
-12
-16
-16
-20
Stops (x10?)
Stops (x]0
Steps (X16?)
Staps (x10
Figure 7: The state visitation percentage (top) and cumulative rewards (bottom) for the highlighted hyperparameter settings in Figure iterative online eigenoption discovery using the SR (CEO) and the DR (RACE), and the random walk (RW). The solid line shows the average over 10 seeds, while the individual seeds are shown in lighter shade.
explore the state space, especially for the larger environments like grid room and grid maze, and does not encounter low-reward regions, causing to obtain even higher cumulative rewards than RACE.
Figure(8shhr shows the state visitation percentage for CEO, RACE, and the random walk in the variations of the environments in Figure [ without any low-reward regions. As the reward function is constant, the top cigenvectors of the SR and the DR are guaranteed by Theorem 3.1 to be identical. It thus does not come as a surprise that we observe very similar state visitation percentages for CEO and RACE,
We now describe details on combining iterative online eigenoption discovery with Q-learning (51). Specifically,
iterative online eigenoption discovery to collect transition data, and then perform offline Q-learning using the collected data. We compare RACE with Q-learning (RACE+Q) with CEO with Q-learning (CEO+Q), and a Q-learning baseline (QL). For the Q-learning bascline, at every iteration, we use the e-greedy policy induced by the current Q-values to collect transitions, and then update the Q-values using the collected transitions.
Rooms
Grid
Grid
Sate SateVisttat vnsiattoo
1.0 ICEO
RACE
1.0
1.0
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
5
5
Steps (x10')
Steps (x100)
Steps (x10')
Steps (x100))
Figure& 8: State visitation percentage for CEO, RACE, and the random walk (RW) in the environments shown in Figure [ without low-reward regions. Shown in lighter shade are the 10 individual seeds.
31

---

CEO
RAGE
CEO
RACE
CE0
RACE
CEO
RAGE
E
Figure 9: Cumulative state visits for CEO and RACE in the environments from Figure averaged over 10 seeds, where darker red indicates more visits. The low-reward regions are enclosed in blue. Although CEO is also applied the problem with low-reward regions, we emphasized them only for RACE because CEO does not see it. As indicated by the lighter red, RACE visits low-reward regions much less than CEO. It is especially clear in grid room (third env. from the left) that RACE takes detours to visit the bottom rooms without passing through low-reward regions.
Figure 10: Left: Larger version of Grid Room (Grid Room (L)); Right: Larger version of Grid Maze (Grid Maze (L)), adapted from prior work [50].
For RACE+Q and CEO+Q. we follow the same hyperparameter search procedure before (see Table 2. For the Q-learning baseline, we perform a grid search over the Q-value initializa- -
tion 10000001100 100, - 10. for e-greedy exploration ((0.01, 0.05, 0.1, 0.15, 0.2)), and step size
(J0.01, 0.03, 1)). We run 10 seeds for each hyperparameter setting, and after identifying the best hyperparameters, re-run 50 seeds avoid maximization bias. As we do not see much performance difference in simpler environments, we introduce two new environments that are larger versions of Grid Room (Grid Room (L)) and Grid Maze (Grid Maze (L)), shown in Figure T0
32

---

Table 3: Hyperparameters for Count-Based Exploration (Sarsa DR).
Environment
1
5
transform(x)
RIVERSWIM
0.5
0.25
100
log([||[12)
SIXARMS
0.5
0.01
0.1
1.5
log([|>|12)
C.4 Count-Based Exploration
To facilitate the learning and application of the DR for exploration in the Riverswim and Sixarms environments, we rescale the environment rewards to lie within the range [-1, 0J. Importantly this rescaling is applied solely for learning the DR; the Sarsa agent continues to operate using the original, unscaled rewards for action-valus estimation. Note that we use the DR defined for state-action-dependeer rewards. Z.
For Sarsa and SRr adopt the code (MIT license) and hyperparameters specified by the origi- nal authors (24] is to be noted that while they use Fimar(s) B3 [ II as the intrinsic reward (24),
we use Fintr(s, a) ,3.. log( log(IZ....:ll)) We empirically find that the logarithmic transformation enhances performance when leveraging the DR for exploration.
For Sarsa+DR, we sweep over different values of 77, ax 3, X, with 7 E {0.01, 0.1, 0.1.0.25,,0, a E {0.01, 0.2, 0.5). 8 E (0.1, 1. 10, 100} and À {1, 2). Here, 77 and denote the step sizes for updating the Q-values and the DR, respectively. /3 is the scaling factor for the intrinsic reward and A controls the relative importance of the deviation cost. We = 0.01 for e-greedy exploration. Additionally, we experimented with different transformations of the DR for computing the intrinsic reward, specifically, transform(X E {[|x| ||x||2, log([|>|]1 log([) (||x|[2)) where x is row of Z. The best hyperparameters, shown in Table 3, were evaluated on 100 independent runs.
C.5 Transfer
We describe the details of the features used for transfer lcarning. We assume that the agent has access to features that can perfectly represent terminal rewards for the DFs, and features that can perfectly represent the reward function for the SFs. For the DR, as there are four terminal states in the environment, and we assume the agent receives the same reward terminal state regardless of which action is chosen, we represent the features as a dimmensiooall one-hot vector, where each entry in the vector corresponds to one terminal state.
For the SR, we have to represent the reward function for all states. There are three types of states in the environment: empty state (indicated by white tile), low-reward state (indicated by red tile), and goal state (indicated by green tilc). All empty states have the same reward of -1, and all low-reward states have the same reward of -20. Therefore, vector, where one entry is activated when the state is an empty state, one entry is activated when the state is low-reward state, and the remaining four entries correspond to the four terminal states.
represent state features
6-dimensional one-hot
Given the features, we compute the DFs and the SFs under randomly sampled configurations of terminal state rewards, where the reward at each terminal state is sampled independently from a normal distribution with 0 mean and 50 standard deviation. We learn the DF's by following the default policy, While we only need to learn the DFs under the default policy, we can learn the SFs with respect to set of policies learned under different reward functions. The better the set of reward functions covers the space of reward functions, the better the transfer policy computed using the SFs. We consider the number of reward functions to be 1, 2, 4, and 8. For each randomly sampled reward function, we compute the SFs with respect to the optimal policy learned under this sampled reward function.
In this work, we ensure that the DFs and the SFs are well learned by allowing the agent to interact with the environment for large number of steps. For the DFs, the agent interacts with the environment while following the default policy for 100K steps. For the SFs, we first use Q-learning [5]] to learn an optimal policy for 100K steps. We then learn the SFs while following the optimal policy for another 100K steps. We use 1.3 and the uniform random policy as the default policy for the DFs, y 0.99 for the SFs, and use step size of 0.1 for both approaches. We perform 50 independent runs.
33

---

Note that the DFs computes optimal policy in the linearly solvable MDP setting that balances the reward function and the cost of deviating from the default policy. The optimal policies computed by the DFs are then softly biased towards the default policy. We mitigate this bias by approximating the optimal transfer policies (see Eq. 10 using deterministic policies greedy over the optimal Q-values.
D Statistical Significance
We report 95% confidence intervals or directly visualize all independent runs if possible. The 95% confidence intervals capture randomness across independent runs, where the primary source of randomness is random exploration, e.g., e-greedy exploration. In transfer learning experiments, randomness also arises from sampling terminal state reward configurations. As we perform large number of independent runs in experiments (N 50), assume the sampling distribution is nor- where ē is the sample standard deviation. mal and compute the 95% confidence interval as 1.96 VN"
E Compute Resources
We use CPUs for all of our experiments. We describe the runtimes for experiments involving the DR. For reward-shaping experiments, each independent takes under 10 minutes. For eigenoption discovery experiments, each independent run takes less than 2 hours in grid task and four rooms, and takes around 10 hours in grid room and grid maze. For count-based exploration experiments, each independent run takes less than one minute. For transfer experiments, each independent run takes less than minutes. Due the large number of independent runs performed for hyperparameter search and preliminary experiments, we estimate the total compute used for the project to be 10.5 CPU core years.
34


Reviews:
[{"review_id": "sBrY24EYkU", "reviewer": "Reviewer_M5W5", "summary": "The authors in this submission studied a similar representation to the popular successor representation (SR), by incorporating the rewards in MDPs. This default representation (DR) was first proposed in a related recent work from neuroscience, and the authors here showed how to extend DR to reinforcement learning. They demonstrated the relationship between SR and DR, and derived the algorithms to learn DR via both dynamic programming and temporal difference learning. Furthermore, they showed the usefulness of DR in different settings through various experiments.", "questions": "* How did you select $\\lambda$ in the DR?\n\n* What's the definition of ``Terminal Reward Configuration'' in Figure 5?\n\n* Is there any explanation on why the norm of the DR can be used as an exploration bonus in count-based exploration?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strength:\n\n* The extension from SR to DR is an interesting direction and this submission showed the promise of using DR in linearly solvable MDPS. \n\n* The authors managed to deriving efficient dynamic programming and TD learning algorithms, and showing the convergence for one of them.\n\n* Through various experiments the authors demonstrated the versatility of the proposed method.\n\n### Weakness:\n\n* The current version focuses on a few simple domains so it would be even more convincing if the authors can show superior performance in more complex environment, along with the extention with neural networks as the authors mentioned.", "was_revised": true, "final_justification": "As the authors addressed my comments and questions, I will keep my score.", "rebuttal": "We thank the reviewer for the thoughtful review. We are happy that the reviewer finds our work interesting and sees it as a promising direction, having only clarifying questions. Performing experiments only in the tabular setting was a conscious choice, as we discuss it below.\n\nWe focused on simple environments in our experiments to fully characterize the properties of the DR, since this paper is the first paper in computational RL on this framework. Including the results for more complicated environments using neural networks could distract from the paper's focus. We believe one has to establish the theoretical foundations and intuitions of an area before focusing on scaling it up, and, importantly, there’s nothing in our framework that suggests it is not scalable. Also, despite focusing on simple environments, as reviewer Ldct noted, the breadth of our experiments is large, exploring several downstream applications of the DR. Nevertheless, we agree with the reviewer that extending the application of the DR to more complex environments using neural networks is an interesting direction that we plan to pursue in the future.\n\nWe now address clarifying questions.\n\n>**Q1**: How did you select $\\lambda$ in the DR?\n\n*Summary*: We select a $\\lambda$ large enough such that our implementation would be numerically stable.\n\nWhile our algorithm is not sensitive to the $\\lambda$ value, in practice, some values of $\\lambda$ can lead to numerical instability. As shown by Equation 7, each entry of the DR involves taking the exponential of the return of trajectories. Since we used negative rewards for all non-terminal states, the exponential of the sum of negative rewards can be very small, leading to numerical issues. As using a larger $\\lambda$ has the effect of reducing the magnitude of the negative rewards in the computation, we increased $\\lambda$ by 0.1 until we settled on the final value of 1.3. Performance was not a consideration in this process.\n\nWe will clarify this $\\lambda$ selection by adding the following paragraph to Appendix C.1:\n\n“As a larger $\\lambda$ has the effect of reducing the magnitude of negative trajectory returns in the DR (see Equation 7), it can alleviate numerical instabilities caused by taking the exponential of very negative trajectory returns. In our experiments, we initially started with $\\lambda=1$, and slowly increased $\\lambda$ by 0.1 until we settled on $\\lambda=1.3$, which, paired with the use of python-flint, resulted in no numerical issues.”\n\n>**Q2**: What's the definition of ``Terminal Reward Configuration'' in Figure 5?\n\n*Summary*: “Terminal reward configuration” here refers to the rewards received upon entering terminal states. \n\nFigure 5 shows the results of transfer learning experiments, where we first learn the default features, before applying them in new scenarios where only the rewards at terminal states change. As shown on the left of Figure 5, the environment has several terminal states, indicated by the color green. When performing transfer, we randomly sampled rewards for each of these terminal states, and we called the set of sampled rewards at the terminal states as “terminal reward configuration”.\n\nWe will make this clearer by adding the phrase “where terminal state reward configuration refers to the set of sampled rewards at the terminal states” in Section 6.4.\n\n>**Q3**: Is there any explanation on why the norm of the DR can be used as an exploration bonus in count-based exploration?\n\n*Summary*: Machado et al. (2020) showed that the norm of the SR captures state visitation, and thus can be used for count-based exploration [1]. The DR similarly captures state visitation, but does so in a way that reflects rewards. \n\nWe first provide an intuition of why the SR can be used as an exploration bonus, before explaining the case for the DR. The SR matrix is usually initialized to all zeros, and the ground-truth SR matrix has all non-negative entries. Every time a transition $(s, a, r, s’)$ is observed, the SR of $s$ is updated to approach the ground-truth value. A state that is visited more is updated more, and its SR is closer to the ground-truth value. This is why the SR captures state visitation, and has been used in count-based exploration.\n\nThe DR can capture state visitation in a similar manner. For the DR, every entry of the ground-truth DR is non-negative. Entries along the diagonal have the largest magnitudes, while off-diagonal entries usually have very small magnitudes, so we chose to focus on the diagonal entries. We initialize the diagonal entries of the DR as the identity matrix (which, in practice, is an optimistic initialization), and every time state $s$ is visited, the corresponding diagonal entry in the DR matrix is updated to approach the ground-truth value from above. Therefore, the DR still captures state visitation and can encourage the agent to visit novel states. However, the DR is different from the SR in that the rate of approaching the ground-truth value is affected by the reward. If a state has a low reward, the corresponding diagonal entry decays faster and approaches the ground-truth value at a faster rate, meaning that the exploration bonus provided to the agent reduces faster. If a state has a high reward, the corresponding diagonal entry decays more slowly, and thus over the course of learning, provides a longer-lasting bonus for the agent. This reward-aware exploration bonus might be the reason why the DR was able to outperform the SR in count-based exploration.\n\n[1] Machado et al. Count-based exploration with the successor representation. AAAI 2020."}, {"review_id": "HKyy1T5lZp", "reviewer": "Reviewer_Ldct", "summary": "Based on the concept of the *default representation*, this paper\ninvestigates how the incorporation of reward signals into\nproto-representations (that is, representations that encode temporal\nsimilarity of states) can influence RL algorithms and tasks\ndownstream.\nFirst, the paper establishes the default representation (DR) as the\nfixed point of a dynamic programming scheme, and presents a\nTD-learning update rule to estimate the DR from samples.\nThis formalism is then generalized to default *features* (DF, much like\nsuccessor features), which can be easily represented / estimated even\nin continuous MDPs.\nIt is shown that optimal value functions can be efficiently extracted\nfrom the DR (and DFs), and additionally, the DR (resp. DFs) can be\nused to zero-shot infer optimal policies for any configuration of\nterminal rewards specified after training.\nThe paper concludes by demonstrating empirically that substituting the\nsuccessor representation for the DR in many RL approaches (e.g., option\ndiscovery, exploration bonuses) tends to lead to more\nefficient learning.", "questions": "I split this section into two parts, the first listing questions /\nconcerns that directly inform my score, and the second which lists\nless pressing questions for clarification.\n\n\n## Critical Questions / Concerns\n\nThere is an important argument missing in Theorem 3.1, which is\npartially addressed in its proof.\nNamely, you claim that the $i$th eigenvectors of the SR and the DR\nare the same, where it is implied that the eigenvectors here are\nranked by the scale of their corresponding eigenvalues.\nHowever, from equation (6), this result is not clear. In the proof of\nthis theorem, you claim that when $\\lambda_{SR, i}$ increases, then\n$\\lambda_{DR, i}$ increases. However, I would expect the relevant\nquantity to be the *magnitude* of the eigenvalue. Is it not possible\nthat, for example, $\\lambda_{SR, i} = -1000$ while $\\lambda_{DR, i} \\approx\n0$, in which case\nthe corresponding eigenvector is highly descriptive of the SR while\nbeing effectively irrelevant for the DR?\nTo be clear, I'm not necessarily convinced that this is possible, but\nI also believe it is not trivially impossible, and it is important to\nestablish.\n\nI have some concerns about equation (7). Particularly I find the\nnotation $\\tau_{s\\to s'}$ imprecise and confusing.\nFirstly, I believe this equation only makes sense, for multiple\nreasons, in the finite horizon case. Maybe most simply put, if you\nhave an MDP with a single state, there are $|\\mathbb{N}|$ trajectories\n$\\tau_{s\\to s}$ (following how this notation is used in the derivation\nof equation 7, where you enumerate over *non-distinct* trajectories),\neach having probability 1. Then the $\\mathbf{Z}(s, s)$ is not\nwell-defined (or rather, it is infinite whenever the reward is\nnonzero).\nIn the case of MDPs with more than one state, it's not immediately\nclear that the proof is correct, for the same reason. The only way I\ncan see the derivation in equation (25) converging generally is if\neither trajectories are finitely long, or there are terminal states,\nbut this isn't specified clearly anywhere.\n\nI also have concerns about the MER. I don't think the $\\mathbf{A}$\nterm is correct. The MER, as you've defined it, it precisely the same\nfor *any* MDP without rewards as long as it is possible to jump between\nany pair of states with nonzero probability. In this case, what\ninformation does the MER provide? Clearly, one could not zero-shot\ninfer optimal policies for instantiations of the terminal reward,\nbecause the optimal policies would be the same for any such MDP.\nI strongly suspect $\\mathbf{A}$ has to include information about the\ntransition probabilities between states.\nIn any case, Definition 5.1 is effectively meaningless on its own. You\nhave not established any useful properties of this object (indeed, as\nI describe just above, the generalization property of the DR doesn't\nappear to hold).\n\nSection 6.2 gives no details at all about how the proposed RACE\nalgorithm works, besides saying it's roughly the same as CEO\nsubstituting the SR for the DR. There should be at least a high level\ndescription of the algorithm, even if the details must be deferred to\nthe appendix.\n\n\n\n## Ancillary Questions\n\n1.  Is there a difference between the objectives of a linearly-solvable\n    MDP with default policy $\\pi_d$ and relative-entropy-regularized RL\n    with reference policy $\\pi_d$ (and, say, finite horizon with\n    $\\gamma=1$)? It would be helpful to explicitly write out the\n    objective.\n2.  In section 6.1, which policy are you estimating the SR from? Is it\n    just the default policy?\n3.  In section 6.1, can you give some intuition about why the\n    potential-based rewards make sense? Particularly in the case of the\n    SR (depending on which policy induces it), why should this\n    potential-based reward be meaningful?\n4.  In section 6.1, are you using the state-only DR leveraging the\n    exact reward function? Under which circumstances would you have\n    access to this DR but not to an optimal policy? Is the idea\n    primarily that you can use the DR to upgrade from a \"soft\" optimal\n    policy (i.e., optimal under the KL penalty) to a \"true\"\n    (un-regularized) optimal policy?", "limitations": "As mentioned above, there seem to be some missing assumptions\nregarding the DR. That said, I believe these assumptions are\nmentioned, just not as explicitly as they should be. Particularly, it\nseems that the DR requires finite episode horizon.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\nFirstly, this paper is visually stunning, and largely well written.\n\nI enjoyed reading about the DR, and I believe the contribution\nof formalizing this concept and establishing algorithmic techniques\nfor estimating it is quite timely.\n\nMoreover, despite the relatively small scale of the experiments, the\n*breadth* of the experiments is quite large, in my opinion. The purpose\nof the empirical section is not to demonstrate state of the art\nperformance, but to explore how the DR can be leveraged downstream,\nand I think the authors have done a great job at this. The results\nwith DR driving exploration bonuses and option discovery were\nsurprising and encouraging to me.\n\n\n## Weaknesses\n\nMy main criticisms of the paper are the following:\n\n1.  The math is quite imprecise, making some of the results difficult\n    to interpret, and in some cases, (I believe) incorrect without\n    further assumptions. I suspect most of this can be fixed by\n    clarifying the statements and assumptions.\n2.  The experimental section, particularly the section about options,\n    does not have enough details. Some of these sections would be\n    effectively impossible to follow / understand without reading the\n    appendix.\n\nMoreover, as I discuss in more detail later, the section on the\nmaximum entropy representation seems basically vacuous, and (in my\nview) it does not add to the paper. At the very least, I highly recommend for this\nsection to be moved to the appendix, at which point space could be\nused to substantially clear up the experimental sections.\n\nThe Questions section outlines more specifically some of the\nissues I highlighted above.\n\n## Minor Issues\n\n1.  Nit: regarding the TD-learning update rule for the SR shown in (3),\n    credit should be also given to the seminal work of Dayan ([7] in\n    your submission's reference list).\n2.  The notation $\\mathrm{KL}(p^\\pi \\| p^{\\pi_d} \\mid S_t)$ is\n    difficult to read and seemingly incorrect (though it's otherwise\n    reasonably clear what you're describing). Why not\n    $\\mathrm{KL}(p^\\pi(\\cdot\\mid S_t)\\| p^{\\pi_d}(\\cdot\\mid S_t))$?\n    It's more verbose, but it correctly conditions both probability\n    kernels on the state. Alternatively you can use kernel notation and\n    write $\\mathrm{KL}(p^\\pi_{S_t} \\| p^{\\pi_d}_{S_t})$, where\n    e.g. $p^\\pi_s\\in\\Delta(\\mathcal{S})$ for any $s\\in\\mathcal{S}$.\n3.  Nit: the notation $\\mathbf{Z}$ should specify $\\lambda$,\n    e.g. $\\mathbf{Z}^\\lambda$.\n4.  You are using $\\lambda$ for both the regularization strength of\n    linearly-solvable MDPs and for eigenvalues.\n5.  On line 174, \"Neurmann series\" should be \"Neumann series\".", "was_revised": true, "final_justification": "While the authors have improved on some of the issues I pointed out in the text, their theoretical results are still a little shaky -- see the discussion I had with them, where they require the introduction of several new assumptions (some of which are more severe than others). I am maintaining my score for now, curious to hear what other reviewers think in the proceeding discussion.", "rebuttal": "We thank the reviewer for the thoughtful and insightful feedback. We are glad that the reviewer finds our work visually appealing and well written, and the experiments broad and encouraging. Given limited space, we here address the critical concerns, and defer the response to remaining questions to a later post.\n\n>**Q1a**: You claim that the ith eigenvectors of the SR and the DR are the same, where it is implied that the eigenvectors here are ranked by the scale of their corresponding eigenvalues. However, from equation (6), this result is not clear.\n\n*Summary: Thank you for pointing this out. We now realize that there’s a minor flaw with Theorem 3.1, and further assumptions need to be made for the order of the eigenvectors to be the same for the SR and DR.*\n\nThe eigenvalues of the SR and DR for the same eigenvector are related by the following:\n\n$\\lambda_{SR}=f(\\lambda_{DR})=[\\gamma (\\lambda_{DR}^{-1}+\\gamma^{-1}-\\exp(-r(s) / \\lambda) ) ]^{-1}$\n\nIf we take the derivative, we have\n$f’(\\lambda_{DR})=(1/\\gamma)[1+\\lambda_{DR}(\\gamma^{-1}-\\exp(-r(s)/\\lambda))]^{-2}$\n\nFor $0 < \\gamma < 1$, this derivative is positive everywhere, except when $\\lambda_{DR}=-1/(\\gamma^{-1}-\\exp(-r(s)/\\lambda))$.\n\nWe overlooked the condition $\\lambda_{DR}=-1/(\\gamma^{-1}-\\exp(-r(s)/\\lambda))$, and falsely concluded that $f$ is a monotonic function, and so the order of the eigenvectors is always preserved. To fix this, we actually need one extra assumption:\nAll eigenvalues of the DR lie on the same side of $-1/(\\gamma^{-1}-\\exp(-r(s)/\\lambda))$.\n\nWhen this holds, the order of the eigenvectors is also preserved. Theorem 3.1 can be easily fixed by including this additional condition, and we just adjust our theorem statement to the following:\n\n“Suppose both the SR and DR are computed with respect to the same policy, i.e., $\\pi = \\pi_d$. When the reward function is constant, i.e., $r(s) = r(s’) \\forall s, s’ \\in \\mathcal{S}$, the SR and DR share the same set of eigenvectors, and the eigenvalues of the SR and DR for the same eigenvector are related as follows:\n\n$\\lambda_{SR} = [\\gamma (\\lambda_{DR}^{-1}-\\exp(-r(s)/\\lambda)+\\gamma^{-1})]^{-1}.$\n\nFurthermore, when $\\gamma^{-1}-\\exp(-r(s)/\\lambda)=0$ or when all eigenvalues of the DR lie on the same side of $-1/(\\gamma^{-1}-\\exp(-r(s)/\\lambda))$, the order of the eigenvectors of the SR and DR by their corresponding eigenvalues is the same.“\n\n>**Q1b**: In the proof of this theorem, you claim that when $\\lambda_{SR,i}$ increases, then $\\lambda_{DR,i}$ increases. Is it not possible that, for example, $\\lambda_{SR,i}=−1000$ while $\\lambda_{DR,i}\\approx0$, in which case the corresponding eigenvector is highly descriptive of the SR while being effectively irrelevant for the DR? \n\n*Summary: We agree with the reviewer, and will make a small adjustment in our argument to support the claim that the DR is a generalization of the SR.*\n\nTheorem 3.1 has been used to claim that the DR is a generalization of the SR: when the reward function is constant, they capture similar information because they share the same eigenvectors, and the order of eigenvectors is also the same. However, the reviewer is right that merely having the same order does not mean that the two matrices capture similar information. Nevertheless, the DR can still be seen as a generalization of the SR, since for the special case of $\\exp(-r(s) / \\lambda) = \\gamma^{-1}$, the order of the eigenvectors is preserved, and the magnitude only differs by a constant $\\gamma$.\n\nWe now provide a full account of how DR is a generalization of the SR:\n1. When the reward function is constant and $\\exp(-r(s) / \\lambda) = \\gamma^{-1}$, the i-th eigenvectors of the DR and SR are identical, and the eigenvalues only differs by $\\gamma$.\n2. When the reward function is constant, and $\\exp(-r(s) / \\lambda) \\neq \\gamma^{-1}$, the DR and SR share the same set of eigenvectors. Furthermore, when all eigenvalues of the DR lie on the same side of $-1 / (\\gamma^{-1} - \\exp(-r(s) / \\lambda))$, the order of the eigenvectors is the same.\n3. When the reward function is not constant, the DR and SR can have different eigenvectors.\n\nWe will highlight the special case of $\\exp(-r(s) / \\lambda) = \\gamma^{-1}$ to better support our claim.\n\n>**Q2**: In equation (7), I find the notation $\\tau_{s \\to s'}$ imprecise and confusing. Firstly, I believe this equation only makes sense, for multiple reasons, in the finite horizon case. Maybe most simply put, if you have an MDP with a single state, there are |N| trajectories $\\tau_{s \\to s'}$ (following how this notation is used in the derivation of equation 7, where you enumerate over non-distinct trajectories), each having probability 1. Then Z(s,s) is not well-defined (or rather, it is infinite whenever the reward is nonzero). In the case of MDPs with more than one state, it's not immediately clear that the proof is correct, for the same reason. The only way I can see the derivation in equation (25) converging generally is if either trajectories are finitely long, or there are terminal states, but this isn't specified clearly anywhere.\n\n*Summary: We will adjust our notation for the trajectories. The reviewer is right that equation 7 relies on an additional assumption, which we only mentioned later in the paper in Theorem 4.1: the assumption that non-terminal state rewards are all negative.*\n\nWe used $\\tau_{s \\to s’}$ to refer to distinct trajectories from $s$ to $s’$, which may be confusing for readers. We will switch to defining distinct trajectories from $s$ to $s’$ by the following:\nLet $\\tau_{s\\to s’}$ denote a trajectory from s to s’.\nLet $\\mathcal{T}\\_{s\\to s’}=${$\\tau\\_{s\\to s’}|\\mathbf{P}^\\pi(\\tau\\_{s\\to s’})>0$} \nAnd when summing over trajectories, we will use $\\sum\\_{\\tau\\in\\mathcal{T\\_{s\\to s'}}}$ instead. Now with the set notation, it will be clear that we are referring to distinct trajectories.\n\nAs for the concern about whether Z(s, s’) is defined, it is defined because we assume that non-terminal state rewards are negative. Therefore, even if a state can transition to itself with probability 1, the contribution of longer trajectories to the summation decays exponentially. While we mentioned this assumption in Theorem 4.1, we agree with the reviewer that we could have mentioned it explicitly earlier. We will add this to Section 2.2. Note that we do not see this assumption as a restriction of the framework (see our response to reviewer M3b3 Q2).\n\n>**Q3**: I also have concerns about the MER. I don't think the A term is correct. The MER, as you've defined it, it precisely the same for any MDP without rewards as long as it is possible to jump between any pair of states with nonzero probability. In this case, what information does the MER provide? Clearly, one could not zero-shot infer optimal policies for instantiations of the terminal reward, because the optimal policies would be the same for any such MDP. \n\n*Summary: The reviewer is correct that the MER would be the same regardless of specific transition probabilities when the graph connectivity stays the same. We claim that this is a property of the MER, rather than an error in our derivation.*\n\nWe first provide intuition in the case of a deterministic transition function where each state-action pair leads to a different next state with probability 1. First, we can show that the linearly solvable MDP formulation can be transformed into a maximum entropy formulation by choosing the default policy to be the uniform random policy:\n\n$r(s)-\\lambda KL(\\pi(\\cdot, s)\\|\\pi_d(\\cdot,s))=r(s)-\\lambda\\mathbb{E}\\_{a\\sim\\pi}[\\log\\pi(a|s)]+\\lambda\\mathbb{E}\\_{a\\sim\\pi}[\\log\\pi_d(a | s)]=r(s)-\\lambda\\log|A|+\\lambda\\mathcal{H}(\\pi(\\cdot|s))$\n\nBy choosing $\\pi_d$ to be the uniform random policy, the linearly solvable MDP formulation is equivalent to a maximum entropy formulation where the rewards are shifted by a constant. In this setting, note that by following the uniform random policy, the agent in fact has a uniform probability of visiting every next state. It then makes intuitive sense for the MER to not capture the specific transition probabilities, since the policy with respect to which the MER is computed assigns a uniform distribution over transitions.\n\nFor stochastic transitions, the MER can be thought of as being computed with respect to a \"generalized\" uniform random policy. An agent following this policy has a uniform transition probability to every next state. Therefore, it is a property of the MER to remain the same when specific transition probabilities change, but the connectivity of the graph remains the same. \n\nWe disagree with the reviewer that one-shot optimal policy inference cannot be performed in this case. It is possible to use the MER, in a similar manner as the DR, to perform transfer. However, the optimal policy we obtain will have the property of disregarding actual transition probabilities, and treating all trajectories as equally likely. This optimal policy may not be useful in the standard settings that we usually consider, but could be informative in certain settings. For example, this optimal policy can possibly be used as an “optimistic” policy that assumes that the agent can transition to a desirable next state with the same probability as any other next state, regardless of how low the actual probability is.  \n\nFinally, we provide a more general view of the SR, DR, and MER: The DR is the most general form that depends on both the transition probabilities and reward function. The SR is a special case when we disregard the reward information. The MER, on the other hand, can be thought of as another special case of the DR that considers rewards but disregards the transition probabilities.\n\nNevertheless, we agree with the reviewer that given the limited exploration of the properties of the MER, including it in the main text can obstruct the flow without providing much value. As the MER is not a central contribution of our paper, we will move it to the Appendix."}, {"review_id": "8UhR4qjHwx", "reviewer": "Reviewer_ksrY", "summary": "The paper presents an extension of the successor feature framework to incorporate the default representation. This representation is obtained by computing the expected future rewards under a reference policy in a similar manner to the state-visitation counts used in the successor feature literature. The authors show several foundational results tying the default representation to the successor representation, present simple but tractable algorithms for computing it, and provide evidence for the usefulness of the representation in simple domains.", "questions": "Are there requirements on the default policy? I assume that it has to be somewhat reasonable for solving the MDP because otherwise no better policy can be learned by deviating from it, correct?\n\nWould it be possible to use the default representation in a policy iteration style where a new improved policy becomes the next iterations default policy? If so, would there be efficient algorithms to shift or update the representations to account for the new ooptimal policies, or would all relevant quantities have to be recomputed?\n\nSeveral equations (espec. the policy update formulation) are highly reminiscent of Relative Entropy Policy Search. It might be good to highlight the connection and cite prior work here, especially as a line of empirically strong relevant algorithms has been developed which could readily be combined with the idea of default features.\n\n[1] Relative Entropy Policy Search, Peters et al., AAAI 2010", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Overall, the paper provides a very clear and interesting account of a novel representation learning scheme which extends the established successor feature regime in a valuable manner. The papers theoretical results are clearly established, well explained, and the proofs are correct as far as I can tell. While the paper only provides small scale experiments and doesn't expand the framework to work with modern deep learning architectures, this limitation was clearly stated in the scope of the paper, and so I do not see this as a weakness. The benefits of the default representation over successor features seems intuitively clear to me, which is also due to the overall high quality of writing which establishes (almost) all concepts very clearly.\n\nWeaknesses\n\nAll weaknesses stated here are minor and should not be seen as barriers to acceptance. I list them in the interest of pointing out potential areas of improvement.\n\nIt would have been useful to add slightly more exposition on the properties of the successor features, for readers not immediately familiar with the field. For example, while the eigenvectors of the successor features are mentioned as \"useful\", some short exposition on what they intuitively characterize would be great. (line 128f)\n\nThe default feature learning approach in Section 5.1 is mentioned to only work with constant non-terminal rewards. However, this limitation is only acknowledged in the expository text of the section, not in the description of the method. It would be helpful to point out how this limitation arises, and maybe (if possible) provide some notes how it could be addressed.\n\nSection 5.2 seems slightly superfluous and somewhat breaks the flow of the paper. While I greatly appreciate that the authors show further extensions of the method, I think this section should either be introduced a bit more clearly, or maybe moved to the appendix? Part of what makes this section slightly worse than the surrounding ones is that relevant concepts, such as the adjacency matrix, are only introduced in the appendix. Alternatively it would be possible to slightly extend the section, but given the tight space, that might be difficult.\n\nMinor nitpicks:\n\n- Please state what distributional assumptions underlie the 95% confidence intervals shaded in the plots? What method was used to calculate the bounds?", "was_revised": true, "final_justification": "I keep my positive score. The answers of the authors to my (minor) questions were satisfactory.", "rebuttal": "We thank the reviewer for the valuable feedback. We are glad that the reviewer finds our paper clear and interesting. We will first address the potential areas of improvements suggested by the reviewer, before moving to the questions.\n\n> **Q1**: It would have been useful to add slightly more exposition on the properties of the successor features, for readers not immediately familiar with the field. For example, while the eigenvectors of the successor features are mentioned as \"useful\", some short exposition on what they intuitively characterize would be great. (line 128f)\n\nIf our paper is accepted, given the extra page, we can make use of the space to add more explanation for the successor representation/features. As an example, we can make the following change to line 128:\n\n“This result is particularly interesting because the eigenvectors of the SR have been shown to be useful in settings such as reward shaping and option discovery. The top eigenvector of the SR captures temporal distance between states, and such distance has been used for reward shaping as it leads to better performance than distances that are based on the coordinates of states [42, 38]. In option discovery, it has been shown that eigenoptions, which are obtained by training a policy to maximize a reward signal induced by the eigenvectors of the SR, allow the agent to reach less explored areas of the state space, facilitating exploration [21]. We will revisit these applications in Section 6.”\n\nApart from this change, we will revise the whole paper to provide more explanation wherever possible.\n\n>**Q2**: The default feature learning approach in Section 5.1 is mentioned to only work with constant non-terminal rewards. However, this limitation is only acknowledged in the expository text of the section, not in the description of the method. It would be helpful to point out how this limitation arises, and maybe (if possible) provide some notes how it could be addressed.\n\n*Summary*: The limitation arises from the fact that default features (DFs) are dependent on non-terminal state rewards. As for ways to address this limitation of the framework, in practice, it may be possible to simply allow the agent to explore the areas around the non-terminal states for which the rewards have changed. \n\nWe will add the following to the end of line 206 of Section 5.1 to better explain how this limitation arises:\n“Note that such a form of optimal value computation is only possible when only the rewards at terminal states change. The reason is because, when only the terminal state rewards change, $\\zeta(s)$ stays the same for all non-terminal $s$. However, when the rewards for non-terminal states change, $\\zeta(s)$ changes since $\\zeta(s) = \\mathbf{Z}_{NN} \\mathbf{P}\\_{NT}^{\\pi\\_d} \\mathbf{\\Phi}$ and $\\mathbf{Z}\\_{NN}$ depends on non-terminal state rewards.”\n\nWithout assuming access to the transition and reward functions, one way to address the change in non-terminal state rewards is to allow the agent to interact with the environment and update the default features to reflect the changes. This method, while seemingly naive, is in fact reasonable since, without assuming access to environment dynamics, the agent has no way to be aware of changes in the environment without experiencing the changes itself. However, since the default feature of every non-terminal state is dependent on other states, even if the reward has not changed in a non-terminal state, the agent still needs to visit it so that changes in other non-terminal states are reflected. In practice, it might be possible to get away with just visiting local regions around the non-terminal states for which the reward has changed, since the contribution of the DFs of a state to those of other states decays exponentially as the distance increases (see Equation 16). Therefore, DFs of states far away from the non-terminal states for which rewards have changed may still be good enough approximations of their ground-truth DFs.\n\nThis way of addressing the changes in non-terminal state rewards for the DFs is speculative, and we have not performed experiments in this particular case, so we do not plan on including this discussion in the paper. Nevertheless, we find this direction very interesting and will pursue this further in our future work.\n\n>**Q3**: Section 5.2 seems slightly superfluous and somewhat breaks the flow of the paper. \n\nWe will move the introduction of the MER to the Appendix.\n\n>**Q4**: Please state what distributional assumptions underlie the 95% confidence intervals shaded in the plots? What method was used to calculate the bounds?\n\nWe discussed the computation of 95% confidence intervals in Appendix D. We now realize that readers might not notice it due to the lack of pointers in the main text, and will add pointers.\n\n>**Q5**: Are there requirements on the default policy? I assume that it has to be somewhat reasonable for solving the MDP because otherwise no better policy can be learned by deviating from it, correct?\n\n*Summary*: To ensure that the deviation cost from the default policy is well-defined for any agent policy, the default policy needs to assign non-zero probability to all actions in all states. Without prior knowledge of the problem, the standard choice is the uniform random policy. \n\nMathematically, in order for the deviation cost from the default policy, which is modeled as the KL divergence between the transition probabilities under the agent’s policy $\\pi$ and the default policy $\\pi_d$, to be well defined, the support of the transition probabilities under $\\pi_d$ should be wider than that under $\\pi$. To ensure that this is true for any $\\pi$, $\\pi_d$ should assign non-zero probability to all actions in all states. We will explicitly state this requirement in line 81, when the default policy is first mentioned. \n\nAs long as the above requirement is satisfied, we can choose any default policy, though, as pointed out by prior work [1], and also noticed by the reviewer, the choice of the default policy affects the quality of the optimal policy. In general, the optimal policy is softly biased towards the default policy. For example, if the default policy assigns large probabilities to $a_1$ in $s_1$, the optimal policy will be inclined to do the same. This is why, without prior knowledge of the problem, the standard practice is to use the uniform random policy.  \n\n>**Q6**: Would it be possible to use the default representation in a policy iteration style where a new improved policy becomes the next iterations default policy?\n\n*Summary*: Our option discovery algorithm, RACE, is related to the algorithm suggested by the reviewer. It is possible to construct such forms of algorithms, and they can be useful.\n\nIn RACE, we in fact compute a DR with respect to a constantly-evolving default policy. This is because, as we compute and follow options, the transitions are not collected strictly under the default policy anymore. Rather, they come from a mixture of the default policy (uniform random) and the computed options. At every iteration, the default representation induces an option, which changes the behavior policy. The new behavior policy can now reach states that are further away from the start state and at the frontier of exploration in a limited number of time steps. The default representation is then updated with respect to this new behavior policy, and induces a new option, and the cycle goes on. Therefore, the idea suggested by the reviewer is possible, and can be useful in certain scenarios. \n\n>**Q7**:  If so, would there be efficient algorithms to shift or update the representations to account for the new optimal policies, or would all relevant quantities have to be recomputed?\n\n*Summary*: When the default policy changes, there are, in general, no efficient algorithms for computing the new DR, unless we assume access to transition dynamics. However, there are two possible ways we can address this. \n\nIn RACE, we do not have an efficient method to directly compute the DR for the new default policy, and only update the DR using transitions and TD learning. While a prior work [1] has demonstrated using matrix inversion lemmas to directly compute the new DR when the default policy changes, such computation is only possible if the transition dynamics is known. For the general case when the environment dynamics is not known, we are currently not aware of efficient algorithms to directly update the DR. \n\nThere are two possible directions to address this: 1) If only the default policy changes, and environment dynamics stays constant, it might be possible to update the default representation using a replay buffer and importance sampling. 2) We can perhaps use the generalization ability of a neural network to help generalize the default representation to new policies, similar to USFA [2]. For example, a neural network can take the state and an encoding of a policy as the input, and output the corresponding default representation. \n\n[1] P. Piray and N. D. Daw. Linear reinforcement learning in planning, grid fields, and cognitive control. Nature Communications, 12(1):4942, 2021.\n\n[2] Borsa et al. Universal successor features approximators. ICLR 2019.\n\n>**Q8**: Several equations (espec. the policy update formulation) are highly reminiscent of Relative Entropy Policy Search.\n\nWe thank the reviewer for pointing out relevant work. Relative entropy policy search constrains the KL divergence between the data distribution induced by the policy and the data distribution of past observed data. This indeed is similar to linearly solvable MDPs, as both frameworks use KL divergence to avoid deviating too much from a reference distribution. One difference is that in linearly solvable MDPs, the KL term is added to the reward function, while in relative entropy policy search, the KL term is treated as a constraint of the optimization. Given the similarities, we will cite this and related papers."}, {"review_id": "yHMBYtlmBn", "reviewer": "Reviewer_M3b3", "summary": "This paper extends the theoretical foundations of the Default Representation (DR), a reward-aware proto-representation in reinforcement learning. Building on prior neuroscience-focused work, the authors develop dynamic programming and temporal-difference learning methods for DR acquisition, characterise its vector space properties relative to the Successor Representation (SR), introduce default features for function approximation, and extend the framework to maximum entropy RL. Empirical evaluations demonstrate DR's advantages over SR in reward shaping, option discovery, exploration, and transfer learning tasks.", "questions": "1. What are the key challenges in applying DR to real-world RL problems compared to established proto-representations? \n2. Does DR's performance advantage persist in environments without penalty regions?", "limitations": "The paper briefly mentions but under-analyses the limitations. Additionally, the biological plausibility claims from earlier DR work are not reconciled with the purely computational focus.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "STRENGTHS:\n* Provides the first complete derivation of DP/TD methods for DR learning, filling a key gap in prior theoretical work. \n* Offers valuable insights into DR's mathematical properties through eigenvectors analysis and trajectory-based interpretations.\n* Introduces default features that enable DR usage without transition dynamics access, enhancing practical applicability.\n* Demonstrates DR's utility across multiple RL paradigms (reward shaping, exploration, transfer) with consistent improvements over SR.\n* Provides functional code implementations supporting claimed results.\n\nWEAKNESSES:\n* The SR's reward-agnostic nature is framed as a weakness, though this design enables flexibility across reward functions - a core SR strength not sufficiently acknowledged.\n* All experiments use negative reward regions, leaving unclear how DR performs in neutral/positive reward landscapes.\n* Focus is exclusively put on SR comparison, while omitting comparisons to modern representation learning approaches (e.g., bisimulation metrics, deep successor features).\n* Maximum entropy extension (MER) receives limited analysis despite being a claimed contribution.", "was_revised": true, "final_justification": "The authors have addressed all of my questions, and concerns from other reviewers as well, hence I will keep a positive score.", "rebuttal": "We thank the reviewer for insightful comments. We are glad that the reviewer finds our work valuable. We will now address the areas of improvements suggested by the reviewer.\n\n> **Q1**: The SR's reward-agnostic nature is framed as a weakness, though this design enables flexibility across reward functions - a core SR strength not sufficiently acknowledged.\n\nWhile we did not intend to frame the reward-agnostic nature of the SR as a weakness, because  our paper explores whether reward-awareness can be beneficial, we have focused extensively on scenarios where being reward-aware is important. However, we agree with the reviewer that the reward-awareness of the DR comes at a price of reduced flexibility, especially in the context of transfer. We will adjust our narrative to better reflect this trade-off. For example, in Section 5.1, when discussing default features, we can add the following after line 218:\n\n“While DFs enable directly computing the optimal policy, they are limited to the scenario where only the terminal state rewards change. On the other hand, while SFs can only compute a policy as good as the ones used to compute the SFs, they can be applied when any state rewards change, allowing more flexibility than DFs.”\n\n> **Q2**: All experiments use negative reward regions, leaving unclear how DR performs in neutral/positive reward landscapes.\n\n*Summary: We assume rewards at non-terminal states to be negative, which is necessary for the DR and dynamic programming/temporal difference learning methods to be well-behaved. For environments with neutral/positive reward landscapes, we can simply transform the rewards to be negative before applying our framework.*\n\nAs shown in Theorem 4.1, the reward for all non-terminal states being negative is a necessary condition for the dynamic programming (DP) method to converge, and for the DR to be well-defined. This is why we use negative rewards for all non-terminal states in our experiments. \n\nHowever, we do not see this as a restriction of this framework, since we can transform a reward function with neutral/positive rewards into one with only negative rewards. For example, in the count-based exploration experiments, the two environments (Riverswim and SixArms) consist of non-negative rewards. In order to apply the DR, denoting the maximum reward by $R_{max}$, we transform every reward by $(r(s, a) - R_{max} - 1) / R_{max}$. Note that the reward transformation is only performed for the DR learning, and for the SARSA part, we use the original reward function. More generally, for a reward function that is non-negative, after applying the transformation, we expect that the DR would treat states (or state-action pairs) having the smallest positive reward as low-reward regions, and states having large positive reward as more desired regions.\n\n> **Q3**: Focus is exclusively put on SR comparison, while omitting comparisons to modern representation learning approaches (e.g., bisimulation metrics, deep successor features).\n\nWe focus on comparing with the SR since our goal is to contrast and characterize the difference between being reward-aware and reward-agnostic. While it would be good to add comparisons, including such comparisons could shift the focus of the paper to comparing the performances of different representation learning methods, rather than establishing the fundamentals of the DR and characterizing the properties of a reward-aware proto-representation. \n\n> **Q4**: Maximum entropy extension (MER) receives limited analysis despite being a claimed contribution.\n\nWe thank the reviewer, and other reviewers, for pointing this out. We will adjust our contribution, and move the discussion of the MER to the Appendix.\n\n> **Q5**: What are the key challenges in applying DR to real-world RL problems compared to established proto-representations?\n\nOne key challenge is that a naive implementation of the framework leads to instability. This is a mechanistic limitation, not a conceptual one.\n\nDue to taking the exponential of negative rewards, when rewards have a large magnitude, the exponential can become very small, causing numerical issues. In our experiments involving small tabular environments, we were able to overcome this using a specialized library for more precise numerical computations. In more complicated environments, this might not be easy. Even if we normalize the reward function, in complicated environments with long horizons, the accumulation of negative rewards over a large number of time steps could still make the problem persist when trying to approximate the DR. \n\n> **Q6**: Does DR's performance advantage persist in environments without penalty regions?\n\nAs we also demonstrated in Figures 6 and 8 in Appendix C, in environments without penalty regions, the DR and the SR have very similar performance. \n\n> **Q7**: The paper briefly mentions but under-analyses the limitations. \n\nWe thank the reviewer for pointing this out. If our paper is accepted, we will use the extra space to better discuss the limitations. We plan to add the following paragraph:\n\n“While the DR provides benefits in environments where being reward-aware is important, there are some limitations. First, naively computing the DR and its eigendecomposition can lead to numerical instabilities due to taking the exponential of negative rewards, as discussed in Appendix C.1. This could be problematic, especially in more complicated environments with long horizons. Second, as we show in Figures 6 and 8 in Appendix C, in environments in which all non-terminal state rewards are the same, the DR will perform as well as the SR. Given the numerical issues associated with the DR, the SR may be a better choice in these environments. Third, while the DR captures the environment dynamics more fully, this comes at the expense of flexibility. While the SR, which is reward-agnostic, can allow transfer learning when any state rewards change, the DR, which is reward-aware, can only perform transfer in settings where only the terminal state rewards change.”\n\nThis paragraph would also address the reviewer’s earlier concern that the flexibility of the SR is not sufficiently acknowledged.\n\n\n> **Q8**:Additionally, the biological plausibility claims from earlier DR work are not reconciled with the purely computational focus.\n\nWhile biological plausibility inspired us to perform this work, our focus is on bridging the gap in the theoretical formulation of the DR so that it can be applied in computational RL. Discussing biological plausibility would fall outside of the scope of the paper."}]

Author rebuttal:

Meta-review: The successor representation (SR) has been widely used in reinforcement learning, for example, for state counting. The paper theoretically and empirically analyzes the default representation which, contrary to SR, takes also reward dynamics into account.

Strengths of the paper include the analysis of DR which is of interest to reinforcement learning research, the theoretical proofs, derived new algorithms/methods, and the writing. 

A weakness of the paper is the focus in experiments on simple domains but this is ok taking into account the focus on analyzing DR.

Overall, the paper provides a valuable contribution that will most likely be used by other research.

During the rebuttal period, there was a very useful discussion mainly between reviewer Ldct and the authors regarding the correctness of the equations in the paper. Due to the discussion the authors promised to modify the assumption of the second part of Theorem 3.1. I encourage the authors to incorporating all the comments that relate to making the theoretical analysis clearer and formally specifying all the required assumptions.
Decision: Accept (poster)