Title: Learning Successor Features the Simple Way
Venue: NeurIPS 2024
Accepted: yes
Keywords: deep reinforcement learning

representation learning

continual reinforcement learning

successor feature

successor representation
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.


Peer reviews:
comments: 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.

STRENGTHS:
The proposed method is simple and easy to implement, which is a significant advantage in practical applications. The effectiveness of the approach is demonstrated in both simple 2D and 3D environments.

WEAKNESSES:
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:
Have the authors tested the proposed method with other RL algorithms and in more complex environments?

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-21

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- The problem is well motivated and the approach offers a simple solution to representational collapse when learning deep successor features. \n- The authors provided a number of insightful ablations. Especially that reconstruction based SF methods have trouble learning a good representation for fully observed settings.\n- The writing is clear and the method is presented in an understandable manner.\n\nWEAKNESSES:\n- 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.\n\n- 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.\n\nQUESTIONS:\n- 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.\n- For the Minigrid environment: Do you learn the successor features from pixels or do you use the built-in symbolic state representations?\n- Since DQN is also robust for the showed environments, I wonder how reward sparsity affects the performance of the different algorithms?\n\nLIMITATIONS:\nThe authors have adequately addressed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- A simple method that is easy to understand, and the presentation is easy to follow\n- Reasonable performance in the experiments\n\nWEAKNESSES:\n- 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.\n- 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.\n- 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.\n\nMinor comment: Eqs.(5)&(6) are for scalars so a norm is unnecessary.\n\nReference:\n\n- Ma, C., Ashley, D.R., Wen, J. and Bengio, Y., 2020. Universal successor features for transfer reinforcement learning. *arXiv preprint arXiv:2001.04025*.\n\nQUESTIONS:\nSee the weakness above.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n# Originality\nThe decoupling of training into separate separate equations is new an intuitive idea. The authors note the inspiration from Liu et. al. (2021) however this algorithm is used as a baseline and clearly out-performed empirically.  Thus, it is clear that the changes made are material and have an impact on model performance.\n\n# Clarity\nThe paper is well written and sections are structured appropriately. Notation is intuitive, consistent and aids understanding. Figure captions are detailed which also aids clarity.\n\n# Quality\nI particularly appreciate some of the additional experiments conducted in support of the algorithm, such as the correlation between the learned SFs and SRs. The core experiments appear sufficiently challenging to separate the proposed algorithm from the baselines, and the baselines which are used are appropriate to challenge the proposed algorithm, The results are interpreted fairly as all algorithms do struggle on at least one domain where the simple SFs do not and consistently perform well.\n\n# Significance\nI do think this work could lead to future work and provide a helpful step in improving SFs and making them more practical. The significance is aided by the originality and simplicity of the approach as it is likely to spur new ideas quickly as a result.\n\nWEAKNESSES:\n# Clarity\nThe 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?\n\nSecondly, 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.\n\n# Quality\nI 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.\n\nQUESTIONS:\nI 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.\n\nIf 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.\n\nLIMITATIONS:\nThe limitations are stated in their own section and a good of consideration given towards the broader impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}