Learning Successor Features the Simple Way

Response to Reviewer nkgZ

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.

Response to Reviewer Tczf

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).

Response to Reviewer Tczf

**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).

Response to Reviewer 2gcv on differences with Ma et al. [1]

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.**

Response to Reviewer 2gcv on Evaluation

# 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.

Response to Reviewer 2gcv on differences with Ma et al. [1]

**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.**

Response to Reviewer 2gcv on Evaluation

**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.

Response to Reviewer 7tgV

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.

Response to Reviewer 7tgV

**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.