Learning Successor Features the Simple Way

Authors: raymond  chua
arna  ghosh
christos  kaplanis
blake a richards
doina  precup
mila ¶ google deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 03:44:37

Abstract: In Deep Reinforcement Learning (RL), it is a challenge to learn representations that do not exhibit catastrophic forgetting or interference in non-stationary environments. Successor Features (SFs) offer a potential solution to this challenge. However, canonical techniques for learning SFs from pixel-level observations often lead to representation collapse, wherein representations degenerate and fail to capture meaningful variations in the data. More recent methods for learning SFs can avoid representation collapse, but they often involve complex losses and multiple learning phases, reducing their efficiency. We introduce a novel, simple method for learning SFs directly from pixels. Our approach uses a combination of a Temporal-difference (TD) loss and a reward prediction loss, which together capture the basic mathematical definition of SFs.  We show that our approach matches or outperforms existing SF learning techniques in both 2D (Minigrid) and 3D (Miniworld) mazes, for both single and continual learning scenarios. As well, our technique is efficient, and can reach higher levels of performance in less time than other approaches. Our work provides a new, streamlined technique for learning SFs directly from pixel observations, with no pretraining required.

Main paper: 
 introduction: 
 Deep reinforcement learning (RL) [Sutton and Barto, 2018 ] is important to modern artificial intelligence (AI), but standard approaches to deep RL can struggle when deployed for continual learning [Parisi et al., 2019 , Hadsell et al., 2020 , Khetarpal et al., 2022] . When either the reward function or the transition dynamics of the environment changes, standard deep RL techniques, such as deep Q-learning, will either struggle to adapt to the changes or they will exhibit catastrophic forgetting [Kirkpatrick et al., 2017 , Kaplanis et al., 2018] . Given that the real-world is often non-stationary, better techniques for deep RL in continual learning are a major goal in AI research [Rusu et al., 2016 , Rolnick et al., 2019 , Powers et al., 2022 , Abbas et al., 2023 , Anand and Precup, 2024] .One potential solution that researchers have explored is the use of Successor Features (SFs). Successor Features, the function approximation variant of Successor Representations (SRs) [Barreto et al., 2017] , decompose the value function into a separate reward function and transition dynamics representation [Dayan, 1993] . In doing so, they make it easier to adapt to changes in the environment, because the network can relearn either the reward function or the transition dynamics separately [Borsa et al., 2018 , Barreto et al., 2018 , 2020 , Hansen et al., 2019 , Lehnert and Littman, 2019 , Liu and Abbeel, 2021] . Furthermore, there are theoretical guarantees that SFs can improve generalization in multi-task Figure 1 : (a) Results from a single task within a 2D two-room environment, illustrating the suboptimal performance of the canonical Successor Features (SF) learning rule (Eq. 4) due to representation collapse. (b) In the canonical SF approach, the average cosine similarity between pairs of SFs converges towards a value of 1, demonstrating representation collapse occurs. (c) The canonical SF learning rule does not develop distinct clusters in its representations, as evidenced by lower silhouette scores and higher Davies-Bouldin scores, which again indicates representation collapse. A mathematical proof can be found in section 3.4. settings [Barreto et al., 2017] . SFs are therefore a promising candidate for deep RL in non-stationary settings.However, learning SFs is non-trivial. The most straightforward solution, which is to use a temporaldifference (TD) error on subsequent observations [Barreto et al., 2018] , can lead to representational collapse, where the artificial neural network maps all inputs to the same point in a high-dimensional representation space. This phenomenon is commonly observed in various deep learning pipelines that end up learning similar or identical latent representations for very different inputs [Aghajanyan et al., 2020] . In RL, representation collapse can lead to different states or state-action pairs being mapped to similar representations, leading to suboptimal policy decisions or inaccurate estimation of values.To solve this problem, a variety of solutions have been proposed. One solution is to use an additional reconstruction loss [Kulkarni et al., 2016 , Zhang et al., 2017 , Machado et al., 2020] to force the network to maintain information about the inputs in its representations. Another solution is to use extensive pretraining coupled with additional loss terms to encourage high-entropy representations [Hansen et al., 2019, Liu and Abbeel, 2021] . More recently, an alternative solution using loss terms to promote orthogonal representations has been put forward [Mahadevan and Maggioni, 2007, Machado et al., 2017b] . Finally, an unconventional approach integrates Q-learning and reward prediction losses with the SF-TD loss, enhancing the learning process by providing additional supervisory signals that improve the robustness and effectiveness of the successor features [Janz et al., 2019] . This method allows the network to simultaneously learn the basis features, successor features, and task encoding vector, with the hope that the learned variables will satisfy their respective constraints.Though these solutions prevent representational collapse, they can impair learning, introduce additional training phases, or add expensive covariance calculations to the loss function [Touati et al., 2022] . Ideally, there would be a way to learn deep SFs directly during task engagement with a simple, easy to calculate loss function.Here, we introduce a simple technique for learning SFs directly during task engagement. We designed a neural network architecture specifically to achieve this training objective. Our approach leverages the mathematical definition of SFs and constructs a loss function with two terms: one that learns the value function with a TD-error, and another that enforces representations that make the rewards linearly predictable. By mathematical definition, this loss is minimized when the system has learned a set of SFs. We show that training with this loss during task engagement, facilitated by our neural network architecture, leads to the learning of deep SFs as well as, or better than, other approaches. It does so with no pretraining required and very minimal computational overhead. As well, we show that our technique improves continual reinforcement learning in dynamic environments, in both 2D grid worlds and 3D mazes. Altogether, our simple deep SF learning approach is an effective way to achieve the benefits of deep SFs without any of the drawbacks. 
related work: 
 Our work builds on an extensive literature on decomposing the value function dating back to the 1990s [Dayan, 1993] . More recent work on learning deep SFs falls broadly into three categories of solutions. The first are solutions that use a reconstruction term in the loss function in order to avoid representational collapse [Kulkarni et al., 2016 , Zhang et al., 2017 , Machado et al., 2020] . This general approach is effective at avoiding collapse, but it can lead to impaired performance on the actual RL task, as we show below. The next set of solutions rely on hand-crafted features [Lehnert et al., 2017 , Barreto et al., 2018 , Borsa et al., 2018 , Madarasz and Behrens, 2019 , Machado et al., 2021 , Emukpere et al., 2021 , Nemecek and Parr, 2021 , Brantley et al., 2021 , McLeod et al., 2022 , Alegre et al., 2022 , Reinke and Alameda-Pineda, 2021] or hand-crafted task knowledge [Hansen et al., 2019 , Filos et al., 2021 , Liu and Abbeel, 2021 , Carvalho et al., 2023a] . In these cases, the networks can learn and generalize well, but hand-crafted solutions cannot scale-up to real-world applications. Another category of solutions uses pretraining of the features in the deep neural network before any engagement with the actual RL task [Fujimoto et al., 2021 , Abdolshah et al., 2021 , Touati et al., 2022 , Carvalho et al., 2023b] . Such solutions are not as applicable for continual RL because they introduce the need to engage in new pretraining when the environment changes, which assumes some form of oracle knowledge of the environment. Finally, there are solutions that rely on additional loss terms to encourage orthogonal representations, since SRs are built off of purely orthogonal tabular inputs [Touati et al., 2022 , Farebrother et al., 2023] . These techniques can improve SF learning, but they require computationally expensive calculations of orthogonality in the basis features.Among these prior approaches, the work most closely related to ours is the application of multiple losses to jointly learn the SFs, a task-encoding vector, and Q-values [Ma et al., 2020] . However, there are several key differences: (1) Our approach does not require the agent to be provided with a goal-it is learned through interaction with the environment; (2) We provide direct evidence that our method works with pixel inputs; (3) We demonstrate that our approach eliminates the need for an SF loss; and (4) By removing the SF loss, we reduce the number of hyperparameters required, thereby simplifying the model.In our results below, we compare our method to these classes of solutions described, namely reconstruction solutions [Machado et al., 2020] , pretraining solutions [Liu and Abbeel, 2021] , and orthogonality solutions [Touati et al., 2022] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We have included a section on the limitations and broader impact of our work in section 9 in the main paper.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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their bestAnswer: [Yes] Justification: All computational experiments are performed over 5 random seeds and all statistical plots include error bars, which represent the standard deviation of the data.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • 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. • It 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 a 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 are 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. 
reinforcement learning: 
 The RL setting is formalized as a Markov Decision Process defined by a tuple (S, A, p, r, ), where S is the set of states, A is the set of actions, r : S ! R is the reward function, p : S ⇥ A ! [0, 1] is the transition probability function and 2 [0, 1) is the discount factor which is being to used to balance the importance of immediate and future rewards [Sutton and Barto, 2018] .At each time step t, the agent observes state S t 2 S and takes an action A t 2 A sampled from a policy ⇡ : S ⇥ A ! [0, 1], resulting in to a transition of next state S t+1 with probability p(S t+1 | S t , A t ) and the reward R t+1 .SFs are defined via a decomposition of the state-action value function (i.e. the expected return), Q, into the reward function and a representation of expected features occupancy for each state S t and action A t of time step t:Q(S t , A t , w) = (S t , A t , w) > w (1)where 2 R n are the SFs that capture expected feature occupancy and w 2 R n is a vector of the task encoding, which can be considered a representation of the reward function [Borsa et al., 2018] .Canonically, the SFs for a state-action pair (s, a) under a policy ⇡ are defined as:⇡ (s, a) ⌘ E ⇡ " 1 X i=t i t i+1 | S t = s, A t = a # (2)where 2 R n is a set of basis features, and ⇡ is the policy [Barreto et al., 2017] .However, as shown by Borsa et al. [2018] , we can treat the task encoding vector w as a way to encode policy ⇡. This results in Universal SFs, (s, a, w), on which we base our work.The task encoding vector w can also be related directly to the rewards themselves via the underlying basis features ( ): R t+1 = (S t+1 ) > w (3)The canonical approach for learning the basis features and successor features for each state S t and action A t of time step t, with respect to policy ⇡, are achieved by optimizing the following SF Temporal-Difference loss:L , = 1 2 k (S t+1 ) + (S t+1 , a, w)) (S t , A t , w)k 2 (4)where action a ⇠ ⇡(S t+1 ). The basis features are typically defined as the normalized output of an encoder, which the SFs learn from concurrently (see Figure 2 for an example).However, when the basis features, , must be learned from high-dimensional, complex observations such as pixels, optimizing Eq. 4 may result in the basis features, , converging to a constant vector. This outcome occurs because it can minimize the loss, as noted by Machado et al. [2020] , which we will also prove mathematically below.Consider the basis features function (•) and the Successor Features (•), omitting the inputs for clarity. The canonical SF-TD loss (Eq. 4) is defined as:EQUATIONUsing proof by contradiction, we aim to show that when both (•) and (•) are constants across all states S, specifically when (•) = c 1 and (•) = c 2 with c 1 = (1 )c 2 , the system satisfies the zero-loss conditions, leading to representation collapse.We start with the assumption that if(•) = c 1 , (•) = c 2 , then L , 6 = 0 8c 1 , c 2 2 R.Substituting (•) = c 1 and (•) = c 2 into the loss function:L , = 1 2 kc 1 + c 2 c 2 k 2 (6)It is trivial to observe that if c 1 = (1 )c 2 , the expression for L , is as follows:L , = 1 2 k(1 )c 2 + c 2 c 2 k 2 = 1 2 k0k 2 = 0 (7)This contradicts our assumption that L , 6 = 0 for a particular relationship between c 1 and c 2 . ⇤ Thus, we have shown that there exist constants c 1 , c 2 such that when (•) = c 1 and (•) = c 2 with c 1 = (1 )c 2 , the system does satisfy the zero-loss conditions, resulting in degenerate solutions for L , , i.e. causing representation collapse. In this collapsed state, (•) loses its ability to distinguish between different states effectively, causing the model to lose critical discriminative information and thus impairing its generalization capabilities.Additionally, we also show empirically in Figure 1(a-c ) of the presence of representation collapse when learning using Eq. 4. In this work, our method aims to mitigate these issues with a novel, simple approach for learning SFs directly from pixels. 
proposed method: 
 The key insight from the proof above (section 3.4) is that preventing representation collapse requires avoiding the scenario where the basis features become a constant vector for all states S, which would minimize the loss without contributing to meaningful learning. Below, we will describe the steps taken in our approach to mitigate these issues causing representation collapse. (S t+1 ) are the normalized output of the encoder using state S t+1 . The task-encoding vector w is learned through the reward prediction loss (Eq. 8). Concatenated with w, the basis features and successor features are learned through computing the Q-values with w and minimizing the Q-SF-TD loss function (Eq. 9). A schematic for continuous actions and previous approaches can be found in Appendix G and H respectively.We note that when the representations form a set of SFs, Eq. ( 1) is satisfied for some w that also satisfies Eq. (3). Therefore, the approach we take to learn SFs is simply to ensure that over the course of the learning and w come to satisfy both of these equations, which can be achieved by using the following loss functions:L w = 1 2 R t+1 (S t+1 ) > w 2 (8) L = 1 2 ŷ (S t , A t , w) > w 2(9) where (S t+1 ) is treated as a constant in Eq. 8 using a stop-gradient operator, and ŷ is the bootstrapped target:ŷ = R t+1 + max a 0 (S t+1 , a 0 , w) > w (10)Here, w is only altered by L w , whereas SF and the basis features are learned via L . Specifically, our proposed approach can overcome representation collapse by treating the basis features as the L2 normalized output from the encoder of the SF network (Figure 2 ), because unlike in Eq. 4, Eq. 8 and Eq. 9 are not minimized by setting to a constant value, given that ŷ and R t+1 are not constants for all states S. Hence, there is nothing encouraging the network to converge to a constant vector, naturally avoiding representational collapse.When the basis features are needed to learn the task encoding vector w through the reward prediction loss (Eq. 8), we apply a stop-gradient operator to treat the basis features as fixed. As we will demonstrate in section 7 "Analysis of Efficiency and Efficacy", this inclusion of a stop-gradient operator is crucial. Without it, learning both the basis features and the task encoding vector w concurrently can lead to learning instability.Next, we will clarify how our approach relates to learning SFs, as they are defined mathematically. Given the straightforward nature of our approach, we refer to the SFs learned as "Simple SFs."In Proposition 1 (Appendix C), we show that our approach ultimately produces true SFs, equivalent to the SFs learned using Eq. 4. Proposition 1 does this by proving that minimizing our losses (Eq. 8 & Eq.9) also minimizes the canonical SF loss used in Universal Successor Features (Eq. 4). Furthermore, Proposition 1 supports the proof above (Section 3.4) that our approach minimizes these losses in a manner such that setting the basis features to a constant is not a solution. Once again, if = c 2 and = c 1 = (1 )c 2 then Eq. 8 & Eq. 9 are not minimized, due to the fact that ŷ and R t+1 in Eq. 10 are not constants for all states S.  
learning successor features the simple way: 
 The architecture for our network is shown in Figure 2 , which is broadly inspired by Liu and Abbeel [2021] . Pixel-level observations, S t , are fed into a convolutional encoder that outputs a latent representation h(S t ), which is used both to construct the basis features and the SFs. To construct the basis features, (S t ), we simply normalize the latent representations h (via L2 normalization, following Machado et al. [2020] ). To calculate the representations (S t , A t , w), the latent representation is combined with the task encoding vector, w, and fed into a multilayer perceptron that generates one set of representations for each possible action, A t . These representations are then combined with the task encoding via a dot product operation to estimate the Q-value function, Q(S t , A t , w) = (S t , A t , w) > w. The policy is then simply an ✏-greedy policy based on the Q-value function.To learn the basis features ( ) and representations ( ), we minimize the losses in Eq. 8 and Eq. 9 using minibatch samples of experience tuples (S t , A t , R t+1 , S t+1 , w), collected while interacting with the environment and sampled from a replay buffer which is similar to Mnih et al. [2015] .Critically, only the task-encoding vector w is learned by optimizing Eq. 8, so a stop gradient operator is applied to the basis features (S t ) (see Figure 2 ). The successor features, , in the bootstrap target, ŷ, actually come from a target network, , which is updated periodically by using the actual network, a common approach in deep RL [Mnih et al., 2015] . The successor features , and all of the upstream network parameters ✓, are learned by minimizing Eq. 9. The full algorithm used for training our network is given in Algorithm 1 in Appendix B. 
experimental results: 
 The environments used in our studies are 10 ⇥ 10 2D grid worlds, 3D Four Rooms environments (Figure 9 in Appendix D) and Mujoco. All studies were conducted exclusively using pixel observations, as the primary motivation for this paper is to address representation collapse when learning with pixel observations.The grid worlds offer both egocentric (partially) and allocentric (fully observable) scenarios while the 3D Four Rooms environments provide exclusively egocentric observations. The rationale behind selecting these environments was threefold: first, to evaluate the agent's learning capabilities across varying levels of environmental visibility, second, to examine its ability to interpret spatial relation-ships and distances, and third, to provide a set of tasks where the transition dynamics are easy to quantify for constructing SRs that can serve as a comparison to evaluate the SFs with.For a more complex setting, we considered the Mujoco environment because it allows for direct manipulation of the reward function and domains switching, such as moving from half-cheetah to walker, given that they both have the same action dimensions.To evaluate the generalization capabilities of the learned SFs, our studies focus on continual learning setting. In the 2D grid worlds and 3D Four Rooms environment, agents are exposed to two cycles of two distinct tasks. These tasks involves changes in reward locations (as shown in Figure 9b & Figure 9d ) and/or changes in environment dynamics (as shown in Figure 9a & Figure 9e ). Additionally, we explored two different scenarios to better simulate real-world conditions. The first scenario involves resetting the replay buffer at each task transition, which emulates drastic distribution shifts typically encountered in real-world applications. The second scenario maintains the replay buffer across task transitions, allowing us to assess the agent's learning continuity in a more stable data setting.In the Mujoco environment, agents are exposed to one cycle of two distinct task as in this setting, we primarily wish to evaluate how well the agents can adapt to new tasks and mitigating interference.In all experiments, we make comparisons with several baselines, namely, a Double Deep Q-Network (DQN) agent [Van Hasselt et al., 2016] and agents learning SFs ( ) with constraints on their basis features ( ), including reconstruction loss [Machado et al., 2020] , orthogonal loss [Touati et al., 2022] , and unlearnable random features [Touati et al., 2022] . Additionally, we compare with an agent that learns SFs through a non-task-engaged pre-training regime [Liu and Abbeel, 2021] . The mathematical definitions of the constraints can be found in Appendix F. To ensure the robustness of our results, all experiments are conducted across 5 different random seeds.The 2D Gridworld environments were developed based on 2D Minigrid [Chevalier-Boisvert et al., 2023] . We created two different layouts of the 2D Gridworld environment, namely Center-Wall (Figure 9a ) and Inverted-LWalls (Figure 9b ). In order to align the setting more closely with the canonical Gridworld environment as described by Sutton and Barto [2018] , we altered the reward function such that it returns a reward of +1 when the agent successfully reaches the goal state and 0 otherwise. For the 2D Gridworld environments, the agents were trained for one million steps per task.Figure 3a presents the cumulative returns for the Center-Wall environment with egocentric observations, while Figure 3b shows the results for allocentric observations.The results show that our agent learns as well as, if not better than, the baseline models. Furthermore, when analysing the cumulative total returns during training, our model, SF Simple, exhibited better transfer that the baseline models. Particularly, SFs that are learned with constraints on the basis features clearly struggle to learn effectively, either due to the additional computational overhead or because representations that fulfill those constraints do not lead to effective policy learning.We developed the 3D Four Rooms environments (Figure 9d ) using Miniworld [Chevalier-Boisvert et al., 2023] . In this environment, the state and action spaces are continuous. In the first task, the agent receives a reward of +1 when it reaches the green box and a reward of -1 when it reaches the yellow box and this alternates for the second task. The agents were trained for five million steps per task. Similarly, the results in Figure 3c show that our agent is able to learn effectively using egocentric pixel observations in a 3D environment.In order to demonstrates our model's capabilities with continuous actions, we consider the Mujoco environment. We followed the established protocol in Yarats et al. [2021] for effective learning with pixels observations in this environment. We started in the half-cheetah domain, rewarding agents for running forward in Task 1. For Task 2, we introduced scenarios with running backwards, running faster, and switching to the walker domain. The results are presented in Figure 4 . Across all scenarios, our model not only maintained high performance but consistently outperformed all baselines in both Task 1 and Task 2, highlighting its superior adaptability and effectiveness in complex environments. This contrasted sharply with other SF-related baseline models, which struggled to adapt under these conditions.7 Analysis of Efficacy and Efficiency Can our SF-learning technique, like traditional SRs, effectively capture the transition dynamics of the environment [Stachenfeld et al., 2017] ? To investigate, we first sought a quantitative measure to compare SFs to SRs. To do this, we trained a simple non-linear decoder to assess which model's SFs can be most effectively decode into SRs. We conducted this evaluation using both allocentric and egocentric observations within the center-wall environment. The results, depicted in Figure 5 , shows that our model demonstrate consistently high accuracy (lower errors) across both settings. This contrasts sharply with SFs developed using reconstruction constraints or random basis features, which, while effective in egocentric settings, perform poorly in allocentric settings where feature sparsity is greater.We next utilized 2D visualizations with geospatial color mapping to differentiate environmental locations, aiming to see if similar SFs that are proximate in neural space are proximate in physical space. Using UMAP [McInnes et al., 2018] for dimension reduction, our results (Figure 6 ) suggest that our simple approach captures environmental statistics comparably to other models, but with less overhead for calculating the loss. Moreover, our technique consistently forms organized spatial clusters across partially, fully, and egocentric observational settings.Additionally, we performed a correlation analysis in 2D Gridworld environments, comparing each spatial position and head direction against analytically computed SRs [Dayan, 1993] , further confirm- Despite this, well-clustered SF models, especially those with orthogonality constraints, may not always translate to effective policy learning, as seen in Figure 3 . In allocentric scenarios, SFs with reconstruction constraints struggle with minimal pixel variations, unlike in the distinct pixel changes in the Four Rooms environment. For more visualizations, see Appendix M.ing the robustness and adaptability of our model's SFs in various observational contexts (Table 6 and Table 7 in Appendix N). 
is stop gradient critical for learning?: 
 With stop grad (ours) Without stop grad (Basis-Rewards) Four Rooms Previous methods that concurrently learn the basis features, , and the task-encoding vector w, often face challenges with learning efficiency and stability, particularly in environments characterized by sparse rewards. This issue is illustrated in Figure 10 in Appendix D of Ma et al. [2020] , where optimizing the reward prediction loss (Eq. 8) can inadvertently drive the basis features towards zero ( ! 0), causing significant representational collapse. Representational collapse not only reduces the discriminative capabilities of but also undermines the agent's ability to differentiate between distinct states, thus severely impacting the overall learning process.As depicted in Figure 2 , our solution involves the strategic use of a stop gradient operator applied to the basis features . This operator prevents the gradient from the reward prediction loss from updating basis features , effectively decoupling the learning of from w, thus ensuring that it retains its critical discriminative statistics, allowing for effective learning as demonstrated in Figure 7 .How robust are the SFs learned using our approach as the environment dynamics become noisier? To explore this question and verify the robustness of our technique, we also created a slippery variant of the Four Rooms environment (Figure 9e ). Specifically, in the top right and bottom left rooms, the agent experiences a "slippery" dynamic: chosen actions have a certain probability of being replaced with random, alternative actions. This design mimics the effects of a low-friction or slippery surface, creating a scenario where the agent's intended movements might lead to unpredictable outcomes. (f) Bar plot showing frames per second achieved by the agent during gradient computation, back-propagation, and interaction with the environment. These metrics provide insights into the computational efficiency and the real-time interaction capabilities of the agent across different tasks or conditions. (g) Bar plot showing the total training duration for completing two exposures of two tasks, demonstrating overall time efficiency. Collectively, these plots reveal that our agent not only learns tasks effectively but also excels in computational efficiency.The results in Figure 8a -d demonstrate that our approach is robust to increasing levels of stochasicity. Notably, when the stochasicity is high (slippery probability >= 0.3), all other SF methods fail to learn effectively in the second task, whereas our approach continues to perform well.How do alternative SF learning methods with extra loss functions, like orthogonality constraints, stack up against our approach in terms of efficiency? We analyzed the number of steps each method takes to learn an effective policy, using a performance threshold defined by the shortest expected episode length from the last 10 episodes. A shorter episode length indicates better performance, as it signifies quicker goal achievement. We noted the timestep when each model first met or exceeded this threshold. Our results, shown in Figure 8e , demonstrate that our method outperforms all baselines in learning efficiency. Additionally, our method leverages simpler compute blocks and loss functions, enhancing computational speed and reducing training duration, as shown by faster frame processing rates (Figure 8f ) and shorter overall training times (Figure 8g ). Therefore, our approach is more efficient than the baseline methods for learning SFs. 
discussion: 
 In this work, we developed a method for learning SFs from pixel-level observations without pretraining or complex auxiliary losses. By applying the mathematical principles of SFs, our system effectively learns during task engagement using standard losses based on typical training returns and rewards. This simplicity and efficiency are key advantages of our approach.Our experiments demonstrate that our method learns SFs effectively under various conditions and surpasses baseline models in continual RL scenarios. It effectively captures environmental transition dynamics and correlates well with analytically computed Successor Representations (SRs), offering a streamlined, efficient strategy for integrating SFs into RL models. Future work could build on this to create more sophisticated models that leverage SFs for enhanced flexibility in RL.Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We provide the information regarding the compute resources in the Appendix L.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or 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 as 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 or failed experiments that didn't make it into the paper). 
limitations and broader impact: 
 The algorithms we developed were evaluated predominantly in simulated environments, which may not capture the diverse complexity of real-world scenarios. A key assumption in our approach is that pixel observations are of good quality. This assumption is critical as poor image quality could substantially degrade the performance and applicability of our algorithms.The use of Successor Features in learning algorithms, as demonstrated in this work, offers significant advantages, particularly in mitigating catastrophic interference. This capability is crucial for the development of machine learning systems that require continuous learning, such as in dynamic environments. For instance, autonomous vehicles operating in ever-changing conditions can retain learned knowledge while adapting to new information, enhancing their safety and reliability.However, the enhanced capabilities of these systems also raise concerns. This research was further enabled by computational resources provided by Calcul Québec2 and the Digital Research Alliance of Canada3 , along with the computational resources support from NVIDIA Corporation.We are grateful to Gheorghe Comanici, Pranshu Malviya, Xing Han Lu, Isabeau Prémont-Schwarz and the anonymous reviewers whose insightful comments and suggestions significantly enhanced the quality of this manuscript. Additionally, our discussions with members of the LiNC lab4 , Mila5 , and early collaborators from Microsoft Research (MSR) have been invaluable in shaping this research. Special thanks to Ida Momennejad, Geoff Gordon and Mehdi Fatemi from MSR for their substantial insights and contributions during the initial phases of this work.1: Initialize task encoding vector w 2: Initialize SF ✓ network, SF ✓ target network 3: for t := 1, T do 4:Receive observation S t from environment 5:A t ✏-greedy using Q(S t , • | w) 6:Send A t to receive S t+1 and R t+1 from environment 7:a 0 2 argmax b ✓ (S t+1 , b, w) > w 8: ŷ = R t+1 + ✓ (S t+1 , a 0 , w) > w 9:L2 normalized output from the encoder of SF network 10:loss ✓ = ( ✓ (S t , A t , w) > w ŷ) 211:loss w = ( > w R t+1 ) 2 12:Gradient descent on ✓ and w 13: end forIn this section, we provide a proof sketch to show that minimizing the Q-SF-TD loss (Eq. 9) will also result in minimizing the canonical universal SF-TD loss [Borsa et al., 2018] for learning the SFs ( (•) 2 R n ). For the sake of brevity, we consider a tabular RL setting, where state s is the current state, s 0 is the next state, a is the current action, and r is the reward of the transition tuple (s, a, s 0 , r) and as per defined in the main text, w 2 R n is the task encoding vector and (•) 2 R n is the set of basis features.Let L SF be the canonical universal SF-TD loss [Borsa et al., 2018] :L SF = 1 2 (s 0 ) + (s 0 , a 0 , w)) (s, a, w) 2 (11)wherea 0 = arg max b Q(s 0 , b, w) = arg max b (s 0 , b) >w and is the discount factor. We treat (s 0 , a 0 , w) as part of the bootstrapped target: ŷSF = (s 0 ) + (s 0 , a 0 , w), which results in semigradient methods [Sutton and Barto, 2018] . Subsequently, the gradient r for L SF (Eq. 4) is defined as:EQUATIONNext, as previously discussed in section 3, the Q-SF-TD loss L which we used to learn the successor features ( ) is defined as:EQUATIONwhere ŷ = r + max a 0 (s 0 , a 0 , w) > w is the bootstrapped target.(Note: Eq. 13 and the bootstrapped target ŷ is the same as Eq. 9 and Eq. 10 respectively, presented in Section 3 of the main text)Following the same reasoning in Eq. 4, the gradient r for L is defined as:r L = r + (s 0 , a 0 , w) > w (s, a, w) > w w (14)Proposition 1 Optimizing r L ' w > r L SF w, where L SF is the canonical loss for universal successor features [Borsa et al., 2018] .Proof. Now, assuming that for any given state s, the reward r for state s can be linearly decomposed into the dot product of the basis features (s) and the task encoding vector w, as suggested by Sutton [1988] , Dayan [1993] , it follows that there exists an optimal set of basis features ⇤ (s). This optimal set ensures that the reward r can be accurately represented as the dot product of ⇤ (•) and the task encoding vector w:r = ⇤ (s 0 ) > w (15) where s 0 is the next state.Thereafter, let us recall that the reward prediction loss L w is defined as:EQUATION(Note: This equation is the same as Eq. 8 presented in Section 3 of the main text.)Substituting the assumption that we made in Eq. 15 into the reward prediction loss (Eq. 16),L w = 1 2 r (s 0 ) > w 2 = 1 2 ⇤ (s 0 ) > w (s 0 ) > w 2 (Subst. r = ⇤ (s 0 ) > w following Eq. 15) = 1 2 ( ⇤ (s 0 ) (s 0 )) > w 2 = 1 2 ✏(s 0 ) > w 2 (17) Where ✏(s 0 ) is the difference between ⇤ (s 0 ) and (s 0 ). Furthermore, if L w ' 0, then ✏(s 0 ) > w = w > ✏(s 0 ) ' 0.Shifting our focus back to the gradient r L of our Q-SF-TD loss function (Eq. 14),r L = r + (s 0 , a 0 , w) > w (s, a, w) > w w (Eq. 14) = w > ( ⇤ (s 0 ) + (s 0 , a 0 , w) (s, a, w))w (Subst. r = ⇤ (s 0 ) > w following Eq. 15)= w > ( (s 0 ) + ✏(s 0 ) + (s 0 , a 0 , w) (s, a, w))w (Subst. ⇤ (s 0 ) = (s 0 ) + ✏(s 0 ) from Eq. 17)= w > ( r L SF + ✏(s 0 )) w (Subst. definition from Eq. 12)= w > r L SF w w > ✏(s 0 )w = w > r L SF w 2 p L w w (Subst.w > ✏(s 0 ) = 2 p L w from Eq. 17) ' w > r L SF w ⇤ (18)In conclusion, this proof demonstrates that the gradients r L computed using our proposed Q-SF-TD loss function (Eq. 13) effectively project the gradients r L SF from the canonical universal SF-TD loss function (Eq. 12) onto the task encoding vector w. This indicates that our loss function maintains the essential characteristics of the canonical form while aligning closely with the specific direction of the task encoding vector w. In this section, we provide more details about the environments used in our experiments.E.1 2D Gridworld Environments.The specific parameters defining the 2D Gridworld environments are detailed in Table 1 . In the Center-Wall environment, a vertical wall splits the area into two distinct regions. Task 1 features a passage from the left to the right side at the bottom, with the goal state located in the top left corner Figure 10a ). In Task 2, the layout is modified: the passage is moved to the top, while the goal state is relocated to the bottom right corner Figure 10b ). These changes are strategically implemented to evaluate the agents' ability to adapt to simultaneous alterations in both the environmental structure and the goal location. To aid in visual analysis, we use a geospatial color mapping initially developed for Task 1 (Figure 10c ). This mapping effectively illustrates the spatial positioning within the environment and is particularly useful in the 2D visualization of the Successor Features and DQN Representations, providing a clearer understanding of how agents interpret and navigate the modified environment (Figures 6, 29 and 30).In the Inverted-Lwalls environment, we placed two L-shaped walls within the environment, one on the left and the other on the right, creating a unique layout. This design results in a single, central path acting as a bottleneck, which the agent must navigate to reach the goal states. Specifically, to access the goal state located on the left side of the environment, the agent is required to traverse this central path while facing north (Figure 11a ). Conversely, to reach the goal state situated on the right, the agent must navigate the same path but facing south (Figure 11b ). This layout ensures that the agent consistently encounters and must maneuver this bottleneck area, regardless of the goal state's location.E.2 3D Miniworld Environments.The actions in this environment consists of moving Forward and Backwards, turning Left and Right.The specific parameters defining the 3D Miniworld environments are detailed in Table 2 . The Four Rooms environment consists of four identical square rooms arranged in a 2x2 grid, with passages connecting the rooms and allowing an agent to move between the rooms (Figure 12a ). Each room in our 3D environment is designed with unique textures, a deliberate choice to reduce the complexity associated with localization ambiguities often encountered in more uniform settings. This variation in textures aids the agent in distinguishing between rooms based solely on visual cues, thereby simulating more realistic navigation scenarios. This setup also allows us to observe how visual diversity impacts the agent's ability to infer its location and navigate to specific goals, providing insights into the interplay between environmental features and SFs learning in a 3D spatial context. Depending on the task, the agent receives a reward of either +1 or -1 when it reaches the yellow or green box. Similar to the Center-Wall environment, we also create a geospatial color mapping for the 2D visualization of the Successor Features and DQN Representations (Figure 12b ).Figure 13 : Slippery Four Rooms (3D) layoutIn the slippery variant of the Four Rooms environment, our goal is to rigorously test the robustness of agents in learning SFs under challenging conditions. Specifically, in the top right and bottom left rooms of this setup, the agent experiences a 'slippery' dynamic: chosen actions have a certain probability of being replaced with random, alternative actions. This design mimics the effects of a low-friction or slippery surface, creating a scenario where the agent's intended movements might lead to unpredictable outcomes. Such a setup is instrumental in assessing the agent's adaptability and the robustness of SF learning in the face of environmental unpredictability. This variant not only challenges the agent to adapt to unexpected changes but also provides valuable insights into the flexibility and resilience of the SFs when navigating environments where control and predictability are compromised.In this work, we only utilised pixels inputs from Mujoco since our focus is on learning SFs directly from pixel observations. For domains, we chose both walker and half-cheetah. We broadly follow the same setup as Yarats et al. [2021] , and included their model as a baseline, which we denote as "DDPG" in our results (Figure 4 ).The codebase from their model is provided in the Unsupervised Reinforcement Learning (URL) Benchmark repository [Laskin et al., 2021] 6 , which we further described in the APS Agent in section F. The specific parameters we used for training in the Mujoco environment are detailed in Table 3 . In this section, we describe how we create our agent as well as the ones we used for comparisons.In addition, we provide the mathematical definitions of the constraints used on the basis features. For all agents, we swept the learning rates for both the SF network and the task encoding (specific for all SFs agents) using a gridsearch. The values ranged from 1e-1 to 1e-6, and the process was repeated using 5 random seeds in both 2D Gridworld and 3D Four Rooms environments. The same was also applied to the Double DQN agent [Van Hasselt et al., 2016] and we took extra care to ensure that the architecture and its number of parameters were as similar as possible to our model. Detailed hyperparameters for learning SFs and the task encoding w for our agent are outlined in Tables 4 and 5 .In our study, we take inspiration from the neural network architecture from Liu and Abbeel [2021] from the Unsupervised Reinforcement Learning (URL) Benchmark repository [Laskin et al., 2021] 7 , which utilizes PyTorch [Paszke et al., 2019] . This repository was chosen for its robust implementation and served as the foundation for all SF-variant agents, including ours. Within the URL Benchmark, the encoder follows the Deep Deterministic Policy Gradient (DDPG) network architecture [Lillicrap et al., 2015] . Notably, there is a discrepancy in the network architecture hyperparameters between the APS paper [Liu and Abbeel, 2021] ) and the URL Benchmark repository. Given the practical implications of these differences, our implementation aligns with the hyperparameters specified in the URL Benchmark.In line with the URL Benchmark's methodology, we initially employed the least squares method to determine the optimal task encoding w. However, we observed that this analytical approach was excessively sensitive in our experimental context, particularly due to its reliance on the mini-batch samples. This sensitivity was especially pronounced in environments with sparse rewards, like those in our study, suggesting that the least squares method might be less suited for such settings. This challenge was not present in the original APS framework [Liu and Abbeel, 2021] , which was structured around distinct pre-training and fine-tuning phases. In contrast, our research focuses exclusively on continuous online learning, introducing unique challenges and dynamics not addressed in the APS paper [Liu and Abbeel, 2021] .At each time step t, the basis features (S t ) are generated from the current state S t using an encoder. Together with the action A t , these features are fed into a reconstruction decoder to predict the next state Ŝt+1 . Both the encoder and decoder are optimized using the reconstruction loss:L recon = ||S t+1 Ŝt+1 || 2 (19)where S t+1 is the ground truth of the next state. The same set of basis features is also utilized in optimizing the Reward Prediction Loss (Eq. 8) and the Q-SF-TD Loss (Eq. 9).At each time step t, the basis features are generated from the current state S t using an encoder.Besides being utilized to optimize the Reward Prediction Loss (Eq. 8) and the Q-SF-TD Loss (Eq. 9), the basis features are also optimized with the orthogonality loss [Koren, 2003 , Mahadevan and Maggioni, 2007 , Machado et al., 2017b,a] :EQUATION)7 https://github.com/rll-research/url_benchmark where states s and s 0 are two different states sampled from the replay buffer D. The first term encourages the basis features (S t ) and (S t+1 ) to be similar and the second term promotes orthogonality by ensuring that the basis features of the different states (s) and (s 0 ) are distinct and decorrelated. Following Touati et al. [2022] , we set the regularization factor = 1.In this agent, the basis features are constrained to be unlearnable random features, which are defined during initialization. The SFs are subsequently learned on top of these predefined basis features.To guarantee that the basis features remain unlearnable throughout the training process, a stop gradient operator is employed.This agent learns Successor Features using a complex learning strategy that integrates three distinct losses: the SF-TD loss (Eq. 4), the reward prediction loss (Eq. 8) and the Q-SF-TD loss (Eq. 9). This multifaceted approach, proposed by Janz et al. [2019] aims to ensure that the learnt SFs satisfy all desired constraints. et al., 2015] and implemented following the URL benchmark [Laskin et al., 2021] . This design modifies our original architecture to accommodate continuous action environments, enabling the model to handle a broader range of control tasks. The model incorporates a linear decomposition of Successor Features and the task encoding vector w to compute Q-value, following Eq.1. Following the DDPG implementation in URL benchmark, actions are sampled from a truncated normal distribution, and LayerNorms are applied to normalize inputs to a unit distribution. [Touati et al., 2022] . A stop gradient operator is also used to prevent the SFs from updating the basis features when optimizing the SF-TD loss. [Kulkarni et al., 2016] I Impact of Learning Rate Variations on Task Encoding Vector Figure 16 : Comparison of learning rates for the task encoding vector in grid worlds and 3D Four Rooms environments. Generally, a lower learning rate is required for the task encoding vector, despite its use of a simple reward prediction loss (Eq. 8), compared to the SF network, which needs more steps to converge due to its involvement in capturing complex environmental dynamics.In this section, we present expanded illustrations of the results initially introduced in the main paper. These larger visual figures provide a clearer and more detailed view to enhance the reader's understanding of our findings. Additionally, we include additional supplementary experimental results that were not featured in the main paper due to space limitations. In the fully-observable scenario, while our agent shows performance comparable to the DQN agent, it is slightly outperformed by the agent employing SFs with orthogonality constraints on its basis features. Notably, despite the superior performance of this latter agent in later tasks during Exposure 2, it initially faces difficulties in developing an effective policy, attributed to the added complexity of adhering to orthogonality constraints. Figure 21 : Evaluation in a slippery Four-Rooms environment with varied slipperiness probabilities, without replay buffer resets at each task transition. This environment features slippery conditions in the top-right and bottom-left rooms for both tasks, Task 1 and Task 2. Both tasks have differing reward structures: In Task 1, rewards are set at +1 for the green box and -1 for the yellow box; in Task 2, this reward scheme is reversed (green box: -1, yellow box: +1). The diagram illustrates the layout of the environment can be found in Figure 9 . Note: The APS Pre-trained agent was tested only at a slippery probability of 0.15; higher probabilities were not evaluated due to performance decline beyond Task 1 of Exposure 1 when the slippery probability is 0.15. The APS Pre-trained agent was tested only at a slippery probability of 0.15; higher probabilities were not evaluated due to performance decline beyond Task 1 of Exposure 1 when the slippery probability is 0.15. K Experimental results of SF + Q-TD + Reward vs SF Simple (Ours)In this section, we present the experimental results of our agent (SF Simple) and the agent which optimizes the three losses (SF + Q-TD + Reward) simultaneously. For more information about this agent, see section F.5.K.1 Continual RL results for Inverted-LWalls environment In both partiallyobservable and fully-observable scenario, the agent, which optimizes three losses simultaneously (SF + Q-TD + Reward), experiences learning instabilities due to the higher complexity involved in managing all constraints.Simple SF (ours) SF + Q-TD + Reward ). The diagram illustrates the layout of the environment can be found in Figure 9 . As observed, simultaneous optimization of three losses (SF + Q-TD + Reward) significantly impedes the agent's ability to learn effectively in a stochastic environment.Figure 28 : Evaluation in a Continual Reinforcement Learning setting across 5 random seeds, with replay buffer resets at each task transition in the 3D Four Rooms environment. Task 1 adheres to the canonical Four Rooms environment dynamics, while Task 2 employs the slippery variant, where chosen actions are altered based on the slippery probability to simulate environmental changes. Throughout both tasks, reward associations remain consistent: +1 for the green box and -1 for the yellow box. The layout of this environment is depicted in Figure 9 . Once again, simultaneous optimization of three losses (SF + Q-TD + Reward) significantly impedes the agent's ability to learn effectively.For our experimental setup, we utilized Python 3 [Van Rossum and Drake, 2009] as the primary programming language. The agent creation and computational components were developed using Jax [Bradbury et al., 2018 , Godwin* et al., 2020] , while Haiku [Hennigan et al., 2020] was employed for implementing the neural network components. For data visualization, we used Matplotlib [Hunter, 2007] and Seaborn [Waskom, 2021] to generate line plots. Additionally, we utilized Plotly8 for creating the violin plots and heat maps used in our correlation analysis, as well as the 2D visualizations of the SFs and DQN Representations. We utilized Scikit-learn [Pedregosa et al., 2011] in our correlation analysis studies as well as the open-source Uniform Manifold Approximation and Projection (UMAP) tool [McInnes et al., 2018] to generate the 2D embeddings of the SFs. The configuration and management of our experiments were facilitated by Hydra [Yadan, 2019] and Weights & Biases [Biewald, 2020] . All experiments, particularly those in the continual learning setting, were conducted using Nvidia V100 GPUs and completed within a maximum of one day. The code used in the study will be released in the near future, following an internal review process.Given that Successor Features (SFs) are action-dependent, and considering the space constraints in the main paper, our visualizations here are more comprehensive. In the main paper, we primarily showcased visualizations for the forward action due to these limitations. However, in this section, we expand our focus to include visual representations for a variety of actions, providing a more holistic view of the SFs' behavior and their influence across different action scenarios. This expanded visualization not only enhances our understanding of the SFs' multidimensional nature but also offers deeper insights into the agent's decision-making process and its interaction with the environment.M.1 Center-wall Environment (Fully-observable) The first column illustrates the initial state of successor features before training, using geospatial color mapping for clear visualization. Subsequent columns correspond to successor features developed for specific actions: Forward, Turn Left, and Turn Right, also visualized using geospatial color mapping. In this scenario, only the agent learning SFs with orthogonality constraints as well as our agent (Simple SF) learned well-clustered representations after training. It's crucial to recognize, however, that while clustered representations may suggest effective learning, they do not automatically equate to successful policy development. These visualizations highlight the varied encoding strategies of agents in response to full observability and different actions. The first column illustrates the initial state of successor features before training, using geospatial color mapping for clear visualization. Subsequent columns correspond to successor features developed for specific actions: Forward, Turn Left, and Turn Right, also visualized using geospatial color mapping. Some agents demonstrate well-clustered representations after training, which typically correlates with improved performance compared to agents with more dispersed or noisy features. It's crucial to recognize, however, that while clustered, color-mapped representations may suggest effective learning, they do not automatically equate to successful policy development. These visualizations highlight the varied encoding strategies of agents in response to partial observability and different actions. Agents operate with solely egocentric observations. Each panel represents the successor features of a different RL agent and action. The first column, using red, green, blue, and yellow to distinguish the four rooms, shows the initial state of successor features pretraining. Subsequent columns depict features for specific actions: Move Forward, Move Backwards, Turn Left, and Turn Right. Except for SF learned using unlearnable random basis features, most agents exhibit well-clustered representations post-training. However, it's important to note that such clustered, color-mapped representations, while indicative of effective learning, do not necessarily translate into successful policy development.Considering that the SRs are not normally distributed [Stachenfeld et al., 2017] , we conduct our correlation analysis in the Grid world environments (Figure 9a and b ) using the Spearman's rank correlation. The SRs were analytically computed using the transition matrix T where T (s 0 | s, a) denotes the probability of transitioning from state s to state s 0 given an action a ⇠ ⇡(• | s):EQUATIONwhere 0  < 1 is the discount factor and I is the identity matrix. The same policy ⇡ was used to generate the transition matrix T and to adjust the final correlations. These adjustments account for less frequently chosen actions and for positions and head directions less likely to be encountered by the agent, as outlined in the main text. Statistics regarding positions and head directions were collected using policy ⇡.In the remaining part of this section, we provide additional detailed violin plots to depict the correlation dynamics in both the Center-wall and Inverted-LWalls environments, covering scenarios that are both partially-observable and fully-observable. These plots are segmented into different stages: before training, after training, and the differences post-training. This segmentation offers a comprehensive view of the agents' learning progression over time. Specifically for the Inverted-LWalls environment, a table is included to provide a summary of mean and standard deviation statistics for these correlations, thus offering a clear quantitative perspective of our findings. Additionally, we present heatmaps that showcase the correlation at each spatial position in the environment for various SF agents. These heatmaps further enrich our analysis by visually representing the spatial distribution of correlation values, highlighting how different agents adapt to the environment. Justification: The abstract and introduction accurately reflect the paper's contributions and scope by clearly outlining our main achievement: the development of an algorithm designed to overcome representation collapse in learning Successor Features from pixel observations efficiently. We provide comprehensive evidence of representation collapse and review various existing approaches to address this issue, highlighting their computational demands and limitations in continual learning settings. Our contributions are precisely these comparative analyses and the introduction of a more efficient algorithm suitable for both single task and continual learning environments.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The data are collected using computational simulation, which does not involve any humans or animals. Most of our studies focus on navigational tasks, which are critical for robotics and self-driving cars.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 a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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: [Yes] Justification: We included a section on the limitations and broader impact of our work in section 9. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 a 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 or 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). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This paper focuses on studying reinforcement learning agents in a computational simulation, therefore we do not foresee such risks with regards to this research paper.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to usage guidelines or 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 a best faith effort. 12. Licenses for existing assets Question: 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: 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 an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: This 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. • Depending on 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. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.https://www.calculquebec.ca/https://alliancecan.ca/enhttps://linclab.mila.quebec/homehttps://mila.quebec/enhttps://github.com/rll-research/url_benchmarkPlotly Technologies Inc. Collaborative data science. Montréal, QC, 2015. https://plot.ly. 
 
 Appendix:  


Decision: Scores 5 5 4 6

This is a solid paper, that is right on the threshold.

Criticism included that the tested environments are too limited and that the paper failed to properly discuss previous work from Ma et al 2020.  The rebuttal added further ablations and experiments, which is good and helps reviewers to increase their scores.  However, reviewer 2gcv wasn't completely convinced about the rebuttal against the concern that a key part of the paper is not completely new, but has been done by Ma et al. 2020 already.  However, 2gcv didn't answer the final rebuttal comment.

Overall the paper is good and the rebuttal did address most concerns, but at the end didn't convince all reviewers in all points.  The resulting score is just at the boundary.  I really appreciate papers trying SIMPLE approaches, so I tend to accept the paper.

Reviews: review: paper_summary: paper_summary: This work introduces a new algorithm for training successor features in deep reinforcement learning. This is achieved by optimizing two separate metrics. The first requires the model to predict the cumulative reward following a full trajectory and optimizes the successor features and the basis feature. The second optimized metric requires the model to predict the reward at the next step using the basis feature and task vector, but in this case only optimizes the task vector. Experiments are conducted in grid worlds and a 3D Four Rooms domain which demonstrates that the proposed algorithm learns more consistently than other SF baselines and also supports faster task switching. Finally, supporting experiments show that the proposed algorithm is faster algorithmically and in terms of wall clock time than SF alternatives and that it also results in more separable SFs which correlate better with successor representations.
strength_weakness: strength_weakness: # Clarity
The figures in this work are laid out poorly and this significantly hinders the readability of the paper. At the least it would help if a reader does not have to look past unseen figures on their way to look at the one being referenced - such as when looking for Figure 6 which comes after Figure 5. Also having the architectures in Figure 1 far from where they are needed and completely out of context is jarring and unhelpful. These architecture diagrams are also very difficult to follow and there is not clear mapping from what is depicted for some of the pieces to any explanation in the text or caption. Most of what is depicted in Figure 1d is not in Section 3, and similarly Section 5 is not detailed enough for me to map onto Figure 1e. I think more detail could be added to the figure itself and to the caption here. Lastly, including the loss functions in the figure, especially ones which have not been explained in the text, like orthogonality loss, is confusing. If these losses are not necessary in the main text then I don't think it is necessary in the figure and so I would remove them. With respect to Proposition 1, it would be better if this was at least in the main text - some kind of proof sketch would be better - but in the interest of space I see why it was omitted. Once again I do then just point to this as a part which could do with more explanation of why this matters and intuition on what it is true. If I am correct, Proposition 1 is the reason why Equations 5 and 6 cannot be optimized with the representation collapse strategy?

Secondly, the Preliminaries section could have more elaboration. Equation 2 in particular is presented without any discussion and $\gamma$ is not introduced at all. A reader with experience on RL and SFs will be fine but less experienced readers will likely be alienated. It would be ideal if this section could set up the ideas to come, and this is done to a degree with it being noted that representation collapse can still optimize Equation 4. More of this insight would just be helpful. Similarly, I would appreciate more discussion on why representation collapse is now not able to optimize Equation 5 and 6. This is merely stated without reason on lines 113 to 115.

# Quality
I am not certain I agree with the assessment from Section 7.2. While it is a worth experiment, the result of simple SFs being more correlated just appears to be due to the fact that it has a more linear latent embedding. This can be seen in Figure 5. SF+Reconstruction has very separable and clear clusters but they are just not organised in a straight line. So correlations - a linear metric - will not work. It seems to me that it would be more appropriate to try decode SRs from the SFs using a simple but nonlinear model and report the final accuracy.
questions: questions: I have asked some questions in the my review above and would appreciate those be answered. I do not have any other questions at this time.

If my question on Proposition 1 is answers and makes sense to me, and the concerns in Quality address I would be likely to advocate for acceptance, with agreement that the clarity would also be improved and figures restructured.
limitations: limitations: The limitations are stated in their own section and a good of consideration given towards the broader impact of this work.
paper_summary: paper_summary: This work presents a model architecture to learn successor features in reinforcement learning. It consists of optimizing Eqs. (5) and (6), i.e., a loss for learning the features and a loss for learning the task specific weights. It claims to avoid representation collapse. Experiments are conducted in common 2D and 3D tasks to show that the proposed method can achieve better performance and higher sample efficiency.
strength_weakness: strength_weakness: - Even though the paper claims that using the reward to train the task weights $w$ is new, this approach has been discussed before (Ma et al., 2020). Specifically, it has shown that using the reward $r$ to train $w$ is inferior to using the $Q$ values (Appendix D, Ma et al., 2020). Of course, the two algorithms are not identical, but it remains unclear why using the reward to learn $w$ is the right choice in the current paper. It is necessary to discuss this prior work.
- Several other design choices need further explanation. L130-139 presents multiple design choices for the model architecture without proper discussion. For example, why the L2 normalization and the layer normalization are required? Why do we need to stop gradient at those specific places? It would be better to have ablation studies to show the importance of these choices.
- Experiment results require further analysis. There is barely any analysis or reasoning for the proposed method. Each subsection in Sec.6 ends with "our method is better" without properly addressing **why** it can achieve better performance. This question is not answered in Sec.7 either. Moreover, the improvement over existing methods is only marginal (see Figs.2&3). More importantly, The quantitive results in Table 1 only show marginal improvement over Orthogonality with a significant performance overlap.

Minor comment: Eqs.(5)&(6) are for scalars so a norm is unnecessary.

Reference:

- Ma, C., Ashley, D.R., Wen, J. and Bengio, Y., 2020. Universal successor features for transfer reinforcement learning. *arXiv preprint arXiv:2001.04025*.
questions: questions: See the weakness above.
limitations: limitations: The authors adequately addressed the limitations
paper_summary: paper_summary: The paper proposes a simpler method to learn Successor Features that avoids representational collapse. For this, the authors decompose the loss function to learn the successor features and task encoding separately. This allows for keeping the basis features fixed while learning the successor features, thus avoiding representational collapse. The experiments involve a continual learning scenario where robustness to task changes is evaluated. The authors show that the method can better adapt to changing tasks.
strength_weakness: strength_weakness: - The tested environments seem to be perhaps too simple for comparison both from the representational and task difficulty perspective, since DQN also has very good relearning capabilities in these environments. Why is it that DQN is better than the successor feature counterparts for the continual learning setting? This seems counterintuitive to me since successor features should be more robust than pure DQN.

- The presentation of the figures has issues. Some Figures are pixelated, i.e., not vector graphics. (e.g., Figure 4 or other environment Figures). Also I think Figure 1 could be split into 2 figures for better readability.
questions: questions: - Should the method not be tested on environments where DQN itself cannot adapt to the new tasks at all? I wonder if the simple approach still holds when the transition dynamics become more complicated or the observations are more noisy.
- For the Minigrid environment: Do you learn the successor features from pixels or do you use the built-in symbolic state representations?
- Since DQN is also robust for the showed environments, I wonder how reward sparsity affects the performance of the different algorithms?
limitations: limitations: The authors have adequately addressed limitations.
paper_summary: paper_summary: This paper presents a method for learning Successor Features (SFs) from pixel-level observations in reinforcement learning (RL) by combining a Temporal-Difference (TD) loss with a reward prediction loss. This approach simplifies the learning process, improves performance, and speeds up learning compared to existing methods.
strength_weakness: strength_weakness: 1. The writing quality of the paper could be improved. The Introduction section reads more like an extensive review of related work rather than setting the context for the proposed method.
2. The first three subplots in Figure 1 are difficult to understand without detailed background information. It is recommended to move these figures to the experimental section.
3. The statement in lines 61 to 62, "without any of the drawbacks," seems too absolute and should be toned down.
4. The experiments only demonstrate effectiveness in a few simple 2D and 3D environments. To further validate the proposed method, it is recommended to test in more complex environments, such as Atari games, similar to APS.
5. The experiments only tested the method with DQN, leaving its effectiveness with other RL methods unknown.
questions: questions: Have the authors tested the proposed method with other RL algorithms and in more complex environments?
limitations: limitations: In addition to the limitations mentioned by the authors in the paper, please refer to my comments in the Weaknesses section for further details.

Comments: title: Concluding Remarks and Acknowledgments
Response to Reviewer 2gcv on Evaluation
Response to Reviewer Tczf
Response to Reviewer 2gcv on differences with Ma et al. [1]
Response to Reviewer 7tgV
Final Day Reminder: Clarifying Concerns and Updating Scores
Final Day Reminder: Clarifying Concerns and Updating Scores
Final Day Reminder: Clarifying Concerns and Updating Scores
Final Day Reminder: Clarifying Concerns and Updating Scores
Response to Reviewer Tczf
Response to Reviewer 2gcv on Evaluation
Response to Reviewer 2gcv on differences with Ma et al. [1]
Response to Reviewer 7tgV
Response to Reviewer nkgZ
Response to Rebuttal by Reviewer 7tgV
comment: Dear all,

As the discussion period comes to a close, we would like to express our sincere gratitude to the reviewers for their valuable feedback, which has greatly guided the clarifications and improvements we plan to include in the final revision of our manuscript.

Among the key updates is the inclusion of results from the complex Mujoco environment using pixel observations, which strongly validate the robustness of our approach. Furthermore, we have clarified that the average episode returns were computed using moving averages and have included cumulative total return figures to demonstrate that our approach exhibits superior transfer across different environments. These results provide strong evidence that our method achieves significant improvements, beyond marginal gains. 

Once again, we would like to emphasize that we presented moving average returns per episode instead of cumulative total return plots in our manuscript to demonstrate that *we allow learning for the first task to converge before introducing the second and subsequent tasks.*

In response to Reviewer 7tgV's suggestion, we have extended our analysis by decoding the Successor Features (SFs) into Successor Representations (SRs) using a simple non-linear decoder. This additional analysis highlights the robustness and versatility of our model's SFs in both egocentric and allocentric observations, which other baseline models did not achieve.

Regarding the concerns raised by Reviewer 2gcv about the similarity between our work and that of Ma et al., 2020 [1], we would like to **reaffirm the novel contributions our work makes: (1) Our approach does not require the agent to be provided with a goal—it is learned; (2) We provide direct evidence that our method works with pixel inputs; (3) We demonstrate that our approach eliminates the need for an SF loss; and (4) By removing the SF loss, we reduce the number of hyperparameters required, thereby simplifying the model.**

In the manuscript (Figures 17 to 21 in Appendix H), we included results comparing our approach to an agent (named "SF + Q-TD + Reward") that combines the Q-learning loss, SF loss, and reward prediction loss, similar to the setup in [1], with the key difference being the inclusion of reward prediction loss, as we do not assume prior knowledge of task specifics. The results showed that our model achieved better learning performance, especially in the continual learning setting (Figure 6a)., and was more computationally efficient (Figure 6b).

Another analysis we included in the general response was to *verify the importance of using a stop-gradient operator on the basis features* when learning the task-encoding vector with the reward prediction loss. We believe that incorporating the stop-gradient operator in our approach enables more effective learning. This is consistent with the findings in [1], which suggest that without the stop-gradient operator on the basis features, concurrently learning both the basis features and the task-encoding vector using the reward prediction loss can lead to learning difficulties.

Although Reviewer Tczf raised concerns about DQN outperforming SFs in the continual RL setting in Figure 2e to 2g, we were unable to directly address this as the specific results being referred to were not clear. However, we invite the reviewers to **refer to the plots generated using the total cumulative return in the same setup as Figures 2e to 2g (in our work) in Figure 2a-2c in general response, where it is clearly evident that our approach performed much better in later tasks.**

More concretely, across all our experimental results in the manuscript, including Figures 2, 3, 6, and 12 to 16, as well as in the complex Mujoco environment, our model, Simple SF, consistently outperforms DQN/DDPG, particularly after the first task.

Finally, we acknowledge the concerns regarding the configuration and density of Figure 1, as well as the need for better insights into Proposition 1. We agree with these observations and, as noted in the comments and general response, we will be making modifications to further improve the clarity and presentation of the manuscript.

The positive feedback on our manuscript's clarity and presentation has been very encouraging. We will incorporate the reviewers' suggestions into the final version of our manuscript. We believe these changes, arising from the discussions, will not only refine our manuscript but also emphasize the simplicity and practicality of our framework, paving the way for future research in the field of Successor Features.

We hope that these revisions address the raised concerns and further solidify the contributions of our manuscript.

[1] Ma, et al., 2020. Universal successor features for transfer reinforcement learning
**This is a re-submit as it seems that our earlier previous response did not notify the reviewers via email.**

# 4. Evaluation
Thank you for your comments, and we apologize for any confusion. To clarify, the average episode return is calculated as a moving average over recent episodes—typically the last 100 episodes experienced by the agent in our case. This metric provides a more immediate snapshot of the agent’s recent performance.

In contrast, the cumulative total return is the sum of all returns accumulated from the moment the agent is first exposed to the current task until the end of the evaluation for that task. This metric reflects the overall performance across the entire evaluation period.

These two metrics can show different trends because the moving average episode return emphasizes recent performance, which may fluctuate, while the cumulative total return captures the long-term accumulation of rewards.

**To emphasize why we presented moving average returns per episode instead of cumulative total return plots in our manuscript, it was to demonstrate that we allow learning for the first task to converge before introducing the second and subsequent tasks.**

We hope this explanation resolves the confusion, and we will ensure that these differences are clearly explained in the manuscript.
**This is a re-submit as it seems that our earlier previous response did not notify the reviewers via email.**

Thank you for taking the time to review our rebuttal. We sincerely appreciate your thoughtful comments and are glad to have the opportunity to provide further clarifications. Please don’t hesitate to reach out if you have any additional questions or concerns.

We appreciate your observation, but **we are unclear which specific plot in Figure 2 you are referring to**, as in all Continual RL plots (Figures 2e to 2g in our paper), our approach (orange) consistently outperforms DQN (blue). Additionally, if you refer to the plots generated using the *total cumulative return* in the same setup as Figures 2e to 2g, as shown in **Figures 2a to 2c in the general response**, it is clearly evident that our approach performed much better in the later tasks.

**To emphasize why we presented (moving) average returns per episode instead of cumulative total return plots in our manuscript, it was to demonstrate that we allow learning for the first task to converge before introducing the second and subsequent tasks.**

Furthermore, we acknowledge that the smaller size of these figures might make the trends less apparent. Therefore, we encourage you to refer to the larger illustrations in Appendix G (Figures 12 to 16), where the replay buffer is not reset to simulate conditions with less interference between task switches. Even under these conditions, our approach (orange) consistently demonstrates superior learning performance compared to DQN (blue).

While the performance improvements in the simpler 2D minigrid environments (Center-Wall and Inverted-LWalls) are less pronounced, they remain significant. In contrast, the more complex 3D Four Rooms environment shows a clearer advantage of our method, as seen in Figures 12 and 13. This trend highlights the robustness of our approach, particularly as task complexity increases, further validating the effectiveness of our method across diverse environments.

Moreover, the newly added results during the rebuttal phase, which utilize the more complex Mujoco environment, also show that our method (orange) outperforms DDPG (blue), a variant of DQN designed for continuous actions.

All these results clearly demonstrate that our method, Simple SF (orange), learns more effectively than DQN and DDPG (blue). This superior performance is due to our method's ability to better generalize and transfer knowledge between tasks, as evidenced by the larger improvements in cumulative total returns when the agent re-encounters the tasks (Exposure 2 in Figure 2 in the General Response).
**This is a re-submit as it seems that our earlier previous response did not notify the reviewers via email.**

Thank you for taking the time to review our rebuttal. We sincerely appreciate your thoughtful comments and are glad to have the opportunity to provide further clarifications. Below, we respond to your excellent points. Please don’t hesitate to reach out if you have any additional questions or concerns.

# 1. Differences
## 1.1 Basis features

First, we see now that we made a mistake in our rebuttal, indeed, in [1] the basis features are learned. As well, the reviewer is correct that, in principle, there is no reason that the approach in [1] could not be applied to pixels. However, to the best of our knowledge, in the paper itself, the authors did not perform experiments and studies involving pixel-based observations. Instead, experiments in [1] were conducted using state inputs, which is likely why their architecture consisting of fully connected networks worked well (Appendix F in [1]). In addition, we ran experiments with the loss from [1] to make a more direct comparison. Please see below for the description of those experiments and the results. 

## 1.3. Direct SF Integration in Q-Learning should eliminate the need for redundant Canonical SF Loss

Thank you for your observation. Indeed, [1] does integrate SFs directly into the Q-value function, similar to our approach, but there is a key difference. **[1] relies on an additional SF loss (Eq. (4) in [1]), known as the Canonical SF-TD loss in our paper, which our method does not require.** Our main contribution lies in the simple (but we believe elegant) architectural design that allows the SFs to be learned directly through the Q-learning loss, eliminating the need for a separate SF loss. 

To highlight the impact of this difference, **we conducted experiments comparing our approach to an agent that combines the Q-learning loss, SF loss, and reward prediction loss, similar to the setup in [1].** Notably, we included the reward prediction loss because, unlike [1], our method does not assume prior knowledge of task specifics, such as goals, which aligns with the expectation in continual learning scenarios. We named this approach “SF + Q-TD + Reward.”

Our results, presented in Figure 6a, demonstrate that the additional SF loss can impair learning efficiency, requiring more time steps to converge to a good policy in the complex 3D Four Rooms environment. Furthermore, this approach is significantly less computationally efficient, as shown by the slower computational speed and longer training duration in Figure 6b. For a detailed comparison of learning performance, **please refer to Figures 17 to 21 in Appendix H.**

We believe these findings underscore the advantages of our method, particularly in terms of efficiency and practicality for continual learning. Moreover, they help to illustrate the key differences between the formulation of our approach from [1]. We agree with the reviewer that our method clearly builds on [1] (which was a seminal paper), but we do feel that what we build on this work represents a novel contribution that can help the field to learn SFs more efficiently, as evidenced by our data. 

## 1.4. Avoiding arbitrary weighting adjustments of Separate SF loss

Thank you for your feedback. We appreciate the opportunity to clarify our position.

On reflection, any claim as to whether or not the SF loss plays a “crucial role” should not hinge on something as basic as the weighting term in the loss. Indeed, the lower weighting used in [1] ($\lambda > 0$) *may be a result of differing scales among the losses or potential conflicts between the SF loss and the Q-learning loss.*

But, from a practical point of view, it is fair to say that [1] requires a careful selection of the weighting coefficient $\lambda$. In contrast, our proposal, through careful algorithmic and architectural design, eliminates the need for a separate SF loss and any concerns about its weighting. In doing so, we directly mitigate the potential problems associated with down-weighting the SF loss, ensuring that SFs meaningfully contribute to the agent's performance.

We hope this addresses your concerns. After reading your reply, we feel that our initial rebuttal did not accurately capture the essence of why our contribution is unique and novel relative to [1]. But, as we described above, **the key advantages of our method are that: 1) We do not need to provide the goal to the agent (it is learned); (2) We provide direct evidence that we can learn off of pixel inputs; (3) We show that we do not need to include the SF loss; and (4) By eliminating the need for the SF loss, we reduce the number of hyperparameters required.**
**This is a re-submit as it seems that our earlier previous response did not notify the reviewers via email.**

Thank you for taking the time to review our rebuttal. We sincerely appreciate your thoughtful feedback and the subsequent adjustment in your evaluation.

In addition, we would like to take this opportunity to further clarify certain points, specifically regarding how our proposed method results in learned Successor Features (SFs) and avoids representation collapse as well as Proposition 1.

# Improving Clarification on Overcoming Representation Collapse
We would like to thank the reviewer for engaging constructively with us, their input has been extremely helpful for improving our paper. Below, we propose two additional modifications to the manuscript to address the points raised by the reviewer in their response to our rebuttal.

First, we will incorporate the proof sketch regarding representation collapse into the main text near line 100, where we initially mention the scenario where the basis features $\phi$ may become a constant vector when the loss is minimized.

Second, at the beginning of Section 4, “Proposed Method,” we will emphasize that the key insight from the proof sketch is that preventing representation collapse requires avoiding the scenario where the basis features $\phi$ become a constant vector for all states, which would minimize the loss without contributing to meaningful learning.

Our approach addresses these constraints by not optimizing the basis features $\phi$ within any loss functions used. Instead, we treat the basis features $\phi$ as the normalized output from the encoder, which is learned using the Q-SF-TD loss (Eq. 5). When the basis features $\phi$ are needed to learn the task encoding vector $w$ through the reward prediction loss (Eq. 6), we apply a stop-gradient operator to treat the basis features $\phi$ as a constant. As we will demonstrate in section 7 “Analysis of Efficiency and Efficacy”, this inclusion of a stop-gradient operator is crucial. Without it, learning both the basis features $\phi$ and the task encoding vector $w$ concurrently can lead to learning instability (as we explained to Reviewer 2gcv). 

# Improving Clarification for Proposition 1

Regarding Proposition 1, in re-reading our own text, we must admit that we agree completely with the reviewer. The way that the text jumps from the discussion of representational collapse in lines 113-115, then brings up Proposition 1, it is natural for a reader to assume that Proposition 1 will deal with representational collapse, and yet it doesn’t. We can see now that this would have caused confusion for readers, potentially suggesting a misleading connection between Proposition 1 and representation collapse.

To improve clarity, we propose the following amendments:

1. Add a concluding sentence after line 115 where we will state, “Next, we will clarify how our approach relates to learning SFs, as they are defined mathematically.”

2. Create a new subsection titled “4.1 Bridging Simple SFs and Universal Successor Features,” where we will expand on the insights related to Proposition 1 (expanding the text currently in lines 116-119).

In terms of expanding on the insights related to Proposition 1, we will highlight the fact that Proposition 1 explains why our approach ultimately produces true SFs. Proposition 1 does this by proving that minimizing our losses (Eqs. 5 & 6) also minimizes the canonical SF loss used in Universal Successor Features (Eq. 4). In order to tie this to the previous section, we will also note that our approach minimizes these losses in a manner such that setting the basis features $\phi$ to a constant is not a solution. 
Specifically, we will note in the text that if one sets $\psi = c_2 $ and $\phi = c_1 = (1 - \gamma) c_2$,  then Eqs. 5 & 6 are not minimized, due to the fact that $\hat{y}$ and $R_{t+1}$ are not constants. 

We believe these revisions will significantly enhance the clarity and robustness of our manuscript. Again, we thank you for your insightful feedback, and please feel free to reach out if any further clarification is needed. We hope you will consider raising your score once more based on these responses.
Dear Reviewer nkgZ,

We hope this message finds you well. As today is the final day for the review discussion, We would like to kindly check in to see if our latest response has addressed your concerns. If the clarifications provided have resolved your questions, we would greatly appreciate it if you could update your score accordingly.

Thank you once again for your time and thoughtful feedback throughout this process. Your input has been invaluable, and we look forward to hearing from you soon.
Dear Reviewer Tczf,

We hope this message finds you well. As today is the final day for the review discussion, We would like to kindly check in to see if our latest response has addressed your concerns. If the clarifications provided have resolved your questions, we would greatly appreciate it if you could update your score accordingly.

Thank you once again for your time and thoughtful feedback throughout this process. Your input has been invaluable, and we look forward to hearing from you soon.
Dear Reviewer 2gcv,

We hope this message finds you well. As today is the final day for the review discussion, We would like to kindly check in to see if our latest response has addressed your concerns. If the clarifications provided have resolved your questions, we would greatly appreciate it if you could update your score accordingly.

Thank you once again for your time and thoughtful feedback throughout this process. Your input has been invaluable, and we look forward to hearing from you soon.
Dear Reviewer 7tgV,

We hope this message finds you well. As today is the final day for the review discussion, We would like to kindly check in to see if our latest response has addressed your concerns. If the clarifications provided have resolved your questions, we would greatly appreciate it if you could update your score accordingly.

Thank you once again for your time and thoughtful feedback throughout this process. Your input has been invaluable, and we look forward to hearing from you soon.
Thank you for taking the time to review our rebuttal. We sincerely appreciate your thoughtful comments and are glad to have the opportunity to provide further clarifications. Please don’t hesitate to reach out if you have any additional questions or concerns.

We appreciate your observation, but **we are unclear which specific plot in Figure 2 you are referring to**, as in all Continual RL plots (Figures 2e to 2g in our paper), our approach (orange) consistently outperforms DQN (blue). Additionally, if you refer to the plots generated using the *total cumulative return* in the same setup as Figures 2e to 2g, as shown in **Figures 2a to 2c in the general response**, it is clearly evident that our approach performed much better in the later tasks.

**To emphasize why we presented (moving) average returns per episode instead of cumulative total return plots in our manuscript, it was to demonstrate that we allow learning for the first task to converge before introducing the second and subsequent tasks.**

Furthermore, we acknowledge that the smaller size of these figures might make the trends less apparent. Therefore, we encourage you to refer to the larger illustrations in Appendix G (Figures 12 to 16), where the replay buffer is not reset to simulate conditions with less interference between task switches. Even under these conditions, our approach (orange) consistently demonstrates superior learning performance compared to DQN (blue).

While the performance improvements in the simpler 2D minigrid environments (Center-Wall and Inverted-LWalls) are less pronounced, they remain significant. In contrast, the more complex 3D Four Rooms environment shows a clearer advantage of our method, as seen in Figures 12 and 13. This trend highlights the robustness of our approach, particularly as task complexity increases, further validating the effectiveness of our method across diverse environments.

Moreover, the newly added results during the rebuttal phase, which utilize the more complex Mujoco environment, also show that our method (orange) outperforms DDPG (blue), a variant of DQN designed for continuous actions.

All these results clearly demonstrate that our method, Simple SF (orange), learns more effectively than DQN and DDPG (blue). This superior performance is due to our method's ability to better generalize and transfer knowledge between tasks, as evidenced by the larger improvements in cumulative total returns when the agent re-encounters the tasks (Exposure 2 in Figure 2 in the General Response).
# 4. Evaluation
Thank you for your comments, and we apologize for any confusion. To clarify, the average episode return is calculated as a moving average over recent episodes—typically the last 100 episodes experienced by the agent in our case. This metric provides a more immediate snapshot of the agent’s recent performance.

In contrast, the cumulative total return is the sum of all returns accumulated from the moment the agent is first exposed to the current task until the end of the evaluation for that task. This metric reflects the overall performance across the entire evaluation period.

These two metrics can show different trends because the moving average episode return emphasizes recent performance, which may fluctuate, while the cumulative total return captures the long-term accumulation of rewards.

**To emphasize why we presented moving average returns per episode instead of cumulative total return plots in our manuscript, it was to demonstrate that we allow learning for the first task to converge before introducing the second and subsequent tasks.**

We hope this explanation resolves the confusion, and we will ensure that these differences are clearly explained in the manuscript.
Thank you for taking the time to review our rebuttal. We sincerely appreciate your thoughtful comments and are glad to have the opportunity to provide further clarifications. Below, we respond to your excellent points. Please don’t hesitate to reach out if you have any additional questions or concerns.

# 1. Differences
## 1.1 Basis features

First, we see now that we made a mistake in our rebuttal, indeed, in [1] the basis features are learned. As well, the reviewer is correct that, in principle, there is no reason that the approach in [1] could not be applied to pixels. However, to the best of our knowledge, in the paper itself, the authors did not perform experiments and studies involving pixel-based observations. Instead, experiments in [1] were conducted using state inputs, which is likely why their architecture consisting of fully connected networks worked well (Appendix F in [1]). In addition, we ran experiments with the loss from [1] to make a more direct comparison. Please see below for the description of those experiments and the results. 

## 1.3. Direct SF Integration in Q-Learning should eliminate the need for redundant Canonical SF Loss

Thank you for your observation. Indeed, [1] does integrate SFs directly into the Q-value function, similar to our approach, but there is a key difference. **[1] relies on an additional SF loss (Eq. (4) in [1]), known as the Canonical SF-TD loss in our paper, which our method does not require.** Our main contribution lies in the simple (but we believe elegant) architectural design that allows the SFs to be learned directly through the Q-learning loss, eliminating the need for a separate SF loss. 

To highlight the impact of this difference, **we conducted experiments comparing our approach to an agent that combines the Q-learning loss, SF loss, and reward prediction loss, similar to the setup in [1].** Notably, we included the reward prediction loss because, unlike [1], our method does not assume prior knowledge of task specifics, such as goals, which aligns with the expectation in continual learning scenarios. We named this approach “SF + Q-TD + Reward.”

Our results, presented in Figure 6a, demonstrate that the additional SF loss can impair learning efficiency, requiring more time steps to converge to a good policy in the complex 3D Four Rooms environment. Furthermore, this approach is significantly less computationally efficient, as shown by the slower computational speed and longer training duration in Figure 6b. For a detailed comparison of learning performance, **please refer to Figures 17 to 21 in Appendix H.**

We believe these findings underscore the advantages of our method, particularly in terms of efficiency and practicality for continual learning. Moreover, they help to illustrate the key differences between the formulation of our approach from [1]. We agree with the reviewer that our method clearly builds on [1] (which was a seminal paper), but we do feel that what we build on this work represents a novel contribution that can help the field to learn SFs more efficiently, as evidenced by our data. 

## 1.4. Avoiding arbitrary weighting adjustments of Separate SF loss

Thank you for your feedback. We appreciate the opportunity to clarify our position.

On reflection, any claim as to whether or not the SF loss plays a “crucial role” should not hinge on something as basic as the weighting term in the loss. Indeed, the lower weighting used in [1] ($\lambda > 0$) *may be a result of differing scales among the losses or potential conflicts between the SF loss and the Q-learning loss.*

But, from a practical point of view, it is fair to say that [1] requires a careful selection of the weighting coefficient $\lambda$. In contrast, our proposal, through careful algorithmic and architectural design, eliminates the need for a separate SF loss and any concerns about its weighting. In doing so, we directly mitigate the potential problems associated with down-weighting the SF loss, ensuring that SFs meaningfully contribute to the agent's performance.

We hope this addresses your concerns. After reading your reply, we feel that our initial rebuttal did not accurately capture the essence of why our contribution is unique and novel relative to [1]. But, as we described above, **the key advantages of our method are that: 1) We do not need to provide the goal to the agent (it is learned); (2) We provide direct evidence that we can learn off of pixel inputs; (3) We show that we do not need to include the SF loss; and (4) By eliminating the need for the SF loss, we reduce the number of hyperparameters required.**
Thank you for taking the time to review our rebuttal. We sincerely appreciate your thoughtful feedback and the subsequent adjustment in your evaluation.

In addition, we would like to take this opportunity to further clarify certain points, specifically regarding how our proposed method results in learned Successor Features (SFs) and avoids representation collapse as well as Proposition 1.

# Improving Clarification on Overcoming Representation Collapse
We would like to thank the reviewer for engaging constructively with us, their input has been extremely helpful for improving our paper. Below, we propose two additional modifications to the manuscript to address the points raised by the reviewer in their response to our rebuttal.

First, we will incorporate the proof sketch regarding representation collapse into the main text near line 100, where we initially mention the scenario where the basis features $\phi$ may become a constant vector when the loss is minimized.

Second, at the beginning of Section 4, “Proposed Method,” we will emphasize that the key insight from the proof sketch is that preventing representation collapse requires avoiding the scenario where the basis features $\phi$ become a constant vector for all states, which would minimize the loss without contributing to meaningful learning.

Our approach addresses these constraints by not optimizing the basis features $\phi$ within any loss functions used. Instead, we treat the basis features $\phi$ as the normalized output from the encoder, which is learned using the Q-SF-TD loss (Eq. 5). When the basis features $\phi$ are needed to learn the task encoding vector $w$ through the reward prediction loss (Eq. 6), we apply a stop-gradient operator to treat the basis features $\phi$ as a constant. As we will demonstrate in section 7 “Analysis of Efficiency and Efficacy”, this inclusion of a stop-gradient operator is crucial. Without it, learning both the basis features $\phi$ and the task encoding vector $w$ concurrently can lead to learning instability (as we explained to Reviewer 2gcv). 

# Improving Clarification for Proposition 1

Regarding Proposition 1, in re-reading our own text, we must admit that we agree completely with the reviewer. The way that the text jumps from the discussion of representational collapse in lines 113-115, then brings up Proposition 1, it is natural for a reader to assume that Proposition 1 will deal with representational collapse, and yet it doesn’t. We can see now that this would have caused confusion for readers, potentially suggesting a misleading connection between Proposition 1 and representation collapse.

To improve clarity, we propose the following amendments:

1. Add a concluding sentence after line 115 where we will state, “Next, we will clarify how our approach relates to learning SFs, as they are defined mathematically.”

2. Create a new subsection titled “4.1 Bridging Simple SFs and Universal Successor Features,” where we will expand on the insights related to Proposition 1 (expanding the text currently in lines 116-119).

In terms of expanding on the insights related to Proposition 1, we will highlight the fact that Proposition 1 explains why our approach ultimately produces true SFs. Proposition 1 does this by proving that minimizing our losses (Eqs. 5 & 6) also minimizes the canonical SF loss used in Universal Successor Features (Eq. 4). In order to tie this to the previous section, we will also note that our approach minimizes these losses in a manner such that setting the basis features $\phi$ to a constant is not a solution. 
Specifically, we will note in the text that if one sets $\psi = c_2 $ and $\phi = c_1 = (1 - \gamma) c_2$,  then Eqs. 5 & 6 are not minimized, due to the fact that $\hat{y}$ and $R_{t+1}$ are not constants. 

We believe these revisions will significantly enhance the clarity and robustness of our manuscript. Again, we thank you for your insightful feedback, and please feel free to reach out if any further clarification is needed. We hope you will consider raising your score once more based on these responses.
We are very pleased to hear that our response helped to answer the reviewer’s questions, and that the reviewer is inclined to accept our paper. Given this, we wonder if the reviewer would be willing to raise their score to reflect the fact that we addressed their questions and put the score more clearly in the “accept” range.
Thank you for your additional ablations and experiments!

* I still feel it is somewhat strange that in Figure 2, DQN is still outperforming other SR methods.  At some point during task change it even outperform your proposed method. I feel the environments don't demonstrate precisely the effectiveness of your method, when DQN is outperforming other SR methods. 

* The results regarding continual RL are encouraging. 

I will keep my score, but increase my confidence to 4.
I thank the authors for their response. I make note of the agreed changes to the figures and their layout in the general comment, this rebuttal and in the discussion with Reviewer Tczf. These seem like appropriate changes and address my concerns with the figures.

For the points on the motivation of Proposition 1 and the proof sketch for representation collapse. I understood how if the $\phi$ and $\psi$ become constant and the loss is minimized. My point is that the insight of the proposition is missing. For example, the sentence in this rebuttal: "This projection is crucial in Continual RL as it aligns the SFs with different tasks, enabling the agent to adapt more rapidly to varying tasks." is key but even then the insight would come from showing how the projection aligns with the task encoding vector and showing or at least arguing why this helps with downstream tasks. I appreciate that the proof and some insight of this sort is in the appendix, but in reality this is the primary contribution and deeper insight here, and is comes in as a passing comment. Perhaps I am missing the point of the work, but the fact that the proposed method behaves in this way seems crucial. In addition, since the representation collapse strategy to the standard SR is noted as being the primary problem, showing how projecting along the task vector fixes this is key. Relying in the main text only on empirics denies the reader this deeper insight, or at least expects them to go and read the proof and obtain the insight themselves from it.

With respect to the RL preliminaries. I appreciate the authors agreeing to this, as it must feel like a nuisance. However, I think it is in the best interest of the broader readership of NeurIPS.

For the correlation analysis. I thank the authors for correcting me, and agree that mention of this would be helpful in the main text. As I mention this was one of my main concerns and I will be raising my score as a result. I also appreciate the new experiments added to the pdf draft and find this to be compelling evidence. I thank the authors for including the decoding experiment as well and find it convincing.

Ultimately, my lingering concern remains with Proposition 1. Essentially lines 113 to 115 explain to us the outcome of the proposed setup, lines 116 to 118 then state literally the Proposition 1 (given it's proximity to lines 113 to 115 leading me to assume it was more related than it is) and then there is a throw forward to the empirical results. So from lines 113 to 119 where the main purpose of the proposed method is being summarize a reader is told what to think, but never shown it. This makes the entire section fall flat, and since this is the technical punchline it makes the paper fall slightly flat. I believe the experimental results (adjusting for my own misunderstandings and the new results) show the intended meaning and I am confident in the correctness of the claims of the work. But a deeper insight into how the proposed method really results in learned SRs and avoid representation collapse still seems missing. Proposition 1, or something of this nature would likely address this. I would raise my score further if this was addressed in the coming days.

I once again thank the authors for their thoughtful response and new experiments. As my quality concerns were due to an error in my understanding, and this has now been corrected, I will raise my score to a 5 and also increase my confidence. I am also raising the soundness and clarity scores in light of the correction and improved figures. I look forward to further discussion on Proposition 1 if the authors are able.
I thank the authors for the additional results and clarifications. Yet, there are some concerns:

# 1. Differences

1\. [1] also uses a NN model to learn the basis features $\phi$ of a state (see Fig.1 of [1]). Whether the feature extractor is pixel-based or not depends on the task.

3\. [1] also integrated the SFs directly into the Q-value function (Eq.(3) of [1]).

4\. There is no connection between "crucial role" and large weighting. There is no guarantee that larger weighting indicates better performance either, as the scales of different losses can be vastly different across tasks. In fact, I found this argument contradicts point 3 above and also defeats the main point of the current paper. The current paper argues that the canonical SF loss is problematic (Fig.1), but now the rebuttal said that one needs to have a larger weighting for the canonical SF loss so that the SFs can play a "crucial role," whatever that means.

# 4. Evaluation

It is unclear why the average episode returns and the cumulative total returns can show different trends. Isn't the former equal to the latter divided by the number of test/evaluation runs?
Thank you to the authors for their answers to my questions and the additional experiments. This has helped me better understand the work. I am currently inclined to accept this paper and maintain my current scores.